Prepare shard files ahead of writes by default

README.md

@@ -304,6 +304,9 @@ the library handles all tiling, padding, and shard assembly internally. See
 [docs/guide.md][docs-guide-md] for a quick orientation to the module structure.
 
 Paced producers can opt into the [bounded input-buffering writer](docs/buffered-writer.md).
+Filesystem streams [prepare one shard generation ahead](docs/shard-preparation.md)
+by default, including cleanup of unused files at shutdown. Set
+`config.disable_shard_preparation = 1` to open files on demand instead.
 
 For writing directly to S3 (or S3-compatible stores), see the
 [S3 storage guide][s3-storage-guide].

docs/shard-preparation.md

@@ -0,0 +1,39 @@
+# Preparing shard files ahead of use
+
+CPU and GPU streams prepare shard files ahead of use by default when the
+sink supports it. The filesystem store supports preparation for Zarr arrays,
+multiarray streams, and NGFF multiscales. Set
+`config.disable_shard_preparation = 1` to open files on demand instead.
+Unsupported sinks, including S3, automatically retain their existing behavior.
+
+Creation waits until the first shard generation is ready. When a stream
+adopts one of those files, a separate worker creates and, where supported,
+sets the capacity of the next file for that slot. There is at most one unused
+prepared file per slot. A layout with 15 spatial shards therefore holds at
+most 15 additional files, independently of the stream's length. A finite
+stream does not prepare files beyond its configured capacity.
+
+Preparation uses one worker separate from the normal write workers. Future
+preparation is excluded from the fences used by writes, metadata publication,
+and footer reuse. If preparation has not finished when a file is needed,
+opening that shard waits for it.
+
+`writer_flush` stops scheduling future preparation, drains accepted input,
+and waits for preparation before removing unused files. `writer_close`
+reports cleanup failures, and destruction also attempts cleanup when explicit
+flush/close was omitted. Files already adopted follow the ordinary shard
+lifecycle. Preparation creates files exclusively: it fails if a target
+already exists, without truncating or deleting that file. Cleanup removes
+only unused files and empty directories created by preparation; it preserves
+pre-existing directories and directories containing used shards.
+
+Store backends expose preparation through the shard pool. `prepare`,
+`wait_prepared`, and `cancel_prepared` form one capability. Supporting
+preparation requires supporting cancellation: wait for outstanding work,
+reclaim unused owned resources, report failures, and allow cancellation to
+be repeated. Opening a prepared shard transfers it to the ordinary writer.
+Preparation and open calls are serialized by the stream; cancellation runs
+after delivery stops using the prepared resources.
+
+Custom sinks need all three `prepare_shards`, `stop_preparing`, and
+`cancel_prepared` hooks. Missing any hook disables preparation for that sink.

src/cpu/stream.body.c

@@ -237,6 +237,8 @@ cpu_stream_flush_batch(struct cpu_stream_view* v)
 struct writer_result
 cpu_stream_flush_body(struct cpu_stream_view* v)
 {
+  if (!v->config->disable_shard_preparation)
+    shard_sink_stop_preparing(v->sink);
   // A create that failed before sizing the layout leaves epoch_elements at 0
   // and the divisions below would fault. Nothing was sized, so nothing to do.
   if (v->layout->epoch_elements == 0)
@@ -338,14 +340,18 @@ cpu_stream_flush_body(struct cpu_stream_view* v)
   // they read.
   if (shard_sink_drain(v->sink))
     failed = 1;
+  if (!v->config->disable_shard_preparation &&
+      shard_sink_cancel_prepared(v->sink))
+    failed = 1;
 
   return failed ? writer_error() : writer_ok();
 }
 
 struct writer_result
 cpu_stream_close_body(struct cpu_stream_view* v)
 {
-  int failed = 0;
+  int failed = !v->config->disable_shard_preparation &&
+               shard_sink_cancel_prepared(v->sink);
 
   // A sink IO error is the one case the shape is withheld: which writes landed
   // is unknowable.

src/cpu/stream.c

@@ -291,6 +291,9 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config,
   }
 
   CHECK(Fail, shard_sink_init_append(sink, &s->cl.dims, s->levels.nlod) == 0);
