diff --git a/CHANGELOG.md b/CHANGELOG.md index 07eebdb..56dca7b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -146,3 +146,14 @@ The first release, an alpha: the plugin API can still change before 0.1.0. every session in one WebAssembly instance that its WebGPU build suspends while it waits. Two models' passes, or a second model built while the first computed, could hang the page and crash the tab. A page's own passes on the same ORT take the turn through `exclusive`. +- `Engine.materialize(view, store, options)` writes a box of a served view (or a whole + level) into a store as an OME-Zarr of its own, zarr v2 with OME 0.4 or zarr v3 with OME + 0.5, gzipped, its empty chunks left out, the metadata the engine serves for the view + (`web/src/layout.ts`, shared with the routes) with the one level as its one dataset; + `directoryStore()` writes into a folder picked with the File System Access API (or the + origin's private file system), `memoryStore()` into memory + ([formats](docs/concepts/formats.md#a-region-written-to-a-store-browser-engine)). Its + chunks share the queue and the kept chunks with the viewer's but run in the background: + behind every chunk a client waits for, in all the queue's slots but one, and never dropped + as abandoned (`submit`, `claim` and `release` take `background`; + [caching](docs/concepts/caching.md#order-of-work-and-requests-given-up-on)). diff --git a/docs/concepts/caching.md b/docs/concepts/caching.md index ffc4426..d44158c 100644 --- a/docs/concepts/caching.md +++ b/docs/concepts/caching.md @@ -208,6 +208,32 @@ one thread, so the first view of a deep chain used one core. The browser page keeps the same rules (`web/src/demand.ts`); it learns of a given-up request from its streamed reply being cancelled, a service worker being told of nothing else. +The browser page has one class of work more: chunks it computes for itself, in the +background, to write a region of a view into a store (`Engine.materialize`; see +[formats](formats.md#a-region-written-to-a-store-browser-engine)). They go through the same +queue as the viewer's chunks, so a chunk the viewer has computed is taken from the kept +chunks and one on its way for the viewer is shared. Three rules keep browsing usable while +a region is written: + +* A background claim is never released because a client gave up, so nothing materialize + asked for is dropped as abandoned; only its own abort gives its waiting chunks up (and + they are dropped if no client wants them either). +* A job only background claims hold waits behind every job a client waits for, whatever + their levels; among themselves, background jobs keep the usual order. +* Background jobs run in all the queue's slots but one (in the one, when there is only + one), so a viewer's request finds a slot free, or the next one to free, rather than + waiting for the region's chunks in line before it. + +A viewer asking for a chunk that waits in the background brings it forward, and it goes +back if the viewer gives up on it. Materialize asks for a few chunks at once (as many as +the queue has slots, by default: the next one is waiting while one is written), not the +whole region, so the queue holds a few of its chunks at a time however large the region. +A chunk only materialize asked for is not kept once written (the kept chunks are for +clients that fetch a chunk again), so a large region does not push what the viewer is +looking at out of the 256 MB; one a client asked for too, while it was on its way, is kept +as any other. A staged view's stage outputs are kept as usual (512 MB), so writing a large +region does cycle through those. + ## What has been computed Zoomed out, a dataset shows its ops run on the coarse level, not a downsample of the diff --git a/docs/concepts/formats.md b/docs/concepts/formats.md index 98abec5..23c80d5 100644 --- a/docs/concepts/formats.md +++ b/docs/concepts/formats.md @@ -73,6 +73,50 @@ view as zarr v2 with OME-Zarr 0.4 (`virtual//zarr2//`), the same chu readers that predate zarr v3's final spec: GDAL 3.8, as the fires demo runs it in the page (gdal3.js), opens a level as a zarr array and translates it to a GeoTIFF. +## A region written to a store (browser engine) + +The browser engine also writes a view, or a box of it, into a store as files +(`Engine.materialize`; see [the design notes](../design.md#another-page-on-the-engine)): a +folder the visitor picks, the origin's private file system or memory. What it writes is +the layout it serves for the view, one level of it, so the result opens wherever the view +did and lands in the same place: the group lists the level written as its one dataset, +`0`, with that level's scale and translation; the level's array has the level's whole +shape (a box keeps its place in it), the chunk shape and dtype served, and chunks +compressed with gzip unless asked not to be; and the chunks of the box, rounded out to +whole chunks. A chunk of nothing but zeros is left out (zarr reads a missing chunk as the +fill value, 0), unless asked to be written. A channel axis the view's ops add is the +first axis and one chunk, as served. As zarr v2 with OME-Zarr 0.4 (the default, which +more tools read): + +```text +.zgroup {"zarr_format": 2} +.zattrs {...extra attributes, "multiscales": [{"name": , "axes": [...], + "datasets": [{"path": "0", "coordinateTransformations": [scale, translation]}], + "version": "0.4"}]} +0/.zarray {"zarr_format": 2, "shape": , "chunks": , "dtype": "// each chunk, gzipped +``` + +As zarr v3 with OME-Zarr 0.5: + +```text +zarr.json {"zarr_format": 3, "node_type": "group", "attributes": {...extra attributes, + "ome": {"version": "0.5", "multiscales": [{... "datasets": [{"path": "0", ...}]}]}}} +0/zarr.json the level's array as served, its codecs [{"name": "bytes", ...}, + {"name": "gzip", "configuration": {"level": 6}}] +0/c/// each chunk, gzipped +``` + +The metadata is written first, so a store whose writing was stopped part way is still a +valid OME-Zarr, zero where no chunk was written. Extra attributes (a consumer's +provenance) sit beside the OME ones in the group's attributes, which they cannot replace. +gzip is what a browser compresses with on its own (`CompressionStream`; level 6 is the +level it writes at, which no reader needs); the Python server's default, blosc, would need +an encoder of its own in the page. + ## Meshes, computed when fetched The `mesh` frontend serves a dataset's surface in Neuroglancer's precomputed mesh formats, diff --git a/docs/design.md b/docs/design.md index d6821d5..e27227e 100644 --- a/docs/design.md +++ b/docs/design.md @@ -516,6 +516,39 @@ as text). What such a page gets, in the order it uses it: so the stage is typed with the rest and run in `npm test` through the WebAssembly backend, which runs in Node as in a browser without WebGPU, on two models of a hundred bytes (`web/test/onnx.test.ts`, the models written by `web/test/fixtures/make_onnx.py`). +* `Engine.materialize(view, store, options)` writes a region of a served view into a store + as an OME-Zarr of its own (`web/src/materialize.ts`), to open again later or copy where + other tools read it: a model's output over a box, computed once in the page and kept. + The store is anything with `put(path, bytes)`; `directoryStore(dir, root)` makes one of + a `FileSystemDirectoryHandle` (a folder the visitor picks with + `showDirectoryPicker({mode: "readwrite"})`, in Chrome and Edge, or the origin's private + file system, `navigator.storage.getDirectory()`), its subdirectories made as needed and + each file written whole through `createWritable()`, and `memoryStore()` one of a `Map`, + for tests and for zipping a small result. The options: the view's `level` (the finest by + default), a voxel box `lo`, `hi` on it (the whole level by default; rounded out to whole + chunks and clipped to the level), `format` (`"zarr2"`, OME-Zarr 0.4, the default, since + more tools read it; or `"zarr3"`, OME-Zarr 0.5), `compression` (`"gzip"`, the default, or + `null`), `skipEmpty` (a chunk of nothing but zeros left out, by default), `concurrency`, + `signal` (an `AbortSignal`), `onProgress` (`done`, `total`, `written`, `bytes`) and + `attributes` (the consumer's provenance, beside the OME metadata). It resolves with + `{chunks, written, skipped, bytes, aborted}`. The layout is the one the engine serves for + the view, from one module (`web/src/layout.ts`) for both, with the one level written as + the group's one dataset at `0` with its own scale and translation, so it lands where the + view does ([formats](concepts/formats.md#a-region-written-to-a-store-browser-engine)). A + single level rather than the pyramid: the engine downsamples nothing itself, and every + level would be the model run again on its own grid. The metadata is written first, so a + store whose writing stopped is still an OME-Zarr, zero where no chunk was written. The + chunks come from the same producer the routes use, through the same queue, but in the + background: they wait behind the viewer's, in all its slots but one, and are never dropped + as abandoned ([caching](concepts/caching.md#order-of-work-and-requests-given-up-on)), so + browsing goes on while a region is written. An abort starts no more chunks, gives up + those waiting for their turn, lets the writes under way end and resolves at once with + `aborted`; a chunk that fails to compute or be written rejects it with the chunk's key + and why, once the writes under way have ended. This is a first piece of + [roadmap](roadmap.md) item 2: an explicit write of a box, asked for by the page, not yet + the cache on disk that browsing fills. The parts that need no engine (the chunks of a + box, the files, gzip, the loop over any producer, the stores) run in `npm test`, the + result read back by zarrita (`web/test/materialize.test.ts`). * `Engine.unserve(view)` stops serving a view: its infos, producer and kept chunks are forgotten and chunks of it waiting to be computed dropped, so a page that serves a model again under a new name (viewers cache by URL, so an edit is a new name) does not keep the diff --git a/docs/roadmap.md b/docs/roadmap.md index a7d0aa2..e123376 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -36,7 +36,12 @@ describe shipped features as future work. viewer becomes the job scheduler. (Ordering the work by what clients ask for, and dropping what they stop waiting for, has shipped; see [caching](concepts/caching.md#order-of-work-and-requests-given-up-on). Inference ops - would join its queues.) + would join its queues. A first piece has shipped in the browser engine: an explicit + write of a box of a view into a store, a folder the visitor picks or memory, as an + OME-Zarr, its chunks computed in the background behind the viewer's + ([formats](concepts/formats.md#a-region-written-to-a-store-browser-engine)). The cache on + disk itself, in either engine, its sharing across workers and the filler that grows it + from what is browsed are still to do.) 3. **Lazy-compute story for non-viewer clients.** Demonstrate dask/tensorstore reading a served pipeline and running downstream analysis with no intermediate written. (The read itself is shown in [getting started](getting-started.md#read-the-result-from-python-instead-of-a-viewer); diff --git a/web/src/demand.ts b/web/src/demand.ts index c06aed7..d770b05 100644 --- a/web/src/demand.ts +++ b/web/src/demand.ts @@ -8,6 +8,14 @@ // placeholder while the fine chunk computes) wants the finer. Within a level the first asked // for goes first. A few run at once, so the GPU works on one while the next ones read, and the // rest cannot flood the network with reads. +// +// One rule more than chunkmirage.demand's: work asked for in the background (a page writing a +// region to a store, Engine.materialize), not for a client waiting on it. A background claim +// is never released by a client giving up, so the work is not dropped, but a job only +// background claims hold waits behind every other, whatever its level, and such jobs take at +// most all the slots but one (one, if there is only one): a viewer's request finds a slot free +// or the next one to free, while the background work goes on in the rest. A client asking for +// a job waiting in the background brings it forward, and it goes back if the client gives up. const SLOTS = 3; // jobs run at once @@ -26,16 +34,21 @@ export class Cancelled extends Error { constructor(why = "no request wants this interface Job { run: () => Promise; level: number; seq: number; // seq: when first asked for - demand: number; // requests holding a claim on it + demand: number; // requests holding a claim on it + background: number; // of them, background claims: a job only they hold waits behind the rest resolve: (v: T) => void; reject: (e: unknown) => void; onDrop: () => void; } /** Jobs waiting, in the order described at the top; `slots` of them run at once. */ export class Queue { - private slots: number; - constructor(slots = SLOTS) { this.slots = slots; } + /** Jobs run at once. */ + readonly slots: number; + /** Of them, jobs only background claims hold (`submit`'s `background`): all but one, so a + * client's request finds a slot free; one, if there is only one. */ + readonly backgroundSlots: number; + constructor(slots = SLOTS) { this.slots = slots; this.backgroundSlots = Math.max(1, slots - 1); } private waiting = new Map>(); - private active = new Map(); // running: key -> level + private active = new Map(); // running, by key dropped = 0; // waiting fits no request wanted any more, never run onChange: () => void = () => {}; // a fit waiting, started, done or dropped private seq = 0; @@ -44,26 +57,38 @@ export class Queue { /** Fits running and waiting, per level. */ byLevel(): { running: Map; waiting: Map } { const count = (levels: Iterable) => { const m = new Map(); for (const l of levels) m.set(l, (m.get(l) ?? 0) + 1); return m; }; - return { running: count(this.active.values()), waiting: count([...this.waiting.values()].map((j) => j.level)) }; + return { running: count([...this.active.values()].map((a) => a.level)), waiting: count([...this.waiting.values()].map((j) => j.level)) }; } - /** A new fit, wanted by one request; `onDrop` runs, synchronously, if it is dropped unrun. */ - submit(key: string, level: number, run: () => Promise, onDrop: () => void): Promise { + /** A new fit, wanted by one request (one in the background, with `background`); `onDrop` + * runs, synchronously, if it is dropped unrun. */ + submit(key: string, level: number, run: () => Promise, onDrop: () => void, background = false): Promise { return new Promise((resolve, reject) => { - const job: Job = { run, level, resolve, reject, onDrop, seq: ++this.seq, demand: 1 }; + const job: Job = { run, level, resolve, reject, onDrop, seq: ++this.seq, demand: 1, background: background ? 1 : 0 }; this.waiting.set(key, job); this.pump(); this.onChange(); }); } - /** Another request wants a fit: if it is still waiting, it counts one more. */ - claim(key: string) { const job = this.waiting.get(key); if (job) job.demand++; } + /** Another request wants a fit (one in the background, with `background`): if it is still + * waiting, it counts one more, and a client's claim brings a background job forward. */ + claim(key: string, background = false) { + const job = this.waiting.get(key); + if (!job) return; + job.demand++; + if (background) job.background++; + else this.pump(); // in the background before: it may take a slot kept for clients + } - /** A request no longer wants a fit: waiting and wanted by no one, it is dropped unrun. */ - release(key: string) { + /** A request no longer wants a fit (`background`: one that claimed it in the background): + * waiting and wanted by no one, it is dropped unrun; wanted in the background alone, it goes + * back behind the rest. */ + release(key: string, background = false) { const job = this.waiting.get(key); - if (!job || --job.demand > 0) return; + if (!job) return; + if (background && job.background > 0) job.background--; + if (--job.demand > 0) return; this.drop(key); } @@ -84,18 +109,26 @@ export class Queue { while (this.active.size < this.slots && this.waiting.size) { let best: [string, Job] | null = null; for (const e of this.waiting) if (!best || before(e[1], best[1])) best = e; - const [key, job] = best!; + const [key, job] = best!, background = inBackground(job); + // the best in the background: so is every job waiting, and they have slots of their own + if (background && [...this.active.values()].filter((a) => a.background).length >= this.backgroundSlots) return; this.waiting.delete(key); - this.active.set(key, job.level); + this.active.set(key, { level: job.level, background }); job.run().then(job.resolve, job.reject).finally(() => { this.active.delete(key); this.pump(); this.onChange(); }); } } } -/** Which of two waiting fits goes first: the finer level (a client asking for one place at - * two levels, as a viewer does to show a coarse placeholder while the fine chunk computes, - * wants the finer one; coarser fits run when nothing finer waits), then the first asked for. */ +/** Whether only background claims hold a job. */ +const inBackground = (job: Job) => job.demand === job.background; + +/** Which of two waiting fits goes first: one a client waits for before one wanted in the + * background alone; then the finer level (a client asking for one place at two levels, as a + * viewer does to show a coarse placeholder while the fine chunk computes, wants the finer + * one; coarser fits run when nothing finer waits), then the first asked for. */ function before(a: Job, b: Job): boolean { + const ba = inBackground(a), bb = inBackground(b); + if (ba !== bb) return bb; return a.level !== b.level ? a.level < b.level : a.seq < b.seq; } diff --git a/web/src/engine.ts b/web/src/engine.ts index 19856fb..d6add92 100644 --- a/web/src/engine.ts +++ b/web/src/engine.ts @@ -10,17 +10,20 @@ // Neuroglancer (virtual///...) and GeoZarr for map clients such as OpenLayers // (virtual//geo////...), the zarr-conventions multiscales, proj: // and spatial: attributes on the group. And zarr v2 with OME 0.4 (virtual//zarr2//), -// for readers that predate zarr v3's final spec, such as GDAL 3.8. +// for readers that predate zarr v3's final spec, such as GDAL 3.8. A region of a view is +// written into a store as the same OME-Zarr by `materialize` (materialize.ts). import { asRead, computedSource, needsPython, TYPED, wholeChunk } from "./asread"; import { Cancelled, Claim, Queue } from "./demand"; import schema from "./generated/chunkmirage.schema.json"; import { Kept } from "./kept"; +import { array2, array3, attrs2, dimensionNames, group3 } from "./layout"; +import { region, writeChunks, type MaterializeOptions, type Materialized, type WritableStore } from "./materialize"; import { route } from "./routes"; import { bytes, fitted, inputBox, levelAt, padded, prod, type Block, type Padding, type Stage, type StageAt } from "./stages"; -import type { Answer, Ask, Later, MeshCall, PipelineView, Plan, Printed, Reply, SourceInfo, ToPyWorker, ToReader, ViewAxis, ViewInfo } from "./types"; +import type { Answer, Ask, Later, MeshCall, PipelineView, Plan, Printed, Reply, SourceInfo, ToPyWorker, ToReader, ViewInfo } from "./types"; -export const TO_SECONDS: Record = { s: 1, second: 1, millisecond: 1e-3, ms: 1e-3, minute: 60, hour: 3600, day: 86400 }; -const OME_UNIT: Record = { nm: "nanometer", um: "micrometer", m: "meter", s: "second" }; +export { isTime, TO_SECONDS } from "./layout"; +export { directoryStore, memoryStore, type MaterializeOptions, type Materialized, type MaterializeProgress, type WritableStore } from "./materialize"; const CONVENTIONS = [ { uuid: "d35379db-88df-4056-af3a-620245f8e347", name: "multiscales" }, { uuid: "f17cb550-5864-4468-aeb7-f3180cfb622f", name: "proj:" }, @@ -28,9 +31,6 @@ const CONVENTIONS = [ ]; const KEPT_BYTES = 256 * 2 ** 20; // computed chunks kept for clients that refetch const STAGED_BYTES = 512 * 2 ** 20; // the outputs of a staged view's stages (stages.ts), kept for the stages after them -const V2_DTYPE: Record = { - uint8: "|u1", int8: "|i1", uint16: " = Object.fromEntries(Object.values(schema.$defs as Record) .flatMap((d) => (d.properties?.op?.const && d["x-packages"] ? [[d.properties.op.const, d["x-packages"]]] : []))); @@ -77,7 +77,6 @@ class Rpc { const asJson = (o: unknown) => ({ status: 200, body: JSON.stringify(o), type: "application/json" }); const notFound = { status: 404, body: "", type: "text/plain" }; -export const isTime = (a: ViewAxis) => a.unit in TO_SECONDS; /** `Engine.start`'s options. `packages`: what Python loads beyond numpy and the ops' own * imports (each op's schema says those). `python`: when the Pyodide workers start; "auto" @@ -107,6 +106,9 @@ export class Engine { private chunks: Queue | null = null; private inflight = new Map>(); private kept = new Kept(KEPT_BYTES); + /** The keys of chunks on their way that a client asked for (`chunk` without `background`): + * only those are kept once computed. */ + private clientWanted = new Set(); private staged = new Kept(STAGED_BYTES); // by stage key, not view: shared by views, kept past an unserve private edits = 0; private producers = new Map Promise>(); @@ -389,20 +391,8 @@ export class Engine { get running(): number { return this.chunks?.running ?? 0; } get waiting(): number { return this.chunks ? [...this.chunks.byLevel().waiting.values()].reduce((a, b) => a + b, 0) : 0; } - private omeGroup(view: string) { - const v = this.infos[view], added = (v.added ?? []).map(() => 0); // channel axes the ops add - return { - zarr_format: 3, node_type: "group", - attributes: { ome: { version: "0.5", multiscales: [{ - name: view, - axes: [...added.map(() => ({ name: "c", type: "channel" })), - ...v.axes.map((a) => ({ name: a.name, type: isTime(a) ? "time" : "space", unit: OME_UNIT[a.unit] ?? (a.unit || undefined) }))], - datasets: v.levels.map((l, i) => ({ path: String(i), coordinateTransformations: [ - { type: "scale", scale: [...added.map(() => 1), ...l.voxel] }, { type: "translation", translation: [...added, ...l.origin] }, - ] })), - }] } }, - }; - } + /** A view's levels, each a dataset of its group at its own index. */ + private datasets(view: string) { return this.infos[view].levels.map((_, level) => ({ path: String(level), level })); } /** The GeoZarr group of a view: the conventions, its levels as the multiscales layout, its * projection (the source's own, which the client must know by `proj`) and bounding box. */ @@ -419,20 +409,6 @@ export class Engine { }; } - /** A level's array; in the GeoZarr layout (`map`) as the map's y, x band, the view's one - * plane along z (a map client reads 2-D bands: with a third axis it would need the - * group's consolidated metadata to know to select it). */ - private array(view: string, level: number, map = false) { - const v = this.infos[view], keep = map ? -2 : 0, added = map ? [] : v.added ?? []; // channels: chunked whole - return { - zarr_format: 3, node_type: "array", shape: [...added, ...v.levels[level].shape.slice(keep)], data_type: v.out, fill_value: 0, - chunk_grid: { name: "regular", configuration: { chunk_shape: [...added, ...this.views[view].chunk.slice(keep)] } }, - chunk_key_encoding: { name: "default", configuration: { separator: "/" } }, - codecs: [{ name: "bytes", configuration: { endian: "little" } }], - dimension_names: [...added.map(() => "c"), ...v.axes.map((a) => a.name).slice(keep)], attributes: {}, - }; - } - /** The level a view's mesh is made from (chunkmirage.meshes.mesh_level). */ meshLevel(view: string): number { const v = this.sources[view], spec = this.views[view].mesh ?? {}; @@ -490,28 +466,74 @@ export class Engine { }, [block.buffer]); } - private chunk(view: string, level: number, index: number[], mesh?: MeshCall): Answered { + /** Chunk `index` of a view's level (or with `mesh`, its mesh): kept, or on its way for + * another request, or queued. `background`: for `materialize`, not a client, so it waits + * behind the clients' work and no client giving up drops it (demand.ts). */ + private chunk(view: string, level: number, index: number[], mesh?: MeshCall, background = false): Answered { const key = `${view}/${mesh ? `mesh${level}` : level}/${index.join(".")}`, done = this.kept.get(key); if (done) return { status: 200, body: done.slice(0), type: "application/octet-stream" }; const chunks = this.chunks!, claim = new Claim(); + if (!background) this.clientWanted.add(key); let pending = this.inflight.get(key); - if (pending) chunks.claim(key); // another request waits for it too + if (pending) chunks.claim(key, background); // another request waits for it too else { - pending = chunks.submit(key, level, () => this.compute(view, level, index, mesh), () => { this.counts.dropped++; }) - .then((b) => { this.counts.computed++; if (this.views[view]) this.kept.put(key, b); return b; }, (e) => { // kept unless unserved meanwhile + pending = chunks.submit(key, level, () => this.compute(view, level, index, mesh), () => { this.counts.dropped++; }, background) + .then((b) => { + // Kept for a client that fetches it again, unless the view was unserved meanwhile or + // only materialize asked for it: a region written chunk by chunk would otherwise push + // what the viewer is looking at out of the 256 MB, for chunks no client may ask for. + this.counts.computed++; + if (this.views[view] && this.clientWanted.has(key)) this.kept.put(key, b); + return b; + }, (e) => { if (!(e instanceof Cancelled)) { this.counts.failed++; this.failure = { view, message: String((e as Error).message ?? e) }; console.error(`chunk ${key}: ${String((e as Error).message ?? e).trim().split("\n").slice(-3).join(" / ")}`); } throw e; }) - .finally(() => { this.inflight.delete(key); this.onChange(); }); + .finally(() => { this.inflight.delete(key); this.clientWanted.delete(key); this.onChange(); }); this.inflight.set(key, pending); } - claim.onCancel(() => chunks.release(key)); + claim.onCancel(() => chunks.release(key, background)); const mine = pending.then((b) => b.slice(0)); // each request its own copy: replies transfer it mine.catch(() => {}); this.onChange(); return { status: 200, type: "application/octet-stream", pending: mine, claim }; } + /** Write a region of `view` into `store` as an OME-Zarr of its own (materialize.ts): one + * level of the view (`options.level`, the finest by default) over a box of its voxels + * (`lo`, `hi`, rounded out to whole chunks; the whole level by default), as zarr v2 with + * OME-Zarr 0.4 or zarr v3 with OME-Zarr 0.5, gzipped, its chunks of nothing but zeros left + * out. The group and the array are written first, so the store reads as the view from the + * start, zero where no chunk is yet; then the chunks, each computed as a client's request + * for it is (kept, or shared with one on its way, and kept after) but in the background: + * behind every chunk a client waits for, in all the queue's slots but one, and never + * dropped because a client gave up on it (demand.ts). `concurrency` chunks are asked for + * at once, as many as the queue has slots by default. Resolves with what was written, or + * on `signal`'s abort with `aborted` (chunks waiting for their turn are dropped, the ones + * written stay); rejects naming the chunk that failed to compute or be written, once the + * writes under way have ended. */ + async materialize(view: string, store: WritableStore, options: MaterializeOptions = {}): Promise { + if (!this.chunks) throw new Error(this.stopped ? "the engine was stopped" : "the engine has not started"); + if (!this.infos[view] && this.sources[view]) await this.ready; // its ops are being planned + const v = this.infos[view]; + if (!v) throw new Error(`${view} is not served`); + const r = region(view, v, this.views[view].chunk, options); + return writeChunks(store, r, (index, signal) => this.backgroundChunk(view, r.level, index, signal), { ...options, concurrency: options.concurrency ?? this.chunks?.slots ?? 1 }); + } + + /** Chunk `index` of a view's level for `materialize`: its whole bytes, as a client's request + * gets them; `signal`'s abort gives it up (dropped unrun if no client wants it either). */ + private async backgroundChunk(view: string, level: number, index: number[], signal: AbortSignal): Promise { + if (!this.chunks) throw new Error("the engine was stopped"); + if (!this.infos[view]) throw new Error(`${view} is no longer served`); + signal.throwIfAborted(); + const a = this.chunk(view, level, index, undefined, true)!; + if (!("pending" in a)) return a.body as ArrayBuffer; // kept + const cancel = () => a.claim.cancel(); + signal.addEventListener("abort", cancel, { once: true }); + try { return await a.pending; } finally { signal.removeEventListener("abort", cancel); } // done: no claim left to give up + } + /** Projection codes of the GeoZarr groups, by view (set by the map page). */ proj: Record = {}; @@ -539,12 +561,12 @@ export class Engine { view = rest[0]; rest = rest.slice(1); if (!this.infos[view]) return notFound; - if (rest.length === 1 && rest[0] === "zarr.json") return asJson(this.omeGroup(view)); + if (rest.length === 1 && rest[0] === "zarr.json") return asJson(group3(view, this.infos[view], this.datasets(view))); if (rest[0] === "mesh" && this.views[view].mesh) return this.mesh(view, rest.slice(1).join("/"), range); } const level = Number(rest[0]); if (!this.infos[view].levels[level]) return notFound; - if (rest.length === 2 && rest[1] === "zarr.json") return asJson(this.array(view, level, map)); + if (rest.length === 2 && rest[1] === "zarr.json") return asJson(array3(this.infos[view], this.views[view].chunk, level, { map })); const added = map ? 0 : this.infos[view].added?.length ?? 0, key = rest.slice(2).map(Number); if (rest[1] !== "c" || key.length !== added + this.infos[view].levels[level].shape.length || key.slice(0, added).some((i) => i !== 0)) return notFound; return this.chunk(view, level, key.slice(added)); // a channel axis is one chunk: the volume's index @@ -555,21 +577,13 @@ export class Engine { private zarr2(view: string, rest: string[]): Answered { const v = this.infos[view]; if (!v) return notFound; - const added = v.added ?? [], names = [...added.map(() => "c"), ...v.axes.map((a) => a.name)]; + const added = v.added ?? []; if (rest.length === 1 && rest[0] === ".zgroup") return asJson({ zarr_format: 2 }); - if (rest.length === 1 && rest[0] === ".zattrs") { - const g = this.omeGroup(view).attributes.ome.multiscales[0]; - return asJson({ multiscales: [{ ...g, version: "0.4" }] }); - } + if (rest.length === 1 && rest[0] === ".zattrs") return asJson(attrs2(view, v, this.datasets(view))); const level = Number(rest[0]), l = v.levels[level]; if (!l) return notFound; - if (rest.length === 2 && rest[1] === ".zattrs") return asJson({ _ARRAY_DIMENSIONS: names }); - if (rest.length === 2 && rest[1] === ".zarray") { - return asJson({ - zarr_format: 2, shape: [...added, ...l.shape], chunks: [...added, ...this.views[view].chunk], dtype: V2_DTYPE[v.out] ?? v.out, - compressor: null, fill_value: 0, order: "C", filters: null, dimension_separator: "/", - }); - } + if (rest.length === 2 && rest[1] === ".zattrs") return asJson({ _ARRAY_DIMENSIONS: dimensionNames(v) }); + if (rest.length === 2 && rest[1] === ".zarray") return asJson(array2(v, this.views[view].chunk, level)); const key = rest.slice(1).map(Number); if (key.length !== added.length + l.shape.length || key.some((i) => !Number.isInteger(i)) || key.slice(0, added.length).some((i) => i !== 0)) return notFound; return this.chunk(view, level, key.slice(added.length)); diff --git a/web/src/layout.ts b/web/src/layout.ts new file mode 100644 index 0000000..ecf24d5 --- /dev/null +++ b/web/src/layout.ts @@ -0,0 +1,89 @@ +// A view's OME-Zarr layouts, as the engine serves them (engine.ts's routes) and writes them +// into a store (materialize.ts), from one place so the two cannot differ: zarr v3 with +// OME-Zarr 0.5, and zarr v2 with OME-Zarr 0.4. The group's metadata, a level's array +// metadata and a chunk's key, as objects ready for JSON. A channel axis the view's ops add +// (`added`) comes first and is one chunk, as the routes serve it. Nothing here imports a +// value, so Node's tests run it. +import type { ViewAxis, ViewInfo } from "./types.ts"; + +export const TO_SECONDS: Record = { s: 1, second: 1, millisecond: 1e-3, ms: 1e-3, minute: 60, hour: 3600, day: 86400 }; +const OME_UNIT: Record = { nm: "nanometer", um: "micrometer", m: "meter", s: "second" }; +/** zarr v2's dtypes, by numpy's name (the engine's dtypes); little endian. */ +export const V2_DTYPE: Record = { + uint8: "|u1", int8: "|i1", uint16: " a.unit in TO_SECONDS; + +/** A level of a view as one of a group's datasets: its path in the group, and which of the + * view's levels it is. */ +export interface Dataset { path: string; level: number } + +/** A zarr v3 codec after `bytes` (a compressor). */ +export interface Codec { name: string; configuration: Record } + +/** The OME multiscales entry of view `name`, whose datasets are the levels `datasets` name: + * the channel axes its ops add, then its own (space, or time by their unit); each level's + * scale its voxel size and its translation its origin, so it lands where the view does. */ +export function multiscale(name: string, v: ViewInfo, datasets: Dataset[]) { + const added = (v.added ?? []).map(() => 0); + return { + name, + axes: [...added.map(() => ({ name: "c", type: "channel" })), + ...v.axes.map((a) => ({ name: a.name, type: isTime(a) ? "time" : "space", unit: OME_UNIT[a.unit] ?? (a.unit || undefined) }))], + datasets: datasets.map(({ path, level }) => ({ path, coordinateTransformations: [ + { type: "scale", scale: [...added.map(() => 1), ...v.levels[level].voxel] }, + { type: "translation", translation: [...added, ...v.levels[level].origin] }, + ] })), + }; +} + +/** A view's zarr v3 group (zarr.json): OME-Zarr 0.5's multiscales in its attributes, beside + * any `extra` attributes (which cannot replace `ome`). */ +export function group3(name: string, v: ViewInfo, datasets: Dataset[], extra: Record = {}) { + return { zarr_format: 3, node_type: "group", attributes: { ...extra, ome: { version: "0.5", multiscales: [multiscale(name, v, datasets)] } } }; +} + +/** A view's zarr v2 group attributes (.zattrs): OME-Zarr 0.4's multiscales, beside any + * `extra` attributes (which cannot replace `multiscales`). The group itself is + * `{"zarr_format": 2}` (.zgroup). */ +export function attrs2(name: string, v: ViewInfo, datasets: Dataset[], extra: Record = {}) { + return { ...extra, multiscales: [{ ...multiscale(name, v, datasets), version: "0.4" }] }; +} + +/** The names of a level's axes: "c" for each channel axis the ops add, then the view's. */ +export const dimensionNames = (v: ViewInfo) => [...(v.added ?? []).map(() => "c"), ...v.axes.map((a) => a.name)]; + +/** A level's zarr v3 array (its zarr.json), chunked `chunk` (the level's own axes), the + * channel axes the ops add one chunk each; its bytes little endian, then `codecs`. In the + * GeoZarr layout (`map`), the map's y, x band alone, the view's one plane along z (a map + * client reads 2-D bands: with a third axis it would need the group's consolidated + * metadata to know to select it). */ +export function array3(v: ViewInfo, chunk: number[], level: number, { map = false, codecs = [] }: { map?: boolean; codecs?: Codec[] } = {}) { + const keep = map ? -2 : 0, added = map ? [] : v.added ?? []; + return { + zarr_format: 3, node_type: "array", shape: [...added, ...v.levels[level].shape.slice(keep)], data_type: v.out, fill_value: 0, + chunk_grid: { name: "regular", configuration: { chunk_shape: [...added, ...chunk.slice(keep)] } }, + chunk_key_encoding: { name: "default", configuration: { separator: "/" } }, + codecs: [{ name: "bytes", configuration: { endian: "little" } }, ...codecs], + dimension_names: map ? v.axes.map((a) => a.name).slice(keep) : dimensionNames(v), attributes: {}, + }; +} + +/** A level's zarr v2 array (.zarray), chunked `chunk`, the channel axes the ops add one + * chunk each, compressed by `compressor` (numcodecs' config; null: as they are). Its axes' + * names are its .zattrs, `{"_ARRAY_DIMENSIONS": dimensionNames(v)}`. */ +export function array2(v: ViewInfo, chunk: number[], level: number, compressor: Record | null = null) { + const added = v.added ?? []; + return { + zarr_format: 2, shape: [...added, ...v.levels[level].shape], chunks: [...added, ...chunk], dtype: V2_DTYPE[v.out] ?? v.out, + compressor, fill_value: 0, order: "C", filters: null, dimension_separator: "/", + }; +} + +/** The key of chunk `index` (on the level's own axes) of the array at `path`: zarr v3's + * default encoding (`c/`, then the indices "/" apart), or zarr v2's with "/" as its + * dimension separator; each of the `added` channel axes is one chunk, index 0. */ +export function chunkKey(format: "zarr3" | "zarr2", path: string, added: number, index: number[]): string { + const at = [...Array(added).fill(0), ...index].join("/"); + return format === "zarr3" ? `${path}/c/${at}` : `${path}/${at}`; +} diff --git a/web/src/materialize.ts b/web/src/materialize.ts new file mode 100644 index 0000000..b16dc8a --- /dev/null +++ b/web/src/materialize.ts @@ -0,0 +1,278 @@ +// A region of a view the engine serves, written into a store as an OME-Zarr of its own +// (Engine.materialize): a box of the view computed once and kept, to open again later or to +// copy where other tools read it. The store is a directory (one the visitor picks with the +// File System Access API's showDirectoryPicker, in Chrome and Edge, or the origin's private +// file system) or memory. What is written is what the engine serves for the view +// (layout.ts), one level of it, its chunks compressed. Here are the parts that need no +// engine, so Node's tests run them: the chunks a box covers, the files written, gzip, and +// the loop that has the chunks made and writes them, from any producer (the engine's queue +// on the page, a fake one in the tests). +import { array2, array3, attrs2, chunkKey, dimensionNames, group3 } from "./layout.ts"; // with its extension: Node's runner resolves no other +import type { ViewInfo } from "./types.ts"; + +/** Where materialize writes: files by "/"-separated path relative to the store's root. */ +export interface WritableStore { + put(path: string, bytes: Uint8Array): Promise; +} + +/** The names of `path`, "/" apart; empty ones (a leading, trailing or doubled "/") left out. + * A store's paths go down from its root, so "." and ".." are refused. */ +function names(path: string): string[] { + const parts = path.split("/").filter(Boolean); + if (parts.some((p) => p === "." || p === "..")) throw new Error(`${path}: a store's paths go down from its root, with no "." or ".."`); + return parts; +} + +/** A directory picked with showDirectoryPicker({mode: "readwrite"}) (or any + * FileSystemDirectoryHandle, e.g. OPFS's navigator.storage.getDirectory()) as a + * WritableStore: subdirectories made as needed (getDirectoryHandle(name, {create: true}), + * each once however many files go in it), each file written whole with createWritable(), + * which writes to a swap file that replaces the file when closed, so a file is never seen + * half written. `root` (optional) is a subdirectory path under it, made as needed (e.g. + * "mito.zarr"). */ +export function directoryStore(dir: FileSystemDirectoryHandle, root?: string): WritableStore { + const base = root ? names(root) : []; + const dirs = new Map>(); + const at = (parts: string[]): Promise => { + if (!parts.length) return Promise.resolve(dir); + const key = parts.join("/"); + let made = dirs.get(key); + if (!made) { + made = at(parts.slice(0, -1)).then((parent) => parent.getDirectoryHandle(parts[parts.length - 1], { create: true })); + dirs.set(key, made); + made.catch(() => dirs.delete(key)); // tried again by the next file in it + } + return made; + }; + return { + async put(path, bytes) { + const parts = [...base, ...names(path)]; + if (parts.length === base.length) throw new Error(`"${path}" is not a file's path`); + const file = await (await at(parts.slice(0, -1))).getFileHandle(parts[parts.length - 1], { create: true }); + const w = await file.createWritable(); + try { + await w.write(bytes as Uint8Array); + await w.close(); + } catch (e) { + await w.abort().catch(() => {}); // the swap file discarded, the file as it was + throw e; + } + }, + }; +} + +/** An in-memory store (Map, by path without a leading "/"), for tests + * and for zipping small results. Each file is a copy of the bytes put. */ +export function memoryStore(): WritableStore & { files: Map } { + const files = new Map(); + return { + files, + async put(path, bytes) { + const parts = names(path); + if (!parts.length) throw new Error(`"${path}" is not a file's path`); + files.set(parts.join("/"), bytes.slice()); + }, + }; +} + +/** How a materialize is going: chunks done (written or skipped as empty) of `total`, and of + * them `written`, in `bytes` as stored. */ +export interface MaterializeProgress { done: number; total: number; written: number; bytes: number } + +/** `Engine.materialize`'s options. */ +export interface MaterializeOptions { + /** The view's level to write (default 0, the finest). */ + level?: number; + /** Voxel box [lo, hi) on that level, in the level's voxel axes (z, y, x...); default the + * whole level. Rounded out to whole chunks, and clipped to the level. */ + lo?: number[]; hi?: number[]; + /** "zarr2" (OME-Zarr 0.4: .zgroup, .zattrs, .zarray, "/" as the dimension separator) or + * "zarr3" (OME-Zarr 0.5, zarr.json); default "zarr2", which more tools read. */ + format?: "zarr3" | "zarr2"; + /** gzip through CompressionStream (zarr v2: the compressor {"id": "gzip", "level": 6}; + * zarr v3: the "gzip" codec after "bytes"), or null for none; default "gzip". */ + compression?: "gzip" | null; + /** A chunk of nothing but zero bytes (the fill value) is not written, zarr reading a + * missing chunk as the fill value; default true. */ + skipEmpty?: boolean; + /** Chunks asked for at once (default: as many as the engine's queue has slots, of which + * all but one run materialize's chunks at most, so the next is waiting while one is + * written). */ + concurrency?: number; + /** Aborted: no chunk is started after it, chunks waiting for their turn are given up, the + * writes under way end, and materialize resolves with `aborted`; chunks written stay. */ + signal?: AbortSignal; + /** Told once the metadata is written (done 0), then after each chunk. */ + onProgress?: (p: MaterializeProgress) => void; + /** Attributes beside the OME multiscales in the group's attributes (the consumer's + * provenance); `ome` (zarr v3) or `multiscales` (zarr v2) among them is the view's. */ + attributes?: Record; +} + +/** What a materialize wrote: of the box's `chunks`, those `written` (in `bytes`, as + * stored) and those `skipped` as empty; `aborted` if its signal stopped it first. */ +export interface Materialized { chunks: number; written: number; skipped: number; bytes: number; aborted: boolean } + +/** The written level's path in the store: its one dataset, "0", as an OME-Zarr's finest + * level is, whichever level of the view it is (the scale and translation say where it lies). */ +export const LEVEL_PATH = "0"; +/** The gzip level written in the metadata: CompressionStream's, zlib's default (a reader + * needs none of it to decompress). */ +export const GZIP_LEVEL = 6; + +/** The chunk grid's indices [lo, hi) covering voxels [lo, hi) of a level of `shape` cut in + * chunks of `chunk`: the box rounded out to whole chunks and clipped to the level; the + * whole level where `lo` or `hi` is left out. A box with no voxel of the level is refused. */ +export function chunkRange(shape: number[], chunk: number[], lo?: number[], hi?: number[]): { lo: number[]; hi: number[] } { + for (const [what, box] of [["lo", lo], ["hi", hi]] as const) { + if (box && (box.length !== shape.length || box.some((x) => typeof x !== "number" || Number.isNaN(x)))) { + throw new Error(`${what} ${JSON.stringify(box)}: one number for each of the level's ${shape.length} axes`); + } + } + const inLo = shape.map((n, a) => Math.min(Math.max(lo?.[a] ?? 0, 0), n)); // the box clipped to the level + const inHi = shape.map((n, a) => Math.min(Math.max(hi?.[a] ?? n, 0), n)); + if (inHi.some((h, a) => h <= inLo[a])) { + throw new Error(`the box from [${(lo ?? shape.map(() => 0)).join(", ")}] to [${(hi ?? shape).join(", ")}] holds no voxel of the level, which is ${shape.join("×")}`); + } + return { lo: inLo.map((x, a) => Math.floor(x / chunk[a])), hi: inHi.map((x, a) => Math.ceil(x / chunk[a])) }; +} + +/** The chunk indices [lo, hi), in C order (the last axis fastest). */ +export function* chunkIndices(lo: number[], hi: number[]): Generator { + if (hi.some((h, a) => h <= lo[a])) return; + const at = [...lo]; + for (;;) { + yield [...at]; + let a = at.length - 1; + while (a >= 0 && ++at[a] >= hi[a]) { at[a] = lo[a]; a--; } + if (a < 0) return; + } +} + +/** `bytes` gzipped, through CompressionStream (the browser's, and Node's). */ +export async function gzip(bytes: Uint8Array): Promise { + const stream = new Blob([bytes as Uint8Array]).stream().pipeThrough(new CompressionStream("gzip")); + return new Uint8Array(await new Response(stream).arrayBuffer()); +} + +/** Whether every byte of `bytes` is zero: an empty chunk (zarr's fill value, 0, for every + * dtype; a float's -0 is not empty, its sign bit set). */ +export function allZero(bytes: Uint8Array): boolean { + let i = 0; + if (bytes.byteOffset % 4 === 0) { // four bytes at a time + const words = new Uint32Array(bytes.buffer, bytes.byteOffset, bytes.byteLength >> 2); + for (const w of words) if (w) return false; + i = words.length * 4; + } + for (; i < bytes.length; i++) if (bytes[i]) return false; + return true; +} + +/** What a materialize writes of view `name`: which `level`, in which `format` and + * `compression`; the metadata `files` (path and bytes, the group's first), the chunk grid's + * indices [lo, hi) the box covers, `total` chunks, and each chunk's `key` in the store. */ +export interface Region { + name: string; level: number; format: "zarr3" | "zarr2"; compression: "gzip" | null; + files: [string, Uint8Array][]; + lo: number[]; hi: number[]; total: number; + key: (index: number[]) => string; +} + +/** What `options` asks to write of view `name` (served as `v`, chunked `chunk`): the group + * listing the one level written as its one dataset, at LEVEL_PATH, with the level's scale and + * translation; the level's array with its whole shape (so the box keeps its place in it), + * the chunks and dtype served, compressed as asked; and the chunks of the box. */ +export function region(name: string, v: ViewInfo, chunk: number[], options: MaterializeOptions = {}): Region { + const { level = 0, format = "zarr2", compression = "gzip", attributes = {} } = options; + if (!Number.isInteger(level) || !v.levels[level]) throw new Error(`${name} has no level ${level}: it has ${v.levels.length}`); + if (format !== "zarr2" && format !== "zarr3") throw new Error(`format ${JSON.stringify(format)}: "zarr2" or "zarr3"`); + if (compression !== "gzip" && compression !== null) throw new Error(`compression ${JSON.stringify(compression)}: "gzip" or null`); + const { lo, hi } = chunkRange(v.levels[level].shape, chunk, options.lo, options.hi); + const datasets = [{ path: LEVEL_PATH, level }], added = (v.added ?? []).length; + const files: [string, unknown][] = format === "zarr3" + ? [ + ["zarr.json", group3(name, v, datasets, attributes)], + [`${LEVEL_PATH}/zarr.json`, array3(v, chunk, level, { codecs: compression ? [{ name: "gzip", configuration: { level: GZIP_LEVEL } }] : [] })], + ] + : [ + [".zgroup", { zarr_format: 2 }], + [".zattrs", attrs2(name, v, datasets, attributes)], + [`${LEVEL_PATH}/.zarray`, array2(v, chunk, level, compression ? { id: "gzip", level: GZIP_LEVEL } : null)], + [`${LEVEL_PATH}/.zattrs`, { _ARRAY_DIMENSIONS: dimensionNames(v) }], + ]; + const json = (o: unknown) => new TextEncoder().encode(`${JSON.stringify(o, null, 2)}\n`); + return { + name, level, format, compression, files: files.map(([path, o]) => [path, json(o)]), + lo, hi, total: lo.reduce((n, l, a) => n * (hi[a] - l), 1), + key: (index) => chunkKey(format, LEVEL_PATH, added, index), + }; +} + +/** Chunk `index` of the level as whole bytes, as served (zarr v3's: little endian, C order, + * the channels first); `signal` aborts when it is no longer wanted. */ +export type Produce = (index: number[], signal: AbortSignal) => Promise; + +const why = (e: unknown) => String((e as Error)?.message ?? e); + +/** Write region `r` into `store`: its metadata, then its chunks, `concurrency` of them asked + * of `produce` at once (one by default), each as it comes, gzipped if `r` says so, those of + * nothing but zeros skipped (`skipEmpty`, the default). Resolves with what was written, or + * on `signal`'s abort with `aborted` (no chunk started after it, the ones asked for given up + * through `produce`'s signal and not waited for, the writes under way ended). A chunk that + * fails to be made or written rejects it, naming the chunk and why, once the writes under + * way have ended; the rest are given up as on an abort. */ +export async function writeChunks(store: WritableStore, r: Region, produce: Produce, + options: Pick = {}): Promise { + const { skipEmpty = true, concurrency = 1, signal, onProgress } = options; + if (!Number.isInteger(concurrency) || concurrency < 1) throw new Error(`concurrency ${concurrency}: a whole number of chunks, 1 or more`); + const got: Materialized = { chunks: r.total, written: 0, skipped: 0, bytes: 0, aborted: false }; + if (signal?.aborted) return { ...got, aborted: true }; + for (const [path, bytes] of r.files) { + try { await store.put(path, bytes); } catch (e) { throw new Error(`writing ${path} of ${r.name}: ${why(e)}`); } + } + // stopped by the caller's abort, or by a chunk that failed + const stop = new AbortController(), onAbort = () => stop.abort(signal?.reason); + signal?.addEventListener("abort", onAbort, { once: true }); + const stopped = new Promise((_, reject) => stop.signal.addEventListener("abort", () => reject(stop.signal.reason), { once: true })); + stopped.catch(() => {}); + let failure: Error | null = null; + const fail = (e: Error) => { failure ??= e; stop.abort(e); }; + const report = () => { + try { onProgress?.({ done: got.written + got.skipped, total: r.total, written: got.written, bytes: got.bytes }); } + catch (e) { console.error(`materialize ${r.name}: onProgress threw: ${why(e)}`); } // the page's own trouble: the writing goes on + }; + report(); + const indices = chunkIndices(r.lo, r.hi); + const lane = async () => { + for (let next = indices.next(); !next.done && !stop.signal.aborted; next = indices.next()) { + const index = next.value, key = r.key(index); + let body: ArrayBuffer; + try { + body = await Promise.race([produce(index, stop.signal), stopped]); // given up at once when stopped + } catch (e) { + if (!stop.signal.aborted) fail(new Error(`chunk ${key} of ${r.name}: ${why(e)}`)); + return; + } + if (stop.signal.aborted) return; // made as it stopped: not written + const bytes = new Uint8Array(body); + if (skipEmpty && allZero(bytes)) { got.skipped++; report(); continue; } + try { + const stored = r.compression === "gzip" ? await gzip(bytes) : bytes; + await store.put(key, stored); + got.written++; got.bytes += stored.byteLength; + } catch (e) { + fail(new Error(`writing chunk ${key} of ${r.name}: ${why(e)}`)); + return; + } + report(); + } + }; + try { + await Promise.all(Array.from({ length: Math.min(concurrency, r.total) }, lane)); + } finally { + signal?.removeEventListener("abort", onAbort); + } + if (signal?.aborted) return { ...got, aborted: true }; + if (failure) throw failure; + return got; +} diff --git a/web/test/demand.test.ts b/web/test/demand.test.ts index 6e93b1f..12658a5 100644 --- a/web/test/demand.test.ts +++ b/web/test/demand.test.ts @@ -3,7 +3,7 @@ // stripped), so it needs nothing installed. import { test } from "node:test"; import assert from "node:assert/strict"; -import { Cancelled, Claim, queue } from "../src/demand.ts"; +import { Cancelled, Claim, Queue, queue } from "../src/demand.ts"; const done = () => Promise.resolve(new Float32Array()); @@ -90,3 +90,92 @@ test("a claim cancelled already runs what is hooked to it at once", () => { claim.onCancel(() => calls++); assert.equal(calls, 2); }); + +// ------------------------------------------------ background work (Engine.materialize's) +/** A job that runs until its gate opens, noting when it starts. */ +function gated(started: string[], key: string) { + let open!: () => void; + const gate = new Promise((resolve) => (open = resolve)); + return { open, run: async () => { started.push(key); await gate; return new Float32Array(); } }; +} +const tick = () => new Promise((r) => setTimeout(r, 0)); + +test("background work waits behind every job a client waits for, whatever its level, and takes all the slots but one", async () => { + const q = new Queue(3), started: string[] = []; + assert.deepEqual([q.slots, q.backgroundSlots], [3, 2]); + const bg = [0, 1, 2].map((i) => gated(started, `background ${i}`)); + const runs = bg.map((b, i) => q.submit(`bg ${i}`, 0, b.run, () => {}, true)); + await tick(); + assert.deepEqual(started, ["background 0", "background 1"]); // the third slot is kept for clients + const client = gated(started, "client, level 3"); + const asked = q.submit("client", 3, client.run, () => {}); + await tick(); + assert.deepEqual(started.slice(2), ["client, level 3"]); // at once, in the slot kept for it + const later = gated(started, "client, later, level 5"); + const asked2 = q.submit("client later", 5, later.run, () => {}); + bg[0].open(); + await tick(); + // a slot freed: the client's job goes before the background one waiting, coarser as it is + assert.deepEqual(started.slice(3), ["client, later, level 5"]); + client.open(); + await tick(); + assert.deepEqual(started.slice(4), ["background 2"]); + for (const g of [...bg, later]) g.open(); + await Promise.all([...runs, asked, asked2]); +}); + +test("a background claim is never released by a client giving up; given up in the background, it is dropped", async () => { + const q = new Queue(1), started: string[] = []; + const busy = gated(started, "busy"); + const running = q.submit("busy", 0, busy.run, () => {}); + let dropped = 0; + const wanted = q.submit("both", 0, async () => { started.push("both"); return new Float32Array(); }, () => dropped++, true); + q.claim("both"); // a client asks for it too + q.release("both"); // and gives up + busy.open(); + await wanted; + assert.deepEqual([started, dropped], [["busy", "both"], 0]); + await running; + const again = gated(started, "busy again"); + const r2 = q.submit("busy again", 0, again.run, () => {}); + const unwanted = q.submit("background alone", 0, async () => new Float32Array(), () => dropped++, true); + q.release("background alone", true); // materialize aborted before its turn + await assert.rejects(unwanted, Cancelled); + assert.equal(dropped, 1); + again.open(); + await r2; +}); + +test("a client asking for a background job brings it forward, and it goes back when the client gives up", async () => { + const q = new Queue(3), started: string[] = []; + const bg = [0, 1, 2, 3].map((i) => gated(started, `background ${i}`)); + const runs = bg.map((b, i) => q.submit(`bg ${i}`, 0, b.run, () => {}, true)); + await tick(); + assert.deepEqual(started, ["background 0", "background 1"]); + q.claim("bg 3"); // a viewer wants the last one: it takes the slot kept for clients at once + await tick(); + assert.deepEqual(started.slice(2), ["background 3"]); + // a fourth waiting in the background, a client's claim on it given up: back behind the rest + const fourth = gated(started, "background 4"); + const r4 = q.submit("bg 4", 0, fourth.run, () => {}, true); + q.claim("bg 4"); + q.release("bg 4"); + const client = gated(started, "client"); + const asked = q.submit("client", 1, client.run, () => {}); + bg[0].open(); + await tick(); + assert.deepEqual(started.slice(3), ["client"]); + for (const g of [...bg, fourth, client]) g.open(); + await Promise.all([...runs, r4, asked]); +}); + +test("with one slot, background work runs when no client waits", async () => { + const q = new Queue(1), started: string[] = []; + assert.equal(q.backgroundSlots, 1); + const bg = gated(started, "background"); + const run = q.submit("bg", 0, bg.run, () => {}, true); + await tick(); + assert.deepEqual(started, ["background"]); + bg.open(); + await run; +}); diff --git a/web/test/materialize.test.ts b/web/test/materialize.test.ts new file mode 100644 index 0000000..08ff27f --- /dev/null +++ b/web/test/materialize.test.ts @@ -0,0 +1,313 @@ +// A region of a view written into a store (src/materialize.ts, Engine.materialize's parts +// that need no engine): the chunks a box covers, rounded out and clipped; the metadata both +// formats write, the layout the engine serves (src/layout.ts); gzip; the loop over a fake +// producer, with empty chunks skipped, progress, abort and failure; the stores; and the +// result read back by zarrita, a zarr implementation of its own, as a reader would. Run by +// `npm test`. +import { test } from "node:test"; +import assert from "node:assert/strict"; +import * as zarr from "zarrita"; +import { array3, chunkKey, group3 } from "../src/layout.ts"; +import { allZero, chunkIndices, chunkRange, directoryStore, gzip, memoryStore, region, writeChunks, type MaterializeProgress, type Produce, type WritableStore } from "../src/materialize.ts"; +import type { ViewInfo } from "../src/types.ts"; + +const nm = (name: string) => ({ name, unit: "nm" }); +/** A view of two levels, served as uint16 in chunks of 2 × 4 × 4. */ +const view: ViewInfo = { + dtype: "uint8", channels: 1, axes: [nm("z"), nm("y"), nm("x")], halo: [0, 0, 0], lead: 0, out: "uint16", added: [], + levels: [{ shape: [5, 6, 7], voxel: [8, 8, 8], origin: [0, 0, 0] }, { shape: [3, 3, 4], voxel: [16, 16, 16], origin: [4, 4, 4] }], +}; +const chunk = [2, 4, 4]; +/** The same view with two channels its ops add (a gradient's, say), as float32. */ +const channels: ViewInfo = { ...view, out: "float32", added: [2] }; + +const json = (b: Uint8Array | undefined) => JSON.parse(new TextDecoder().decode(b)); +const files = (s: { files: Map }) => [...s.files.keys()].sort(); +async function gunzip(b: Uint8Array): Promise { + const stream = new Blob([b as Uint8Array]).stream().pipeThrough(new DecompressionStream("gzip")); + return new Uint8Array(await new Response(stream).arrayBuffer()); +} +const tick = () => new Promise((r) => setTimeout(r, 0)); + +/** Level 0's voxel (z, y, x) as the fake producer makes it: never zero inside the level. */ +const value = (z: number, y: number, x: number) => 1 + 100 * z + 10 * y + x; +/** Chunk `index` of level 0 as the engine gives it: whole, zero past the level's edge; the + * chunks of `empty` all zero, as a model's output where there is nothing. */ +function fakeChunk(index: number[], empty: string[] = []): ArrayBuffer { + const out = new Uint16Array(chunk[0] * chunk[1] * chunk[2]), [Z, Y, X] = view.levels[0].shape; + if (empty.includes(index.join("."))) return out.buffer; + for (let z = 0; z < chunk[0]; z++) for (let y = 0; y < chunk[1]; y++) for (let x = 0; x < chunk[2]; x++) { + const g = [index[0] * chunk[0] + z, index[1] * chunk[1] + y, index[2] * chunk[2] + x]; + if (g[0] < Z && g[1] < Y && g[2] < X) out[(z * chunk[1] + y) * chunk[2] + x] = value(g[0], g[1], g[2]); + } + return out.buffer; +} +const producer = (empty: string[] = []): Produce => async (index) => fakeChunk(index, empty); + +/** `store`'s files as zarrita reads a store (keys with a leading "/"). */ +const readable = (s: { files: Map }) => ({ async get(key: string) { return s.files.get(key.replace(/^\//, "")); } }); + +test("a box covers the chunks it touches, rounded out and clipped to the level; the whole level by default", () => { + const shape = view.levels[0].shape; // 5 × 6 × 7 in chunks of 2 × 4 × 4: a grid of 3 × 2 × 2 + assert.deepEqual(chunkRange(shape, chunk), { lo: [0, 0, 0], hi: [3, 2, 2] }); + assert.deepEqual(chunkRange(shape, chunk, [1, 3, 5], [3, 5, 6]), { lo: [0, 0, 1], hi: [2, 2, 2] }); + assert.deepEqual(chunkRange(shape, chunk, [2, 4, 4], [4, 8, 8]), { lo: [1, 1, 1], hi: [2, 2, 2] }); // on chunk edges: no more + assert.deepEqual(chunkRange(shape, chunk, [-10, 2, 3], [100, 6, 100]), { lo: [0, 0, 0], hi: [3, 2, 2] }); // past the edges: clipped + assert.deepEqual(chunkRange(shape, chunk, [4, 0, 0]), { lo: [2, 0, 0], hi: [3, 2, 2] }); // lo alone: to the level's end + assert.deepEqual(chunkRange(shape, chunk, [0.5, 0, 0], [1.5, 1, 1]), { lo: [0, 0, 0], hi: [1, 1, 1] }); // fractions: rounded out too + assert.throws(() => chunkRange(shape, chunk, [5, 0, 0]), /holds no voxel of the level, which is 5×6×7/); // starts at the end + assert.throws(() => chunkRange(shape, chunk, [0, 0, 0], [0, 6, 7]), /holds no voxel/); + assert.throws(() => chunkRange(shape, chunk, [0, 0], [1, 1]), /lo \[0,0\]: one number for each of the level's 3 axes/); + assert.throws(() => chunkRange(shape, chunk, undefined, [1, Number.NaN, 1]), /hi /); +}); + +test("the chunks of a range in C order, the last axis fastest", () => { + assert.deepEqual([...chunkIndices([0, 1], [2, 3])], [[0, 1], [0, 2], [1, 1], [1, 2]]); + assert.deepEqual([...chunkIndices([1, 1, 1], [2, 2, 2])], [[1, 1, 1]]); + assert.deepEqual([...chunkIndices([0, 0], [0, 3])], []); +}); + +test("zarr v2: .zgroup, .zattrs with OME 0.4 multiscales listing the one level at 0, its .zarray gzipped and .zattrs", () => { + const r = region("mito", view, chunk, { level: 1, attributes: { provenance: { model: "mito-v3" } } }); + assert.deepEqual(r.files.map(([p]) => p), [".zgroup", ".zattrs", "0/.zarray", "0/.zattrs"]); + const [group, attrs, zarray, dims] = r.files.map(([, b]) => json(b)); + assert.deepEqual(group, { zarr_format: 2 }); + assert.deepEqual(attrs, { + provenance: { model: "mito-v3" }, + multiscales: [{ + name: "mito", + axes: [{ name: "z", type: "space", unit: "nanometer" }, { name: "y", type: "space", unit: "nanometer" }, { name: "x", type: "space", unit: "nanometer" }], + // level 1's own scale and translation, so it lands where the view's level 1 does + datasets: [{ path: "0", coordinateTransformations: [{ type: "scale", scale: [16, 16, 16] }, { type: "translation", translation: [4, 4, 4] }] }], + version: "0.4", + }], + }); + assert.deepEqual(zarray, { + zarr_format: 2, shape: [3, 3, 4], chunks: [2, 4, 4], dtype: " { + const r = region("mito", view, chunk, { format: "zarr3", lo: [2, 0, 4], hi: [5, 3, 7], attributes: { provenance: 1 } }); + assert.deepEqual(r.files.map(([p]) => p), ["zarr.json", "0/zarr.json"]); + const [group, array] = r.files.map(([, b]) => json(b)); + assert.deepEqual(group, { + zarr_format: 3, node_type: "group", attributes: { + provenance: 1, + ome: { version: "0.5", multiscales: [{ + name: "mito", + axes: [{ name: "z", type: "space", unit: "nanometer" }, { name: "y", type: "space", unit: "nanometer" }, { name: "x", type: "space", unit: "nanometer" }], + datasets: [{ path: "0", coordinateTransformations: [{ type: "scale", scale: [8, 8, 8] }, { type: "translation", translation: [0, 0, 0] }] }], + }] }, + }, + }); + assert.deepEqual(array, { + zarr_format: 3, node_type: "array", shape: [5, 6, 7], data_type: "uint16", fill_value: 0, // the whole level: the box keeps its place + chunk_grid: { name: "regular", configuration: { chunk_shape: [2, 4, 4] } }, + chunk_key_encoding: { name: "default", configuration: { separator: "/" } }, + codecs: [{ name: "bytes", configuration: { endian: "little" } }, { name: "gzip", configuration: { level: 6 } }], + dimension_names: ["z", "y", "x"], attributes: {}, + }); + assert.deepEqual([r.lo, r.hi, r.total, r.key([2, 0, 1])], [[1, 0, 1], [3, 1, 2], 2, "0/c/2/0/1"]); + // the metadata is the engine's as served, but for the one dataset and the codec + assert.deepEqual(array.codecs.slice(0, 1), array3(view, chunk, 0).codecs); + assert.deepEqual(group3("mito", view, [{ path: "0", level: 0 }, { path: "1", level: 1 }]).attributes.ome.multiscales[0].datasets[1].path, "1"); +}); + +test("uncompressed when asked; a channel axis the ops add is one chunk, index 0, as the engine serves it", () => { + const two = region("grad", channels, chunk, { compression: null }); + const zarray = json(two.files.find(([p]) => p === "0/.zarray")![1]); + assert.deepEqual([zarray.shape, zarray.chunks, zarray.dtype, zarray.compressor], [[2, 5, 6, 7], [2, 2, 4, 4], " p === "0/.zattrs")![1]), { _ARRAY_DIMENSIONS: ["c", "z", "y", "x"] }); + const ms = json(two.files.find(([p]) => p === ".zattrs")![1]).multiscales[0]; + assert.deepEqual(ms.axes[0], { name: "c", type: "channel" }); + assert.deepEqual(ms.datasets[0].coordinateTransformations, [{ type: "scale", scale: [1, 8, 8, 8] }, { type: "translation", translation: [0, 0, 0, 0] }]); + assert.equal(two.key([1, 0, 1]), "0/0/1/0/1"); + const three = region("grad", channels, chunk, { format: "zarr3", compression: null }); + const array = json(three.files[1][1]); + assert.deepEqual([array.shape, array.chunk_grid.configuration.chunk_shape, array.codecs.length], [[2, 5, 6, 7], [2, 2, 4, 4], 1]); + assert.equal(three.key([1, 0, 1]), "0/c/0/1/0/1"); + assert.equal(chunkKey("zarr3", "0", 0, [3, 2, 1]), "0/c/3/2/1"); + assert.throws(() => region("grad", channels, chunk, { level: 2 }), /grad has no level 2: it has 2/); + assert.throws(() => region("grad", channels, chunk, { format: "n5" as "zarr2" }), /format "n5"/); +}); + +test("the group's attributes keep the view's OME metadata whatever the extra ones are called", () => { + const two = region("v", view, chunk, { attributes: { multiscales: "mine", note: "kept" } }); + const attrs = json(two.files[1][1]); + assert.equal(attrs.note, "kept"); + assert.equal(attrs.multiscales[0].name, "v"); + const three = region("v", view, chunk, { format: "zarr3", attributes: { ome: "mine" } }); + assert.equal(json(three.files[0][1]).attributes.ome.version, "0.5"); +}); + +test("gzip round trips through CompressionStream and DecompressionStream; zeros are told from the rest", async () => { + const chunk0 = new Uint8Array(fakeChunk([0, 0, 0])); + assert.deepEqual(await gunzip(await gzip(chunk0)), chunk0); + const bytes = new Uint8Array(4096).fill(7, 100, 3000), packed = await gzip(bytes); // a model's output: much of it alike + assert.deepEqual([packed[0], packed[1]], [0x1f, 0x8b]); // gzip's magic + assert.ok(packed.length < bytes.length / 10, `${packed.length} bytes`); + assert.deepEqual(await gunzip(packed), bytes); + assert.equal(allZero(new Uint8Array(64)), true); + assert.equal(allZero(new Uint8Array(new ArrayBuffer(67), 1, 65)), true); // not on a word: byte by byte + const one = new Uint8Array(67); + one[66] = 1; + assert.equal(allZero(one), false); // past the last whole word + assert.equal(allZero(new Uint8Array(new Float32Array([0, -0]).buffer)), false); // -0 has its sign bit +}); + +for (const format of ["zarr2", "zarr3"] as const) { + test(`${format}: the box's chunks written gzipped, the empty one skipped, and read back by zarrita as the view`, async () => { + const store = memoryStore(), seen: MaterializeProgress[] = []; + // z voxels 1 to 3 of the whole y, x: chunks z 0 and 1, the second of them empty + const got = await writeChunks(store, region("v", view, chunk, { format, lo: [1, 0, 0], hi: [3, 6, 7] }), producer(["1.0.1"]), + { concurrency: 3, onProgress: (p) => seen.push(p) }); + const keys = files(store).filter((f) => !f.endsWith("json") && !f.includes(".z")); + const at = (i: string) => (format === "zarr3" ? `0/c/${i}` : `0/${i}`); + assert.deepEqual(keys, ["0.0.0", "0.0.1", "0.1.0", "0.1.1", "1.0.0", "1.1.0", "1.1.1"].map((i) => at(i.replaceAll(".", "/"))).sort()); + const bytes = keys.reduce((n, k) => n + store.files.get(k)!.length, 0); + assert.deepEqual(got, { chunks: 8, written: 7, skipped: 1, bytes, aborted: false }); + assert.deepEqual(await gunzip(store.files.get(at("1/1/1"))!), new Uint8Array(fakeChunk([1, 1, 1]))); + // progress: once with the metadata written, then after each chunk, done = written + skipped + assert.deepEqual(seen[0], { done: 0, total: 8, written: 0, bytes: 0 }); + assert.deepEqual(seen.at(-1), { done: 8, total: 8, written: 7, bytes }); + const steps = seen.slice(1).map((p, i) => [p.done - seen[i].done, p.written - seen[i].written, p.bytes - seen[i].bytes]); + assert.deepEqual(steps.map(([d]) => d), [1, 1, 1, 1, 1, 1, 1, 1]); + assert.deepEqual(steps.filter(([, w, b]) => w === 0 && b === 0).length, 1); // the one skipped + // a reader: the box's chunks as made, the skipped one and those outside the box zero + const root = zarr.root(readable(store)); + const group = format === "zarr3" ? await zarr.open.v3(root, { kind: "group" }) : await zarr.open.v2(root, { kind: "group" }); + assert.ok(group.attrs.ome ?? group.attrs.multiscales); + const arr = format === "zarr3" ? await zarr.open.v3(root.resolve("0"), { kind: "array" }) : await zarr.open.v2(root.resolve("0"), { kind: "array" }); + assert.deepEqual([arr.shape, arr.dtype], [[5, 6, 7], "uint16"]); + const data = (await zarr.get(arr)).data as Uint16Array; + const [, Y, X] = arr.shape; + for (let z = 0; z < 5; z++) for (let y = 0; y < Y; y++) for (let x = 0; x < X; x++) { + const written = z < 4 && !(z >= 2 && y < 4 && x >= 4); // chunks z 0 and 1, but for chunk (1, 0, 1) + assert.equal(data[(z * Y + y) * X + x], written ? value(z, y, x) : 0, `voxel ${z}, ${y}, ${x}`); + } + }); +} + +test("empty chunks are written too when asked, and uncompressed ones as they are", async () => { + const store = memoryStore(); + const got = await writeChunks(store, region("v", view, chunk, { compression: null, lo: [2, 0, 4], hi: [4, 4, 8] }), producer(["1.0.1"]), { skipEmpty: false }); + assert.deepEqual(got, { chunks: 1, written: 1, skipped: 0, bytes: 64, aborted: false }); + assert.deepEqual(store.files.get("0/1/0/1"), new Uint8Array(32 * 2)); +}); + +test("no more than `concurrency` chunks are asked for at once", async () => { + let asked = 0, most = 0; + const produce: Produce = async (index) => { + most = Math.max(most, ++asked); + await tick(); + asked--; + return fakeChunk(index); + }; + const got = await writeChunks(memoryStore(), region("v", view, chunk), produce, { concurrency: 2 }); + assert.deepEqual([got.written, most], [12, 2]); + await assert.rejects(writeChunks(memoryStore(), region("v", view, chunk), produce, { concurrency: 0 }), /concurrency 0/); +}); + +test("an abort stops starting chunks, gives up those asked for and resolves at once; the chunks written stay", { timeout: 5000 }, async () => { + const store = memoryStore(), stop = new AbortController(), asked: number[][] = [], given: AbortSignal[] = []; + let atAbort = -1; + // the first two chunks come at once, the rest never (a model still running, say) + const produce: Produce = (index, signal) => { + asked.push(index); + given.push(signal); + return asked.length <= 2 ? Promise.resolve(fakeChunk(index)) : new Promise(() => {}); + }; + const got = await writeChunks(store, region("v", view, chunk), produce, { + concurrency: 2, signal: stop.signal, onProgress: (p) => { if (p.written === 2) { stop.abort(); atAbort = asked.length; } }, + }); + assert.deepEqual(got, { chunks: 12, written: 2, skipped: 0, bytes: got.bytes, aborted: true }); + assert.equal(asked.length, atAbort); // none asked for after the abort + assert.ok(atAbort > 2); // and the one asked for when it came, never made, was not waited for + assert.ok(given.every((s) => s.aborted)); + assert.deepEqual(files(store), [".zattrs", ".zgroup", "0/.zarray", "0/.zattrs", "0/0/0/0", "0/0/0/1"]); + // aborted before it starts: nothing is written + const none = memoryStore(); + assert.deepEqual(await writeChunks(none, region("v", view, chunk), producer(), { signal: AbortSignal.abort() }), + { chunks: 12, written: 0, skipped: 0, bytes: 0, aborted: true }); + assert.equal(none.files.size, 0); +}); + +test("a chunk that fails rejects naming it and why, once the writes under way have ended", async () => { + let slowWritten = false; + const memory = memoryStore(); + const store: WritableStore = { + async put(path, bytes) { + if (path === "0/0/0/0") { await new Promise((r) => setTimeout(r, 30)); slowWritten = true; } + return memory.put(path, bytes); + }, + }; + const produce: Produce = async (index) => { + if (index.join(".") === "0.0.1") { await tick(); throw new Error("the model ran out of GPU memory"); } + return fakeChunk(index); + }; + await assert.rejects(writeChunks(store, region("v", view, chunk), produce, { concurrency: 2 }), (e: Error) => { + assert.equal(e.message, "chunk 0/0/0/1 of v: the model ran out of GPU memory"); + assert.ok(slowWritten, "the chunk being written when the other failed was written first"); + return true; + }); + assert.ok(memory.files.has("0/0/0/0") && !memory.files.has("0/0/1/0")); // nothing started after the failure + const full: WritableStore = { async put(path) { if (path.startsWith("0/c/1/")) throw new Error("the disk is full"); } }; + await assert.rejects(writeChunks(full, region("v", view, chunk, { format: "zarr3" }), producer()), /^Error: writing chunk 0\/c\/1\/0\/0 of v: the disk is full$/); + const refused: WritableStore = { async put() { throw new Error("permission revoked"); } }; + await assert.rejects(writeChunks(refused, region("v", view, chunk), producer()), /^Error: writing \.zgroup of v: permission revoked$/); +}); + +test("the memory store keeps a copy of each file under its path, without empty names", async () => { + const store = memoryStore(), bytes = new Uint8Array([1, 2, 3]); + await store.put("/a//b/c", bytes); + bytes[0] = 9; + assert.deepEqual([...store.files], [["a/b/c", new Uint8Array([1, 2, 3])]]); + await assert.rejects(store.put("a/../b", bytes), /with no "\." or "\.\."/); + await assert.rejects(store.put("/", bytes), /not a file's path/); +}); + +/** A FileSystemDirectoryHandle's few methods directoryStore calls, over a tree of Maps. */ +interface Dir { dirs: Map; files: Map } +function fakeDirectory(calls: string[], failWrite?: string) { + const make = (path: string, d: Dir): FileSystemDirectoryHandle => ({ + async getDirectoryHandle(name: string, { create = false } = {}) { + calls.push(`dir ${path}${name}`); + await tick(); // as a real handle is: concurrent puts must not make one directory twice + if (!d.dirs.has(name)) { assert.ok(create); d.dirs.set(name, { dirs: new Map(), files: new Map() }); } + return make(`${path}${name}/`, d.dirs.get(name)!); + }, + async getFileHandle(name: string) { + return { + async createWritable() { + const parts: Uint8Array[] = []; + return { + async write(b: Uint8Array) { if (name === failWrite) throw new Error("quota exceeded"); parts.push(b.slice()); }, + async close() { d.files.set(name, new Uint8Array(parts.flatMap((p) => [...p]))); calls.push(`closed ${path}${name}`); }, + async abort() { calls.push(`aborted ${path}${name}`); }, + }; + }, + }; + }, + }) as unknown as FileSystemDirectoryHandle; + const tree: Dir = { dirs: new Map(), files: new Map() }; + return { tree, handle: make("", tree) }; +} + +test("a directory as a store: the root and subdirectories made once each, each file written whole", async () => { + const calls: string[] = [], { tree, handle } = fakeDirectory(calls); + const store = directoryStore(handle, "runs/mito.zarr"); + await Promise.all([store.put("0/c/0/0/0", new Uint8Array([1])), store.put("0/c/0/0/1", new Uint8Array([2])), store.put("zarr.json", new Uint8Array([3]))]); + assert.deepEqual(calls.filter((c) => c.startsWith("dir")), ["dir runs", "dir runs/mito.zarr", "dir runs/mito.zarr/0", "dir runs/mito.zarr/0/c", "dir runs/mito.zarr/0/c/0", "dir runs/mito.zarr/0/c/0/0"]); + const zarrDir = tree.dirs.get("runs")!.dirs.get("mito.zarr")!; + assert.deepEqual(zarrDir.files.get("zarr.json"), new Uint8Array([3])); + assert.deepEqual([...zarrDir.dirs.get("0")!.dirs.get("c")!.dirs.get("0")!.dirs.get("0")!.files], [["0", new Uint8Array([1])], ["1", new Uint8Array([2])]]); + // a write that fails is aborted (the swap file dropped) and rejects + const failing = fakeDirectory(calls, "bad"); + await assert.rejects(directoryStore(failing.handle).put("bad", new Uint8Array([1])), /quota exceeded/); + assert.ok(calls.includes("aborted bad")); + assert.equal(failing.tree.files.size, 0); +});