fix: sse connection

app/client/hooks/__tests__/use-server-events.test.tsx

@@ -0,0 +1,132 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { cleanup, render, screen } from "@testing-library/react";
+import { useEffect, useState } from "react";
+import { useServerEvents } from "../use-server-events";
+
+class MockEventSource {
+	static instances: MockEventSource[] = [];
+
+	public onerror: ((event: Event) => void) | null = null;
+	public close = mock(() => {});
+	private listeners = new Map<string, Set<(event: Event) => void>>();
+
+	constructor(public url: string) {
+		MockEventSource.instances.push(this);
+	}
+
+	addEventListener(type: string, listener: EventListenerOrEventListenerObject) {
+		const listeners = this.listeners.get(type) ?? new Set<(event: Event) => void>();
+		const callback = typeof listener === "function" ? listener : (event: Event) => listener.handleEvent(event);
+		listeners.add(callback);
+		this.listeners.set(type, listeners);
+	}
+
+	emit(type: string, data: unknown) {
+		const event = new MessageEvent(type, {
+			data: JSON.stringify(data),
+		});
+		for (const listener of this.listeners.get(type) ?? []) {
+			listener(event);
+		}
+	}
+
+	static reset() {
+		MockEventSource.instances = [];
+	}
+}
+
+const originalEventSource = globalThis.EventSource;
+const originalConsoleInfo = console.info;
+const originalConsoleError = console.error;
+
+const createTestQueryClient = () =>
+	new QueryClient({
+		defaultOptions: {
+			queries: {
+				retry: false,
+				gcTime: Infinity,
+			},
+			mutations: {
+				gcTime: Infinity,
+			},
+		},
+	});
+
+const ConnectionConsumer = () => {
+	useServerEvents();
+	return null;
+};
+
+const BackupCompletedListener = ({ scheduleId }: { scheduleId: string }) => {
+	const { addEventListener } = useServerEvents();
+	const [status, setStatus] = useState("pending");
+
+	useEffect(() => {
+		const abortController = new AbortController();
+
+		addEventListener(
+			"backup:completed",
+			(event) => {
+				if (event.scheduleId === scheduleId) {
+					setStatus(event.status);
+				}
+			},
+			{ signal: abortController.signal },
+		);
+
+		return () => abortController.abort();
+	}, [addEventListener, scheduleId]);
+
+	return <div>{status}</div>;
+};
+
+describe("useServerEvents", () => {
+	beforeEach(() => {
+		MockEventSource.reset();
+		globalThis.EventSource = MockEventSource as unknown as typeof EventSource;
+		console.info = mock(() => {});
+		console.error = mock(() => {});
+	});
+
+	afterEach(() => {
+		cleanup();
+		globalThis.EventSource = originalEventSource;
+		console.info = originalConsoleInfo;
+		console.error = originalConsoleError;
+		MockEventSource.reset();
+	});
+
+	test("shares one EventSource across consumers and invalidates queries once on backup completion", async () => {
+		const queryClient = createTestQueryClient();
+		const invalidateQueries = mock(async () => undefined);
+		const refetchQueries = mock(async () => undefined);
+		queryClient.invalidateQueries = invalidateQueries as typeof queryClient.invalidateQueries;
+		queryClient.refetchQueries = refetchQueries as typeof queryClient.refetchQueries;
+
+		render(
+			<QueryClientProvider client={queryClient}>
+				<ConnectionConsumer />
+				<BackupCompletedListener scheduleId="0b9c940b" />
+			</QueryClientProvider>,
+		);
+
+		expect(MockEventSource.instances).toHaveLength(1);
+
+		MockEventSource.instances[0]?.emit("backup:completed", {
+			organizationId: "default-org",
+			scheduleId: "0b9c940b",
+			volumeName: "synology",
+			repositoryName: "swiss-backup",
+			status: "success",
+		});
+
+		expect(await screen.findByText("success")).toBeTruthy();
+		expect(invalidateQueries).toHaveBeenCalledTimes(1);
+		expect(refetchQueries).not.toHaveBeenCalled();
+
+		cleanup();
+
+		expect(MockEventSource.instances[0]?.close).toHaveBeenCalledTimes(1);
+	});
+});

app/client/hooks/use-server-events.ts

