[Collections] testPopulateTargetTrie sometimes gets stuck (#3522)

Motivation:
`PackageCollectionsStorageTests/testPopulateTargetTrie` gets stuck sometimes because `targetTrieReady` is already set and `callback` is only invoked in the body of `memoize`.

https://ci.swift.org/job/oss-swift-package-ubuntu-20_04-aarch64/79/console

Modifications:
Rearrange logic in `populateTargetTrie` such that `callback` is invoked outside of `memoize` body.

Also, add `populateTargetTrieLock` to prevent more than one thread from calling `memoize` at the same time.

Author: yim-lee

Sources/PackageCollections/Storage/SQLitePackageCollectionsStorage.swift

@@ -45,7 +45,8 @@ final class SQLitePackageCollectionsStorage: PackageCollectionsStorage, Closable
 
     // Targets have in-memory trie in addition to SQLite FTS as optimization
     private let targetTrie = Trie<CollectionPackage>()
-    private var targetTrieReady = ThreadSafeBox<Bool>()
+    private var targetTrieReady: Bool?
+    private let populateTargetTrieLock = Lock()
 
     init(location: SQLite.Location? = nil, configuration: Configuration = .init(), diagnosticsEngine: DiagnosticsEngine? = nil) {
         self.location = location ?? .path(localFileSystem.swiftPMCacheDirectory.appending(components: "package-collection.db"))
@@ -503,7 +504,7 @@ final class SQLitePackageCollectionsStorage: PackageCollectionsStorage, Closable
             var matchingCollections = Set<Model.CollectionIdentifier>()
 
             // Trie is more performant for target search; use it if available
-            if self.targetTrieReady.get() ?? false {
+            if self.populateTargetTrieLock.withLock({ self.targetTrieReady }) ?? false {
                 do {
                     switch type {
                     case .exactMatch:
@@ -761,13 +762,18 @@ final class SQLitePackageCollectionsStorage: PackageCollectionsStorage, Closable
 
     internal func populateTargetTrie(callback: @escaping (Result<Void, Error>) -> Void = { _ in }) {
         DispatchQueue.sharedConcurrent.async(group: nil, qos: .background, flags: .assignCurrentContext) {
-            self.targetTrieReady.memoize {
-                do {
+            do {
+                try self.populateTargetTrieLock.withLock { // Prevent race to populate targetTrie
+                    // Exit early if we've already done the computation before
+                    guard self.targetTrieReady == nil else {
+                        return
+                    }
+
                     // since running on low priority thread make sure the database has not already gone away
                     switch (try self.withStateLock { self.state }) {
                     case .disconnected, .disconnecting:
-                        callback(.success(()))
-                        return false
+                        self.targetTrieReady = false
+                        return
                     default:
                         break
                     }
@@ -796,12 +802,11 @@ final class SQLitePackageCollectionsStorage: PackageCollectionsStorage, Closable
                             }
                         }
                     }
-                    callback(.success(()))
-                    return true
-                } catch {
-                    callback(.failure(error))
-                    return false
+                    self.targetTrieReady = true
                 }
+                callback(.success(()))
+            } catch {
+                callback(.failure(error))
             }
         }
     }

Tests/PackageCollectionsTests/PackageCollectionsStorageTests.swift

@@ -1,7 +1,7 @@
 /*
  This source file is part of the Swift.org open source project
 
- Copyright (c) 2020 Apple Inc. and the Swift project authors
+ Copyright (c) 2020-2021 Apple Inc. and the Swift project authors
  Licensed under Apache License v2.0 with Runtime Library Exception
 
  See http://swift.org/LICENSE.txt for license information
@@ -231,10 +231,8 @@ class PackageCollectionsStorageTests: XCTestCase {
             do {
                 try tsc_await { callback in storage2.populateTargetTrie(callback: callback) }
 
-                do {
-                    let searchResult = try tsc_await { callback in storage2.searchTargets(query: targetName, type: .exactMatch, callback: callback) }
-                    XCTAssert(searchResult.items.count > 0, "should get results")
-                }
+                let searchResult = try tsc_await { callback in storage2.searchTargets(query: targetName, type: .exactMatch, callback: callback) }
+                XCTAssert(searchResult.items.count > 0, "should get results")
             } catch {
                 // It's possible that some platforms don't have support FTS
                 XCTAssertEqual(false, storage2.useSearchIndices.get(), "populateTargetTrie should fail only if FTS is not available")