console: keep the redesigned cluster page's live data in TanStack DB collections

console/src/api/materialize/cluster/largestClusterReplica.ts

@@ -38,6 +38,7 @@ export function buildLargestClusterReplicaQuery(clusterId: string) {
       )
       .leftJoin("mz_cluster_replica_sizes as crs", "crs.size", "cr.size")
       .select([
+        "cr.id",
         "cr.name",
         "cr.size",
         sql<string | null>`${replicaHeapLimit}::text`.as("heapLimit"),

console/src/api/materialize/cluster/largestMaintainedQueries.ts

@@ -15,6 +15,7 @@ import {
   executeSqlV2,
   queryBuilder,
 } from "~/api/materialize";
+import { buildSubscribeQuery } from "~/api/materialize/buildSubscribeQuery";
 
 export type LargestMaintainedQueriesParams = {
   replicaHeapLimit: number;
@@ -162,6 +163,29 @@ export function buildLargestMaintainedObjectSizesQuery({
   );
 }
 
+/** One object's arrangement size on a replica, in bytes. */
+export interface ReplicaObjectSize {
+  objectId: string;
+  size: number | null;
+}
+
+/**
+ * Live arrangement sizes of every object on a replica. Filtering on the
+ * literal replica id makes this a lookup on `mz_object_arrangement_sizes_ind`.
+ */
+export function buildReplicaObjectSizesSubscribe(replicaId: string) {
+  return buildSubscribeQuery<ReplicaObjectSize>(
+    queryBuilder
+      .selectFrom("mz_object_arrangement_sizes")
+      .select([
+        "object_id as objectId",
+        sql<number | null>`size::float8`.as("size"),
+      ])
+      .where("replica_id", "=", replicaId),
+    { upsertKey: "objectId" },
+  );
+}
+
 export async function fetchLargestMaintainedObjectSizes({
   params,
   queryKey,

console/src/api/materialize/cluster/replicaUtilizationBinning.ts

@@ -475,3 +475,17 @@ export const hasUtilizationReading = (point: DataPoint) =>
 export const hasNoUtilizationMetrics = (graphData: ReplicaData[]) =>
   graphData.length > 0 &&
   graphData.every(({ data }) => !data.some(hasUtilizationReading));
+
+/**
+ * Server-binned rows inside a window that starts at `startDate`, in time order.
+ * Held rows from a wider window would otherwise stretch the chart past the
+ * selected range, and ENVELOPE UPSERT yields them unordered.
+ */
+export function clipBinnedRowsToWindow<T extends { bucketStart: Date }>(
+  rows: T[],
+  startDate: Date,
+): T[] {
+  return rows
+    .filter((row) => row.bucketStart.getTime() >= startDate.getTime())
+    .sort((a, b) => a.bucketStart.getTime() - b.bucketStart.getTime());
+}

console/src/api/materialize/subscribeCollection.test.ts

@@ -10,6 +10,7 @@
 import { getStore } from "~/jotai";
 
 import {
+  createScopedSubscribeCollections,
   createSubscribeCollection,
   syncEngineCacheKey,
 } from "./subscribeCollection";
@@ -257,6 +258,22 @@ describe("createSubscribeCollection", () => {
     expect(getStore().get(statusAtom).snapshotComplete).toBe(false);
   });
 
+  it("marks held rows as refreshing until a new subscribe's snapshot lands", () => {
+    const { applySnapshot, statusAtom } = freshCollection();
+    applySnapshot(state([{ id: "1", name: "a" }]));
+    expect(getStore().get(statusAtom).isRefreshing).toBe(false);
+
+    // A new subscribe opens with an empty pre-snapshot.
+    applySnapshot(state([], false));
+    expect(getStore().get(statusAtom)).toMatchObject({
+      snapshotComplete: true,
+      isRefreshing: true,
+    });
+
+    applySnapshot(state([{ id: "1", name: "b" }]));
+    expect(getStore().get(statusAtom).isRefreshing).toBe(false);
+  });
+
   it("surfaces error and snapshotComplete on the status atom", () => {
     const { applySnapshot, statusAtom } = freshCollection();
 
@@ -302,3 +319,18 @@ describe("createSubscribeCollection", () => {
     }
   });
 });
+
+describe("createScopedSubscribeCollections", () => {
+  it("keeps one collection per scope and key, dropping other scopes' on a scope change", () => {
+    const collections = createScopedSubscribeCollections<Row>({
+      id: "scoped-test",
+      getKey: (row) => row.id,
+    });
+    const first = collections.get("org|regionA", "u1");
+    expect(collections.get("org|regionA", "u1")).toBe(first);
+    expect(collections.get("org|regionA", "u2")).not.toBe(first);
+
+    expect(collections.get("org|regionB", "u1")).not.toBe(first);
+    expect(collections.get("org|regionA", "u1")).not.toBe(first);
+  });
+});

console/src/api/materialize/subscribeCollection.ts

@@ -57,6 +57,9 @@ function writePersistedRows<T>(key: string, rows: T[]) {
 export interface SubscribeCollectionStatus {
   error: SubscribeError | undefined;
   snapshotComplete: boolean;
+  /** Rows from an earlier snapshot are served, and no running subscribe has
+   * delivered a newer one yet. */
+  isRefreshing: boolean;
 }
 
 export interface SubscribeCollection<T extends object> {
@@ -106,6 +109,7 @@ export function createSubscribeCollection<T extends object>(options: {
   const statusAtom = atom<SubscribeCollectionStatus>({
     error: undefined,
     snapshotComplete: false,
+    isRefreshing: false,
   });
 
   // Captured when the collection starts syncing; cleared when it pauses.
@@ -214,7 +218,8 @@ export function createSubscribeCollection<T extends object>(options: {
     const current = store.get(statusAtom);
     if (
       current.error !== next.error ||
-      current.snapshotComplete !== next.snapshotComplete
+      current.snapshotComplete !== next.snapshotComplete ||
+      current.isRefreshing !== next.isRefreshing
     ) {
       store.set(statusAtom, next);
     }
@@ -227,7 +232,11 @@ export function createSubscribeCollection<T extends object>(options: {
     const isEmptyPreload =
       !state.snapshotComplete && state.data.length === 0 && !state.error;
     if (isEmptyPreload) {
-      writeStatus({ error: undefined, snapshotComplete: lastSnapshotComplete });
+      writeStatus({
+        error: undefined,
+        snapshotComplete: lastSnapshotComplete,
+        isRefreshing: lastSnapshotComplete,
+      });
       return;
     }
 
@@ -240,6 +249,7 @@ export function createSubscribeCollection<T extends object>(options: {
     writeStatus({
       error: state.error,
       snapshotComplete: state.snapshotComplete,
+      isRefreshing: false,
     });
   };
 
@@ -296,7 +306,12 @@ export function createSubscribeCollection<T extends object>(options: {
     // On a scope change this also flushes the emptied set, so the collection
     // stops serving the previous scope's rows while the new snapshot loads.
     materialize();
-    writeStatus({ error: undefined, snapshotComplete: lastSnapshotComplete });
+    // Cached rows stand in until the live snapshot arrives.
+    writeStatus({
+      error: undefined,
+      snapshotComplete: lastSnapshotComplete,
+      isRefreshing: lastSnapshotComplete,
+    });
   };
 
   const clearPendingPersist = () => {
@@ -312,7 +327,11 @@ export function createSubscribeCollection<T extends object>(options: {
     lastSnapshotComplete = false;
     liveSnapshotApplied = false;
     materialize();
-    writeStatus({ error: undefined, snapshotComplete: false });
+    writeStatus({
+      error: undefined,
+      snapshotComplete: false,
+      isRefreshing: false,
+    });
   };
 
   const suspendPersistence = () => {
@@ -332,3 +351,51 @@ export function createSubscribeCollection<T extends object>(options: {
     suspendPersistence,
   };
 }
+
+export interface ScopedSubscribeCollections<T extends object> {
+  /** The collection for `key` under `scope`, created on first use. */
+  get: (scope: string, key: string) => SubscribeCollection<T>;
+  /** A session target while the scope resolves. It is shared by every key, so
+   * nothing may subscribe into it. */
+  placeholder: SubscribeCollection<T>;
+}
+
+/**
+ * In-memory collections, one per tenant scope (`organizationId|regionId`) and
+ * key, such as a cluster id. Keys repeat across environments, so a collection
+ * never serves another scope's rows. Rows stay in memory until the scope
+ * changes, which makes returning to a page instant. They are not persisted,
+ * since an entry per key would have no bound in localStorage.
+ */
+export function createScopedSubscribeCollections<T extends object>(options: {
+  /** Prefix of each collection's id. */
+  id: string;
+  getKey: (row: T) => string;
+}): ScopedSubscribeCollections<T> {
+  const { id, getKey } = options;
+  const byKey = new Map<string, SubscribeCollection<T>>();
+  let currentScope: string | undefined;
+  return {
+    get: (scope, key) => {
+      if (scope !== currentScope) {
+        // A return to an earlier scope gets fresh collections, so the dropped
+        // ones are never served again.
+        byKey.clear();
+        currentScope = scope;
+      }
+      let collection = byKey.get(key);
+      if (!collection) {
+        collection = createSubscribeCollection<T>({
+          id: `${id}|${scope}|${key}`,
+          getKey,
+        });
+        byKey.set(key, collection);
+      }
+      return collection;
+    },
+    placeholder: createSubscribeCollection<T>({
+      id: `${id}|placeholder`,
+      getKey,
+    }),
+  };
+}

console/src/api/materialize/subscribeSession.ts

@@ -185,6 +185,14 @@ export function createCollectionSubscribeSession<
       unsubscribeScope?.();
       keepAlive.unsubscribe();
       session.destroy();
+      // No subscribe backs the held rows any more, so they read as refreshing
+      // until a later session's snapshot lands, as behind a new subscribe's
+      // empty pre-snapshot.
+      target.applySnapshot({
+        data: [],
+        snapshotComplete: false,
+        error: undefined,
+      });
     },
   };
 }

console/src/api/materialize/useSubscribe.test.tsx

@@ -159,6 +159,34 @@ describe("subscribeSession", () => {
     }
   });
 
+  it("marks a collection's held rows as refreshing once its session ends", async () => {
+    const store = setTwoRegions();
+    const target = createSubscribeCollection<Row>({
+      id: "session-test-destroy",
+      getKey: (row) => row.id,
+    });
+    const session = createCollectionSubscribeSession<Row, Row>({
+      store,
+      request,
+      upsertKey,
+      select: (row) => row.data,
+      target,
+    });
+    target.applySnapshot({
+      data: [{ id: "1" }],
+      snapshotComplete: true,
+      error: undefined,
+    });
+    session.destroy();
+
+    await flush();
+    expect(target.collection.size).toBe(1);
+    expect(store.get(target.statusAtom)).toMatchObject({
+      snapshotComplete: true,
+      isRefreshing: true,
+    });
+  });
+
   it("clears an atom-fed collection when the region switches", async () => {
     const store = setTwoRegions();
     const sourceAtom = atom<SubscribeState<Row>>(loadingState);