@@ -1,4 +1,4 @@
-import { useQueryClient } from "@tanstack/react-query";
+import { type QueryClient, useQueryClient } from "@tanstack/react-query";
 import { useCallback, useEffect, useRef } from "react";
 import { serverEventNames, type ServerEventPayloadMap } from "~/schemas/server-events";
 
@@ -16,6 +16,13 @@ type EventHandlerMap = {
 	[K in ServerEventType]?: EventHandlerSet<K>;
 };
 
+type SharedServerEventsState = {
+	eventSource: EventSource | null;
+	handlers: EventHandlerMap;
+	queryClient: QueryClient | null;
+	subscribers: number;
+};
+
 const invalidatingEvents = new Set<ServerEventType>([
 	"backup:completed",
 	"volume:updated",
@@ -31,103 +38,151 @@ export type RestoreProgressEvent = ServerEventsPayloadMap["restore:progress"];
 export type RestoreCompletedEvent = ServerEventsPayloadMap["restore:completed"];
 export type BackupProgressEvent = ServerEventsPayloadMap["backup:progress"];
 
+const sharedState: SharedServerEventsState = {
+	eventSource: null,
+	handlers: {},
+	queryClient: null,
+	subscribers: 0,
+};
+
 const parseEventData = <T extends ServerEventType>(event: Event): ServerEventsPayloadMap[T] =>
 	JSON.parse((event as MessageEvent<string>).data) as ServerEventsPayloadMap[T];
 
+const isAbortError = (error: unknown): error is Error => error instanceof Error && error.name === "AbortError";
+
+const emit = <T extends ServerEventType>(eventName: T, data: ServerEventsPayloadMap[T]) => {
+	const handlers = sharedState.handlers[eventName] as EventHandlerSet<T> | undefined;
+	handlers?.forEach((handler) => {
+		handler(data);
+	});
+};
+
+const refreshQueriesForEvent = (eventName: ServerEventType) => {
+	if (!invalidatingEvents.has(eventName) || !sharedState.queryClient) {
+		return;
+	}
+
+	void sharedState.queryClient.invalidateQueries().catch((error) => {
+		if (!isAbortError(error)) {
+			console.error(`[SSE] Failed to refresh queries after ${eventName}:`, error);
+		}
+	});
+};
+
+const connectEventSource = (queryClient: QueryClient) => {
+	sharedState.queryClient = queryClient;
+	if (sharedState.eventSource) {
+		return;
+	}
+
+	const eventSource = new EventSource("/api/v1/events");
+	sharedState.eventSource = eventSource;
+
+	eventSource.addEventListener("connected", (event) => {
+		const data = parseEventData<"connected">(event);
+		console.info("[SSE] Connected to server events");
+		emit("connected", data);
+	});
+
+	eventSource.addEventListener("heartbeat", (event) => {
+		emit("heartbeat", parseEventData<"heartbeat">(event));
+	});
+
+	for (const eventName of serverEventNames) {
+		eventSource.addEventListener(eventName, (event) => {
+			const data = parseEventData<typeof eventName>(event);
+			console.info(`[SSE] ${eventName}:`, data);
+
+			refreshQueriesForEvent(eventName);
+
+			if (eventName === "volume:status_changed") {
+				const statusData = data as ServerEventsPayloadMap["volume:status_changed"];
+				emit("volume:status_changed", statusData);
+				emit("volume:updated", statusData);
+				return;
+			}
+
+			emit(eventName, data);
+		});
+	}
+
+	eventSource.onerror = (error) => {
+		console.error("[SSE] Connection error:", error);
+	};
+};
+
+const disconnectEventSource = () => {
+	if (!sharedState.eventSource) {
+		return;
+	}
+
+	console.info("[SSE] Disconnecting from server events");
+	sharedState.eventSource.close();
+	sharedState.eventSource = null;
+	sharedState.queryClient = null;
+	sharedState.handlers = {};
+};
+
+const addSharedEventListener = <T extends ServerEventType>(
+	eventName: T,
+	handler: EventHandler<T>,
+	options?: { signal?: AbortSignal },
+) => {
+	if (options?.signal?.aborted) {
+		return () => {};
+	}
+
+	const existingHandlers = sharedState.handlers[eventName] as EventHandlerSet<T> | undefined;
+	const eventHandlers = existingHandlers ?? new Set<EventHandler<T>>();
+	eventHandlers.add(handler);
+	sharedState.handlers[eventName] = eventHandlers as EventHandlerMap[T];
+
+	const unsubscribe = () => {
+		const handlers = sharedState.handlers[eventName] as EventHandlerSet<T> | undefined;
+		handlers?.delete(handler);
+		if (handlers && handlers.size === 0) {
+			delete sharedState.handlers[eventName];
+		}
+		if (options?.signal) {
+			options.signal.removeEventListener("abort", unsubscribe);
+		}
+	};
+
+	if (options?.signal) {
+		options.signal.addEventListener("abort", unsubscribe, { once: true });
+	}
+
+	return unsubscribe;
+};
+
 /**
  * Hook to listen to Server-Sent Events (SSE) from the backend
  * Automatically handles cache invalidation for backup and volume events
  */
 export function useServerEvents() {
 	const queryClient = useQueryClient();
-	const eventSourceRef = useRef<EventSource | null>(null);
+	const addEventListener = useCallback(addSharedEventListener, []);
-	const handlersRef = useRef<EventHandlerMap>({});
+	const hasMountedRef = useRef(false);
-	const emit = useCallback(<T extends ServerEventType>(eventName: T, data: ServerEventsPayloadMap[T]) => {
-		const handlers = handlersRef.current[eventName] as EventHandlerSet<T> | undefined;
-		handlers?.forEach((handler) => {
-			handler(data);
-		});
-	}, []);
 
 	useEffect(() => {
-		const eventSource = new EventSource("/api/v1/events");
+		connectEventSource(queryClient);
-		eventSourceRef.current = eventSource;
+		if (!hasMountedRef.current) {
-
+			sharedState.subscribers += 1;
-		eventSource.addEventListener("connected", (event) => {
+			hasMountedRef.current = true;
-			const data = parseEventData<"connected">(event);
-			console.info("[SSE] Connected to server events");
-			emit("connected", data);
-		});
-
-		eventSource.addEventListener("heartbeat", (event) => {
-			emit("heartbeat", parseEventData<"heartbeat">(event));
-		});
-
-		for (const eventName of serverEventNames) {
-			eventSource.addEventListener(eventName, (event) => {
-				const data = parseEventData<typeof eventName>(event);
-				console.info(`[SSE] ${eventName}:`, data);
-
-				if (invalidatingEvents.has(eventName)) {
-					void queryClient.invalidateQueries();
-				}
-
-				if (eventName === "backup:completed") {
-					void queryClient.refetchQueries();
-				}
-
-				if (eventName === "volume:status_changed") {
-					const statusData = data as ServerEventsPayloadMap["volume:status_changed"];
-					emit("volume:status_changed", statusData);
-					emit("volume:updated", statusData);
-					return;
-				}
-
-				emit(eventName, data);
-			});
 		}
 
-		eventSource.onerror = (error) => {
-			console.error("[SSE] Connection error:", error);
-		};
-
 		return () => {
-			console.info("[SSE] Disconnecting from server events");
+			if (!hasMountedRef.current) {
-			eventSource.close();
+				return;
-			eventSourceRef.current = null;
+			}
+
+			hasMountedRef.current = false;
+			sharedState.subscribers = Math.max(0, sharedState.subscribers - 1);
+			if (sharedState.subscribers === 0) {
+				disconnectEventSource();
+			}
 		};
-	}, [emit, queryClient]);
+	}, [queryClient]);
-
-	const addEventListener = useCallback(
-		<T extends ServerEventType>(eventName: T, handler: EventHandler<T>, options?: { signal?: AbortSignal }) => {
-			if (options?.signal?.aborted) {
-				return () => {};
-			}
-
-			const existingHandlers = handlersRef.current[eventName] as EventHandlerSet<T> | undefined;
-			const eventHandlers = existingHandlers ?? new Set<EventHandler<T>>();
-			eventHandlers.add(handler);
-			handlersRef.current[eventName] = eventHandlers as EventHandlerMap[T];
-
-			const unsubscribe = () => {
-				const handlers = handlersRef.current[eventName] as EventHandlerSet<T> | undefined;
-				handlers?.delete(handler);
-				if (handlers && handlers.size === 0) {
-					delete handlersRef.current[eventName];
-				}
-				if (options?.signal) {
-					options.signal.removeEventListener("abort", unsubscribe);
-				}
-			};
-
-			if (options?.signal) {
-				options.signal.addEventListener("abort", unsubscribe, { once: true });
-			}
-
-			return unsubscribe;
-		},
-		[],
-	);
 
 	return { addEventListener };
 }