perf(connectors): singleflight ensureWarehouseRunning per warehouse (#420)

* perf(connectors): singleflight ensureWarehouseRunning per warehouse

Concurrent analytics queries on a cold warehouse now share one poll loop
and a single warehouses.start call. Joiners receive broadcast onStatus
updates with last-status replay; ref-counted abort cancels only when all
waiters leave.

Closes #419

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>

* refactor(appkit): simplify warehouse readiness singleflight helpers

Consolidate subscribe/join/spawn into focused helpers and tighten the
abort-race handler without changing behavior.

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>

* fix(connectors): grace-period abort for warehouse readiness singleflight

Delay shared poll cancellation by 100ms so React StrictMode unmount→remount
can rejoin the in-flight warmup instead of surfacing UPSTREAM_ERROR on first render.

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>

* fix(appkit): run-to-completion readiness and StrictMode abort handling

Let shared warehouse warmup finish when waiters disconnect instead of a
brittle grace timer. Classify cancellations as STREAM_ABORTED, suppress
error SSE frames for them, and ignore late envelopes on aborted hooks.

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>

* fix(connectors): hybrid orphan abort for warehouse readiness

Abort shared readiness polls on the next microtask when all waiters
leave before warehouses.start, so true orphans do not poll until timeout.
After start is issued, the poll runs to completion so the warm-path cache
stays primed. Synchronous StrictMode remounts rejoin before the microtask.

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>

---------

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>
Co-authored-by: MarioCadenas <MarioCadenas@users.noreply.github.com>

Author: MarioCadenas

packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts

@@ -1,16 +1,37 @@
 import { act, renderHook, waitFor } from "@testing-library/react";
 import { afterEach, describe, expect, test, vi } from "vitest";
 
+let capturedCallbacks: {
+  onMessage?: (msg: { data: string }) => void;
+  onError?: (err: Error) => void;
+  signal?: AbortSignal;
+} = {};
+
+const mockFetchArrow = vi.fn();
+const mockProcessArrowBuffer = vi.fn();
+
 // Mock connectSSE so the hook does not attempt a real network request.
-const mockConnectSSE = vi.fn().mockImplementation((_opts: unknown) => {
-  // Return a never-resolving promise; tests don't need the result.
-  return new Promise<void>(() => {});
-});
+const mockConnectSSE = vi
+  .fn()
+  .mockImplementation(
+    (opts: {
+      onMessage?: (msg: { data: string }) => void;
+      onError?: (err: Error) => void;
+      signal?: AbortSignal;
+    }) => {
+      capturedCallbacks = {
+        onMessage: opts.onMessage,
+        onError: opts.onError,
+        signal: opts.signal,
+      };
+      return new Promise<void>(() => {});
+    },
+  );
 
 vi.mock("@/js", () => ({
   ArrowClient: {
-    fetchArrow: vi.fn(),
-    processArrowBuffer: vi.fn(),
+    fetchArrow: (...args: unknown[]) => mockFetchArrow(...args),
+    processArrowBuffer: (...args: unknown[]) => mockProcessArrowBuffer(...args),
   },
   connectSSE: (...args: unknown[]) => mockConnectSSE(...args),
 }));
@@ -22,29 +43,32 @@ vi.mock("../use-query-hmr", () => ({
 
 import { useAnalyticsQuery } from "../use-analytics-query";
 
+function markAborted() {
+  const sig = capturedCallbacks.signal;
+  if (!sig) throw new Error("signal not captured yet");
+  Object.defineProperty(sig, "aborted", { value: true, configurable: true });
+}
+
 describe("useAnalyticsQuery", () => {
   afterEach(() => {
+    capturedCallbacks = {};
     vi.clearAllMocks();
   });
 
   test("does not refetch when params object is structurally equal across renders", () => {
-    // Each render passes a fresh object literal — the common footgun.
     const { rerender } = renderHook(
       ({ limit }: { limit: number }) =>
         // biome-ignore lint/suspicious/noExplicitAny: typed registry not available in tests
         useAnalyticsQuery("test_query" as any, { limit } as any),
       { initialProps: { limit: 10 } },
     );
 
-    // Initial render triggers exactly one connection.
     expect(mockConnectSSE).toHaveBeenCalledTimes(1);
 
-    // Re-render with structurally-equal-but-new-reference params.
     rerender({ limit: 10 });
     rerender({ limit: 10 });
     rerender({ limit: 10 });
 
-    // Should NOT have refetched — the hook stabilized the params reference.
     expect(mockConnectSSE).toHaveBeenCalledTimes(1);
   });
 
@@ -93,27 +117,26 @@ describe("useAnalyticsQuery", () => {
 
   describe("warehouse_status", () => {
     test("surfaces warehouseStatus while waiting and clears loading on result", async () => {
-      // Capture the connectSSE options so we can drive onMessage manually.
       let capturedOnMessage:
         | ((msg: { id: string; data: string }) => void)
         | null = null;
-      mockConnectSSE.mockImplementationOnce((opts: any) => {
-        capturedOnMessage = opts.onMessage;
-        return new Promise<void>(() => {});
-      });
+      mockConnectSSE.mockImplementationOnce(
+        (opts: { onMessage?: (msg: { id: string; data: string }) => void }) => {
+          capturedOnMessage = opts.onMessage ?? null;
+          return new Promise<void>(() => {});
+        },
+      );
 
       const { result } = renderHook(() =>
         // biome-ignore lint/suspicious/noExplicitAny: typed registry not available in tests
         useAnalyticsQuery("test_query" as any),
       );
 
-      // Initially: loading, no status, no data.
       expect(result.current.loading).toBe(true);
       expect(result.current.warehouseStatus).toBeNull();
       expect(result.current.data).toBeNull();
       expect(capturedOnMessage).toBeTruthy();
 
-      // Server emits a STARTING status — UI should show progress, still loading.
       act(() => {
         capturedOnMessage?.({
           id: "1",
@@ -133,7 +156,6 @@ describe("useAnalyticsQuery", () => {
       expect(result.current.loading).toBe(true);
       expect(result.current.data).toBeNull();
 
-      // Then RUNNING.
       act(() => {
         capturedOnMessage?.({
           id: "2",
@@ -148,7 +170,6 @@ describe("useAnalyticsQuery", () => {
         expect(result.current.warehouseStatus?.state).toBe("RUNNING");
       });
 
-      // Finally the SQL result lands.
       act(() => {
         capturedOnMessage?.({
           id: "3",
@@ -164,21 +185,19 @@ describe("useAnalyticsQuery", () => {
       });
       expect(result.current.data).toEqual([{ id: 1, name: "row1" }]);
       expect(result.current.error).toBeNull();
-      // warehouseStatus is left at its last observed value (RUNNING) so
-      // consumers that gated on `state !== "RUNNING"` flip back to data.
       expect(result.current.warehouseStatus?.state).toBe("RUNNING");
     });
 
     test("surfaces an error when a warehouse_status event has no status payload", async () => {
-      // A malformed frame must terminate the stream so the hook doesn't
-      // stay stuck in `loading: true` after a clean stream close.
       let capturedOnMessage:
         | ((msg: { id: string; data: string }) => void)
         | null = null;
-      mockConnectSSE.mockImplementationOnce((opts: any) => {
-        capturedOnMessage = opts.onMessage;
-        return new Promise<void>(() => {});
-      });
+      mockConnectSSE.mockImplementationOnce(
+        (opts: { onMessage?: (msg: { id: string; data: string }) => void }) => {
+          capturedOnMessage = opts.onMessage ?? null;
+          return new Promise<void>(() => {});
+        },
+      );
 
       const consoleError = vi
         .spyOn(console, "error")
@@ -203,4 +222,68 @@ describe("useAnalyticsQuery", () => {
       consoleError.mockRestore();
     });
   });
+
+  describe("aborted controller", () => {
+    test("ignores late error envelope after the controller was aborted", async () => {
+      const { result } = renderHook(() =>
+        // biome-ignore lint/suspicious/noExplicitAny: typed registry not available in tests
+        useAnalyticsQuery("test_query" as any),
+      );
+
+      await waitFor(() => expect(capturedCallbacks.signal).toBeDefined());
+
+      markAborted();
+
+      await act(async () => {
+        await capturedCallbacks.onMessage?.({
+          data: JSON.stringify({
+            type: "error",
+            error: "Statement failed: The operation was aborted.",
+            code: "UPSTREAM_ERROR",
+          }),
+        });
+      });
+
+      expect(result.current.error).toBeNull();
+    });
+
+    test("ignores late result envelope after the controller was aborted", async () => {
+      const { result } = renderHook(() =>
+        // biome-ignore lint/suspicious/noExplicitAny: typed registry not available in tests
+        useAnalyticsQuery("test_query" as any),
+      );
+
+      await waitFor(() => expect(capturedCallbacks.signal).toBeDefined());
+
+      markAborted();
+
+      await act(async () => {
+        await capturedCallbacks.onMessage?.({
+          data: JSON.stringify({ type: "result", data: [{ id: 99 }] }),
+        });
+      });
+
+      expect(result.current.data).toBeNull();
+    });
+
+    test("ignores late arrow envelope after the controller was aborted", async () => {
+      const { result } = renderHook(() =>
+        useAnalyticsQuery("test_query", null, { format: "ARROW_STREAM" }),
+      );
+
+      await waitFor(() => expect(capturedCallbacks.signal).toBeDefined());
+
+      markAborted();
+
+      await act(async () => {
+        await capturedCallbacks.onMessage?.({
+          data: JSON.stringify({ type: "arrow", statement_id: "stmt-123" }),
+        });
+      });
+
+      expect(mockFetchArrow).not.toHaveBeenCalled();
+      expect(result.current.data).toBeNull();
+      expect(result.current.error).toBeNull();
+    });
+  });
 });

packages/appkit-ui/src/react/hooks/use-analytics-query.ts

@@ -263,6 +263,9 @@ export function useAnalyticsQuery<
       payload,
       signal: abortController.signal,
       onMessage: async (message) => {
+        // Drop late envelopes from a stream whose controller was already
+        // aborted (React StrictMode unmount→remount). Mirrors onError below.
+        if (abortController.signal.aborted) return;
         try {
           const parsed = JSON.parse(message.data) as Record<string, unknown>;
           await handleAnalyticsSseMessage(parsed, sseContext);