From 43dfe409e1528395ff9746d86caeed05dfc104ca Mon Sep 17 00:00:00 2001
From: Axel Bocciarelli <axelboc@gmail.com>
Date: Thu, 8 Aug 2024 09:03:27 +0200
Subject: [PATCH 1/2] Move code out of `MockApi` and into utilities

---
 packages/app/src/providers/mock/mock-api.ts | 53 ++++-----------------
 packages/app/src/providers/mock/utils.ts    | 42 ++++++++++++++++
 2 files changed, 52 insertions(+), 43 deletions(-)

diff --git a/packages/app/src/providers/mock/mock-api.ts b/packages/app/src/providers/mock/mock-api.ts
index 31b9622e8..e191a7106 100644
--- a/packages/app/src/providers/mock/mock-api.ts
+++ b/packages/app/src/providers/mock/mock-api.ts
@@ -2,7 +2,6 @@ import {
   assertArrayShape,
   assertDefined,
   hasNumericType,
-  isGroup,
 } from '@h5web/shared/guards';
 import type {
   ArrayShape,
@@ -22,9 +21,13 @@ import {
 import { DataProviderApi } from '../api';
 import type { ExportFormat, ExportURL, ValuesStoreParams } from '../models';
 import { makeMockFile } from './mock-file';
-import { findMockEntity, sliceValue } from './utils';
-
-export const SLOW_TIMEOUT = 3000;
+import {
+  cancellableDelay,
+  findMockEntity,
+  getChildrenPaths,
+  sliceValue,
+  SLOW_TIMEOUT,
+} from './utils';
 
 export class MockApi extends DataProviderApi {
   private readonly mockFile: GroupWithChildren;
@@ -33,6 +36,7 @@ export class MockApi extends DataProviderApi {
     private readonly _getExportURL?: DataProviderApi['getExportURL'],
   ) {
     const mockFile = makeMockFile();
+
     super(mockFile.name);
     this.mockFile = mockFile;
   }
@@ -65,7 +69,7 @@ export class MockApi extends DataProviderApi {
     }
 
     if (dataset.name.startsWith('slow')) {
-      await this.cancellableDelay(signal);
+      await cancellableDelay(signal);
     }
 
     const { value } = dataset;
@@ -132,43 +136,6 @@ export class MockApi extends DataProviderApi {
   }
 
   public override async getSearchablePaths(path: string): Promise<string[]> {
-    return this.getEntityPaths(path);
-  }
-
-  private getEntityPaths(entityPath: string): string[] {
-    const entity = findMockEntity(this.mockFile, entityPath);
-    if (!entity) {
-      return [];
-    }
-
-    if (!isGroup(entity)) {
-      return [entity.path];
-    }
-
-    return entity.children.reduce<string[]>(
-      (acc, child) => [...acc, ...this.getEntityPaths(child.path)],
-      [entity.path],
-    );
-  }
-
-  private async cancellableDelay(signal?: AbortSignal) {
-    await new Promise<void>((resolve, reject) => {
-      const timeout = setTimeout(() => {
-        signal?.removeEventListener('abort', handleAbort);
-        resolve();
-      }, SLOW_TIMEOUT);
-
-      function handleAbort() {
-        clearTimeout(timeout);
-        signal?.removeEventListener('abort', handleAbort);
-        reject(
-          new Error(
-            typeof signal?.reason === 'string' ? signal.reason : 'cancelled',
-          ),
-        );
-      }
-
-      signal?.addEventListener('abort', handleAbort);
-    });
+    return getChildrenPaths(this.mockFile, path);
   }
 }
diff --git a/packages/app/src/providers/mock/utils.ts b/packages/app/src/providers/mock/utils.ts
index 756745c14..09c406781 100644
--- a/packages/app/src/providers/mock/utils.ts
+++ b/packages/app/src/providers/mock/utils.ts
@@ -18,6 +18,8 @@ import ndarray from 'ndarray';
 
 import { applyMapping } from '../../vis-packs/core/utils';
 
+export const SLOW_TIMEOUT = 3000;
+
 export function findMockEntity(
   group: GroupWithChildren,
   path: string,
@@ -58,3 +60,43 @@ export function sliceValue<T extends DType>(
 
   return mappedArray.data;
 }
+
+export function getChildrenPaths(
+  mockFile: GroupWithChildren,
+  entityPath: string,
+): string[] {
+  const entity = findMockEntity(mockFile, entityPath);
+  if (!entity) {
+    return [];
+  }
+
+  if (!isGroup(entity)) {
+    return [entity.path];
+  }
+
+  return entity.children.reduce<string[]>(
+    (acc, child) => [...acc, ...getChildrenPaths(mockFile, child.path)],
+    [entity.path],
+  );
+}
+
+export async function cancellableDelay(signal?: AbortSignal) {
+  await new Promise<void>((resolve, reject) => {
+    const timeout = setTimeout(() => {
+      signal?.removeEventListener('abort', handleAbort);
+      resolve();
+    }, SLOW_TIMEOUT);
+
+    function handleAbort() {
+      clearTimeout(timeout);
+      signal?.removeEventListener('abort', handleAbort);
+      reject(
+        new Error(
+          typeof signal?.reason === 'string' ? signal.reason : 'cancelled',
+        ),
+      );
+    }
+
+    signal?.addEventListener('abort', handleAbort);
+  });
+}

From 992975c0caea8f383a230a4059d1e99910749ae1 Mon Sep 17 00:00:00 2001
From: Axel Bocciarelli <axelboc@gmail.com>
Date: Thu, 8 Aug 2024 12:04:53 +0200
Subject: [PATCH 2/2] Use zustand store inside suspense fetch store to track
 progress

---
 packages/app/src/__tests__/Explorer.test.tsx  |  2 +-
 packages/app/src/__tests__/NexusPack.test.tsx |  2 +-
 .../app/src/__tests__/Visualizer.test.tsx     |  2 +-
 packages/app/src/providers/DataProvider.tsx   |  5 --
 packages/app/src/providers/api.ts             | 50 +++++--------------
 .../app/src/providers/h5grove/h5grove-api.ts  | 21 ++++++--
 packages/app/src/providers/hsds/hsds-api.ts   |  7 ++-
 packages/app/src/vis-packs/ValueLoader.tsx    | 31 ++++++------
 packages/shared/package.json                  |  3 ++
 packages/shared/src/react-suspense-fetch.ts   | 41 +++++++++++++--
 pnpm-lock.yaml                                |  4 ++
 11 files changed, 97 insertions(+), 71 deletions(-)

diff --git a/packages/app/src/__tests__/Explorer.test.tsx b/packages/app/src/__tests__/Explorer.test.tsx
index dd9f7f115..0e9c628c5 100644
--- a/packages/app/src/__tests__/Explorer.test.tsx
+++ b/packages/app/src/__tests__/Explorer.test.tsx
@@ -1,7 +1,7 @@
 import { screen } from '@testing-library/react';
 import { expect, test } from 'vitest';
 
-import { SLOW_TIMEOUT } from '../providers/mock/mock-api';
+import { SLOW_TIMEOUT } from '../providers/mock/utils';
 import { renderApp } from '../test-utils';
 
 test('select root group by default', async () => {
diff --git a/packages/app/src/__tests__/NexusPack.test.tsx b/packages/app/src/__tests__/NexusPack.test.tsx
index b169fd788..24c4e103d 100644
--- a/packages/app/src/__tests__/NexusPack.test.tsx
+++ b/packages/app/src/__tests__/NexusPack.test.tsx
@@ -1,7 +1,7 @@
 import { screen } from '@testing-library/react';
 import { expect, test } from 'vitest';
 
-import { SLOW_TIMEOUT } from '../providers/mock/mock-api';
+import { SLOW_TIMEOUT } from '../providers/mock/utils';
 import {
   getSelectedVisTab,
   getVisTabs,
diff --git a/packages/app/src/__tests__/Visualizer.test.tsx b/packages/app/src/__tests__/Visualizer.test.tsx
index d0094e1ac..eb3656794 100644
--- a/packages/app/src/__tests__/Visualizer.test.tsx
+++ b/packages/app/src/__tests__/Visualizer.test.tsx
@@ -1,7 +1,7 @@
 import { screen } from '@testing-library/react';
 import { expect, test } from 'vitest';
 
-import { SLOW_TIMEOUT } from '../providers/mock/mock-api';
+import { SLOW_TIMEOUT } from '../providers/mock/utils';
 import { mockConsoleMethod, renderApp } from '../test-utils';
 import { Vis } from '../vis-packs/core/visualizations';
 
diff --git a/packages/app/src/providers/DataProvider.tsx b/packages/app/src/providers/DataProvider.tsx
index cafdbfe6a..3f48e8ebe 100644
--- a/packages/app/src/providers/DataProvider.tsx
+++ b/packages/app/src/providers/DataProvider.tsx
@@ -11,7 +11,6 @@ import type {
   AttrName,
   AttrValuesStore,
   EntitiesStore,
-  ProgressCallback,
   ValuesStore,
 } from './models';
 
@@ -24,8 +23,6 @@ export interface DataContextValue {
 
   // Undocumented
   getExportURL?: DataProviderApi['getExportURL'];
-  addProgressListener: (cb: ProgressCallback) => void;
-  removeProgressListener: (cb: ProgressCallback) => void;
   getSearchablePaths?: DataProviderApi['getSearchablePaths'];
 }
 
@@ -92,8 +89,6 @@ function DataProvider(props: PropsWithChildren<Props>) {
         valuesStore,
         attrValuesStore,
         getExportURL: api.getExportURL?.bind(api),
-        addProgressListener: api.addProgressListener.bind(api),
-        removeProgressListener: api.removeProgressListener.bind(api),
         getSearchablePaths: api.getSearchablePaths?.bind(api),
       }}
     >
diff --git a/packages/app/src/providers/api.ts b/packages/app/src/providers/api.ts
index ed9892514..03b11080f 100644
--- a/packages/app/src/providers/api.ts
+++ b/packages/app/src/providers/api.ts
@@ -1,4 +1,3 @@
-/* eslint-disable promise/prefer-await-to-callbacks */
 import type {
   ArrayShape,
   AttributeValues,
@@ -7,6 +6,7 @@ import type {
   ProvidedEntity,
   Value,
 } from '@h5web/shared/hdf5-models';
+import type { OnProgress } from '@h5web/shared/react-suspense-fetch';
 import type {
   AxiosInstance,
   AxiosProgressEvent,
@@ -16,18 +16,10 @@ import type {
 } from 'axios';
 import axios from 'axios';
 
-import type {
-  ExportFormat,
-  ExportURL,
-  ProgressCallback,
-  ValuesStoreParams,
-} from './models';
+import type { ExportFormat, ExportURL, ValuesStoreParams } from './models';
 
 export abstract class DataProviderApi {
   protected readonly client: AxiosInstance;
-  protected progress = new Map<ValuesStoreParams, number>();
-
-  private readonly progressListeners = new Set<ProgressCallback>();
 
   public constructor(
     public readonly filepath: string,
@@ -53,36 +45,25 @@ export abstract class DataProviderApi {
 
   public getSearchablePaths?(path: string): Promise<string[]>;
 
-  public addProgressListener(cb: ProgressCallback): void {
-    this.progressListeners.add(cb);
-    cb([...this.progress.values()]); // notify once
-  }
-
-  public removeProgressListener(cb: ProgressCallback): void {
-    this.progressListeners.delete(cb);
-  }
-
   protected async cancellableFetchValue(
     endpoint: string,
-    storeParams: ValuesStoreParams,
     queryParams: Record<string, string | boolean | undefined>,
     signal?: AbortSignal,
+    onProgress?: OnProgress,
     responseType?: ResponseType,
   ): Promise<AxiosResponse> {
-    this.progress.set(storeParams, 0);
-    this.notifyProgressChange();
-
     try {
       return await this.client.get(endpoint, {
         signal,
         params: queryParams,
         responseType,
-        onDownloadProgress: (evt: AxiosProgressEvent) => {
-          if (evt.total !== undefined && evt.total > 0) {
-            this.progress.set(storeParams, evt.loaded / evt.total);
-            this.notifyProgressChange();
-          }
-        },
+        onDownloadProgress:
+          onProgress &&
+          ((evt: AxiosProgressEvent) => {
+            if (evt.total !== undefined && evt.total > 0) {
+              onProgress(evt.loaded / evt.total);
+            }
+          }),
       });
     } catch (error) {
       if (axios.isCancel(error)) {
@@ -93,21 +74,16 @@ export abstract class DataProviderApi {
         );
       }
       throw error;
-    } finally {
-      // Remove progress when request fulfills
-      this.progress.delete(storeParams);
-      this.notifyProgressChange();
     }
   }
 
-  private notifyProgressChange() {
-    this.progressListeners.forEach((cb) => cb([...this.progress.values()]));
-  }
-
   public abstract getEntity(path: string): Promise<ProvidedEntity>;
+
   public abstract getValue(
     params: ValuesStoreParams,
     signal?: AbortSignal,
+    onProgress?: OnProgress,
   ): Promise<unknown>;
+
   public abstract getAttrValues(entity: Entity): Promise<AttributeValues>;
 }
diff --git a/packages/app/src/providers/h5grove/h5grove-api.ts b/packages/app/src/providers/h5grove/h5grove-api.ts
index 7ca6522d5..355f7758e 100644
--- a/packages/app/src/providers/h5grove/h5grove-api.ts
+++ b/packages/app/src/providers/h5grove/h5grove-api.ts
@@ -8,6 +8,7 @@ import type {
   Value,
 } from '@h5web/shared/hdf5-models';
 import { DTypeClass } from '@h5web/shared/hdf5-models';
+import type { OnProgress } from '@h5web/shared/react-suspense-fetch';
 import type { AxiosRequestConfig } from 'axios';
 
 import { DataProviderApi } from '../api';
@@ -44,21 +45,29 @@ export class H5GroveApi extends DataProviderApi {
   public override async getValue(
     params: ValuesStoreParams,
     signal?: AbortSignal,
+    onProgress?: OnProgress,
   ): Promise<H5GroveDataResponse> {
     const { dataset } = params;
 
     if (dataset.type.class === DTypeClass.Opaque) {
-      return new Uint8Array(await this.fetchBinaryData(params, signal));
+      return new Uint8Array(
+        await this.fetchBinaryData(params, signal, onProgress),
+      );
     }
 
     const DTypedArray = h5groveTypedArrayFromDType(dataset.type);
     if (DTypedArray) {
-      const buffer = await this.fetchBinaryData(params, signal, true);
+      const buffer = await this.fetchBinaryData(
+        params,
+        signal,
+        onProgress,
+        true,
+      );
       const array = new DTypedArray(buffer);
       return hasScalarShape(dataset) ? array[0] : array;
     }
 
-    return this.fetchData(params, signal);
+    return this.fetchData(params, signal, onProgress);
   }
 
   public override async getAttrValues(
@@ -146,16 +155,17 @@ export class H5GroveApi extends DataProviderApi {
   private async fetchData(
     params: ValuesStoreParams,
     signal?: AbortSignal,
+    onProgress?: OnProgress,
   ): Promise<H5GroveDataResponse> {
     const { data } = await this.cancellableFetchValue(
       `/data/`,
-      params,
       {
         path: params.dataset.path,
         selection: params.selection,
         flatten: true,
       },
       signal,
+      onProgress,
     );
 
     return data;
@@ -164,11 +174,11 @@ export class H5GroveApi extends DataProviderApi {
   private async fetchBinaryData(
     params: ValuesStoreParams,
     signal?: AbortSignal,
+    onProgress?: OnProgress,
     safe = false,
   ): Promise<ArrayBuffer> {
     const { data } = await this.cancellableFetchValue(
       '/data/',
-      params,
       {
         path: params.dataset.path,
         selection: params.selection,
@@ -176,6 +186,7 @@ export class H5GroveApi extends DataProviderApi {
         dtype: safe ? 'safe' : undefined,
       },
       signal,
+      onProgress,
       'arraybuffer',
     );
 
diff --git a/packages/app/src/providers/hsds/hsds-api.ts b/packages/app/src/providers/hsds/hsds-api.ts
index 10b1eb4da..4580a831e 100644
--- a/packages/app/src/providers/hsds/hsds-api.ts
+++ b/packages/app/src/providers/hsds/hsds-api.ts
@@ -17,6 +17,7 @@ import type {
 } from '@h5web/shared/hdf5-models';
 import { EntityKind } from '@h5web/shared/hdf5-models';
 import { buildEntityPath, getChildEntity } from '@h5web/shared/hdf5-utils';
+import type { OnProgress } from '@h5web/shared/react-suspense-fetch';
 
 import { DataProviderApi } from '../api';
 import type { ExportFormat, ExportURL, ValuesStoreParams } from '../models';
@@ -127,11 +128,12 @@ export class HsdsApi extends DataProviderApi {
   public override async getValue(
     params: ValuesStoreParams,
     signal?: AbortSignal,
+    onProgress?: OnProgress,
   ): Promise<unknown> {
     const { dataset } = params;
     assertHsdsDataset(dataset);
 
-    const value = await this.fetchValue(dataset.id, params, signal);
+    const value = await this.fetchValue(dataset.id, params, signal, onProgress);
 
     // https://github.com/HDFGroup/hsds/issues/88
     // HSDS does not reduce the number of dimensions when selecting indices
@@ -217,13 +219,14 @@ export class HsdsApi extends DataProviderApi {
     entityId: HsdsId,
     params: ValuesStoreParams,
     signal?: AbortSignal,
+    onProgress?: OnProgress,
   ): Promise<HsdsValueResponse> {
     const { selection } = params;
     const { data } = await this.cancellableFetchValue(
       `/datasets/${entityId}/value`,
-      params,
       { select: selection && `[${selection}]` },
       signal,
+      onProgress,
     );
     return data.value;
   }
diff --git a/packages/app/src/vis-packs/ValueLoader.tsx b/packages/app/src/vis-packs/ValueLoader.tsx
index 16c27d084..23553b44c 100644
--- a/packages/app/src/vis-packs/ValueLoader.tsx
+++ b/packages/app/src/vis-packs/ValueLoader.tsx
@@ -1,5 +1,5 @@
 import { useTimeoutEffect, useToggle } from '@react-hookz/web';
-import { useEffect, useState } from 'react';
+import { useStore } from 'zustand';
 
 import { useDataContext } from '../providers/DataProvider';
 import { CANCELLED_ERROR_MSG } from '../providers/utils';
@@ -13,20 +13,15 @@ interface Props {
 
 function ValueLoader(props: Props) {
   const { isSlice = false } = props;
-  const { valuesStore, addProgressListener, removeProgressListener } =
-    useDataContext();
-
-  const [progress, setProgress] = useState<number[]>();
-
-  useEffect(() => {
-    addProgressListener(setProgress);
-    return () => removeProgressListener(setProgress);
-  }, [addProgressListener, removeProgressListener, setProgress]);
+  const { valuesStore } = useDataContext();
 
   // Wait a bit before showing loader to avoid flash
   const [isReady, toggleReady] = useToggle();
   useTimeoutEffect(toggleReady, 100);
 
+  // Track progress
+  const { ongoing } = useStore(valuesStore.progressStore);
+
   return (
     <div className={styles.loader} data-testid="LoadingDatasetValue">
       {isReady && (
@@ -42,13 +37,17 @@ function ValueLoader(props: Props) {
             <div />
             <div />
           </div>
-          {progress && (
-            <div className={styles.progressBars}>
-              {progress.slice(0, MAX_PROGRESS_BARS).map((val, index) => (
-                <progress className={styles.progress} key={index} value={val} /> // eslint-disable-line react/no-array-index-key
+          <div className={styles.progressBars}>
+            {[...ongoing.entries()]
+              .slice(0, MAX_PROGRESS_BARS)
+              .map(([key, val]) => (
+                <progress
+                  className={styles.progress}
+                  key={`${key.dataset.path}_${key.selection || ''}`}
+                  value={val}
+                />
               ))}
-            </div>
-          )}
+          </div>
           <p>{isSlice ? 'Loading current slice' : 'Loading data'}...</p>
           <p>
             <button
diff --git a/packages/shared/package.json b/packages/shared/package.json
index a50fbabb8..bce0d3737 100644
--- a/packages/shared/package.json
+++ b/packages/shared/package.json
@@ -35,6 +35,9 @@
     "ndarray-ops": "1.2.2",
     "react": ">=18"
   },
+  "optionalDependencies": {
+    "zustand": "4.5.4"
+  },
   "devDependencies": {
     "@types/d3-array": "~3.2.1",
     "@types/d3-format": "~3.0.4",
diff --git a/packages/shared/src/react-suspense-fetch.ts b/packages/shared/src/react-suspense-fetch.ts
index 5604e9de3..e20ea8670 100644
--- a/packages/shared/src/react-suspense-fetch.ts
+++ b/packages/shared/src/react-suspense-fetch.ts
@@ -1,9 +1,14 @@
+import type { StoreApi } from 'zustand/vanilla';
+import { createStore } from 'zustand/vanilla';
+
 type FetchFunc<Input, Result> = (
   input: Input,
   signal: AbortSignal,
+  onProgress: OnProgress,
 ) => Promise<Result>;
 
 type AreEqual<Input> = (a: Input, b: Input) => boolean;
+export type OnProgress = (value: number) => void;
 
 interface Instance<Result> {
   get: () => Result;
@@ -20,6 +25,13 @@ export interface FetchStore<Input, Result> {
   evictErrors: () => void;
   abort: (input: Input, reason?: string) => void;
   abortAll: (reason?: string) => void;
+  get progressStore(): StoreApi<ProgressState<Input>>;
+}
+
+interface ProgressState<Input> {
+  ongoing: Map<Input, number>;
+  setProgress: (input: Input, value: number) => void;
+  clearProgress: (input: Input) => void;
 }
 
 export function createFetchStore<Input, Result>(
@@ -28,15 +40,30 @@ export function createFetchStore<Input, Result>(
 ) {
   const cache = createCache<Input, Instance<Result>>(areEqual);
 
+  const progressStore = createStore<ProgressState<Input>>((set, get) => ({
+    ongoing: new Map(),
+    setProgress: (input, value) => {
+      const ongoing = new Map(get().ongoing);
+      ongoing.set(input, value);
+      set({ ongoing });
+    },
+    clearProgress: (input) => {
+      const ongoing = new Map(get().ongoing);
+      ongoing.delete(input);
+      set({ ongoing });
+    },
+  }));
+
   return {
     has: (input: Input): boolean => cache.has(input),
     prefetch: (input: Input): void => {
       if (!cache.has(input)) {
-        cache.set(input, createInstance(input, fetchFunc));
+        cache.set(input, createInstance(input, fetchFunc, progressStore));
       }
     },
     get: (input: Input): Result => {
-      const instance = cache.get(input) || createInstance(input, fetchFunc);
+      const instance =
+        cache.get(input) || createInstance(input, fetchFunc, progressStore);
       cache.set(input, instance);
       return instance.get();
     },
@@ -63,6 +90,9 @@ export function createFetchStore<Input, Result>(
     abortAll: (reason?: string): void => {
       cache.values().forEach((instance) => instance.abort(reason));
     },
+    get progressStore() {
+      return progressStore;
+    },
   };
 }
 
@@ -104,6 +134,7 @@ function createCache<K, V>(areEqual: AreEqual<K>) {
 function createInstance<Input, Result>(
   input: Input,
   fetchFunc: FetchFunc<Input, Result>,
+  progressStore: StoreApi<ProgressState<Input>>,
 ): Instance<Result> {
   let promise: Promise<void> | undefined;
   let result: Result | undefined;
@@ -112,10 +143,14 @@ function createInstance<Input, Result>(
 
   promise = (async () => {
     try {
-      result = await fetchFunc(input, controller.signal);
+      progressStore.getState().setProgress(input, 0);
+      result = await fetchFunc(input, controller.signal, (value) => {
+        progressStore.getState().setProgress(input, value);
+      });
     } catch (error_) {
       error = error_;
     } finally {
+      progressStore.getState().clearProgress(input);
       promise = undefined;
     }
   })();
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index d5caecc76..875f9290e 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -579,6 +579,10 @@ importers:
       d3-array:
         specifier: 3.2.4
         version: 3.2.4
+    optionalDependencies:
+      zustand:
+        specifier: 4.5.4
+        version: 4.5.4(@types/react@18.3.3)(react@18.3.1)
     devDependencies:
       '@types/d3-array':
         specifier: ~3.2.1