+  CHECK(Fail,
+        shard_sink_prepare_first(
+          sink, s->shard, s->levels.nlod, s->shard_alignment, config) == 0);
 
   s->writer.append = cpu_append;
   s->writer.flush = cpu_flush_final;
@@ -303,6 +306,8 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config,
   return s;
 
 Fail:
+  if (!config->disable_shard_preparation && shard_sink_cancel_prepared(sink))
+    log_error("CPU stream preparation cleanup failed during rollback");
   // Construction rollback releases resources without finalizing the sink.
   tile_stream_cpu_release_resources(s);
   return NULL;

src/gpu/stream.c

@@ -397,6 +397,8 @@ publish_array_shape(struct compress_agg_array* ar, struct stream_context* ctx)
 struct writer_result
 stream_flush_body(struct stream_engine* e, struct stream_context* ctx)
 {
+  if (!ctx->config.disable_shard_preparation)
+    shard_sink_stop_preparing(ctx->sink);
   // A create that fails before sizing the layout leaves epoch_elements at 0;
   // the divisions below would then fault. Nothing was ever sized, so there is
   // nothing to flush.
@@ -429,13 +431,21 @@ stream_flush_body(struct stream_engine* e, struct stream_context* ctx)
       r = finalize_all_levels(e, ctx);
   }
 
+  // Finishing a partial batch can fail after handing it to delivery. Join
+  // through the scheduler before draining writes or cancelling preparation.
+  if (schedule_deliver_kicked(e, ctx).error)
+    r = writer_error();
+
   // Writes queued anywhere above can still fail, and destroy frees the buffers
   // they read.
   if (shard_sink_drain(ctx->sink)) {
     r = writer_error();
   }
   schedule_quiesce_output(e, ctx->sink);
 
+  if (!ctx->config.disable_shard_preparation &&
+      shard_sink_cancel_prepared(ctx->sink))
+    r = writer_error();
   collect_ingest_timing(e);
 
   return r;
@@ -449,6 +459,10 @@ stream_close_body(struct compress_agg_array* ar, struct stream_context* ctx)
 
   struct writer_result r = writer_ok();
 
+  if (!ctx->config.disable_shard_preparation &&
+      shard_sink_cancel_prepared(ctx->sink))
+    r = writer_error();
+
   // The shape is written after this, so it never names data still queued.
   const int sink_failed = shard_sink_drain(ctx->sink);
   if (sink_failed)

src/gpu/stream.init.c

@@ -350,6 +350,9 @@ tile_stream_gpu_rollback(struct tile_stream_gpu* s)
   const int pushed = cu_ctx_push(s->engine.cuda);
   gpu_delivery_stop_join(&s->engine.delivery);
   gpu_streams_sync(&s->engine.streams);
+  if (!s->ctx.config.disable_shard_preparation &&
+      shard_sink_cancel_prepared(s->ctx.sink))
+    log_error("GPU stream preparation cleanup failed during rollback");
   tile_stream_gpu_release_resources(s, pushed);
 }
 
@@ -451,6 +454,13 @@ tile_stream_gpu_create(const struct tile_stream_configuration* config,
         shard_sink_init_append(sink, &out->ctx.dims, out->ctx.levels.nlod) ==
           0);
 
+  CHECK(FailPhase2,
+        shard_sink_prepare_first(sink,
+                                 out->engine.compress_agg.ar.shard,
+                                 out->ctx.levels.nlod,
+                                 out->ctx.shard_alignment,
+                                 config) == 0);
+
   CHECK(FailPhase2, gpu_placement_leave(&placement_scope) == 0);
   out->flushed = 0;
   out->closed = 0;

src/multiarray/stream.c

