Fix ResolvedTopicReference memory leak

Sources/SwiftDocC/Model/Identifier.swift

@@ -1,7 +1,7 @@
 /*
  This source file is part of the Swift.org open source project
 
- Copyright (c) 2021 Apple Inc. and the Swift project authors
+ Copyright (c) 2021-2022 Apple Inc. and the Swift project authors
  Licensed under Apache License v2.0 with Runtime Library Exception
 
  See https://swift.org/LICENSE.txt for license information
@@ -89,7 +89,7 @@ public struct ResolvedTopicReference: Hashable, Codable, Equatable, CustomString
     typealias ReferenceKey = String
 
     /// A synchronized reference cache to store resolved references.
-    static var sharedPool = Synchronized([ReferenceBundleIdentifier: [ReferenceKey: ResolvedTopicReference]]())
+    static var sharedPool = Synchronized([ReferenceBundleIdentifier: [ReferenceKey: Weak<ResolvedTopicReference.Storage>]]())
 
     /// Clears cached references belonging to the bundle with the given identifier.
     /// - Parameter bundleIdentifier: The identifier of the bundle to which the method should clear belonging references.
@@ -161,8 +161,8 @@ public struct ResolvedTopicReference: Hashable, Codable, Equatable, CustomString
             sourceLanguages: sourceLanguages
         )
         let cached = Self.sharedPool.sync { $0[bundleIdentifier]?[key] }
-        if let resolved = cached {
-            self = resolved
+        if let resolved = cached?.value {
+            self = .init(storage: resolved)
             return
         }
 
@@ -174,10 +174,16 @@ public struct ResolvedTopicReference: Hashable, Codable, Equatable, CustomString
         )
 
         // Cache the reference
-        Self.sharedPool.sync { $0[bundleIdentifier, default: [:]][key] = self }
+        Self.sharedPool.sync { sharedPool in
+            sharedPool[bundleIdentifier, default: [:]][key] = Weak(value: self._storage)
+        }
+    }
+
+    fileprivate init(storage: Storage) {
+        self._storage = storage
     }
 
-    private static func cacheKey(
+    fileprivate static func cacheKey(
         urlReadablePath path: String,
         urlReadableFragment fragment: String?,
         sourceLanguages: Set<SourceLanguage>
@@ -427,6 +433,21 @@ public struct ResolvedTopicReference: Hashable, Codable, Equatable, CustomString
             self.pathComponents = self.url.pathComponents
             self.absoluteString = self.url.absoluteString
         }
+
+        deinit {
+            ResolvedTopicReference.sharedPool.sync { sharedPool in
+                let key = ResolvedTopicReference.cacheKey(
+                    urlReadablePath: self.path,
+                    urlReadableFragment: self.fragment,
+                    sourceLanguages: self.sourceLanguages
+                )
+                sharedPool[bundleIdentifier]?.removeValue(forKey: key)
+
+                if sharedPool[bundleIdentifier]?.isEmpty ?? false {
+                    sharedPool.removeValue(forKey: bundleIdentifier)
+                }
+            }
+        }
     }
 }
 

Sources/SwiftDocC/Utility/Weak.swift

@@ -0,0 +1,14 @@
+/*
+ This source file is part of the Swift.org open source project
+
+ Copyright (c) 2022 Apple Inc. and the Swift project authors
+ Licensed under Apache License v2.0 with Runtime Library Exception
+
+ See https://swift.org/LICENSE.txt for license information
+ See https://swift.org/CONTRIBUTORS.txt for Swift project authors
+*/
+
+/// A wrapper that provides weak ownership of a value.
+struct Weak<T: AnyObject> {
+    weak var value: T?
+}

Tests/SwiftDocCTests/Infrastructure/DocumentationContext/DocumentationContextTests.swift

