Fix a race condition.

The reader lock must be acquired before retrieving a scratch context,
otherwise another thread could still tear it down between the time it
was created and the creation of the reader object.

There are other shortcomings in GetSwiftScratchContext() that are
documented, but not addressed, in this patch to make it easier to review.

Author: adrian-prantl

lldb/include/lldb/Core/SwiftScratchContextReader.h

@@ -52,7 +52,7 @@ class SwiftScratchContextReader {
   TypeSystemSwiftTypeRefForExpressions *m_ts;
 
 public:
-  SwiftScratchContextReader(std::shared_mutex &mutex,
+  SwiftScratchContextReader(std::shared_lock<std::shared_mutex> &&lock,
                             TypeSystemSwiftTypeRefForExpressions &ts);
   SwiftScratchContextReader(const SwiftScratchContextReader &) = delete;
   SwiftScratchContextReader(SwiftScratchContextReader &&other) = default;

lldb/source/Target/Target.cpp

@@ -2716,62 +2716,69 @@ llvm::Optional<SwiftScratchContextReader> Target::GetSwiftScratchContext(
       }
     }
   }
- 
-  auto get_or_create_fallback_context =
-      [&]() -> TypeSystemSwiftTypeRefForExpressions * {
-    ModuleLanguage idx = {lldb_module, lldb::eLanguageTypeSwift};
-    auto cached = m_scratch_typesystem_for_module.find(idx);
-    if (cached != m_scratch_typesystem_for_module.end()) {
-      auto *cached_ts =
-          llvm::cast<TypeSystemSwiftTypeRefForExpressions>(cached->second.get());
-      if (!cached_ts)
-        return nullptr;
+
+  auto get_cached_module_ts =
+      [&](Module *lldb_module) -> TypeSystemSwiftTypeRefForExpressions * {
+    ModuleLanguage key = {lldb_module, lldb::eLanguageTypeSwift};
+    auto cached = m_scratch_typesystem_for_module.find(key);
+    if (cached != m_scratch_typesystem_for_module.end())
+      return llvm::cast<TypeSystemSwiftTypeRefForExpressions>(cached->second.get());
+    return nullptr;
+  };
+
+  auto maybe_create_fallback_context = [&]() {
+    ModuleLanguage key = {lldb_module, lldb::eLanguageTypeSwift};
+    if (auto *cached_ts = get_cached_module_ts(lldb_module)) {
       auto *cached_ast_ctx =
           llvm::dyn_cast_or_null<SwiftASTContextForExpressions>(
               cached_ts->GetSwiftASTContext());
       if (cached_ast_ctx && cached_ast_ctx->HasFatalErrors() &&
           !m_cant_make_scratch_type_system.count(lldb::eLanguageTypeSwift)) {
         DisplayFallbackSwiftContextErrors(cached_ast_ctx);
         // Try again.
-        m_scratch_typesystem_for_module.erase(cached);
-        return nullptr;
+        // FIXME: Shouldn't this continue rather than return?
+        m_scratch_typesystem_for_module.erase(key);
+        if (log)
+          log->Printf("erased module-wide scratch context with errors\n");
+        return;
       }
       if (log)
         log->Printf("returned cached module-wide scratch context\n");
-      return cached_ts;
+      return;
     }
 
     if (!create_on_demand) {
       if (log)
         log->Printf("not allowed to create a new context\n");
-      return nullptr;
+      return;
     }
 
-    // Call for its side effects of establishing the Swift scratch type system.
+    // Call for its side effects of establishing the Swift scratch type
+    // system.
     auto type_system_or_err =
         GetScratchTypeSystemForLanguage(eLanguageTypeSwift, false);
     if (!type_system_or_err) {
       llvm::consumeError(type_system_or_err.takeError());
-      return nullptr;
+      return;
     }
 
     auto &lock = GetSwiftScratchContextLock();
     if (!lock.try_lock()) {
       if (log)
         log->Printf("couldn't acquire scratch context lock\n");
-      return nullptr;
+      return;
     }
     auto unlock = llvm::make_scope_exit([&lock] { lock.unlock(); });
 
-    // With the lock held, get the current scratch type system. This ensures the
-    // current instance is used even in the unlikely event it was changed during
-    // the brief window between the call to `GetScratchTypeSystemForLanguage`
-    // and taking the lock.
+    // With the lock held, get the current scratch type system. This ensures
+    // the current instance is used even in the unlikely event it was changed
+    // during the brief window between the call to
+    // `GetScratchTypeSystemForLanguage` and taking the lock.
     type_system_or_err = m_scratch_type_system_map.GetTypeSystemForLanguage(
         eLanguageTypeSwift, this, false);
     if (!type_system_or_err) {
       llvm::consumeError(type_system_or_err.takeError());
-      return nullptr;
+      return;
     }
 
     if (auto *global_scratch_ctx =
@@ -2785,41 +2792,52 @@ llvm::Optional<SwiftScratchContextReader> Target::GetSwiftScratchContext(
     auto typesystem_sp = std::make_shared<TypeSystemSwiftTypeRefForExpressions>(
         lldb::eLanguageTypeSwift, *this, *lldb_module);
     typesystem_sp->GetSwiftASTContext();
-    m_scratch_typesystem_for_module.insert({idx, typesystem_sp});
+    m_scratch_typesystem_for_module.insert({key, typesystem_sp});
     if (log)
       log->Printf("created module-wide scratch context\n");
-    return typesystem_sp.get();
+    return;
   };
 
-  TypeSystemSwiftTypeRefForExpressions *swift_scratch_ctx = nullptr;
-  if (lldb_module && m_use_scratch_typesystem_per_module)
-    swift_scratch_ctx = get_or_create_fallback_context();
-  if (!swift_scratch_ctx) {
-    if (log)
-      log->Printf("returned project-wide scratch context\n");
+  llvm::Optional<SwiftScratchContextReader> reader;
+  if (lldb_module && m_use_scratch_typesystem_per_module) {
+    maybe_create_fallback_context();
+    std::shared_lock<std::shared_mutex> lock(GetSwiftScratchContextLock());
+    if (auto *cached_ts = get_cached_module_ts(lldb_module)) {
+      reader = SwiftScratchContextReader(std::move(lock), *cached_ts);
+      if (log)
+        log->Printf("returned module-wide scratch context\n");
+    }
+  }
+  // FIXME: Don't return the project-wide context after requesting the
+  // module-wide one.
+  if (!reader) {
+    std::shared_lock<std::shared_mutex> lock(GetSwiftScratchContextLock());
     auto type_system_or_err =
         GetScratchTypeSystemForLanguage(eLanguageTypeSwift, create_on_demand);
-    if (type_system_or_err)
-      swift_scratch_ctx =
-          llvm::cast_or_null<TypeSystemSwiftTypeRefForExpressions>(
-              type_system_or_err->get());
-    else
+    if (type_system_or_err) {
+      if (auto *ts = llvm::cast_or_null<TypeSystemSwiftTypeRefForExpressions>(
+              type_system_or_err->get())) {
+        reader = SwiftScratchContextReader(std::move(lock), *ts);
+        if (log)
+          log->Printf("returned project-wide scratch context\n");
+      }
+    } else
       llvm::consumeError(type_system_or_err.takeError());
   }
 
-  StackFrameSP frame_sp = exe_scope.CalculateStackFrame();
-  if (frame_sp && frame_sp.get() && swift_scratch_ctx) {
-    SymbolContext sc =
-        frame_sp->GetSymbolContext(lldb::eSymbolContextEverything);
-    Status status = swift_scratch_ctx->PerformCompileUnitImports(sc);
-    if (status.Fail())
-      Debugger::ReportError(status.AsCString(), GetDebugger().GetID());
+  if (reader) {
+    // Perform compile unit imports.
+    assert(reader->get());
+    StackFrameSP frame_sp = exe_scope.CalculateStackFrame();
+    if (frame_sp && frame_sp.get()) {
+      SymbolContext sc =
+          frame_sp->GetSymbolContext(lldb::eSymbolContextEverything);
+      Status status = reader->get()->PerformCompileUnitImports(sc);
+      if (status.Fail())
+        Debugger::ReportError(status.AsCString(), GetDebugger().GetID());
+    }
   }
-
-  if (!swift_scratch_ctx)
-    return llvm::None;
-  return SwiftScratchContextReader(GetSwiftScratchContextLock(),
-                                   *swift_scratch_ctx);
+  return reader;
 }
 
 static std::shared_mutex *
@@ -2833,8 +2851,9 @@ GetSwiftScratchContextMutex(const ExecutionContext *exe_ctx) {
 }
 
 SwiftScratchContextReader::SwiftScratchContextReader(
-    std::shared_mutex &mutex, TypeSystemSwiftTypeRefForExpressions &ts)
-    : m_lock(mutex), m_ts(&ts) {}
+    std::shared_lock<std::shared_mutex> &&lock,
+    TypeSystemSwiftTypeRefForExpressions &ts)
+    : m_lock(std::move(lock)), m_ts(&ts) {}
 
 SwiftScratchContextLock::SwiftScratchContextLock(
     const ExecutionContext *exe_ctx) {