Protect DWARFCompileUnit::m_die_array by a new mutex

[jankratochvil]

Summary

I find it is unrelated to DW_TAG_partial_unit. For example LLDB Python Reference says:

consider that lldb can run in a multithreaded environment, and another thread might call the "script" command, changing the value out from under you.

Therefore if BuildAddressRangeTable called ExtractDIEsIfNeeded(false), then another thread started processing data from m_die_array and then the first thread called final ClearDIEs() the second thread would crash, wouldn't it?

DW_TAG_partial_unit for DWZ would get solved even by this patch.

This patch requires D46810.

Some notes of mine:

ExtractDIEsIfNeeded ro/rw, 0->1, 1->n, 0->n
AddUnitDIE rw, called with lock by: DWARFUnit::ExtractDIEsIfNeeded
AppendDIEsWithTag ro, called by ParseCompileUnitFunctions,
                      called by Module::ParseAllDebugSymbols after ParseVariablesForContext using ExtractDIEsIfNeeded(false)
ClearDIEs ro/rw, called by BuildAddressRangeTable, SymbolFileDWARF::Index
GetDIE expand+ro, called by MANY, only using m_die_array after: ExtractDIEsIfNeeded(false)
IndexPrivate ro, called by Index, Index called only after: ExtractDIEsIfNeeded(false)
GetUnitDIEPtrOnly expand+ro, only using m_die_array after: ExtractDIEsIfNeeded(true)
                             called by BuildAddressRangeTable, ParseProducerInfo, GetLanguageType, GetIsOptimized
DIEPtr            expand+ro, only using m_die_array after: ExtractDIEsIfNeeded(false)
                             called by BuildAddressRangeTable, GetFunctionAranges

[clayborg]

source/Plugins/SymbolFile/DWARF/DWARFCompileUnit.cpp

123 ↗

