diff --git a/packages/protocol/src/broker.ts b/packages/protocol/src/broker.ts index 0eb8e77..64e7b07 100644 --- a/packages/protocol/src/broker.ts +++ b/packages/protocol/src/broker.ts @@ -53,6 +53,7 @@ interface PendingRequest { interface OwnedToken { sessionId: string; localToken: string; + upstreamToken: string; releaseMethod: string; } @@ -108,8 +109,12 @@ interface BrokerConnection { const TOKEN_METHODS = new Map([ ["transaction_v1_broadcast", "transaction_v1_stop"], + ["transactionWatch_v1_submitAndWatch", "transactionWatch_v1_unwatch"], ["statement_subscribeStatement", "statement_unsubscribeStatement"], ]); +const RELEASE_METHODS = new Set(TOKEN_METHODS.values()); +const MAX_EARLY_SUBSCRIPTION_TOKENS = 32; +const MAX_EARLY_SUBSCRIPTION_EVENTS_PER_TOKEN = 16; function isJsonRpcObject( value: unknown, @@ -189,6 +194,10 @@ function cloneWithRewrittenFirstParam( return { ...request, params }; } +function releaseResultFor(method: string): unknown { + return method === "statement_unsubscribeStatement" ? true : null; +} + export interface ChainBrokerManager { connectRemote( genesisHash: string, @@ -227,7 +236,11 @@ class ChainBroker { private readonly sessions = new Map(); private readonly pending = new Map(); private readonly localToOwned = new Map(); - private readonly upstreamToOwned = new Map(); + private readonly upstreamToOwned = new Map>(); + private readonly earlySubscriptions = new Map< + string, + SubscriptionMessage[] + >(); private readonly localFollowTokens = new Map< string, { sessionId: string; followKey: string } @@ -375,10 +388,16 @@ class ChainBroker { /** Rewrite a session-owned token to its upstream token and forward. */ private routeGenericRequest(session: Session, request: JsonRpcRequest): void { + const method = request.method as string; + if (RELEASE_METHODS.has(method)) { + this.routeOwnedReleaseRequest(session, request, method); + return; + } + const rewritten = this.rewriteOwnedToken(session, request); if (rewritten === null) { brokerLog( - `routeGenericRequest: unknown token for session ${session.id}, method=${request.method as string}`, + `routeGenericRequest: unknown token for session ${session.id}, method=${method}`, ); this.sendToSession( session, @@ -397,7 +416,64 @@ class ChainBroker { this.pending.set(upstreamId, { sessionId: session.id, clientId: request.id ?? null, - method: request.method as string, + method, + }); + this.sendUpstream({ ...rewritten, id: upstreamId }); + } + + private routeOwnedReleaseRequest( + session: Session, + request: JsonRpcRequest, + method: string, + ): void { + const params = Array.isArray(request.params) ? request.params : []; + const localToken = typeof params[0] === "string" ? params[0] : null; + const owned = + localToken !== null ? this.localToOwned.get(localToken) : undefined; + if ( + localToken === null || + owned?.sessionId !== session.id || + owned.releaseMethod !== method + ) { + this.sendToSession( + session, + buildJsonRpcError(request.id ?? null, "Unknown subscription/token"), + ); + return; + } + + const upstreamToken = owned.upstreamToken; + const released = this.releaseOwnedToken(localToken, false); + if (released === null) { + this.sendToSession( + session, + buildJsonRpcError(request.id ?? null, "Unknown subscription/token"), + ); + return; + } + + if (!released.lastOwner) { + if (request.id !== undefined) { + this.sendToSession( + session, + buildJsonRpcResult(request.id ?? null, releaseResultFor(method)), + ); + } + return; + } + + const rewritten = cloneWithRewrittenFirstParam(request, upstreamToken); + if (request.id === undefined) { + this.sendUpstream(rewritten); + return; + } + + const upstreamId = `broker:${this.requestCounter.toString(36)}:${session.id}`; + this.requestCounter += 1; + this.pending.set(upstreamId, { + sessionId: session.id, + clientId: request.id ?? null, + method, }); this.sendUpstream({ ...rewritten, id: upstreamId }); } @@ -569,21 +645,7 @@ class ChainBroker { return null; } - const upstreamToken = this.getUpstreamToken(firstParam); - if (upstreamToken === null) { - return null; - } - - return cloneWithRewrittenFirstParam(request, upstreamToken); - } - - private getUpstreamToken(localToken: string): string | null { - for (const [upstreamToken, owned] of this.upstreamToOwned.entries()) { - if (owned.localToken === localToken) { - return upstreamToken; - } - } - return null; + return cloneWithRewrittenFirstParam(request, owned.upstreamToken); } private handleUpstreamMessage(message: unknown): void { @@ -641,9 +703,16 @@ class ChainBroker { const event = result.event; const rawSub = parsed.params?.subscription; const token = typeof rawSub === "string" ? rawSub : "?"; - // Find which session owns this token - const owned = this.upstreamToOwned.get(token); - const sessionTag = owned ? owned.sessionId : "unknown"; + const ownedLocals = this.upstreamToOwned.get(token); + const sessionTag = + ownedLocals !== undefined && ownedLocals.size > 0 + ? [...ownedLocals] + .map( + (localToken) => + this.localToOwned.get(localToken)?.sessionId ?? "?", + ) + .join(",") + : "unknown"; if (event === "newBlock") { brokerLog( `← raw newBlock [${sessionTag}] hash=${String(result.blockHash).slice(0, 18)}… parent=${String(result.parentBlockHash).slice(0, 18)}… token=${token.slice(0, 12)}…`, @@ -716,6 +785,7 @@ class ChainBroker { buildJsonRpcResult(pendingLocal.requestId, pendingLocal.localToken), ); } + this.flushEarlySubscriptions(response.result); return; } @@ -735,10 +805,16 @@ class ChainBroker { const owned: OwnedToken = { sessionId: pending.sessionId, localToken, + upstreamToken: response.result, releaseMethod, }; this.localToOwned.set(localToken, owned); - this.upstreamToOwned.set(response.result, owned); + let localTokens = this.upstreamToOwned.get(response.result); + if (localTokens === undefined) { + localTokens = new Set(); + this.upstreamToOwned.set(response.result, localTokens); + } + localTokens.add(localToken); session.ownedTokens.add(localToken); brokerLog( `Token mapped: ${localToken} ↔ ${response.result} (${pending.method})`, @@ -754,6 +830,9 @@ class ChainBroker { rewritten.result = result; } this.sendToSession(session, rewritten); + if (releaseMethod !== undefined && typeof response.result === "string") { + this.flushEarlySubscriptions(response.result); + } } private handleUpstreamSubscription(message: SubscriptionMessage): void { @@ -818,41 +897,93 @@ class ChainBroker { return; } - const owned = this.upstreamToOwned.get(upstreamToken); - if (!owned) { + const ownedLocals = this.upstreamToOwned.get(upstreamToken); + if (ownedLocals === undefined || ownedLocals.size === 0) { + if (this.hasPendingSubscriptionRequest()) { + this.bufferEarlySubscription(upstreamToken, message); + return; + } brokerLog(`← upstream subscription for unknown token: ${upstreamToken}`); return; } - const session = this.sessions.get(owned.sessionId); - if (session?.connected !== true) { - brokerLog( - `← upstream subscription for disconnected session: ${owned.sessionId}`, - ); - return; - } - const eventResult = message.params?.result; const eventType = isJsonRpcObject(eventResult) ? typeof eventResult.event === "string" ? eventResult.event : "unknown" : "?"; - brokerLog( - `← subscription [${owned.sessionId}] event=${eventType} method=${String(message.method)}`, - ); + const localTokens = [...ownedLocals]; + for (const localToken of localTokens) { + const owned = this.localToOwned.get(localToken); + if (owned === undefined) { + continue; + } - this.sendToSession(session, { - ...message, - params: { - ...message.params, - subscription: owned.localToken, - }, - }); + const session = this.sessions.get(owned.sessionId); + if (session?.connected !== true) { + brokerLog( + `← upstream subscription for disconnected session: ${owned.sessionId}`, + ); + continue; + } + + brokerLog( + `← subscription [${owned.sessionId}] event=${eventType} method=${String(message.method)}`, + ); + + this.sendToSession(session, { + ...message, + params: { + ...message.params, + subscription: owned.localToken, + }, + }); + } if (isJsonRpcObject(eventResult) && eventResult.event === "stop") { - brokerLog(`Token stopped by upstream: ${owned.localToken}`); - this.releaseOwnedToken(owned.localToken, false); + brokerLog(`Token stopped by upstream: ${upstreamToken}`); + for (const localToken of localTokens) { + this.releaseOwnedToken(localToken, false); + } + } + } + + private hasPendingSubscriptionRequest(): boolean { + return [...this.pending.values()].some( + ({ method }) => + method === "chainHead_v1_follow" || TOKEN_METHODS.has(method), + ); + } + + private bufferEarlySubscription( + upstreamToken: string, + message: SubscriptionMessage, + ): void { + let events = this.earlySubscriptions.get(upstreamToken); + if (events === undefined) { + if (this.earlySubscriptions.size >= MAX_EARLY_SUBSCRIPTION_TOKENS) { + const oldestToken = this.earlySubscriptions.keys().next().value; + if (oldestToken !== undefined) { + this.earlySubscriptions.delete(oldestToken); + } + } + events = []; + this.earlySubscriptions.set(upstreamToken, events); + } + if (events.length < MAX_EARLY_SUBSCRIPTION_EVENTS_PER_TOKEN) { + events.push(message); + } + } + + private flushEarlySubscriptions(upstreamToken: string): void { + const events = this.earlySubscriptions.get(upstreamToken); + if (events === undefined) { + return; + } + this.earlySubscriptions.delete(upstreamToken); + for (const event of events) { + this.handleUpstreamSubscription(event); } } @@ -889,46 +1020,49 @@ class ChainBroker { } } - private releaseOwnedToken(localToken: string, notifyUpstream: boolean): void { + private releaseOwnedToken( + localToken: string, + notifyUpstream: boolean, + ): { lastOwner: boolean } | null { const owned = this.localToOwned.get(localToken); if (!owned) { - return; + return null; } this.localToOwned.delete(localToken); const session = this.sessions.get(owned.sessionId); session?.ownedTokens.delete(localToken); - let upstreamTokenToDelete: string | null = null; - for (const [upstreamToken, candidate] of this.upstreamToOwned.entries()) { - if (candidate.localToken === localToken) { - upstreamTokenToDelete = upstreamToken; - break; - } + const localTokens = this.upstreamToOwned.get(owned.upstreamToken); + if (localTokens?.delete(localToken) !== true) { + return null; } - if (upstreamTokenToDelete === null) { - return; + + if (localTokens.size > 0) { + return { lastOwner: false }; } - this.upstreamToOwned.delete(upstreamTokenToDelete); + this.upstreamToOwned.delete(owned.upstreamToken); if (!notifyUpstream) { - return; + return { lastOwner: true }; } this.sendUpstream({ jsonrpc: "2.0", id: `broker-release:${this.requestCounter.toString(36)}`, method: owned.releaseMethod, - params: [upstreamTokenToDelete], + params: [owned.upstreamToken], }); this.requestCounter += 1; + return { lastOwner: true }; } private disconnectUpstream(): void { this.pending.clear(); this.localToOwned.clear(); this.upstreamToOwned.clear(); + this.earlySubscriptions.clear(); this.localFollowTokens.clear(); this.sharedFollows.clear(); this.upstreamFollowTokens.clear(); diff --git a/packages/protocol/tests/broker.test.ts b/packages/protocol/tests/broker.test.ts index c3b5a9e..42b7038 100644 --- a/packages/protocol/tests/broker.test.ts +++ b/packages/protocol/tests/broker.test.ts @@ -201,6 +201,282 @@ describe("createChainBrokerManager", () => { expect(harness.disconnect).toHaveBeenCalledTimes(1); }); + it("releases transactionWatch subscriptions on disconnect", () => { + const harness = createProviderHarness(); + const manager = createChainBrokerManager(() => harness.provider); + const connection = manager.connectRemote("asset-hub", "conn-a", () => {}); + + connection?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 1, + method: "transactionWatch_v1_submitAndWatch", + params: ["0x0102"], + }), + ); + + const upstreamRequest = harness.sent[0] as { id: string; method: string }; + expect(upstreamRequest.method).toBe("transactionWatch_v1_submitAndWatch"); + harness.emit({ jsonrpc: "2.0", id: upstreamRequest.id, result: "up-tx" }); + + connection?.disconnect(); + + const release = harness.sent[1] as { method: string; params: string[] }; + expect(release.method).toBe("transactionWatch_v1_unwatch"); + expect(release.params[0]).toBe("up-tx"); + }); + + it("delivers transactionWatch events received before the subscribe response", () => { + const harness = createProviderHarness(); + const manager = createChainBrokerManager(() => harness.provider); + const messages: string[] = []; + const connection = manager.connectRemote("asset-hub", "conn-a", (message) => + messages.push(message), + ); + + connection?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 1, + method: "transactionWatch_v1_submitAndWatch", + params: ["0x0102"], + }), + ); + + const upstreamRequest = harness.sent[0] as { id: string }; + harness.emit({ + jsonrpc: "2.0", + method: "transactionWatch_v1_watchEvent", + params: { + subscription: "up-tx", + result: { event: "finalized", block: { hash: "0xabc" } }, + }, + }); + expect(messages).toEqual([]); + + harness.emit({ + jsonrpc: "2.0", + id: upstreamRequest.id, + result: "up-tx", + }); + + const response = JSON.parse(messages[0] ?? "{}") as { result: string }; + expect(JSON.parse(messages[1] ?? "{}")).toEqual({ + jsonrpc: "2.0", + method: "transactionWatch_v1_watchEvent", + params: { + subscription: response.result, + result: { event: "finalized", block: { hash: "0xabc" } }, + }, + }); + }); + + it("fans out same-token statement notifications to every local owner", () => { + const harness = createProviderHarness(); + const manager = createChainBrokerManager(() => harness.provider); + const messagesA: string[] = []; + const messagesB: string[] = []; + const connectionA = manager.connectRemote("asset-hub", "conn-a", (m) => + messagesA.push(m), + ); + const connectionB = manager.connectRemote("asset-hub", "conn-b", (m) => + messagesB.push(m), + ); + + connectionA?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 1, + method: "statement_subscribeStatement", + params: [{ matchAll: ["0xtopic"] }], + }), + ); + connectionB?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 2, + method: "statement_subscribeStatement", + params: [{ matchAll: ["0xtopic"] }], + }), + ); + + harness.emit({ + jsonrpc: "2.0", + id: (harness.sent[0] as { id: string }).id, + result: "up-stmt", + }); + harness.emit({ + jsonrpc: "2.0", + id: (harness.sent[1] as { id: string }).id, + result: "up-stmt", + }); + + const localTokenA = (JSON.parse(messagesA[0] ?? "{}") as { result: string }) + .result; + const localTokenB = (JSON.parse(messagesB[0] ?? "{}") as { result: string }) + .result; + expect(localTokenA).not.toBe(localTokenB); + + harness.emit({ + jsonrpc: "2.0", + method: "statement_statement", + params: { + subscription: "up-stmt", + result: { event: "newStatements", data: { statements: [] } }, + }, + }); + + expect(messagesA[1]).toBe( + JSON.stringify({ + jsonrpc: "2.0", + method: "statement_statement", + params: { + subscription: localTokenA, + result: { event: "newStatements", data: { statements: [] } }, + }, + }), + ); + expect(messagesB[1]).toBe( + JSON.stringify({ + jsonrpc: "2.0", + method: "statement_statement", + params: { + subscription: localTokenB, + result: { event: "newStatements", data: { statements: [] } }, + }, + }), + ); + + connectionA?.disconnect(); + expect( + harness.sent.filter( + (message) => + (message as JsonRpcRequest).method === + "statement_unsubscribeStatement", + ), + ).toHaveLength(0); + + connectionB?.disconnect(); + const releases = harness.sent.filter( + (message) => + (message as JsonRpcRequest).method === "statement_unsubscribeStatement", + ); + expect(releases).toHaveLength(1); + expect((releases[0]?.params as unknown[])[0]).toBe("up-stmt"); + }); + + it("ref-counts same-token statement unsubscribe requests", () => { + const harness = createProviderHarness(); + const manager = createChainBrokerManager(() => harness.provider); + const messagesA: string[] = []; + const messagesB: string[] = []; + const connectionA = manager.connectRemote("asset-hub", "conn-a", (m) => + messagesA.push(m), + ); + const connectionB = manager.connectRemote("asset-hub", "conn-b", (m) => + messagesB.push(m), + ); + + connectionA?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 1, + method: "statement_subscribeStatement", + params: [{ matchAll: ["0xtopic"] }], + }), + ); + connectionB?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 2, + method: "statement_subscribeStatement", + params: [{ matchAll: ["0xtopic"] }], + }), + ); + harness.emit({ + jsonrpc: "2.0", + id: (harness.sent[0] as { id: string }).id, + result: "up-stmt", + }); + harness.emit({ + jsonrpc: "2.0", + id: (harness.sent[1] as { id: string }).id, + result: "up-stmt", + }); + const localTokenA = (JSON.parse(messagesA[0] ?? "{}") as { result: string }) + .result; + const localTokenB = (JSON.parse(messagesB[0] ?? "{}") as { result: string }) + .result; + + connectionA?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 10, + method: "statement_unsubscribeStatement", + params: [localTokenA], + }), + ); + expect(JSON.parse(messagesA.at(-1) ?? "{}")).toEqual({ + jsonrpc: "2.0", + id: 10, + result: true, + }); + expect( + harness.sent.filter( + (message) => + (message as JsonRpcRequest).method === + "statement_unsubscribeStatement", + ), + ).toHaveLength(0); + + harness.emit({ + jsonrpc: "2.0", + method: "statement_statement", + params: { + subscription: "up-stmt", + result: { event: "newStatements", data: { statements: [] } }, + }, + }); + expect(messagesA).toHaveLength(2); + expect(messagesB.at(-1)).toBe( + JSON.stringify({ + jsonrpc: "2.0", + method: "statement_statement", + params: { + subscription: localTokenB, + result: { event: "newStatements", data: { statements: [] } }, + }, + }), + ); + + connectionB?.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 11, + method: "statement_unsubscribeStatement", + params: [localTokenB], + }), + ); + const release = harness.sent.at(-1) as { + id: string; + method: string; + params: string[]; + }; + expect(release.method).toBe("statement_unsubscribeStatement"); + expect(release.params[0]).toBe("up-stmt"); + + harness.emit({ + jsonrpc: "2.0", + id: release.id, + result: true, + } as JsonRpcMessage); + expect(JSON.parse(messagesB.at(-1) ?? "{}")).toEqual({ + jsonrpc: "2.0", + id: 11, + result: true, + }); + }); + it("reuses the warm upstream when a new session attaches after every previous one disconnected", () => { const harness = createProviderHarness(); const manager = createChainBrokerManager(() => harness.provider);