feat(datasource/graphene) added ability to choose a mesh bounding box

src/chunk_manager/backend.ts

@@ -450,8 +450,11 @@ function startChunkDownload(chunk: Chunk) {
     (error: any) => {
       if (chunk.downloadAbortController === downloadAbortController) {
         chunk.downloadAbortController = undefined;
-        chunk.downloadFailed(error);
-        console.log(`Error retrieving chunk ${chunk}: ${error}`);
+        if (error === "retry") {
+        } else {
+          chunk.downloadFailed(error);
+          console.log(`Error retrieving chunk ${chunk}: ${error}`);
+        }
       }
     },
   );

src/chunk_manager/base.ts

@@ -110,6 +110,10 @@ export interface ChunkSourceParametersConstructor<T> {
   RPC_ID: string;
 }
 
+export interface ChunkSourceStateConstructor<T> {
+  new (): T;
+}
+
 export class LayerChunkProgressInfo {
   numVisibleChunksNeeded = 0;
   numVisibleChunksAvailable = 0;

src/chunk_manager/frontend.ts

@@ -16,6 +16,7 @@
 
 import type {
   ChunkSourceParametersConstructor,
+  ChunkSourceStateConstructor,
   LayerChunkProgressInfo,
 } from "#src/chunk_manager/base.js";
 import {
@@ -493,22 +494,30 @@ export interface ChunkSource {
 
 export function WithParameters<
   Parameters,
+  State,
   TBase extends ChunkSourceConstructor,
 >(
   Base: TBase,
   parametersConstructor: ChunkSourceParametersConstructor<Parameters>,
+  state?: ChunkSourceStateConstructor<State>,
 ) {
+  state;
   type WithParametersOptions = InstanceType<TBase>["OPTIONS"] & {
     parameters: Parameters;
+    state?: State;
   };
   @registerSharedObjectOwner(parametersConstructor.RPC_ID)
   class C extends Base {
     declare OPTIONS: WithParametersOptions;
     parameters: Parameters;
+    state: State;
     constructor(...args: any[]) {
       super(...args);
       const options: WithParametersOptions = args[1];
       this.parameters = options.parameters;
+      if (options.state) {
+        this.state = options.state;
+      }
     }
     initializeCounterpart(rpc: RPC, options: any) {
       options.parameters = this.parameters;

src/datasource/graphene/backend.ts

@@ -27,13 +27,14 @@ import {
   getGrapheneFragmentKey,
   GRAPHENE_MESH_NEW_SEGMENT_RPC_ID,
   ChunkedGraphSourceParameters,
-  MeshSourceParameters,
+  MeshSourceParametersWithFocus,
   CHUNKED_GRAPH_LAYER_RPC_ID,
   CHUNKED_GRAPH_RENDER_LAYER_UPDATE_SOURCES_RPC_ID,
   RENDER_RATIO_LIMIT,
   isBaseSegmentId,
   parseGrapheneError,
   getHttpSource,
+  startLayerForBBox,
 } from "#src/datasource/graphene/base.js";
 import { decodeManifestChunk } from "#src/datasource/precomputed/backend.js";
 import { WithSharedKvStoreContextCounterpart } from "#src/kvstore/backend.js";
@@ -98,7 +99,7 @@ function downloadFragmentWithSharding(
 function downloadFragment(
   fragmentKvStore: KvStoreWithPath,
   fragmentId: string,
-  parameters: MeshSourceParameters,
+  parameters: MeshSourceParametersWithFocus,
   signal: AbortSignal,
 ): Promise<ReadResponse> {
   if (parameters.sharding) {
@@ -123,7 +124,7 @@ async function decodeDracoFragmentChunk(
 @registerSharedObject()
 export class GrapheneMeshSource extends WithParameters(
   WithSharedKvStoreContextCounterpart(MeshSource),
-  MeshSourceParameters,
+  MeshSourceParametersWithFocus,
 ) {
   manifestRequestCount = new Map<string, number>();
   newSegments = new Uint64Set();
@@ -136,6 +137,13 @@ export class GrapheneMeshSource extends WithParameters(
     this.parameters.fragmentUrl,
   );
 
+  focusBoundingBox: SharedWatchableValue<Float32Array | undefined>;
+
+  constructor(rpc: RPC, options: MeshSourceParametersWithFocus) {
+    super(rpc, options);
+    this.focusBoundingBox = rpc.get(options.focusBoundingBox);
+  }
+
   addNewSegment(segment: bigint) {
     const { newSegments } = this;
     newSegments.add(segment);
@@ -146,12 +154,41 @@ export class GrapheneMeshSource extends WithParameters(
   }
 
   async download(chunk: ManifestChunk, signal: AbortSignal) {
+    const {
+      focusBoundingBox: { value: focusBoundingBox },
+    } = this;
+    const { chunkSize, nBitsForLayerId } = this.parameters;
+
+    const unregister = this.registerDisposer(
+      this.focusBoundingBox.changed.add(() => {
+        chunk.downloadAbortController?.abort("retry");
+        unregister();
+        if (chunk.newRequestedState !== ChunkState.NEW) {
+          // re-download manifest
+          this.chunkManager.queueManager.updateChunkState(
+            chunk,
+            ChunkState.QUEUED,
+          );
+        }
+      }),
+    );
+
     const { parameters, newSegments, manifestRequestCount } = this;
-    if (isBaseSegmentId(chunk.objectId, parameters.nBitsForLayerId)) {
+    if (isBaseSegmentId(chunk.objectId, nBitsForLayerId)) {
       return decodeManifestChunk(chunk, { fragments: [] });
     }
+
     const { fetchOkImpl, baseUrl } = this.manifestHttpSource;
-    const manifestPath = `/manifest/${chunk.objectId}:${parameters.lod}?verify=1&prepend_seg_ids=1`;
+    let manifestPath = `/manifest/${chunk.objectId}:${parameters.lod}?verify=1&prepend_seg_ids=1`;
+    if (focusBoundingBox) {
+      const rank = focusBoundingBox.length / 2;
+      const startLayer = startLayerForBBox(focusBoundingBox, chunkSize);
+      const boundsStr = Array.from(
+        new Array(rank),
+        (_, i) => `${focusBoundingBox[i]}-${focusBoundingBox[i + rank]}`,
+      ).join("_");
+      manifestPath += `&bounds=${boundsStr}&start_layer=${startLayer}`;
+    }
     const response = await (
       await fetchOkImpl(baseUrl + manifestPath, { signal })
     ).json();

src/datasource/graphene/base.ts

@@ -27,7 +27,7 @@ import type {
   DataType,
 } from "#src/sliceview/base.js";
 import { makeSliceViewChunkSpecification } from "#src/sliceview/base.js";
-import type { mat4 } from "#src/util/geom.js";
+import type { mat4, vec3 } from "#src/util/geom.js";
 import type { FetchOk, HttpError } from "#src/util/http_request.js";
 
 export const PYCG_APP_VERSION = 1;
@@ -59,10 +59,14 @@ export class MeshSourceParameters {
   lod: number;
   sharding: Array<ShardingParameters> | undefined;
   nBitsForLayerId: number;
-
+  chunkSize: vec3;
   static RPC_ID = "graphene/MeshSource";
 }
 
+export class MeshSourceParametersWithFocus extends MeshSourceParameters {
+  focusBoundingBox: number;
+}
+
 export class MultiscaleMeshMetadata {
   transform: mat4;
   lodScaleMultiplier: number;
@@ -172,3 +176,17 @@ export function getHttpSource(
   }
   return { fetchOkImpl, baseUrl: joinBaseUrlAndPath(baseUrl, path) };
 }
+
+export const startLayerForBBox = (
+  focusBoundingBox: Float32Array<ArrayBufferLike>,
+  chunkSize: vec3,
+) => {
+  const rank = 3; // TODO need to change vec3 otherwise it doesn't make sense to make this a variable
+  let minChunks = Number.POSITIVE_INFINITY;
+  for (let i = 0; i < rank; i++) {
+    const length = focusBoundingBox[i + rank] - focusBoundingBox[i];
+    const numChunks = Math.ceil(length / chunkSize[i]);
+    minChunks = Math.min(minChunks, numChunks);
+  }
+  return 2 + Math.max(0, Math.floor(Math.log2(minChunks / 1)));
+};