@@ -478,6 +478,12 @@ multiarray_tile_stream_cpu_create(
     CHECK(Fail,
           shard_sink_init_append(
             desc->sink, &desc->cl.dims, desc->levels.nlod) == 0);
+    CHECK(Fail,
+          shard_sink_prepare_first(desc->sink,
+                                   desc->shard,
+                                   desc->levels.nlod,
+                                   desc->shard_alignment,
+                                   &desc->config) == 0);
   }
   for (int i = 0; i < n_arrays; ++i) {
     ms->arrays[i].flushed = 0;
@@ -486,6 +492,11 @@ multiarray_tile_stream_cpu_create(
   return ms;
 
 Fail:
+  if (ms && ms->arrays)
+    for (int i = 0; i < ms->n_arrays; ++i)
+      if (!ms->arrays[i].config.disable_shard_preparation &&
+          shard_sink_cancel_prepared(ms->arrays[i].sink))
+        log_error("CPU multiarray preparation cleanup failed during rollback");
   // Construction rollback releases resources without finalizing any sink.
   multiarray_tile_stream_cpu_release_resources(ms);
   return NULL;
@@ -733,14 +744,20 @@ flush_impl(struct multiarray_writer* self)
   struct multiarray_tile_stream_cpu* ms =
     container_of(self, struct multiarray_tile_stream_cpu, writer);
 
+  // Every array must finish cleanup even if an earlier array fails.
+  int failed = 0;
   for (int a = 0; a < ms->n_arrays; ++a) {
     struct array_descriptor* desc = &ms->arrays[a];
     // Finalizing twice would re-finalize an already-closed sink, which
     // deadlocks on Windows.
     if (desc->flushed)
       continue;
     if (desc->cursor_elements == 0 && desc->batch_accumulated == 0) {
+      if (!desc->config.disable_shard_preparation &&
+          shard_sink_cancel_prepared(desc->sink))
+        failed = 1;
       desc->flushed = 1;
+      desc->closed = 0;
       continue;
     }
 
@@ -757,16 +774,11 @@ flush_impl(struct multiarray_writer* self)
     desc->flushed = 1;
     desc->closed = 0;
     if (r.error)
-      goto Error;
+      failed = 1;
   }
 
   return (struct multiarray_writer_result){
-    .error = multiarray_writer_ok,
-  };
-
-Error:
-  return (struct multiarray_writer_result){
-    .error = multiarray_writer_fail,
+    .error = failed ? multiarray_writer_fail : multiarray_writer_ok,
   };
 }
 

src/multiarray/stream.gpu.c

@@ -285,6 +285,9 @@ flush_impl(struct multiarray_writer* self)
     if (desc->flushed)
       continue;
     if (desc->ctx.cursor_elements == 0 && desc->st.sched.accumulated == 0) {
+      if (!desc->ctx.config.disable_shard_preparation &&
+          shard_sink_cancel_prepared(desc->ctx.sink))
+        failed = 1;
       desc->flushed = 1;
       continue;
     }
@@ -379,6 +382,11 @@ multiarray_tile_stream_gpu_rollback(struct multiarray_tile_stream_gpu* ms)
   const int pushed = cu_ctx_push(ms->engine.cuda);
   gpu_delivery_stop_join(&ms->engine.delivery);
   gpu_streams_sync(&ms->engine.streams);
+  if (ms->arrays)
+    for (int a = 0; a < ms->n_arrays; ++a)
+      if (!ms->arrays[a].ctx.config.disable_shard_preparation &&
+          shard_sink_cancel_prepared(ms->arrays[a].ctx.sink))
+        log_error("GPU multiarray preparation cleanup failed during rollback");
   multiarray_tile_stream_gpu_release_resources(ms, pushed);
 }
 
