Keep market subscriptions alive until cancellation

Author: outerlook

src/handlers/subscribe/handler.ts

@@ -320,7 +320,8 @@ export function createSubscribeHandler(deps: SubscribeDeps) {
 		const markStreamClosed = () => {
 			streamClosed = true;
 		};
-		const isStreamClosed = () => streamClosed || call.destroyed;
+		const isStreamClosed = () =>
+			streamClosed || call.cancelled || call.writableEnded;
 		const closeOwnedBroker = (): Promise<void> => {
 			if (ownedBrokerClosePromise) {
 				return ownedBrokerClosePromise;
@@ -337,9 +338,7 @@ export function createSubscribeHandler(deps: SubscribeDeps) {
 			void closeOwnedBroker();
 		};
 
-		call.once("close", markStreamClosed);
 		call.once("cancelled", markStreamClosed);
-		call.once("close", closeOwnedBrokerOnCallEnd);
 		call.once("cancelled", closeOwnedBrokerOnCallEnd);
 		call.once("error", closeOwnedBrokerOnCallEnd);
 		call.once("end", () => {
@@ -800,7 +799,6 @@ export function createSubscribeHandler(deps: SubscribeDeps) {
 				type: subscriptionType,
 			});
 		} finally {
-			call.off("close", closeOwnedBrokerOnCallEnd);
 			call.off("cancelled", closeOwnedBrokerOnCallEnd);
 			call.off("error", closeOwnedBrokerOnCallEnd);
 			await closeOwnedBroker();

test/fixtures/ohlcv-collector-fake-exchange.ts

@@ -6,10 +6,16 @@ class ShutdownTestExchange {
 	enableRateLimit = false;
 	timeout = 0;
 	#keepAlive: ReturnType<typeof setInterval> | undefined;
+	#fetchCount = 0;
 
 	extendExchangeOptions(): void {}
 
 	async fetchOHLCVWs(): Promise<number[][]> {
+		this.#fetchCount += 1;
+		const countPath = process.env.OHLCV_TEST_EXCHANGE_COUNT_PATH;
+		if (countPath) {
+			await Bun.write(countPath, String(this.#fetchCount));
+		}
 		if (!this.#keepAlive) {
 			this.#keepAlive = setInterval(() => {}, 1_000);
 			const activePath = process.env.OHLCV_TEST_EXCHANGE_ACTIVE_PATH;

test/ohlcv-collector-shutdown.test.ts

@@ -11,14 +11,33 @@ async function waitForFile(filePath: string): Promise<void> {
 	throw new Error(`Timed out waiting for ${filePath}`);
 }
 
+async function waitForFetchCount(
+	filePath: string,
+	minimum: number,
+): Promise<number> {
+	for (let attempt = 0; attempt < 200; attempt += 1) {
+		if (await Bun.file(filePath).exists()) {
+			const count = Number(await Bun.file(filePath).text());
+			if (count >= minimum) {
+				return count;
+			}
+		}
+		await Bun.sleep(10);
+	}
+	throw new Error(`Timed out waiting for ${minimum} fetches in ${filePath}`);
+}
+
 async function runShutdownCase(exchangeCloseHangs: boolean): Promise<{
 	exitCode: number;
 	closeMarker: string;
+	countBeforeShutdown: number;
+	countAtShutdown: number;
 	output: string;
 }> {
 	const fixtureId = crypto.randomUUID();
 	const configPath = `/tmp/ohlcv-shutdown-${fixtureId}.json`;
 	const activePath = `/tmp/ohlcv-shutdown-${fixtureId}.active`;
+	const countPath = `/tmp/ohlcv-shutdown-${fixtureId}.count`;
 	const closedPath = `/tmp/ohlcv-shutdown-${fixtureId}.closed`;
 	await Bun.write(
 		configPath,
@@ -38,6 +57,7 @@ async function runShutdownCase(exchangeCloseHangs: boolean): Promise<{
 			CEX_BROKER_OHLCV_COLLECTOR_CONFIG: configPath,
 			CEX_BROKER_OHLCV_ARCHIVE_BOOTSTRAP_LIMIT: "0",
 			OHLCV_TEST_EXCHANGE_ACTIVE_PATH: activePath,
+			OHLCV_TEST_EXCHANGE_COUNT_PATH: countPath,
 			OHLCV_TEST_EXCHANGE_CLOSED_PATH: closedPath,
 			OHLCV_TEST_EXCHANGE_CLOSE_HANG: String(exchangeCloseHangs),
 		},
@@ -49,6 +69,14 @@ async function runShutdownCase(exchangeCloseHangs: boolean): Promise<{
 
 	try {
 		await waitForFile(activePath);
+		const countBeforeShutdown = await waitForFetchCount(countPath, 5);
+		expect(await Bun.file(closedPath).exists()).toBe(false);
+		await Bun.sleep(100);
+		const countAtShutdown = await waitForFetchCount(
+			countPath,
+			countBeforeShutdown + 1,
+		);
+		expect(await Bun.file(closedPath).exists()).toBe(false);
 		child.kill("SIGTERM");
 		const result = await Promise.race([
 			child.exited.then((exitCode) => ({ exitCode })),
@@ -64,6 +92,8 @@ async function runShutdownCase(exchangeCloseHangs: boolean): Promise<{
 		return {
 			exitCode: result.exitCode,
 			closeMarker: await Bun.file(closedPath).text(),
+			countBeforeShutdown,
+			countAtShutdown,
 			output: `${await stdout}\n${await stderr}`,
 		};
 	} finally {
@@ -72,7 +102,7 @@ async function runShutdownCase(exchangeCloseHangs: boolean): Promise<{
 			await child.exited;
 		}
 		await Promise.all(
-			[configPath, activePath, closedPath].map(async (filePath) => {
+			[configPath, activePath, countPath, closedPath].map(async (filePath) => {
 				if (await Bun.file(filePath).exists()) {
 					await Bun.file(filePath).delete();
 				}
@@ -85,6 +115,7 @@ test("entrypoint exits promptly on SIGTERM after an exchange stream opens", asyn
 	const result = await runShutdownCase(false);
 	expect(result.exitCode).toBe(0);
 	expect(result.closeMarker).toBe("closed");
+	expect(result.countAtShutdown).toBeGreaterThan(result.countBeforeShutdown);
 });
 
 test("entrypoint bounds shutdown when an exchange close does not resolve", async () => {

test/subscribe-handler.test.ts

@@ -33,6 +33,7 @@ function createSubscribeCall(
 	};
 	const writeResults = [...(options.writeResults ?? [])];
 	const call = Object.assign(emitter, {
+		cancelled: false,
 		metadata: new grpc.Metadata(),
 		request,
 		getPeer: () => "127.0.0.1:1234",
@@ -65,6 +66,14 @@ function createSubscribeCall(
 	};
 }
 
+function cancelSubscribeCall(
+	call: grpc.ServerWritableStream<SubscribeRequest, SubscribeResponse>,
+): void {
+	call.cancelled = true;
+	call.emit("cancelled", "cancelled");
+	call.destroy();
+}
+
 function nextTick(): Promise<void> {
 	return new Promise((resolve) => setTimeout(resolve, 0));
 }
@@ -140,7 +149,7 @@ async function expectBackpressureWaitsForDrain({
 
 	call.emit("drain");
 	await waitFor(() => controlledWatch.calls.length === 2);
-	call.emit("close");
+	cancelSubscribeCall(call);
 	controlledWatch.resolvers[1]?.(secondValue);
 	await handlerPromise;
 
@@ -179,6 +188,32 @@ function createThrowingExchange(
 }
 
 describe("subscribe handler", () => {
+	test("keeps a subscription active when close fires without cancellation", async () => {
+		const controlledWatch = createControlledWatch();
+		const exchange = {
+			watchTrades: controlledWatch.watch,
+		} as unknown as Exchange;
+		const { call } = createSubscribeCall({
+			cex: "binance",
+			symbol: "BTC/USDT",
+			type: SubscriptionType.TRADES,
+		});
+		const handler = createSubscribeHandler({
+			brokers: createPool(exchange),
+			whitelistIps: ["*"],
+		});
+		const handlerPromise = handler(call);
+
+		await waitFor(() => controlledWatch.calls.length === 1);
+		call.emit("close");
+		controlledWatch.resolvers[0]?.([{ id: "trade-1" }]);
+		await waitFor(() => controlledWatch.calls.length === 2);
+
+		cancelSubscribeCall(call);
+		controlledWatch.resolvers[1]?.([{ id: "trade-2" }]);
+		await handlerPromise;
+	});
+
 	test.each([
 		{
 			type: SubscriptionType.TRADES,
@@ -360,7 +395,7 @@ describe("subscribe handler", () => {
 			// forwarder over the real transport rather than racing the round trip.
 			await waitFor(() => archiver.getStats().flushed >= 1);
 			await waitFor(() => controlledWatch.calls.length >= 2);
-			call.emit("close");
+			cancelSubscribeCall(call);
 			controlledWatch.resolvers[1]?.({
 				bids: [[100, 1.5]],
 				asks: [[101, 2]],
@@ -449,7 +484,7 @@ describe("subscribe handler", () => {
 			// forwarder over the real transport rather than racing the round trip.
 			await waitFor(() => archiver.getStats().flushed >= 1);
 			await waitFor(() => controlledWatch.calls.length >= 2);
-			call.emit("close");
+			cancelSubscribeCall(call);
 			controlledWatch.resolvers[1]?.([[1_700_000_000_000, 1, 2, 0.5, 1.5, 10]]);
 			await handlerPromise;