diff --git a/src/main/runtime/rpc/core.ts b/src/main/runtime/rpc/core.ts index c59c8da64d3..0eae5f48bb3 100644 --- a/src/main/runtime/rpc/core.ts +++ b/src/main/runtime/rpc/core.ts @@ -67,6 +67,8 @@ export type RpcContext = { signal?: AbortSignal // Why: per-WebSocket key so the server reaps a closing socket's subscriptions without touching sibling sockets sharing the deviceToken. connectionId?: string + // An unsubscribe cannot retire a registration created after its dispatch began. + subscriptionRegistrationVersion?: number // Why: shared-control multiplexes many logical streams over one socket; the frame id lets handlers register cleanup per logical stream. requestId?: string // Why: paired mobile device token; state-owning handlers use it to clean up when that device disconnects. diff --git a/src/main/runtime/rpc/dispatcher.ts b/src/main/runtime/rpc/dispatcher.ts index 2c407207197..6ed2f1fd9eb 100644 --- a/src/main/runtime/rpc/dispatcher.ts +++ b/src/main/runtime/rpc/dispatcher.ts @@ -96,6 +96,10 @@ export class RpcDispatcher { runtime: this.runtime, signal: options?.signal, connectionId: options?.connectionId, + subscriptionRegistrationVersion: + request.method === 'terminal.unsubscribe' + ? this.runtime.getSubscriptionRegistrationVersion() + : undefined, requestId: request.id, clientId: options?.clientId, clientKind: options?.clientKind, diff --git a/src/main/runtime/rpc/methods/terminal/terminal-viewport-methods.ts b/src/main/runtime/rpc/methods/terminal/terminal-viewport-methods.ts index 91bdf25d84d..da1121c89df 100644 --- a/src/main/runtime/rpc/methods/terminal/terminal-viewport-methods.ts +++ b/src/main/runtime/rpc/methods/terminal/terminal-viewport-methods.ts @@ -82,12 +82,12 @@ export const TERMINAL_VIEWPORT_METHODS_AFTER_STREAMS = [ defineMethod({ name: 'terminal.unsubscribe', params: TerminalUnsubscribe, - handler: async (params, { runtime, connectionId }) => { - // Why: only the connection that owns the subscription may retire it — a stale - // unsubscribe from the pre-reconnect socket names an id the replacement now owns. + handler: async (params, { runtime, connectionId, subscriptionRegistrationVersion }) => { + // Fence both socket replacement and a newer subscription on the same socket. let unsubscribed = runtime.cleanupSubscriptionIfOwnedByConnection( params.subscriptionId, - connectionId + connectionId, + subscriptionRegistrationVersion ) // Why: older builds send a bare-handle subscriptionId, so also try the reconstructed `${terminal}:${clientId}` composite key. // Why AND over the calls that ran: a clientless stream registers under the bare id @@ -97,7 +97,8 @@ export const TERMINAL_VIEWPORT_METHODS_AFTER_STREAMS = [ unsubscribed = runtime.cleanupSubscriptionIfOwnedByConnection( `${params.subscriptionId}:${params.client.id}`, - connectionId + connectionId, + subscriptionRegistrationVersion ) && unsubscribed } return { unsubscribed } diff --git a/src/main/runtime/rpc/rpc-streaming-dispatcher.ts b/src/main/runtime/rpc/rpc-streaming-dispatcher.ts index 6eddcb748be..1428b72d188 100644 --- a/src/main/runtime/rpc/rpc-streaming-dispatcher.ts +++ b/src/main/runtime/rpc/rpc-streaming-dispatcher.ts @@ -65,6 +65,11 @@ export class RpcStreamingDispatcher { if (!isStreamingMethod(method)) { try { + // Capture before middleware yields to a replacement subscribe on the same connection. + const subscriptionRegistrationVersion = + request.method === 'terminal.unsubscribe' + ? runtime.getSubscriptionRegistrationVersion() + : undefined const clientHostedBrowser = await routeDispatcherClientHostedBrowserRpc( runtime, request.method, @@ -109,6 +114,7 @@ export class RpcStreamingDispatcher { signal: options?.signal, requestId: request.id, connectionId: options?.connectionId, + subscriptionRegistrationVersion, clientId: options?.clientId, pairedDeviceId: options?.pairedDeviceId, clientKind: options?.clientKind, diff --git a/src/main/runtime/rpc/subscription-registry-test-double.ts b/src/main/runtime/rpc/subscription-registry-test-double.ts index b2618f6ddc3..d74e8bc329f 100644 --- a/src/main/runtime/rpc/subscription-registry-test-double.ts +++ b/src/main/runtime/rpc/subscription-registry-test-double.ts @@ -1,6 +1,7 @@ import type { SubscriptionRegistration } from '../orca-runtime' type Cleanup = () => void | Promise +type Entry = { cleanup: Cleanup; version: number } export type SubscriptionRegistryDouble = { registerSubscriptionCleanup: (id: string, cleanup: Cleanup, connectionId?: string) => void @@ -10,7 +11,12 @@ export type SubscriptionRegistryDouble = { connectionId?: string ) => SubscriptionRegistration cleanupSubscription: (id: string) => void - cleanupSubscriptionIfOwnedByConnection: (id: string, connectionId: string | undefined) => boolean + cleanupSubscriptionIfOwnedByConnection: ( + id: string, + connectionId: string | undefined, + throughVersion?: number + ) => boolean + getSubscriptionRegistrationVersion: () => number cleanupSubscriptionsForConnection: (connectionId: string) => void /** Test-only inspection; the runtime deliberately exposes no such accessor. */ peekCleanup: (id: string) => Cleanup | undefined @@ -32,8 +38,9 @@ export type SubscriptionRegistryDouble = { * subscribe/teardown must use this instead of a bare Map. */ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble { - const cleanups = new Map() - const inFlight = new Map }>() + const cleanups = new Map() + const inFlight = new Map }>() + let registrationVersion = 0 const byConnection = new Map>() const connectionByEntry = new Map() @@ -54,25 +61,25 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble { } const cleanupAndWait = (id: string): Promise => { - const cleanup = cleanups.get(id) - if (!cleanup) { + const entry = cleanups.get(id) + if (!entry) { return Promise.resolve() } // Mirrors cleanupSubscriptionAndWait: join an in-flight attempt for this exact owner. const existing = inFlight.get(id) - if (existing?.cleanup === cleanup) { + if (existing?.entry === entry) { return existing.promise } let result: void | Promise try { - result = cleanup() + result = entry.cleanup() } catch (error) { result = Promise.reject(error) } const promise = Promise.resolve(result) .then(() => { // Only the generation that registered this callback may retire it. - if (cleanups.get(id) !== cleanup) { + if (cleanups.get(id) !== entry) { return } cleanups.delete(id) @@ -83,7 +90,7 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble { inFlight.delete(id) } }) - inFlight.set(id, { cleanup, promise }) + inFlight.set(id, { entry, promise }) return promise } @@ -92,7 +99,7 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble { void cleanupAndWait(id).catch(() => undefined) } - const cleanupOwned = (id: string, expected: Cleanup): void => { + const cleanupOwned = (id: string, expected: Entry): void => { if (cleanups.get(id) !== expected) { return } @@ -103,15 +110,16 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble { id: string, cleanup: Cleanup, connectionId?: string - ): void => { + ): Entry => { const existing = cleanups.get(id) if (existing) { removeIndex(id) cleanupOwned(id, existing) } - cleanups.set(id, cleanup) + const entry = { cleanup, version: ++registrationVersion } + cleanups.set(id, entry) if (!connectionId) { - return + return entry } let set = byConnection.get(connectionId) if (!set) { @@ -120,27 +128,29 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble { } set.add(id) connectionByEntry.set(id, connectionId) + return entry } return { registerSubscriptionCleanup, registerOwnedSubscriptionCleanup: (id, cleanup, connectionId) => { - registerSubscriptionCleanup(id, cleanup, connectionId) + const entry = registerSubscriptionCleanup(id, cleanup, connectionId) return { - releaseIfCurrent: () => cleanupOwned(id, cleanup) + releaseIfCurrent: () => cleanupOwned(id, entry) } }, cleanupSubscription, - cleanupSubscriptionIfOwnedByConnection: (id, connectionId) => { - if (!connectionId) { - cleanupSubscription(id) - return true - } + getSubscriptionRegistrationVersion: () => registrationVersion, + cleanupSubscriptionIfOwnedByConnection: (id, connectionId, throughVersion) => { + const entry = cleanups.get(id) // Mirrors the production early-out: an unregistered id is already gone, not refused. - if (!cleanups.has(id)) { + if (!entry) { return true } - if (connectionByEntry.get(id) !== connectionId) { + if (throughVersion !== undefined && entry.version > throughVersion) { + return false + } + if (connectionId && connectionByEntry.get(id) !== connectionId) { return false } cleanupSubscription(id) @@ -163,6 +173,6 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble { byConnection.delete(connectionId) } }, - peekCleanup: (id) => cleanups.get(id) + peekCleanup: (id) => cleanups.get(id)?.cleanup } } diff --git a/src/main/runtime/rpc/terminal-subscribe-ownership.test.ts b/src/main/runtime/rpc/terminal-subscribe-ownership.test.ts index 9e036b879ce..d59287639b6 100644 --- a/src/main/runtime/rpc/terminal-subscribe-ownership.test.ts +++ b/src/main/runtime/rpc/terminal-subscribe-ownership.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from 'vitest' import type { OrcaRuntimeService } from '../orca-runtime' +import { RuntimeSubscriptionRegistry } from '../runtime-subscription-registry' import type { RpcRequest } from './core' import { RpcDispatcher } from './dispatcher' import { TERMINAL_METHODS } from './methods/terminal' @@ -14,6 +15,7 @@ function stubRuntime( waiters: Waiter[], overrides: Record = {} ): OrcaRuntimeService { + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: This partial runtime supplies the terminal RPC methods these tests invoke. return { getRuntimeId: () => 'test-runtime', registerRemoteTerminalViewSubscriber: () => () => {}, @@ -36,6 +38,7 @@ function stubRuntime( registerOwnedSubscriptionCleanup: vi.fn(registry.registerOwnedSubscriptionCleanup), cleanupSubscription: vi.fn(registry.cleanupSubscription), cleanupSubscriptionIfOwnedByConnection: vi.fn(registry.cleanupSubscriptionIfOwnedByConnection), + getSubscriptionRegistrationVersion: registry.getSubscriptionRegistrationVersion, subscribeToPtyExit: vi.fn((_ptyId: string, listener: () => void) => { waiters.push({ resolve: listener }) return vi.fn() @@ -291,6 +294,87 @@ describe('terminal.unsubscribe connection ownership', () => { expect(JSON.parse(replies[0]!).result).toEqual({ unsubscribed: false }) }) + it.each( + [binaryParams, leaseOnlyParams].flatMap((params) => + ['streaming', 'unary'].flatMap((transport) => + [SUBSCRIPTION_ID, 'terminal-1'].map((subscriptionId) => ({ + params, + transport, + subscriptionId + })) + ) + ) + )( + 'keeps the replacement when $transport unsubscribe of $subscriptionId yields', + async ({ params, transport, subscriptionId }) => { + const registry = new RuntimeSubscriptionRegistry() + const runtime = stubRuntime(createSubscriptionRegistryDouble(), [], { + registerOwnedSubscriptionCleanup: registry.registerOwned.bind(registry), + cleanupSubscriptionIfOwnedByConnection: registry.cleanupIfOwnedByConnection.bind(registry), + getSubscriptionRegistrationVersion: registry.getRegistrationVersion.bind(registry) + }) + const dispatcher = new RpcDispatcher({ runtime, methods: TERMINAL_METHODS }) + const connection = streamOptions('conn-a') + const original = dispatcher.dispatchStreaming(makeRequest(params), vi.fn(), connection) + await vi.waitFor(() => expect(runtime.handleMobileSubscribe).toHaveBeenCalledTimes(1)) + + const replies: string[] = [] + const retiring = + transport === 'streaming' + ? dispatcher.dispatchStreaming( + unsubscribeRequest(subscriptionId), + (reply) => replies.push(reply), + { connectionId: 'conn-a' } + ) + : dispatcher + .dispatch(unsubscribeRequest(subscriptionId), { connectionId: 'conn-a' }) + .then((reply) => replies.push(JSON.stringify(reply))) + const replacementMessages: string[] = [] + const replacement = dispatcher.dispatchStreaming( + { ...makeRequest(params), id: 'req-replacement' }, + (reply) => replacementMessages.push(reply), + connection + ) + await retiring + + try { + expect(runtime.handleMobileSubscribe).toHaveBeenCalledTimes(2) + expect(runtime.handleMobileUnsubscribe).toHaveBeenCalledTimes(1) + expect(replacementMessages.some((reply) => JSON.parse(reply).result?.type === 'end')).toBe( + false + ) + expect(JSON.parse(replies[0]!).result).toEqual({ unsubscribed: false }) + } finally { + registry.cleanupForConnection('conn-a') + await Promise.all([original, replacement]) + } + } + ) + + it('rejects malformed unsubscribe params before capturing registration state', async () => { + const registry = createSubscriptionRegistryDouble() + const runtime = stubRuntime(registry, []) + const captureVersion = vi.spyOn(runtime, 'getSubscriptionRegistrationVersion') + const cleanup = vi.fn() + registry.registerSubscriptionCleanup(SUBSCRIPTION_ID, cleanup, 'conn-a') + const dispatcher = new RpcDispatcher({ runtime, methods: TERMINAL_METHODS }) + const replies: string[] = [] + + await dispatcher.dispatchStreaming( + { ...unsubscribeRequest(SUBSCRIPTION_ID), params: { subscriptionId: 42 } }, + (reply) => replies.push(reply), + { connectionId: 'conn-a' } + ) + + expect(JSON.parse(replies[0]!)).toMatchObject({ + ok: false, + error: { code: 'invalid_argument' } + }) + expect(captureVersion).not.toHaveBeenCalled() + expect(cleanup).not.toHaveBeenCalled() + registry.cleanupSubscriptionsForConnection('conn-a') + }) + // T5 same-connection: preservation — the owning connection may still unsubscribe. it('honors an unsubscribe from the connection that owns the subscription', async () => { const registry = createSubscriptionRegistryDouble() diff --git a/src/main/runtime/runtime-service-command-surface.ts b/src/main/runtime/runtime-service-command-surface.ts index 55a1d020a7d..27a5ebdbdf9 100644 --- a/src/main/runtime/runtime-service-command-surface.ts +++ b/src/main/runtime/runtime-service-command-surface.ts @@ -25,6 +25,7 @@ export type RuntimeServiceCommandSurface = { cleanupSubscriptionsByPrefix: RuntimeSubscriptionRegistry['cleanupByPrefix'] cleanupSubscriptionsForConnection: RuntimeSubscriptionRegistry['cleanupForConnection'] cleanupSubscriptionIfOwnedByConnection: RuntimeSubscriptionRegistry['cleanupIfOwnedByConnection'] + getSubscriptionRegistrationVersion: RuntimeSubscriptionRegistry['getRegistrationVersion'] onNotificationDispatched: RuntimeMobileNotificationController['onDispatched'] getMobileNotificationListenerCount: RuntimeMobileNotificationController['getListenerCount'] dispatchMobileNotification: RuntimeMobileNotificationController['dispatch'] @@ -114,6 +115,7 @@ export function installRuntimeServiceCommandSurface( cleanupSubscriptionsForConnection: subscriptions.cleanupForConnection.bind(subscriptions), cleanupSubscriptionIfOwnedByConnection: subscriptions.cleanupIfOwnedByConnection.bind(subscriptions), + getSubscriptionRegistrationVersion: subscriptions.getRegistrationVersion.bind(subscriptions), onNotificationDispatched: notifications.onDispatched.bind(notifications), getMobileNotificationListenerCount: notifications.getListenerCount.bind(notifications), dispatchMobileNotification: notifications.dispatch.bind(notifications), diff --git a/src/main/runtime/runtime-subscription-registry.test.ts b/src/main/runtime/runtime-subscription-registry.test.ts new file mode 100644 index 00000000000..093b855d483 --- /dev/null +++ b/src/main/runtime/runtime-subscription-registry.test.ts @@ -0,0 +1,125 @@ +import { describe, expect, it, vi } from 'vitest' +import { RuntimeSubscriptionRegistry } from './runtime-subscription-registry' + +describe('subscription registration versions', () => { + it.each(['conn-a', undefined])( + 'fences a delayed unsubscribe across replacement (%s)', + async (connectionId) => { + const registry = new RuntimeSubscriptionRegistry() + const originalCleanup = vi.fn() + const replacementCleanup = vi.fn() + registry.register('terminal:generation', originalCleanup, 'conn-a') + const admittedVersion = registry.getRegistrationVersion() + registry.register('terminal:generation', replacementCleanup, 'conn-a') + + expect( + registry.cleanupIfOwnedByConnection('terminal:generation', connectionId, admittedVersion) + ).toBe(false) + await Promise.resolve() + expect(originalCleanup).toHaveBeenCalledTimes(1) + expect(replacementCleanup).not.toHaveBeenCalled() + + expect( + registry.cleanupIfOwnedByConnection( + 'terminal:generation', + connectionId, + registry.getRegistrationVersion() + ) + ).toBe(true) + await registry.cleanupAndWait('terminal:generation') + expect(replacementCleanup).toHaveBeenCalledTimes(1) + } + ) + + it('does not let a missing target at admission cancel a later registration', async () => { + const registry = new RuntimeSubscriptionRegistry() + const admittedVersion = registry.getRegistrationVersion() + const cleanup = vi.fn() + registry.register('terminal:later', cleanup, 'conn-a') + expect(registry.cleanupIfOwnedByConnection('terminal:later', 'conn-a', admittedVersion)).toBe( + false + ) + expect(cleanup).not.toHaveBeenCalled() + await registry.cleanupAndWait('terminal:later') + }) + + it('allows admitted cleanup after unrelated subscriptions register', async () => { + const registry = new RuntimeSubscriptionRegistry() + const cleanup = vi.fn() + registry.register('terminal:first', cleanup, 'conn-a') + const admittedVersion = registry.getRegistrationVersion() + registry.register('terminal:other', vi.fn(), 'conn-a') + expect(registry.cleanupIfOwnedByConnection('terminal:first', 'conn-a', admittedVersion)).toBe( + true + ) + await registry.cleanupAndWait('terminal:first') + expect(cleanup).toHaveBeenCalledTimes(1) + await registry.cleanupAndWait('terminal:other') + }) + + it('still refuses a different connection with a current registration version', async () => { + const registry = new RuntimeSubscriptionRegistry() + const cleanup = vi.fn() + registry.register('terminal:owned', cleanup, 'conn-a') + expect( + registry.cleanupIfOwnedByConnection( + 'terminal:owned', + 'conn-b', + registry.getRegistrationVersion() + ) + ).toBe(false) + expect(cleanup).not.toHaveBeenCalled() + await registry.cleanupAndWait('terminal:owned') + }) + + it('reports an already-retired subscription as gone', async () => { + const registry = new RuntimeSubscriptionRegistry() + registry.register('terminal:retired', vi.fn(), 'conn-a') + const admittedVersion = registry.getRegistrationVersion() + await registry.cleanupAndWait('terminal:retired') + expect(registry.cleanupIfOwnedByConnection('terminal:retired', 'conn-a', admittedVersion)).toBe( + true + ) + }) + + it('keeps registration ownership even when a cleanup callback is reused', async () => { + const registry = new RuntimeSubscriptionRegistry() + const cleanup = vi.fn() + const original = registry.registerOwned('terminal:reused', cleanup, 'conn') + const replacement = registry.registerOwned('terminal:reused', cleanup, 'conn') + original.releaseIfCurrent() + await Promise.resolve() + expect(cleanup).toHaveBeenCalledTimes(1) + replacement.releaseIfCurrent() + await registry.cleanupAndWait('terminal:reused') + expect(cleanup).toHaveBeenCalledTimes(2) + }) + + it('does not merge different registrations of the same async cleanup callback', async () => { + const registry = new RuntimeSubscriptionRegistry() + const gate = Promise.withResolvers() + const cleanup = vi.fn(() => gate.promise) + registry.register('terminal:async', cleanup, 'conn') + const originalCleanup = registry.cleanupAndWait('terminal:async') + registry.register('terminal:async', cleanup, 'conn') + const replacementCleanup = registry.cleanupAndWait('terminal:async') + expect(cleanup).toHaveBeenCalledTimes(2) + gate.resolve() + await Promise.all([originalCleanup, replacementCleanup]) + }) + + it('does not retry a replaced registration that reused the cleanup callback', async () => { + const registry = new RuntimeSubscriptionRegistry() + const gate = Promise.withResolvers() + const cleanup = vi.fn() + registry.register('terminal:retry', cleanup, 'conn') + registry.retryAfter('terminal:retry', cleanup, gate.promise) + registry.register('terminal:retry', cleanup, 'conn') + gate.resolve() + await Promise.resolve() + await Promise.resolve() + expect(cleanup).toHaveBeenCalledTimes(1) + await registry.cleanupAndWait('terminal:retry') + expect(cleanup).toHaveBeenCalledTimes(2) + }) +}) diff --git a/src/main/runtime/runtime-subscription-registry.ts b/src/main/runtime/runtime-subscription-registry.ts index febab3f85ff..93e04619618 100644 --- a/src/main/runtime/runtime-subscription-registry.ts +++ b/src/main/runtime/runtime-subscription-registry.ts @@ -1,17 +1,24 @@ type SubscriptionCleanup = () => void | Promise +type SubscriptionEntry = { cleanup: SubscriptionCleanup; version: number } + export type SubscriptionRegistration = { releaseIfCurrent(): void } export class RuntimeSubscriptionRegistry { - private readonly cleanups = new Map() + private readonly cleanups = new Map() private readonly cleanupPromises = new Map< string, - { cleanup: SubscriptionCleanup; promise: Promise } + { entry: SubscriptionEntry; promise: Promise } >() private readonly subscriptionsByConnection = new Map>() private readonly connectionBySubscription = new Map() + private registrationVersion = 0 + + getRegistrationVersion(): number { + return this.registrationVersion + } register(subscriptionId: string, cleanup: SubscriptionCleanup, connectionId?: string): void { const existing = this.cleanups.get(subscriptionId) @@ -19,7 +26,7 @@ export class RuntimeSubscriptionRegistry { this.removeConnectionIndex(subscriptionId) this.cleanup(subscriptionId) } - this.cleanups.set(subscriptionId, cleanup) + this.cleanups.set(subscriptionId, { cleanup, version: ++this.registrationVersion }) if (!connectionId) { return } @@ -38,18 +45,23 @@ export class RuntimeSubscriptionRegistry { connectionId?: string ): SubscriptionRegistration { this.register(subscriptionId, cleanup, connectionId) - return { releaseIfCurrent: () => this.cleanupOwned(subscriptionId, cleanup) } + const version = this.registrationVersion + return { releaseIfCurrent: () => this.cleanupOwned(subscriptionId, version) } } - cleanupIfOwnedByConnection(subscriptionId: string, connectionId?: string): boolean { - if (!connectionId) { - this.cleanup(subscriptionId) + cleanupIfOwnedByConnection( + subscriptionId: string, + connectionId?: string, + throughVersion?: number + ): boolean { + const entry = this.cleanups.get(subscriptionId) + if (!entry) { return true } - if (!this.cleanups.has(subscriptionId)) { - return true + if (throughVersion !== undefined && entry.version > throughVersion) { + return false } - if (this.connectionBySubscription.get(subscriptionId) !== connectionId) { + if (connectionId && this.connectionBySubscription.get(subscriptionId) !== connectionId) { return false } this.cleanup(subscriptionId) @@ -63,15 +75,19 @@ export class RuntimeSubscriptionRegistry { } retryAfter(subscriptionId: string, cleanupOwner: SubscriptionCleanup, gate: Promise): void { + const entry = this.cleanups.get(subscriptionId) const failedGeneration = this.cleanupPromises.get(subscriptionId) void gate.then( async () => { - await (failedGeneration?.cleanup === cleanupOwner + if (entry?.cleanup !== cleanupOwner) { + return + } + await (failedGeneration?.entry === entry ? failedGeneration.promise.catch(() => undefined) : undefined) - while (this.cleanups.get(subscriptionId) === cleanupOwner) { + while (this.cleanups.get(subscriptionId) === entry) { const newerGeneration = this.cleanupPromises.get(subscriptionId) - if (newerGeneration?.cleanup === cleanupOwner) { + if (newerGeneration?.entry === entry) { await newerGeneration.promise.catch(() => undefined) continue } @@ -84,23 +100,23 @@ export class RuntimeSubscriptionRegistry { } async cleanupAndWait(subscriptionId: string): Promise { - const cleanup = this.cleanups.get(subscriptionId) - if (!cleanup) { + const entry = this.cleanups.get(subscriptionId) + if (!entry) { return } const inFlight = this.cleanupPromises.get(subscriptionId) - if (inFlight?.cleanup === cleanup) { + if (inFlight?.entry === entry) { return inFlight.promise } let cleanupResult: void | Promise try { - cleanupResult = cleanup() + cleanupResult = entry.cleanup() } catch (error) { cleanupResult = Promise.reject(error) } const promise = Promise.resolve(cleanupResult) .then(() => { - if (this.cleanups.get(subscriptionId) !== cleanup) { + if (this.cleanups.get(subscriptionId) !== entry) { return } this.cleanups.delete(subscriptionId) @@ -111,7 +127,7 @@ export class RuntimeSubscriptionRegistry { this.cleanupPromises.delete(subscriptionId) } }) - this.cleanupPromises.set(subscriptionId, { cleanup, promise }) + this.cleanupPromises.set(subscriptionId, { entry, promise }) return promise } @@ -139,8 +155,8 @@ export class RuntimeSubscriptionRegistry { } } - private cleanupOwned(subscriptionId: string, expectedCleanup: SubscriptionCleanup): void { - if (this.cleanups.get(subscriptionId) !== expectedCleanup) { + private cleanupOwned(subscriptionId: string, expectedVersion: number): void { + if (this.cleanups.get(subscriptionId)?.version !== expectedVersion) { return } this.cleanup(subscriptionId)