Yielding webworker

app/web/package.json

@@ -73,6 +73,7 @@
     "cm6-theme-gruvbox-dark": "^0.2.0",
     "codemirror": "^6.0.1",
     "comlink": "^4.4.2",
+    "comlink-async-generator": "^0.0.1",
     "d3": "^7.9.0",
     "date-fns": "^2.29.2",
     "elkjs": "^0.10.0",

app/web/src/newhotness/Explore.vue

@@ -1,5 +1,5 @@
 <template>
-  <DelayedLoader v-if="componentListRaw.isLoading.value" :size="'full'" />
+  <DelayedLoader v-if="componentListRaw.isPending.value" :size="'full'" />
   <section v-else :class="clsx('grid h-full', showGrid ? 'explore' : 'map')">
     <!-- Left column -->
     <!-- 12 pixel padding to align with the SI logo -->
@@ -215,7 +215,7 @@ import { useQuery } from "@tanstack/vue-query";
 import { useVirtualizer } from "@tanstack/vue-virtual";
 import { Fzf } from "fzf";
 import { tw } from "@si/vue-lib";
-import { bifrost, useMakeArgs, useMakeKey } from "@/store/realtime/heimdall";
+import { bifrost, bifrostGenerator, useMakeArgs, useMakeKey } from "@/store/realtime/heimdall";
 import {
   BifrostActionViewList,
   BifrostComponentList,
@@ -332,22 +332,64 @@ const kind = computed(() =>
 const id = computed(() =>
   selectedView.value ? selectedView.value : ctx.workspacePk.value,
 );
-const componentQueryKey = key(kind, id);
+const componentQueryKey = key(EntityKind.ComponentList, ctx.workspacePk.value);
 
 const componentListRaw = useQuery<
-  BifrostComponentList | ViewComponentList | null
+  BifrostComponentInList[]
 >({
   queryKey: componentQueryKey,
-  queryFn: async () => {
-    const arg = selectedView.value
-      ? args(EntityKind.ViewComponentList, selectedView.value)
-      : args(EntityKind.ComponentList);
-    return await bifrost<BifrostComponentList | ViewComponentList>(arg);
+  queryFn: async (context) => {
+    // PSA: there is an experimental `streamedQuery` in tanstack/react_query
+    // I've stolen most of it here
+    const query = context.client
+      .getQueryCache()
+      .find({ queryKey: context.queryKey, exact: true })
+
+    const isRefetch = !!query && query.state.data !== undefined
+
+    if (isRefetch) {
+      query.setState({
+        status: 'pending',
+        data: undefined,
+        error: null,
+        fetchStatus: 'fetching',
+      })
+    }
+
+    const result: Array<BifrostComponentInList> = [];
+
+    const arg = args(EntityKind.ComponentList);
+    const stream = bifrostGenerator<BifrostComponentInList>(arg);
+    for await (const chunk of stream) {
+      if (!chunk)
+        break;
+      if (context.signal.aborted) {
+        break;
+      }
+
+      // don't append to the cache directly when replace-refetching
+      if (!isRefetch ) {
+        context.client.setQueryData<Array<BifrostComponentInList>>(
+          context.queryKey,
+          (prev = []) => {
+            return [...prev, chunk]
+          },
+        )
+      }
+      result.push(chunk)
+    }
+
+    // finalize result: replace-refetching needs to write to the cache
+    if (isRefetch && !context.signal.aborted) {
+      context.client.setQueryData<Array<BifrostComponentInList>>(context.queryKey, result)
+    }
+
+    return context.client.getQueryData(context.queryKey)!
   },
 });
 
 const componentList = computed(
-  () => componentListRaw.data.value?.components ?? [],
+  () => componentListRaw.data.value ?? [],
 );
 
 const scrollRef = ref<HTMLDivElement>();

app/web/src/store/realtime/heimdall.ts

@@ -1,4 +1,5 @@
 import * as Comlink from "comlink";
+import { asyncGeneratorTransferHandler } from 'comlink-async-generator';
 import { computed, reactive, Reactive, inject, ComputedRef, unref } from "vue";
 import { QueryClient } from "@tanstack/vue-query";
 import {
@@ -19,6 +20,8 @@ import router from "@/router";
 import { useChangeSetsStore } from "../change_sets.store";
 import { useWorkspacesStore } from "../workspaces.store";
 
+Comlink.transferHandlers.set('asyncGenerator', asyncGeneratorTransferHandler);
+
 let token: string | undefined;
 let queryClient: QueryClient;
 export const init = async (bearerToken: string, _queryClient: QueryClient) => {
@@ -124,17 +127,50 @@ export const bifrost = async <T>(args: {
 }): Promise<Reactive<T> | null> => {
   if (!initCompleted.value) throw new Error("bifrost not initiated");
   const start = Date.now();
-  const maybeAtomDoc = await db.get(
+  const generator = await db.get(
     args.workspaceId,
     args.changeSetId,
     args.kind,
     args.id,
   );
+  const maybeAtomDoc = await generator.next();
   const end = Date.now();
   // eslint-disable-next-line no-console
   console.log("🌈 bifrost query", args.kind, args.id, end - start, "ms");
-  if (maybeAtomDoc === -1) return null;
-  return reactive(maybeAtomDoc);
+  if (maybeAtomDoc.value === -1) return null;
+  return reactive(maybeAtomDoc.value);
+};
+
+export async function* bifrostGenerator<T>(args: {
+  workspaceId: string;
+  changeSetId: ChangeSetId;
+  kind: EntityKind;
+  id: Id;
+}): AsyncGenerator<T | null> {
+  if (!initCompleted.value) throw new Error("bifrost not initiated");
+  const start = Date.now();
+  const generator = await db.get(
+    args.workspaceId,
+    args.changeSetId,
+    args.kind,
+    args.id,
+  );
+  let first = true 
+  let firstTime;
+  for await (const maybeAtomDoc of await generator) {
+    if (first) {
+      first = false;
+      firstTime = Date.now();
+    }
+    if (maybeAtomDoc === -1) yield null;
+    else {
+      console.log("yielding...")
+      yield maybeAtomDoc as T;
+    }
+  }
+  const end = Date.now();
+  // eslint-disable-next-line no-console
+  console.log("🌈 bifrost query", args.kind, args.id, firstTime ? firstTime-start : Infinity, "ms", end - start, "ms");
 };
 
 export const getPossibleConnections = async (args: {

app/web/src/workers/types/dbinterface.ts

@@ -70,7 +70,7 @@ export interface DBInterface {
     changeSetId: ChangeSetId,
     kind: EntityKind,
     id: Id,
-  ): Promise<typeof NOROW | AtomDocument>;
+  ): AsyncGenerator<object | -1>;
   mjolnirBulk(
     workspaceId: string,
     changeSetId: ChangeSetId,

app/web/src/workers/webworker.ts

@@ -1,4 +1,5 @@
 import * as Comlink from "comlink";
+import { asyncGeneratorTransferHandler } from 'comlink-async-generator';
 import { applyPatch as applyOperations } from "fast-json-patch";
 import sqlite3InitModule, {
   Database,
@@ -89,6 +90,8 @@ import {
   bustQueueAdd,
 } from "./mjolnir_queue";
 
+Comlink.transferHandlers.set('asyncGenerator', asyncGeneratorTransferHandler);
+
 let otelEndpoint = import.meta.env.VITE_OTEL_EXPORTER_OTLP_ENDPOINT;
 if (!otelEndpoint) otelEndpoint = "http://localhost:8080";
 const exporter = new OTLPTraceExporter({
@@ -1582,12 +1585,13 @@ const allOutgoingConns = new DefaultMap<
 >(() => new DefaultMap(() => ({})));
 
 const coldStartComputed = async (workspaceId: string, changeSetId: string) => {
-  const data = (await get(
+  const dataGenerator = get(
     workspaceId,
     changeSetId,
     EntityKind.ComponentList,
     workspaceId,
-  )) as BifrostComponentList | -1;
+  )
+  const data = (await dataGenerator.next()).value as BifrostComponentList | -1;
 
   if (data === -1) return;
 
@@ -1613,15 +1617,17 @@ const coldStartComputed = async (workspaceId: string, changeSetId: string) => {
     true,
   );
 
-  const list = (await get(
+  const listGenerator =  get(
     workspaceId,
     changeSetId,
     EntityKind.IncomingConnectionsList,
     workspaceId,
     undefined,
     undefined,
     false, // don't compute
-  )) as BifrostIncomingConnectionsList | -1;
+  );
+
+  const list = (await listGenerator.next()).value as BifrostIncomingConnectionsList | -1;
 
   if (list === -1) return;
 
@@ -1998,15 +2004,16 @@ const getReferences = async (
     return [bifrost, hasReferenceError];
   } else if (kind === EntityKind.Component) {
     const data = atomDoc as EddaComponent;
-    const sv = (await get(
+    const svGenerator = get(
       workspaceId,
       changeSetId,
       data.schemaVariantId.kind,
       data.schemaVariantId.id,
       undefined,
       indexChecksum,
       followComputed,
-    )) as SchemaVariant | -1;
+    )
+    const sv = (await svGenerator.next()).value as SchemaVariant | -1;
 
     if (sv === -1) {
       hasReferenceError = true;
@@ -2034,15 +2041,16 @@ const getReferences = async (
       );
     }
 
-    const sm = (await get(
+    const smGenerator = get(
       workspaceId,
       changeSetId,
       data.schemaMembers.kind,
       data.schemaMembers.id,
       undefined,
       indexChecksum,
       followComputed,
-    )) as SchemaMembers | -1;
+    )
+    const sm = (await smGenerator.next()).value as SchemaMembers | -1;
 
     if (sm === -1) {
       hasReferenceError = true;
@@ -2182,7 +2190,7 @@ const getReferences = async (
     return [list, hasReferenceError];
   } else if (kind === EntityKind.IncomingConnections) {
     const raw = atomDoc as EddaIncomingConnections;
-    const component = (await get(
+    const componentGenerator = get(
       workspaceId,
       changeSetId,
       EntityKind.Component,
@@ -2191,7 +2199,8 @@ const getReferences = async (
       indexChecksum,
       false,
       false,
-    )) as BifrostComponent | -1;
+    );
+    const component = (await componentGenerator.next()).value as BifrostComponent | -1;
 
     clearWeakReferences(changeSetId, {
       kind: EntityKind.IncomingConnections,
@@ -2375,7 +2384,7 @@ const getReferences = async (
   }
 };
 
-const get = async (
+async function* get (
   workspaceId: string,
   changeSetId: ChangeSetId,
   kind: EntityKind,
@@ -2384,7 +2393,7 @@ const get = async (
   indexChecksum?: string,
   followComputed = true,
   followReferences = true,
-): Promise<-1 | object> => {
+): AsyncGenerator<object | -1, void> {
   const sql = `
     select
       data
@@ -2422,13 +2431,17 @@ const get = async (
   );
   if (data === NOROW) {
     mjolnir(workspaceId, changeSetId, kind, id, checksum);
-    return -1;
+    yield -1;
+    return;
   }
   const atomDoc = decodeDocumentFromDB(data as ArrayBuffer);
   // debug("📄 atom doc", atomDoc);
 
   // THIS GETS REPLACED WITH AUTO-GEN CODE
-  if (!followReferences) return atomDoc;
+  if (!followReferences) {
+    yield atomDoc;
+    return;
+  }
 
   try {
     const [docAndRefs, hasReferenceError] = await getReferences(
@@ -2443,16 +2456,28 @@ const get = async (
     // this is a choice, we could send through objects that don't match the types
     // and potentially have something drawn on the screen—but that seems worse
     // for the possible side-effects
-    if (hasReferenceError) return -1;
+    if (hasReferenceError) {
+      yield -1;
+      return;
+    }
 
     if (followComputed) {
-      return await getComputed(docAndRefs, workspaceId, changeSetId, kind, id);
+      const completedData = await getComputed(docAndRefs, workspaceId, changeSetId, kind, id);
+      if (kind === EntityKind.ComponentList) {
+        for (const c of (completedData as BifrostComponentList).components) {
+          yield c;
+        };
+      } else
+        yield completedData;
+      return;
     }
-    return docAndRefs;
+    yield docAndRefs;
+    return;
   } catch (err) {
     // eslint-disable-next-line no-console
     console.error(err);
-    return -1;
+    yield -1;
+    return;
   }
 };