chore -cached loading (V3 omezarr support) step one (#205)

Author: froyo-np

packages/core/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@alleninstitute/vis-core",
-    "version": "0.0.4",
+    "version": "0.0.5",
     "contributors": [
         {
             "name": "Lane Sawyer",

packages/core/src/workers/worker-pool.ts

@@ -26,7 +26,7 @@ export enum WorkerStatus {
 
 export class WorkerPool {
     #workers: Worker[];
-    #promises: Map<number, MessagePromise>;
+    #promises: Map<string, MessagePromise>;
     #timeOfPreviousHeartbeat: Map<number, number>;
     #which: number;
 
@@ -36,25 +36,37 @@ export class WorkerPool {
         for (let i = 0; i < size; i++) {
             this.#workers[i] = new Worker(workerModule, { type: 'module' });
             this.#workers[i].onmessage = (msg) => this.#handleMessage(i, msg);
-            this.#timeOfPreviousHeartbeat.set(i, Date.now());
+            this.#timeOfPreviousHeartbeat.set(i, 0);
         }
         this.#promises = new Map();
         this.#which = 0;
     }
-
+    /**
+     * Warning - nothing in this class should be considered useable after
+     * calling this method - any/all methods called should be expected to be
+     * completely unreliable. dont call me unless you're about to dispose of all references to this object
+     */
+    destroy() {
+        for (let i = 0; i < this.#workers.length; i++) {
+            this.#workers[i].terminate();
+        }
+        this.#workers = [];
+    }
     #handleMessage(workerIndex: number, msg: MessageEvent<unknown>) {
         const { data } = msg;
-        const messagePromise = this.#promises.get(workerIndex);
         if (isHeartbeatMessage(data)) {
             this.#timeOfPreviousHeartbeat.set(workerIndex, Date.now());
             return;
         }
-        if (messagePromise === undefined) {
-            logger.warn('unexpected message from worker');
-            return;
-        }
         if (isWorkerMessageWithId(data)) {
-            this.#promises.delete(workerIndex);
+            logger.debug(`worker ${workerIndex} responded to a message`);
+            const { id } = data;
+            const messagePromise = this.#promises.get(id);
+            if (messagePromise === undefined) {
+                logger.warn('unexpected message from worker');
+                return;
+            }
+            this.#promises.delete(id);
             if (!messagePromise.validator(data)) {
                 const reason = 'invalid response from worker: message type did not match expected type';
                 logger.error(reason);
@@ -63,28 +75,32 @@ export class WorkerPool {
             }
             messagePromise.resolve(data);
         } else {
+            logger.debug(`worker ${workerIndex} received an invalid message`);
             const reason = 'encountered an invalid message; skipping';
-            logger.error(reason);
-            messagePromise.reject(new Error(reason));
+            logger.warn(reason);
         }
     }
 
     #roundRobin() {
         this.#which = (this.#which + 1) % this.#workers.length;
     }
 
-    submitRequest(
+    async submitRequest(
         message: WorkerMessage,
         responseValidator: MessageValidator<WorkerMessageWithId>,
         transfers: Transferable[],
         signal?: AbortSignal | undefined,
     ): Promise<WorkerMessageWithId> {
+        if (this.#workers.length < 1) {
+            return Promise.reject('this woorker pool has been disposed');
+        }
         const reqId = `rq${uuidv4()}`;
         const workerIndex = this.#which;
         const messageWithId = { ...message, id: reqId };
         const messagePromise = this.#createMessagePromise(responseValidator);
+        logger.debug(`worker ${workerIndex} being handed a request`);
 
-        this.#promises.set(workerIndex, messagePromise);
+        this.#promises.set(reqId, messagePromise);
 
         if (signal) {
             signal.addEventListener('abort', () => {
@@ -126,7 +142,11 @@ export class WorkerPool {
         if (!this.#isValidIndex(workerIndex)) {
             throw new Error('invalid worker index');
         }
-        const delta = Date.now() - (this.#timeOfPreviousHeartbeat.get(workerIndex) ?? 0);
+        const lastHeartbeat = this.#timeOfPreviousHeartbeat.get(workerIndex) ?? 0;
+        if (lastHeartbeat === 0) {
+            return WorkerStatus.Unresponsive;
+        }
+        const delta = Date.now() - lastHeartbeat;
         if (delta && delta > 1500) {
             return WorkerStatus.Unresponsive;
         }

packages/omezarr/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@alleninstitute/vis-omezarr",
-    "version": "0.0.14",
+    "version": "0.0.15",
     "contributors": [
         {
             "name": "Lane Sawyer",

packages/omezarr/src/index.ts

@@ -55,3 +55,16 @@ export {
 } from './zarr/loading';
 
 export { type CancelRequest, type ZarrSliceRequest, makeOmeZarrSliceLoaderWorker } from './sliceview/worker-loader';
+export { decoderFactory } from './zarr/cache-lower';
+export { setupFetchDataWorker } from './zarr/cached-loading/fetch-data.worker-loader';
+export {
+    type TransferrableRequestInit,
+    type FetchMessagePayload,
+    type FetchMessage,
+    type FetchResponseMessage,
+    type CancelMessage,
+    isFetchMessage,
+    isFetchResponseMessage,
+    isCancelMessage,
+    isCancellationError,
+} from './zarr/cached-loading/fetch-data.interface';

packages/omezarr/src/zarr/cache-lower.ts

@@ -0,0 +1,52 @@
+import type { OmeZarrShapedDataset, OmeZarrMetadata } from './types';
+import { type ZarrRequest, buildQuery, loadZarrArrayFileFromStore } from './loading';
+import { VisZarrDataError } from '../errors';
+import * as zarr from 'zarrita';
+import { logger } from '@alleninstitute/vis-core';
+import { ZarrFetchStore, type CachingMultithreadedFetchStoreOptions } from './cached-loading/store';
+
+export function decoderFactory(url: string, workerModule: URL, options?: CachingMultithreadedFetchStoreOptions) {
+    const store = new ZarrFetchStore(url, workerModule, options);
+    const getSlice = async (
+        metadata: OmeZarrMetadata,
+        req: ZarrRequest,
+        level: OmeZarrShapedDataset,
+        signal?: AbortSignal,
+    ) => {
+        if (metadata.url !== url) {
+            throw new Error(
+                'trying to use a decoder from a different store - we cant do that yet, although we could build a map of url->stores here if we wanted later - TODO',
+            );
+        }
+        const scene = metadata.attrs.multiscales[0];
+        const { axes } = scene;
+        if (!level) {
+            const message = 'invalid Zarr data: no datasets found';
+            logger.error(message);
+            throw new VisZarrDataError(message);
+        }
+        const arr = metadata.arrays.find((a) => a.path === level.path);
+        if (!arr) {
+            const message = `cannot load slice: no array found for path [${level.path}]`;
+            logger.error(message);
+            throw new VisZarrDataError(message);
+        }
+        const { raw } = await loadZarrArrayFileFromStore(store, arr.path, metadata.zarrVersion, false);
+        const result = await zarr.get(raw, buildQuery(req, axes, level.shape), { opts: { signal: signal ?? null } });
+        if (typeof result === 'number') {
+            throw new Error('oh noes, slice came back all weird');
+        }
+        const { shape, data } = result;
+        if (typeof data !== 'object' || !('buffer' in data)) {
+            throw new Error('slice was malformed, array-buffer response required');
+        }
+        // biome-ignore lint/suspicious/noExplicitAny: <hard to prove - but the typeof check above is sufficient for this to be safe>
+        return { shape, data: new Float32Array(data as any) };
+    };
+    return {
+        decoder: getSlice,
+        destroy: () => {
+            store.destroy();
+        },
+    };
+}

packages/omezarr/src/zarr/cached-loading/fetch-data.interface.ts

@@ -0,0 +1,90 @@
+import type { AbsolutePath, RangeQuery } from 'zarrita';
+import z from 'zod';
+
+export type TransferrableRequestInit = Omit<RequestInit, 'body' | 'headers' | 'signal'> & {
+    body?: string;
+    headers?: [string, string][] | Record<string, string>;
+};
+
+export type FetchMessagePayload = {
+    rootUrl: string;
+    path: AbsolutePath;
+    range?: RangeQuery | undefined;
+    options?: TransferrableRequestInit | undefined;
+};
+
+export const FETCH_MESSAGE_TYPE = 'fetch' as const;
+export const FETCH_RESPONSE_MESSAGE_TYPE = 'fetch-response' as const;
+export const CANCEL_MESSAGE_TYPE = 'cancel' as const;
+
+export type FetchMessage = {
+    type: typeof FETCH_MESSAGE_TYPE;
+    id: string;
+    payload: FetchMessagePayload;
+};
+
+export type FetchResponseMessage = {
+    type: typeof FETCH_RESPONSE_MESSAGE_TYPE;
+    id: string;
+    payload: ArrayBufferLike | undefined;
+};
+
+export type CancelMessage = {
+    type: typeof CANCEL_MESSAGE_TYPE;
+    id: string;
+};
+
+const FetchMessagePayloadSchema = z.object({
+    rootUrl: z.string().nonempty(),
+    path: z.string().nonempty().startsWith('/'),
+    range: z
+        .union([
+            z.object({
+                offset: z.number(),
+                length: z.number(),
+            }),
+            z.object({
+                suffixLength: z.number(),
+            }),
+        ])
+        .optional(),
+    options: z.unknown().optional(), // being "lazy" for now; doing a full schema for this could be complex and fragile
+});
+
+const FetchMessageSchema = z.object({
+    type: z.literal(FETCH_MESSAGE_TYPE),
+    id: z.string().nonempty(),
+    payload: FetchMessagePayloadSchema,
+});
+
+const FetchResponseMessageSchema = z.object({
+    type: z.literal(FETCH_RESPONSE_MESSAGE_TYPE),
+    id: z.string().nonempty(),
+    payload: z.unknown().optional(), // unclear if it's feasible/wise to define a schema for this one
+});
+
+const CancelMessageSchema = z.object({
+    type: z.literal(CANCEL_MESSAGE_TYPE),
+    id: z.string().nonempty(),
+});
+
+export function isFetchMessage(val: unknown): val is FetchMessage {
+    return FetchMessageSchema.safeParse(val).success;
+}
+
+export function isFetchResponseMessage(val: unknown): val is FetchResponseMessage {
+    return FetchResponseMessageSchema.safeParse(val).success;
+}
+
+export function isCancelMessage(val: unknown): val is CancelMessage {
+    return CancelMessageSchema.safeParse(val).success;
+}
+
+export function isCancellationError(err: unknown): boolean {
+    return (
+        err === 'cancelled' ||
+        (typeof err === 'object' &&
+            err !== null &&
+            (('name' in err && err.name === 'AbortError') || ('code' in err && err.code === 20)))
+    );
+}