@@ -496,6 +504,12 @@ multiarray_tile_stream_gpu_create(
     const struct stream_context* ctx = &ms->arrays[a].ctx;
     CHECK(Fail,
           shard_sink_init_append(ctx->sink, &ctx->dims, ctx->levels.nlod) == 0);
+    CHECK(Fail,
+          shard_sink_prepare_first(ctx->sink,
+                                   ms->arrays[a].st.agg.shard,
+                                   ctx->levels.nlod,
+                                   ctx->shard_alignment,
+                                   &ctx->config) == 0);
   }
   for (int a = 0; a < n_arrays; ++a) {
     ms->arrays[a].flushed = 0;

src/ngff/ngff_multiscale.c

@@ -104,6 +104,36 @@ ngff_multiscale_open(struct shard_sink* self,
   return NULL;
 }
 
+static int
+ngff_multiscale_prepare_shards(struct shard_sink* self,
+                               uint8_t level,
+                               uint64_t capacity)
+{
+  struct ngff_multiscale* ms = container_of(self, struct ngff_multiscale, base);
+  if (level >= ms->nlod)
+    return 1;
+  struct shard_sink* child = zarr_array_as_shard_sink(ms->levels[level]);
+  return child->prepare_shards(child, level, capacity);
+}
+
+static void
+ngff_multiscale_stop_preparing(struct shard_sink* self)
+{
+  struct ngff_multiscale* ms = container_of(self, struct ngff_multiscale, base);
+  for (int lv = 0; lv < ms->nlod; ++lv)
+    shard_sink_stop_preparing(zarr_array_as_shard_sink(ms->levels[lv]));
+}
+
+static int
+ngff_multiscale_cancel_prepared(struct shard_sink* self)
+{
+  struct ngff_multiscale* ms = container_of(self, struct ngff_multiscale, base);
+  int rc = 0;
+  for (int lv = 0; lv < ms->nlod; ++lv)
+    rc |= shard_sink_cancel_prepared(zarr_array_as_shard_sink(ms->levels[lv]));
+  return rc;
+}
+
 static int
 submit_append(struct ngff_multiscale* ms,
               uint8_t level,
@@ -235,6 +265,11 @@ ngff_multiscale_init(struct store* store,
     CHECK(Fail_ms, ngff_axes_copy(ms->axes, cfg->axes, cfg->rank) == 0);
 
   ms->base.open = ngff_multiscale_open;
+  if (pool->prepare && pool->wait_prepared && pool->cancel_prepared) {
+    ms->base.prepare_shards = ngff_multiscale_prepare_shards;
+    ms->base.stop_preparing = ngff_multiscale_stop_preparing;
+    ms->base.cancel_prepared = ngff_multiscale_cancel_prepared;
+  }
   ms->base.update_append = ngff_multiscale_update_append;
   ms->base.queue_append =
     pool->queue_metadata ? ngff_multiscale_queue_append : NULL;

src/platform/platform_io.darwin.c

@@ -17,6 +17,24 @@ platform_mkdir(const char* path)
   return 0;
 }
 
+int
+platform_mkdir_new(const char* path)
+{
+  if (mkdir(path, 0755) == 0)
+    return 1;
+  struct stat st;
+  return errno == EEXIST && stat(path, &st) == 0 && S_ISDIR(st.st_mode) ? 0
+                                                                        : -1;
+}
+
+int
+platform_remove_empty_directory(const char* path)
+{
+  if (rmdir(path) == 0 || errno == ENOENT)
+    return 0;
+  return errno == ENOTEMPTY || errno == EEXIST ? 1 : -1;
+}
+
 int
 platform_mkdirp(const char* path)
 {
@@ -40,13 +58,16 @@ platform_mkdirp(const char* path)
 platform_fd
 platform_open_write(const char* path, int flags)
 {
-  int oflags = O_WRONLY | O_CREAT | O_TRUNC;
+  int oflags = O_WRONLY | O_CREAT;
+  oflags |= flags & PLATFORM_OPEN_EXCLUSIVE ? O_EXCL : O_TRUNC;
   int fd = open(path, oflags, 0644);
   if (fd < 0)
     return fd;
   if (flags & PLATFORM_OPEN_UNBUFFERED) {
     if (fcntl(fd, F_NOCACHE, 1) != 0) {
       close(fd);
+      if (flags & PLATFORM_OPEN_EXCLUSIVE)
+        unlink(path);
       return -1;
     }
   }

src/platform/platform_io.h

@@ -17,14 +17,23 @@ typedef int platform_fd;
 int
 platform_mkdir(const char* path);
 
+// Returns 1 when created, 0 when already a directory, -1 on error.
+int
+platform_mkdir_new(const char* path);
+
+// Returns 0 if removed or missing, 1 if nonempty, -1 on other failures.
+int
+platform_remove_empty_directory(const char* path);
+
 // Create a directory and all parent directories. Returns 0 on success.
 int
 platform_mkdirp(const char* path);
 
 // Flags for platform_open_write.
 enum
 {
-  PLATFORM_OPEN_UNBUFFERED = 1
+  PLATFORM_OPEN_UNBUFFERED = 1,
+  PLATFORM_OPEN_EXCLUSIVE = 2
 };
 
 // Open a file for writing (create/truncate). Returns PLATFORM_FD_INVALID on