Fix memory leak in cache's reference counting (#1296)

.changeset/breezy-waves-cover.md

@@ -0,0 +1,5 @@
+---
+'houdini': patch
+---
+
+Fix memory leak in cache's reference counting

.github/workflows/canary.yml

@@ -6,7 +6,7 @@ on:
     push:
         branches:
             - next
-            - plugin-runtime-search
+            - error-memory-leak
 
 env:
     CI: true

packages/houdini/src/runtime/cache/cache.ts

@@ -106,6 +106,11 @@ export class Cache {
 
 	// register the provided callbacks with the fields specified by the selection
 	subscribe(spec: SubscriptionSpec, variables: {} = {}) {
+		// if the cache is disabled, dont do anything
+		if (this._internal_unstable.disabled) {
+			return
+		}
+
 		// add the subscribers to every field in the specification
 		return this._internal_unstable.subscriptions.add({
 			parent: spec.parentID || rootID,
@@ -334,7 +339,7 @@ export class Cache {
 
 class CacheInternal {
 	// for server-side requests we need to be able to flag the cache as disabled so we dont write to it
-	private _disabled = false
+	disabled = false
 
 	_config?: ConfigFile
 	storage: InMemoryStorage
@@ -380,10 +385,10 @@ class CacheInternal {
 		this.createComponent = createComponent ?? (() => ({}))
 
 		// the cache should always be disabled on the server, unless we're testing
-		this._disabled = disabled
+		this.disabled = disabled
 		try {
 			if (process.env.HOUDINI_TEST === 'true') {
-				this._disabled = false
+				this.disabled = false
 			}
 		} catch {
 			// if process.env doesn't exist, that's okay just use the normal value
@@ -421,7 +426,7 @@ class CacheInternal {
 		forceStale?: boolean
 	}): FieldSelection[] {
 		// if the cache is disabled, dont do anything
-		if (this._disabled) {
+		if (this.disabled) {
 			return []
 		}
 

packages/houdini/src/runtime/cache/subscription.ts

@@ -375,10 +375,13 @@ export class InMemorySubscriptions {
 		let targets: SubscriptionSpec['set'][] = []
 
 		const subscriber = this.subscribers.get(id)
-		const subscriberField = subscriber?.get(fieldName)
+		if (!subscriber) {
+			return
+		}
+		const subscriberField = subscriber.get(fieldName)
 
 		for (const spec of specs) {
-			const counts = subscriber?.get(fieldName)?.referenceCounts
+			const counts = subscriber.get(fieldName)?.referenceCounts
 
 			// if we dont know this field/set combo, there's nothing to do (probably a bug somewhere)
 			if (!counts?.has(spec.set)) {
@@ -394,6 +397,11 @@ export class InMemorySubscriptions {
 				// remove the reference to the set function
 				counts.delete(spec.set)
 			}
+
+			// if we have no more references to the field, we need to remove it from the map
+			if (counts.size === 0) {
+				subscriber.delete(fieldName)
+			}
 		}
 
 		// we do need to remove the set from the list
@@ -402,6 +410,11 @@ export class InMemorySubscriptions {
 				([{ set }]) => !targets.includes(set)
 			)
 		}
+
+		// if we got this far and there are no subscribers on the field, we need to clean things up
+		if (subscriber.size === 0) {
+			this.subscribers.delete(id)
+		}
 	}
 
 	removeAllSubscribers(id: string, targets?: SubscriptionSpec[], visited: string[] = []) {
@@ -441,4 +454,15 @@ export class InMemorySubscriptions {
 			}
 		}
 	}
+
+	get size() {
+		let size = 0
+		for (const [, nodeCounts] of this.subscribers) {
+			for (const [, { referenceCounts }] of nodeCounts) {
+				size += [...referenceCounts.values()].reduce((size, count) => size + count, 0)
+			}
+		}
+
+		return size
+	}
 }

packages/houdini/src/runtime/cache/tests/subscriptions.test.ts

@@ -2939,3 +2939,144 @@ test('overwrite null value with list', function () {
 		friends: [],
 	})
 })
+
+test('removing all subscribers of a field cleans up reference count object', function () {
+	// instantiate the cache
+	const cache = new Cache(config)
+
+	const selection: SubscriptionSelection = {
+		fields: {
+			viewer: {
+				type: 'User',
+				visible: true,
+				keyRaw: 'viewer',
+				selection: {
+					fields: {
+						id: {
+							type: 'ID',
+							visible: true,
+							keyRaw: 'id',
+						},
+						firstName: {
+							type: 'String',
+							visible: true,
+							keyRaw: 'firstName',
+						},
+					},
+				},
+			},
+		},
+	}
+
+	// add some data to the cache
+	cache.write({
+		selection,
+		data: {
+			viewer: {
+				id: '1',
+				firstName: 'bob',
+			},
+		},
+	})
+
+	// subscribe to the list
+	const spec = {
+		set: vi.fn(),
+		selection,
+		rootType: 'Query',
+	}
+	cache.subscribe(spec)
+
+	// sanity check
+	expect(cache._internal_unstable.subscriptions.size).toEqual(3)
+
+	// remove the subscription
+	cache.unsubscribe(spec)
+
+	// make sure the subscribers object is empty
+	expect(cache._internal_unstable.subscriptions.size).toEqual(0)
+})
+
+test('reference count garbage collection requires totally empty garbage', function () {
+	// instantiate the cache
+	const cache = new Cache(config)
+
+	const selection1: SubscriptionSelection = {
+		fields: {
+			viewer: {
+				type: 'User',
+				visible: true,
+				keyRaw: 'viewer',
+				selection: {
+					fields: {
+						id: {
+							type: 'ID',
+							visible: true,
+							keyRaw: 'id',
+						},
+						firstName: {
+							type: 'String',
+							visible: true,
+							keyRaw: 'firstName',
+						},
+					},
+				},
+			},
+		},
+	}
+
+	const selection2: SubscriptionSelection = {
+		fields: {
+			viewer: {
+				type: 'User',
+				visible: true,
+				keyRaw: 'viewer',
+				selection: {
+					fields: {
+						firstName: {
+							type: 'String',
+							visible: true,
+							keyRaw: 'firstName',
+						},
+					},
+				},
+			},
+		},
+	}
+
+	// add some data to the cache
+	cache.write({
+		selection: selection1,
+		data: {
+			viewer: {
+				id: '1',
+				firstName: 'bob',
+			},
+		},
+	})
+
+	// subscribe to selection 1
+	const spec1 = {
+		set: vi.fn(),
+		selection: selection1,
+		rootType: 'Query',
+	}
+	cache.subscribe(spec1)
+
+	// subscribe to selection 2
+	const spec2 = {
+		set: vi.fn(),
+		selection: selection2,
+		rootType: 'Query',
+	}
+	cache.subscribe(spec2)
+
+	// sanity check
+	expect(cache._internal_unstable.subscriptions.size).toEqual(5)
+
+	// remove the subscription from spec 1 which should clear the subscription on id, but not first name
+	cache.unsubscribe(spec1)
+
+	// make sure the subscribers object is empty
+	expect(cache._internal_unstable.subscriptions.size).toEqual(2)
+})