(On Diff #124296)

This isn't initialized and it is captured and used in already_parsed below?

123–137 ↗

(On Diff #124296)

I really can't tell what any of the above changed code is doing besides adding the mutex. Seems like you are trying to emulate a condition variable to halt other threads trying to get access to the DIEs. But m_die_array_finished isn't initialized in the constructor or inlined in the header file and initial_die_array_size isn't initialized so I can't really tell what this is supposed to actually do.

128 ↗

(On Diff #124296)

initial_die_array_size is used uninitialized when calling already_parsed()

131 ↗

(On Diff #124296)

initial_die_array_size is used uninitialized when calling already_parsed()

source/Plugins/SymbolFile/DWARF/DWARFCompileUnit.h

74 ↗

(On Diff #124296)

This isn't initialized anywhere.

[jankratochvil]

source/Plugins/SymbolFile/DWARF/DWARFCompileUnit.cpp

123 ↗

(On Diff #124296)

I am sorry but the missing initialization was just due to a wrong patch splitting. Next patch in the series did initialize it. Anyway it has been reworked now (after dropping whole ExtractDIEsIfNeededKind for its AllDiesButCuDieOnlyForPartialUnits).

123–137 ↗

(On Diff #124296)

Condition variable is not required here I think. m_die_array_finished is initialized before any time it gets read.

[jankratochvil]

It is now reworked without FileOffset and the remapping to unique DIE offsets for each DW_TAG_partial_unit inclusion by DW_TAG_imported_unit, as discussed in: https://lists.llvm.org/pipermail/lldb-commits/Week-of-Mon-20180409/040324.html

[clayborg]

source/Plugins/SymbolFile/DWARF/DWARFUnit.cpp

56 ↗

(On Diff #144576)

Remove the "m_" prefix on this local variable as that indicates it is a member variable.

Might be better to use CleanUp utility from "lldb/Utility/CleanUp.h" instead of a shared pointer to void? Took me a bit to figure out what this was doing. If this is a very accepted C++ way to run a function that is scoped that is fine, but it seems a little unclear what this was doing initially.

[jankratochvil]

Removed the m_die_array_size_atomic_set m_ prefix, used CleanUp.

[jankratochvil]

source/Plugins/SymbolFile/DWARF/DWARFUnit.cpp

56 ↗

(On Diff #144576)

Done, done. I have googled the shared_ptr cleanup for this purpose at StackOverflow. Updated, thanks.

[clayborg]

source/Plugins/SymbolFile/DWARF/DWARFUnit.cpp

335 ↗

(On Diff #144589)

You can't really play with the m_die_array_size_atomic without locking the mutex, which really means m_die_array_size_atomic is not very useful. Seems like there is still a race condition:

• thread 1 calls this function and executes "m_die_array.swap(tmp_array);" but not " m_die_array_size_atomic = 0;"
• thread 2 calls DWARFUnit::ExtractDIEsIfNeeded() and calls "already_parsed()" which returns true because m_die_array_size_atomic is still some valid value
• thread 1 clears the DIE array and sets m_die_array_size_atomic to zero
• thread 2 has no DIEs when it tries to use them

[clayborg]

Much better.

[jankratochvil]

source/Plugins/SymbolFile/DWARF/DWARFUnit.cpp

335 ↗

(On Diff #144589)

I see I did screw up even that, sorry for that. I believe there still could be an atomic_t speedup but given you are fine with the performance of a mutex I have implemented it that way.

[labath]

lldb/trunk/source/Plugins/SymbolFile/DWARF/DWARFUnit.cpp

405–407

You're accessing m_die_array without holding a lock.

[clayborg]

Good catch

[jankratochvil]

Revert "Protect DWARFCompileUnit::m_die_array by a new mutex"
Pavel Labath found this patch is incomplete and racy. I think there needs to be some more mutexes even before considering DW_TAG_partial_unit.
This reverts commit 331229 which was: https://reviews.llvm.org/D40470
GIT commit cc242f2a5e10a8e6f8e11a268d376a76324c1d23

[clayborg]

I think we should expose the llvm::sys::RWMutex and lets clients like BuildAddressRangeTable and SymbolFileDWARF::Index use it and get rid of m_die_array_usecount all together.

source/Plugins/SymbolFile/DWARF/DWARFUnit.h

181 ↗

(On Diff #149037)

Why don't we expose the llvm::sys::RWMutex for BuildAddressRangeTable and SymbolFileDWARF::Index? Then they can just grab the read lock?

[jankratochvil]

source/Plugins/SymbolFile/DWARF/DWARFUnit.h

181 ↗

(On Diff #149037)

Wouldn't you prefer an RAII lock (DWARFUnit::ScopedExtractDIEs) instead of the mutex exposure?

There are two mutexes (m_die_array_mutex and m_die_array_scoped_mutex), I tried it using a single mutex (as you IMO suggest) but then one needs to downgrade exclusive->shared lock there which requires a retry label there so I did not like the single lock (mutex2.patch).

[jankratochvil]

RAII lock DWARFUnit::ScopedExtractDIEs with mutexes m_die_array_mutex, m_die_array_scoped_mutex and m_first_die_mutex.

[jankratochvil]

Provided variants with a new RAII lock guard vs. unchanged caller bool clear_dies variable. And also 1 vs. 2 mutexes for m_die_array. Please choose one, I do not mind any.

• mutexes2raii.patch - RAII lock guard, 2 mutexes - preferred by me, posted here
• mutexes2var.patch - caller variable, 2 mutexes
• mutexes1raii.patch - RAII lock guard, 1 mutex (lock retrying)
• mutexes1var.patch - caller variable, 1 mutex (lock retrying) - IMO this was suggested by you

[jankratochvil]

FYI I have reordered m_cu->m_die_array_scoped_mutex.unlock_shared(); in DWARFUnit::ScopedExtractDIEs::~ScopedExtractDIEs() during commit (it was written this new way in some other of the 4 offered variants).

[jankratochvil]

Thanks for rL333996 fixups.

lldb/trunk/source/Plugins/SymbolFile/DWARF/DWARFUnit.h

 //===-- DWARFUnit.h ---------------------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SymbolFileDWARF_DWARFUnit_h_
 #define SymbolFileDWARF_DWARFUnit_h_
 
 #include "DWARFDIE.h"
 #include "DWARFDebugInfoEntry.h"
 #include "lldb/lldb-enumerations.h"
+#include "llvm/Support/RWMutex.h"
 
 class DWARFUnit;
 class DWARFCompileUnit;
 class NameToDIE;
 class SymbolFileDWARF;
 class SymbolFileDWARFDwo;
 
 typedef std::shared_ptr<DWARFUnit> DWARFUnitSP;
class DWARFUnit {
 
   using die_iterator_range =
       llvm::iterator_range<DWARFDebugInfoEntry::collection::iterator>;
 
 public:
   virtual ~DWARFUnit();
 
   void ExtractUnitDIEIfNeeded();
-  bool ExtractDIEsIfNeeded();
+  void ExtractDIEsIfNeeded();
+
+  class ScopedExtractDIEs {
+    DWARFUnit *m_cu;
+  public:
+    bool m_clear_dies = false;
+    ScopedExtractDIEs(DWARFUnit *cu);
+    ~ScopedExtractDIEs();
+    DISALLOW_COPY_AND_ASSIGN(ScopedExtractDIEs);
+    ScopedExtractDIEs(ScopedExtractDIEs &&rhs);
+    ScopedExtractDIEs &operator=(ScopedExtractDIEs &&rhs);
+  };
+  ScopedExtractDIEs ExtractDIEsScoped();
+
   DWARFDIE LookupAddress(const dw_addr_t address);
   size_t AppendDIEsWithTag(const dw_tag_t tag,
                            DWARFDIECollection &matching_dies,
                            uint32_t depth = UINT32_MAX) const;
   bool Verify(lldb_private::Stream *s) const;
   virtual void Dump(lldb_private::Stream *s) const = 0;
   //------------------------------------------------------------------
   /// Get the data that contains the DIE information for this unit.
public:
   const DWARFAbbreviationDeclarationSet *GetAbbreviations() const;
   dw_offset_t GetAbbrevOffset() const;
   uint8_t GetAddressByteSize() const { return m_addr_size; }
   dw_addr_t GetBaseAddress() const { return m_base_addr; }
   dw_addr_t GetAddrBase() const { return m_addr_base; }
   dw_addr_t GetRangesBase() const { return m_ranges_base; }
   void SetAddrBase(dw_addr_t addr_base, dw_addr_t ranges_base,
                    dw_offset_t base_obj_offset);
-  void ClearDIEs();
   void BuildAddressRangeTable(SymbolFileDWARF *dwarf,
                               DWARFDebugAranges *debug_aranges);
 
   lldb::ByteOrder GetByteOrder() const;
 
   lldb_private::TypeSystem *GetTypeSystem();
 
   const DWARFDebugAranges &GetFunctionAranges();
protected:
   DWARFUnit(SymbolFileDWARF *dwarf);
 
   SymbolFileDWARF *m_dwarf = nullptr;
   std::unique_ptr<SymbolFileDWARFDwo> m_dwo_symbol_file;
   const DWARFAbbreviationDeclarationSet *m_abbrevs = nullptr;
   void *m_user_data = nullptr;
   // The compile unit debug information entry item
   DWARFDebugInfoEntry::collection m_die_array;
+  mutable llvm::sys::RWMutex m_die_array_mutex;
+  // It is used for tracking of ScopedExtractDIEs instances.
+  mutable llvm::sys::RWMutex m_die_array_scoped_mutex;
+  // ScopedExtractDIEs instances should not call ClearDIEsRWLocked()
+  // as someone called ExtractDIEsIfNeeded().
+  std::atomic<bool> m_cancel_scopes;
   // GetUnitDIEPtrOnly() needs to return pointer to the first DIE.
   // But the first element of m_die_array after ExtractUnitDIEIfNeeded()
   // would possibly move in memory after later ExtractDIEsIfNeeded().
   DWARFDebugInfoEntry m_first_die;
+  llvm::sys::RWMutex m_first_die_mutex;
   // A table similar to the .debug_aranges table, but this one points to the
   // exact DW_TAG_subprogram DIEs
   std::unique_ptr<DWARFDebugAranges> m_func_aranges_ap;
   dw_addr_t m_base_addr = 0;
   dw_offset_t m_length = 0;
   uint16_t m_version = 0;
   uint8_t m_addr_size = 0;
   DWARFProducer m_producer = eProducerInvalid;
protected:
   // in the main object file
   dw_offset_t m_base_obj_offset = DW_INVALID_OFFSET;
 
   // Offset of the initial length field.
   dw_offset_t m_offset;
 
 private:
   void ParseProducerInfo();
+  void ExtractDIEsRWLocked();
+  void ClearDIEsRWLocked();
 
   // Get the DWARF unit DWARF debug informration entry. Parse the single DIE
   // if needed.
   const DWARFDebugInfoEntry *GetUnitDIEPtrOnly() {
     ExtractUnitDIEIfNeeded();
+    // m_first_die_mutex is not required as m_first_die is never cleared.
     if (!m_first_die)
       return NULL;
     return &m_first_die;
   }
 
   // Get all DWARF debug informration entries. Parse all DIEs if needed.
   const DWARFDebugInfoEntry *DIEPtr() {
     ExtractDIEsIfNeeded();

lldb/trunk/source/Plugins/SymbolFile/DWARF/DWARFUnit.cpp

 #include "SymbolFileDWARFDwo.h"
 
 using namespace lldb;
 using namespace lldb_private;
 using namespace std;
 
 extern int g_verbose;
 
-DWARFUnit::DWARFUnit(SymbolFileDWARF *dwarf) : m_dwarf(dwarf) {}
+DWARFUnit::DWARFUnit(SymbolFileDWARF *dwarf)
+    : m_dwarf(dwarf), m_cancel_scopes(false) {}
 
 DWARFUnit::~DWARFUnit() {}
 
 //----------------------------------------------------------------------
 // Parses first DIE of a compile unit.
 //----------------------------------------------------------------------
 void DWARFUnit::ExtractUnitDIEIfNeeded() {
+  {
+    llvm::sys::ScopedReader lock(m_first_die_mutex);
+    if (m_first_die)
+      return; // Already parsed
+  }
+  llvm::sys::ScopedWriter lock(m_first_die_mutex);
   if (m_first_die)
     return; // Already parsed
 
   static Timer::Category func_cat(LLVM_PRETTY_FUNCTION);
   Timer scoped_timer(
       func_cat, "%8.8x: DWARFUnit::ExtractUnitDIEIfNeeded()", m_offset);
 
   // Set the offset to that of the first DIE and calculate the start of the
if (offset < GetNextCompileUnitOffset() &&
     return;
   }
 
   ExtractDIEsEndCheck(offset);
 }
 
 //----------------------------------------------------------------------
 // Parses a compile unit and indexes its DIEs if it hasn't already been done.
+// It will leave this compile unit extracted forever.
+//----------------------------------------------------------------------
+void DWARFUnit::ExtractDIEsIfNeeded() {
+  m_cancel_scopes = true;
+
+  {
+    llvm::sys::ScopedReader lock(m_die_array_mutex);
+    if (!m_die_array.empty())
+      return; // Already parsed
+  }
+  llvm::sys::ScopedWriter lock(m_die_array_mutex);
+  if (!m_die_array.empty())
+    return; // Already parsed
+
+  ExtractDIEsRWLocked();
+}
+
+//----------------------------------------------------------------------
+// Parses a compile unit and indexes its DIEs if it hasn't already been done.
+// It will clear this compile unit after returned instance gets out of scope,
+// no other ScopedExtractDIEs instance is running for this compile unit
+// and no ExtractDIEsIfNeeded() has been executed during this ScopedExtractDIEs
+// lifetime.
 //----------------------------------------------------------------------
-bool DWARFUnit::ExtractDIEsIfNeeded() {
+DWARFUnit::ScopedExtractDIEs DWARFUnit::ExtractDIEsScoped() {
+  ScopedExtractDIEs scoped(this);
+
+  {
+    llvm::sys::ScopedReader lock(m_die_array_mutex);
+    if (!m_die_array.empty())
+      return std::move(scoped); // Already parsed
+  }
+  llvm::sys::ScopedWriter lock(m_die_array_mutex);
   if (!m_die_array.empty())
-    return 0; // Already parsed
+    return std::move(scoped); // Already parsed
+
+  // Otherwise m_die_array would be already populated.
+  lldbassert(!m_cancel_scopes);
+
+  ExtractDIEsRWLocked();
+  scoped.m_clear_dies = true;
+  return scoped;
+}
+
+DWARFUnit::ScopedExtractDIEs::ScopedExtractDIEs(DWARFUnit *cu) : m_cu(cu) {
+  lldbassert(m_cu);
+  m_cu->m_die_array_scoped_mutex.lock_shared();
+}
+
+DWARFUnit::ScopedExtractDIEs::~ScopedExtractDIEs() {
+  if (!m_cu)
+    return;
+  m_cu->m_die_array_scoped_mutex.unlock_shared();
+  if (!m_clear_dies || m_cu->m_cancel_scopes)
+    return;
+  // Be sure no other ScopedExtractDIEs is running anymore.
+  llvm::sys::ScopedWriter lock_scoped(m_cu->m_die_array_scoped_mutex);
+  llvm::sys::ScopedWriter lock(m_cu->m_die_array_mutex);
+  if (m_cu->m_cancel_scopes)
+    return;
+  m_cu->ClearDIEsRWLocked();
+}
+
+DWARFUnit::ScopedExtractDIEs::ScopedExtractDIEs(ScopedExtractDIEs &&rhs)
+    : m_cu(rhs.m_cu), m_clear_dies(rhs.m_clear_dies) {
+  rhs.m_cu = nullptr;
+}
+
+DWARFUnit::ScopedExtractDIEs &DWARFUnit::ScopedExtractDIEs::operator=(
+    DWARFUnit::ScopedExtractDIEs &&rhs) {
+  m_cu = rhs.m_cu;
+  rhs.m_cu = nullptr;
+  m_clear_dies = rhs.m_clear_dies;
+  return *this;
+}
+
+//----------------------------------------------------------------------
+// Parses a compile unit and indexes its DIEs, m_die_array_mutex must be
+// held R/W and m_die_array must be empty.
+//----------------------------------------------------------------------
+void DWARFUnit::ExtractDIEsRWLocked() {
+  llvm::sys::ScopedWriter first_die_lock(m_first_die_mutex);
 
   static Timer::Category func_cat(LLVM_PRETTY_FUNCTION);
   Timer scoped_timer(
       func_cat, "%8.8x: DWARFUnit::ExtractDIEsIfNeeded()", m_offset);
 
   // Set the offset to that of the first DIE and calculate the start of the
   // next compilation unit header.
   lldb::offset_t offset = GetFirstDIEOffset();
void DWARFUnit::ExtractDIEsRWLocked() {
   m_die_array.shrink_to_fit();
 
   ExtractDIEsEndCheck(offset);
 
   if (m_dwo_symbol_file) {
     DWARFUnit *dwo_cu = m_dwo_symbol_file->GetCompileUnit();
     dwo_cu->ExtractDIEsIfNeeded();
   }
-
-  return true;
 }
 
 //--------------------------------------------------------------------------
 // Final checks for both ExtractUnitDIEIfNeeded() and ExtractDIEsIfNeeded().
 //--------------------------------------------------------------------------
 void DWARFUnit::ExtractDIEsEndCheck(lldb::offset_t offset) const {
   // Give a little bit of info if we encounter corrupt DWARF (our offset should
   // always terminate at or before the start of the next compilation unit
if (log && log->GetVerbose()) {
     if (m_die_array.empty())
       strm.Printf("error: no DIE for compile unit");
     else
       m_die_array[0].Dump(m_dwarf, this, strm, UINT32_MAX);
     log->PutString(strm.GetString());
   }
 }
 
+// m_die_array_mutex must be already held as read/write.
 void DWARFUnit::AddUnitDIE(const DWARFDebugInfoEntry &cu_die) {
   uint64_t base_addr = cu_die.GetAttributeValueAsAddress(
       m_dwarf, this, DW_AT_low_pc, LLDB_INVALID_ADDRESS);
   if (base_addr == LLDB_INVALID_ADDRESS)
     base_addr = cu_die.GetAttributeValueAsAddress(
         m_dwarf, this, DW_AT_entry_pc, 0);
   SetBaseAddress(base_addr);
 
DWARFDIE DWARFUnit::LookupAddress(const dw_addr_t address) {
   }
   return DWARFDIE();
 }
 
 size_t DWARFUnit::AppendDIEsWithTag(const dw_tag_t tag,
             DWARFDIECollection &dies,
             uint32_t depth) const {
   size_t old_size = dies.Size();
+  {
+    llvm::sys::ScopedReader lock(m_die_array_mutex);
     DWARFDebugInfoEntry::const_iterator pos;
     DWARFDebugInfoEntry::const_iterator end = m_die_array.end();
     for (pos = m_die_array.begin(); pos != end; ++pos) {
       if (pos->Tag() == tag)
         dies.Append(DWARFDIE(this, &(*pos)));
     }
+  }
 
   // Return the number of DIEs added to the collection
   return dies.Size() - old_size;
 }
 
 
 lldb::user_id_t DWARFUnit::GetID() const {
   dw_offset_t local_id =

 
 void DWARFUnit::SetAddrBase(dw_addr_t addr_base,
                             dw_addr_t ranges_base,
                             dw_offset_t base_obj_offset) {
   m_addr_base = addr_base;
   m_ranges_base = ranges_base;
   m_base_obj_offset = base_obj_offset;
 }
 
-void DWARFUnit::ClearDIEs() {
+// It may be called only with m_die_array_mutex held R/W.
+void DWARFUnit::ClearDIEsRWLocked() {
   m_die_array.clear();
   m_die_array.shrink_to_fit();
 
   if (m_dwo_symbol_file)
-    m_dwo_symbol_file->GetCompileUnit()->ClearDIEs();
+    m_dwo_symbol_file->GetCompileUnit()->ClearDIEsRWLocked();
 }
 
 void DWARFUnit::BuildAddressRangeTable(SymbolFileDWARF *dwarf,
                                        DWARFDebugAranges *debug_aranges) {
   // This function is usually called if there in no .debug_aranges section in
   // order to produce a compile unit level set of address ranges that is
   // accurate.
 
if (num_ranges > 0) {
       return; // We got all of our ranges from the DW_AT_ranges attribute
     }
   }
   // We don't have a DW_AT_ranges attribute, so we need to parse the DWARF
 
   // If the DIEs weren't parsed, then we don't want all dies for all compile
   // units to stay loaded when they weren't needed. So we can end up parsing
   // the DWARF and then throwing them all away to keep memory usage down.
-  const bool clear_dies = ExtractDIEsIfNeeded();
+  ScopedExtractDIEs clear_dies(ExtractDIEsScoped());
 
   die = DIEPtr();
   if (die)
     die->BuildAddressRangeTable(dwarf, this, debug_aranges);
 
   if (debug_aranges->GetNumRanges() == num_debug_aranges) {
     // We got nothing from the functions, maybe we have a line tables only
     // situation. Check the line tables and build the arange table from this.
if (sc.comp_unit) {
           const LineTable::FileAddressRanges::Entry &range =
               file_ranges.GetEntryRef(idx);
           debug_aranges->AppendRange(GetOffset(), range.GetRangeBase(),
                                      range.GetRangeEnd());
         }
       }
     }
   }
-
-  // Keep memory down by clearing DIEs if this generate function caused them to
-  // be parsed
-  if (clear_dies)
-    ClearDIEs();
 }
 
 lldb::ByteOrder DWARFUnit::GetByteOrder() const {
   return m_dwarf->GetObjectFile()->GetByteOrder();
 }
 
 TypeSystem *DWARFUnit::GetTypeSystem() {
   if (m_dwarf)

lldb/trunk/source/Plugins/SymbolFile/DWARF/ManualDWARFIndex.cpp

void ManualDWARFIndex::Index() {
   Timer scoped_timer(func_cat, "%p", static_cast<void *>(&debug_info));
 
   const uint32_t num_compile_units = debug_info.GetNumCompileUnits();
   if (num_compile_units == 0)
     return;
 
   std::vector<IndexSet> sets(num_compile_units);
 
-  // std::vector<bool> might be implemented using bit test-and-set, so use
-  // uint8_t instead.
-  std::vector<uint8_t> clear_cu_dies(num_compile_units, false);
+  //----------------------------------------------------------------------
+  // Keep memory down by clearing DIEs for any compile units if indexing
+  // caused us to load the compile unit's DIEs.
+  //----------------------------------------------------------------------
+  std::vector<llvm::Optional<DWARFUnit::ScopedExtractDIEs>>
+      clear_cu_dies(num_compile_units);
   auto parser_fn = [&](size_t cu_idx) {
     DWARFUnit *dwarf_cu = debug_info.GetCompileUnitAtIndex(cu_idx);
     if (dwarf_cu)
       IndexUnit(*dwarf_cu, sets[cu_idx]);
   };
 
   auto extract_fn = [&debug_info, &clear_cu_dies](size_t cu_idx) {
     DWARFUnit *dwarf_cu = debug_info.GetCompileUnitAtIndex(cu_idx);
-    if (dwarf_cu) {
-      // dwarf_cu->ExtractDIEsIfNeeded() will return false if the DIEs
-      // for a compile unit have already been parsed.
-      if (dwarf_cu->ExtractDIEsIfNeeded())
-        clear_cu_dies[cu_idx] = true;
-    }
+    if (dwarf_cu)
+      clear_cu_dies[cu_idx] = dwarf_cu->ExtractDIEsScoped();
   };
 
   // Create a task runner that extracts dies for each DWARF compile unit in a
   // separate thread
   //----------------------------------------------------------------------
   // First figure out which compile units didn't have their DIEs already
   // parsed and remember this.  If no DIEs were parsed prior to this index
   // function call, we are going to want to clear the CU dies after we are
void ManualDWARFIndex::Index() {
   TaskPool::RunTasks([&]() { finalize_fn(&IndexSet::function_basenames); },
                      [&]() { finalize_fn(&IndexSet::function_fullnames); },
                      [&]() { finalize_fn(&IndexSet::function_methods); },
                      [&]() { finalize_fn(&IndexSet::function_selectors); },
                      [&]() { finalize_fn(&IndexSet::objc_class_selectors); },
                      [&]() { finalize_fn(&IndexSet::globals); },
                      [&]() { finalize_fn(&IndexSet::types); },
                      [&]() { finalize_fn(&IndexSet::namespaces); });
-
-  //----------------------------------------------------------------------
-  // Keep memory down by clearing DIEs for any compile units if indexing
-  // caused us to load the compile unit's DIEs.
-  //----------------------------------------------------------------------
-  for (uint32_t cu_idx = 0; cu_idx < num_compile_units; ++cu_idx) {
-    if (clear_cu_dies[cu_idx])
-      debug_info.GetCompileUnitAtIndex(cu_idx)->ClearDIEs();
-  }
 }
 
 void ManualDWARFIndex::IndexUnit(DWARFUnit &unit, IndexSet &set) {
   Log *log = LogChannelDWARF::GetLogIfAll(DWARF_LOG_LOOKUPS);
 
   if (log) {
     m_module.LogMessage(
         log, "ManualDWARFIndex::IndexUnit for compile unit at .debug_info[0x%8.8x]",