[SourceKit] Ensure ASTCache is guarded by mutex

There were a couple of accesses not guarded by
`CacheMtx`, introduce a couple of methods that
guard them, renaming `getASTProducer` while here.

Also make sure we don't ever insert a producer
after it has been purposefully removed by e.g a
close that removes the cached AST.

Author: hamishknight

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -631,7 +631,16 @@ struct SwiftASTManager::Implementation {
     });
   }
 
-  ASTProducerRef getASTProducer(SwiftInvocationRef InvokRef);
+  /// Retrieve the ASTProducer for a given invocation, creating one if needed.
+  ASTProducerRef getOrCreateASTProducer(SwiftInvocationRef InvokRef);
+
+  /// Retrieve the ASTProducer for a given invocation, returning \c nullopt if
+  /// not present.
+  std::optional<ASTProducerRef> getASTProducer(SwiftInvocationRef Invok);
+
+  /// Updates the cache entry to account for any changes to the ASTProducer
+  /// for the given invocation.
+  void updateASTProducer(SwiftInvocationRef Invok);
 
   FileContent
   getFileContent(StringRef FilePath, bool IsPrimary,
@@ -780,7 +789,7 @@ void SwiftASTManager::processASTAsync(
     const void *OncePerASTToken, SourceKitCancellationToken CancellationToken,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fileSystem) {
   assert(fileSystem);
-  ASTProducerRef Producer = Impl.getASTProducer(InvokRef);
+  ASTProducerRef Producer = Impl.getOrCreateASTProducer(InvokRef);
 
   Impl.cleanDeletedConsumers();
   {
@@ -816,12 +825,31 @@ void SwiftASTManager::processASTAsync(
       });
 }
 
+std::optional<ASTProducerRef>
+SwiftASTManager::Implementation::getASTProducer(SwiftInvocationRef Invok) {
+  llvm::sys::ScopedLock L(CacheMtx);
+  return ASTCache.get(Invok->Impl.Key);
+}
+
+void SwiftASTManager::Implementation::updateASTProducer(
+    SwiftInvocationRef Invok) {
+  llvm::sys::ScopedLock L(CacheMtx);
+
+  // Get and set the producer to update its cost in the cache. If we don't
+  // have a value, then this is a race where we've removed the cached AST, but
+  // still have a build waiting to complete after cancellation, we don't need
+  // to do anything in that case.
+  if (auto Producer = ASTCache.get(Invok->Impl.Key))
+    ASTCache.set(Invok->Impl.Key, *Producer);
+}
+
 void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
+  llvm::sys::ScopedLock L(Impl.CacheMtx);
   Impl.ASTCache.remove(Invok->Impl.Key);
 }
 
-ASTProducerRef
-SwiftASTManager::Implementation::getASTProducer(SwiftInvocationRef InvokRef) {
+ASTProducerRef SwiftASTManager::Implementation::getOrCreateASTProducer(
+    SwiftInvocationRef InvokRef) {
   llvm::sys::ScopedLock L(CacheMtx);
   std::optional<ASTProducerRef> OptProducer = ASTCache.get(InvokRef->Impl.Key);
   if (OptProducer.has_value())
@@ -1342,7 +1370,7 @@ void ASTProducer::enqueueConsumer(
               [This]() { This->cleanBuildOperations(); });
           // Re-register the object with the cache to update its memory
           // cost.
-          Mgr->Impl.ASTCache.set(This->InvokRef->Impl.Key, This);
+          Mgr->Impl.updateASTProducer(This->InvokRef);
         }
       };