@@ -1,7 +1,7 @@
 /*
  This source file is part of the Swift.org open source project
 
- Copyright (c) 2021 Apple Inc. and the Swift project authors
+ Copyright (c) 2021-2022 Apple Inc. and the Swift project authors
  Licensed under Apache License v2.0 with Runtime Library Exception
 
  See https://swift.org/LICENSE.txt for license information
@@ -2159,11 +2159,35 @@ let expected = """
         XCTAssertNotNil(try context.entity(with: referenceForPath("/Collisions/SharedStruct/iOSVar")))
     }
 
+    func testContextDoesNotLeakReferences() throws {
+        #if os(Linux) || os(Android)
+            throw XCTSkip("'autoreleasepool' is not supported on this platform so this test can not be run reliably.")
+        #endif
+
+        // Verify there is no pool bucket for the bundle we're about to test
+        XCTAssertNil(ResolvedTopicReference.sharedPool.sync({ $0[#function] }))
+
+        try autoreleasepool {
+            let (url, _, _) = try testBundleAndContext(copying: "TestBundle", excludingPaths: [], codeListings: [:], configureBundle: { rootURL in
+                let infoPlistURL = rootURL.appendingPathComponent("Info.plist", isDirectory: false)
+                try! String(contentsOf: infoPlistURL)
+                    .replacingOccurrences(of: "org.swift.docc.example", with: #function)
+                    .write(to: infoPlistURL, atomically: true, encoding: .utf8)
+            })
+
+            try FileManager.default.removeItem(at: url)
+        }
+
+
+        // Verify there is no pool bucket for the bundle after we've released all references to it
+        XCTAssertNil(ResolvedTopicReference.sharedPool.sync({ $0[#function] }))
+    }
+
     func testContextCachesReferences() throws {
         // Verify there is no pool bucket for the bundle we're about to test
         XCTAssertNil(ResolvedTopicReference.sharedPool.sync({ $0[#function] }))
 
-        let (url, _, _) = try testBundleAndContext(copying: "TestBundle", excludingPaths: [], codeListings: [:], configureBundle: { rootURL in
+        let (url, bundle, context) = try testBundleAndContext(copying: "TestBundle", excludingPaths: [], codeListings: [:], configureBundle: { rootURL in
             let infoPlistURL = rootURL.appendingPathComponent("Info.plist", isDirectory: false)
             try! String(contentsOf: infoPlistURL)
                 .replacingOccurrences(of: "org.swift.docc.example", with: #function)
@@ -2195,6 +2219,11 @@ let expected = """
         // Verify creating a new reference added to the ones loaded with the context
         XCTAssertNotEqual(beforeCount, ResolvedTopicReference.sharedPool.sync({ $0[#function]!.count }))
 
+        // Access the bundle and the context so that they're kept in memory and the resolved topic
+        // references are not released from the cache
+        XCTAssertEqual(bundle.identifier, #function)
+        XCTAssertEqual(context.knownPages.count, 39)
+
         // Purge the pool
         ResolvedTopicReference.purgePool(for: #function)
     }

Tests/SwiftDocCTests/Infrastructure/ResolvedTopicReferenceTests.swift

@@ -1,7 +1,7 @@
 /*
  This source file is part of the Swift.org open source project
 
- Copyright (c) 2021 Apple Inc. and the Swift project authors
+ Copyright (c) 2021-2022 Apple Inc. and the Swift project authors
  Licensed under Apache License v2.0 with Runtime Library Exception
 
  See https://swift.org/LICENSE.txt for license information
@@ -78,4 +78,49 @@ class ResolvedTopicReferenceTests: XCTestCase {
             _ = topicReference.absoluteString
         }
     }
+
+    func testResolvedTopicReferenceDoesNotLeakMemory() throws {
+        #if os(Linux) || os(Android)
+            throw XCTSkip("'autoreleasepool' is not supported on this platform so this test can not be run reliably.")
+        #endif
+
+        // Use the name of the test as a unique bundle identifier for this test. Otherwise, when
+        // these tests are run sequentially, this test could fail because of stale topic references
+        // that still exist in other pools of the cache.
+        //
+        // This isn't a concern when running the tests in parallel because each test is run in its
+        // own process.
+        let bundleIdentifier = #function
+
+        autoreleasepool {
+            var topicReference: ResolvedTopicReference? = ResolvedTopicReference(
+                bundleIdentifier: bundleIdentifier,
+                path: "/documentation/path/sub-path",
+                fragment: nil,
+                sourceLanguage: .swift
+            )
+
+            ResolvedTopicReference.sharedPool.sync { sharedPool in
+                XCTAssertNotNil(sharedPool[bundleIdentifier])
+                XCTAssertEqual(
+                    sharedPool[bundleIdentifier]?.values.first?.value?.path,
+                    "/documentation/path/sub-path"
+                )
+            }
+
+            topicReference = nil
+
+            // The below assertion just adds a use of the 'topicReference' variable. Otherwise
+            // we end up with a warning:
+            //
+            //    Variable 'topicReference' was written to, but never read
+
+            XCTAssertNil(topicReference)
+        }
+
+        ResolvedTopicReference.sharedPool.sync { sharedPool in
+            XCTAssertNil(sharedPool[bundleIdentifier]?.values.first?.value)
+            XCTAssertNil(sharedPool[bundleIdentifier])
+        }
+    }
 }