From 67e33356f53819bd379bb03e33ed44d999f5adbd Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Fri, 18 Sep 2026 21:59:28 +0000 Subject: [PATCH 1/2] Prepare shard files ahead of writes --- README.md | 2 + docs/shard-preparation.md | 39 ++++ src/cpu/stream.body.c | 6 +- src/cpu/stream.c | 10 + src/gpu/stream.c | 13 ++ src/gpu/stream.init.c | 10 + src/multiarray/stream.c | 14 ++ src/multiarray/stream.gpu.c | 15 ++ src/ngff/ngff_multiscale.c | 35 ++++ src/platform/platform_io.darwin.c | 24 ++- src/platform/platform_io.h | 11 +- src/platform/platform_io.posix.c | 22 ++- src/platform/platform_io.win32.c | 36 +++- src/types.stream.h | 1 + src/writer.c | 14 ++ src/writer.h | 15 ++ src/zarr/io_backend.fs.c | 16 ++ src/zarr/io_backend.fs.h | 5 + src/zarr/shard_delivery.c | 29 +++ src/zarr/shard_delivery.h | 7 + src/zarr/shard_pool.h | 12 ++ src/zarr/shard_pool_fs.c | 212 ++++++++++++++++++++- src/zarr/shard_pool_fs.h | 2 + src/zarr/zarr_array.c | 110 +++++++++-- tests/CMakeLists.txt | 11 ++ tests/test_shard_preparation.c | 228 ++++++++++++++++++++++ tests/test_stream_preparation.c | 301 ++++++++++++++++++++++++++++++ tests/test_zarr_readback.c | 1 + tests/test_zarr_readback_gpu.c | 1 + 29 files changed, 1177 insertions(+), 25 deletions(-) create mode 100644 docs/shard-preparation.md create mode 100644 tests/test_shard_preparation.c create mode 100644 tests/test_stream_preparation.c diff --git a/README.md b/README.md index 21be57a4..91c8fb48 100644 --- a/README.md +++ b/README.md @@ -304,6 +304,8 @@ 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 can also [prepare one shard generation ahead](docs/shard-preparation.md) +with `config.prepare_shards = 1`, including cleanup of unused files at shutdown. For writing directly to S3 (or S3-compatible stores), see the [S3 storage guide][s3-storage-guide]. diff --git a/docs/shard-preparation.md b/docs/shard-preparation.md new file mode 100644 index 00000000..68fa1d81 --- /dev/null +++ b/docs/shard-preparation.md @@ -0,0 +1,39 @@ +# Preparing shard files ahead of use + +Set `config.prepare_shards = 1` before creating a CPU or GPU stream to +prepare shard files ahead of use. This is opt-in and currently supported by +the filesystem store, including Zarr arrays and NGFF multiscales. + +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 the matching `prepare_shards`, `stop_preparing`, and +`cancel_prepared` hooks. Enabling the option on an unsupported sink, including +the current S3 backend, fails stream creation. With the option left at zero, +existing sinks retain their ordinary open-on-demand behavior. diff --git a/src/cpu/stream.body.c b/src/cpu/stream.body.c index da1a84af..67ae4a24 100644 --- a/src/cpu/stream.body.c +++ b/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->prepare_shards) + 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,6 +340,8 @@ cpu_stream_flush_body(struct cpu_stream_view* v) // they read. if (shard_sink_drain(v->sink)) failed = 1; + if (v->config->prepare_shards && shard_sink_cancel_prepared(v->sink)) + failed = 1; return failed ? writer_error() : writer_ok(); } @@ -345,7 +349,7 @@ cpu_stream_flush_body(struct cpu_stream_view* v) struct writer_result cpu_stream_close_body(struct cpu_stream_view* v) { - int failed = 0; + int failed = v->config->prepare_shards && shard_sink_cancel_prepared(v->sink); // A sink IO error is the one case the shape is withheld: which writes landed // is unknowable. diff --git a/src/cpu/stream.c b/src/cpu/stream.c index 7cfa6a9b..37b787da 100644 --- a/src/cpu/stream.c +++ b/src/cpu/stream.c @@ -291,6 +291,14 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, } CHECK(Fail, shard_sink_init_append(sink, &s->cl.dims, s->levels.nlod) == 0); + if (config->prepare_shards) + CHECK(Fail, + shard_sink_prepare_first( + sink, + s->shard, + s->levels.nlod, + s->shard_alignment, + config->codec.id != CODEC_NONE) == 0); s->writer.append = cpu_append; s->writer.flush = cpu_flush_final; @@ -303,6 +311,8 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, return s; Fail: + if (config->prepare_shards && 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; diff --git a/src/gpu/stream.c b/src/gpu/stream.c index 4e066ccd..a7a5873a 100644 --- a/src/gpu/stream.c +++ b/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.prepare_shards) + 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. @@ -436,6 +438,14 @@ stream_flush_body(struct stream_engine* e, struct stream_context* ctx) } schedule_quiesce_output(e, ctx->sink); + if (ctx->config.prepare_shards) { + for (int fc = 0; fc < 2; ++fc) + if (gpu_delivery_pending(&e->delivery, fc) && + gpu_delivery_join(&e->delivery, fc).error) + r = writer_error(); + if (shard_sink_cancel_prepared(ctx->sink)) + r = writer_error(); + } collect_ingest_timing(e); return r; @@ -449,6 +459,9 @@ stream_close_body(struct compress_agg_array* ar, struct stream_context* ctx) struct writer_result r = writer_ok(); + if (ctx->config.prepare_shards && 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) diff --git a/src/gpu/stream.init.c b/src/gpu/stream.init.c index 0f8286a4..0524acac 100644 --- a/src/gpu/stream.init.c +++ b/src/gpu/stream.init.c @@ -350,6 +350,8 @@ 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.prepare_shards && 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 +453,14 @@ tile_stream_gpu_create(const struct tile_stream_configuration* config, shard_sink_init_append(sink, &out->ctx.dims, out->ctx.levels.nlod) == 0); + if (config->prepare_shards) + CHECK(FailPhase2, + shard_sink_prepare_first(sink, + out->engine.compress_agg.ar.shard, + out->ctx.levels.nlod, + out->ctx.shard_alignment, + config->codec.id != CODEC_NONE) == 0); + CHECK(FailPhase2, gpu_placement_leave(&placement_scope) == 0); out->flushed = 0; out->closed = 0; diff --git a/src/multiarray/stream.c b/src/multiarray/stream.c index 2ebd3ff5..cb25285d 100644 --- a/src/multiarray/stream.c +++ b/src/multiarray/stream.c @@ -478,6 +478,13 @@ multiarray_tile_stream_cpu_create( CHECK(Fail, shard_sink_init_append( desc->sink, &desc->cl.dims, desc->levels.nlod) == 0); + if (desc->config.prepare_shards) + CHECK(Fail, + shard_sink_prepare_first(desc->sink, + desc->shard, + desc->levels.nlod, + desc->shard_alignment, + desc->config.codec.id != CODEC_NONE) == 0); } for (int i = 0; i < n_arrays; ++i) { ms->arrays[i].flushed = 0; @@ -486,6 +493,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.prepare_shards && + 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; @@ -740,6 +752,8 @@ flush_impl(struct multiarray_writer* self) if (desc->flushed) continue; if (desc->cursor_elements == 0 && desc->batch_accumulated == 0) { + if (desc->config.prepare_shards && shard_sink_cancel_prepared(desc->sink)) + goto Error; desc->flushed = 1; continue; } diff --git a/src/multiarray/stream.gpu.c b/src/multiarray/stream.gpu.c index 2bb93afe..36bdd1d4 100644 --- a/src/multiarray/stream.gpu.c +++ b/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.prepare_shards && + 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.prepare_shards && + 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,13 @@ 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); + if (ctx->config.prepare_shards) + CHECK(Fail, + shard_sink_prepare_first(ctx->sink, + ms->arrays[a].st.agg.shard, + ctx->levels.nlod, + ctx->shard_alignment, + ctx->config.codec.id != CODEC_NONE) == 0); } for (int a = 0; a < n_arrays; ++a) { ms->arrays[a].flushed = 0; diff --git a/src/ngff/ngff_multiscale.c b/src/ngff/ngff_multiscale.c index 40405298..55f4a80b 100644 --- a/src/ngff/ngff_multiscale.c +++ b/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; diff --git a/src/platform/platform_io.darwin.c b/src/platform/platform_io.darwin.c index 627e5212..b741bc6b 100644 --- a/src/platform/platform_io.darwin.c +++ b/src/platform/platform_io.darwin.c @@ -17,6 +17,25 @@ 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 || errno == ENOTEMPTY || + errno == EEXIST) + return 0; + return -1; +} + int platform_mkdirp(const char* path) { @@ -40,13 +59,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; } } diff --git a/src/platform/platform_io.h b/src/platform/platform_io.h index 17fc1a4e..d4b0524f 100644 --- a/src/platform/platform_io.h +++ b/src/platform/platform_io.h @@ -17,6 +17,14 @@ 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); + +// Missing and nonempty directories are left alone; other failures return -1. +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); @@ -24,7 +32,8 @@ 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 diff --git a/src/platform/platform_io.posix.c b/src/platform/platform_io.posix.c index 614af883..78ead5be 100644 --- a/src/platform/platform_io.posix.c +++ b/src/platform/platform_io.posix.c @@ -19,6 +19,25 @@ 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 || errno == ENOTEMPTY || + errno == EEXIST) + return 0; + return -1; +} + int platform_mkdirp(const char* path) { @@ -42,7 +61,8 @@ 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; if (flags & PLATFORM_OPEN_UNBUFFERED) oflags |= O_DIRECT; return open(path, oflags, 0644); diff --git a/src/platform/platform_io.win32.c b/src/platform/platform_io.win32.c index 853e8b54..ad451110 100644 --- a/src/platform/platform_io.win32.c +++ b/src/platform/platform_io.win32.c @@ -13,6 +13,31 @@ platform_mkdir(const char* path) return -1; } +int +platform_mkdir_new(const char* path) +{ + if (CreateDirectoryA(path, NULL)) + return 1; + if (GetLastError() != ERROR_ALREADY_EXISTS) + return -1; + DWORD attrs = GetFileAttributesA(path); + return attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY) + ? 0 + : -1; +} + +int +platform_remove_empty_directory(const char* path) +{ + if (RemoveDirectoryA(path)) + return 0; + DWORD error = GetLastError(); + return error == ERROR_FILE_NOT_FOUND || error == ERROR_PATH_NOT_FOUND || + error == ERROR_DIR_NOT_EMPTY + ? 0 + : -1; +} + int platform_mkdirp(const char* path) { @@ -44,9 +69,14 @@ platform_open_write(const char* path, int flags) { // PLATFORM_OPEN_UNBUFFERED is a no-op on Windows: NO_BUFFERING + sub-sector // EOF corrupts the trailing partial sector on 4K-logical NTFS volumes. - (void)flags; - return CreateFileA( - path, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL); + return CreateFileA(path, + GENERIC_WRITE, + 0, + NULL, + flags & PLATFORM_OPEN_EXCLUSIVE ? CREATE_NEW + : CREATE_ALWAYS, + FILE_ATTRIBUTE_NORMAL, + NULL); } int diff --git a/src/types.stream.h b/src/types.stream.h index aa60b45a..210a523f 100644 --- a/src/types.stream.h +++ b/src/types.stream.h @@ -154,6 +154,7 @@ struct tile_stream_configuration // reports more pending than this int max_threads; // 0 = OpenMP default int full_memcpy_timing; + int prepare_shards; // opt-in; requires preparation and cleanup from the sink }; struct tile_stream_status diff --git a/src/writer.c b/src/writer.c index dd67494c..6541259f 100644 --- a/src/writer.c +++ b/src/writer.c @@ -88,6 +88,20 @@ writer_finished_at(const void* beg, const void* end) return r; } +void +shard_sink_stop_preparing(struct shard_sink* s) +{ + if (s && s->stop_preparing) + s->stop_preparing(s); +} + +int +shard_sink_cancel_prepared(struct shard_sink* s) +{ + shard_sink_stop_preparing(s); + return s && s->cancel_prepared ? s->cancel_prepared(s) : 0; +} + uint64_t shard_sink_pending_bytes(const struct shard_sink* s) { diff --git a/src/writer.h b/src/writer.h index ab40b124..0aa993d6 100644 --- a/src/writer.h +++ b/src/writer.h @@ -91,6 +91,15 @@ struct shard_sink uint8_t level, uint64_t shard_index); + // Paired capability: prepare the first generation before returning, then + // one generation ahead on open. Stop may race with delivery; cancellation + // runs after delivery joins and waits for, and reclaims, unused resources. + int (*prepare_shards)(struct shard_sink* self, + uint8_t level, + uint64_t capacity); + void (*stop_preparing)(struct shard_sink* self); + int (*cancel_prepared)(struct shard_sink* self); + // Optional: update append dim extents in metadata (e.g. zarr.json shape). // Called with the empty extent at stream creation, periodically during // streaming for synchronous sinks, and from the writer's close. @@ -132,6 +141,12 @@ struct shard_sink size_t (*required_shard_alignment)(const struct shard_sink* self); }; +void +shard_sink_stop_preparing(struct shard_sink* s); + +int +shard_sink_cancel_prepared(struct shard_sink* s); + uint64_t shard_sink_pending_bytes(const struct shard_sink* s); diff --git a/src/zarr/io_backend.fs.c b/src/zarr/io_backend.fs.c index aa925b1e..04390d41 100644 --- a/src/zarr/io_backend.fs.c +++ b/src/zarr/io_backend.fs.c @@ -103,6 +103,22 @@ io_backend_fs_reserve_file(struct io_backend_fs* b) return (struct io_file_token){ 0 }; } +struct io_file_token +io_backend_fs_adopt_file(struct io_backend_fs* b, platform_fd fd) +{ + if (fd == PLATFORM_FD_INVALID) + return (struct io_file_token){ 0 }; + struct io_file_token token = io_backend_fs_reserve_file(b); + if (token.generation == 0) + return token; + platform_mutex_lock(b->mutex); + b->files[token.index].fd = fd; + if (++b->handle_count > b->peak_handle_count) + b->peak_handle_count = b->handle_count; + platform_mutex_unlock(b->mutex); + return token; +} + static struct file_entry* find_locked(struct io_backend_fs* b, struct io_file_token file) { diff --git a/src/zarr/io_backend.fs.h b/src/zarr/io_backend.fs.h index 847a9aa9..ce7bdf0a 100644 --- a/src/zarr/io_backend.fs.h +++ b/src/zarr/io_backend.fs.h @@ -1,5 +1,6 @@ #pragma once +#include "platform/platform_io.h" #include "zarr/io_backend.h" #include @@ -30,6 +31,10 @@ io_backend_fs_as_backend(struct io_backend_fs* b); struct io_file_token io_backend_fs_reserve_file(struct io_backend_fs* b); +// Takes ownership of fd on success. No scheduled request may name it yet. +struct io_file_token +io_backend_fs_adopt_file(struct io_backend_fs* b, platform_fd fd); + // Only a reservation whose open was never posted may be cancelled. void io_backend_fs_cancel_file(struct io_backend_fs* b, struct io_file_token file); diff --git a/src/zarr/shard_delivery.c b/src/zarr/shard_delivery.c index 202a1fe8..59e1c060 100644 --- a/src/zarr/shard_delivery.c +++ b/src/zarr/shard_delivery.c @@ -243,6 +243,35 @@ record_finalized(struct shard_state* ss, struct shard_sink* sink) } } +int +shard_sink_prepare_first(struct shard_sink* sink, + const struct shard_state* shards, + int nlod, + size_t alignment, + int padded) +{ + if (!sink || !sink->prepare_shards || !sink->stop_preparing || + !sink->cancel_prepared) { + log_error("sink does not support shard preparation with cleanup"); + return 1; + } + for (int lv = 0; lv < nlod; ++lv) { + uint64_t capacity = shards[lv].shard_file_capacity; + if (capacity && padded && alignment > 1) { + const uint64_t gaps = shards[lv].chunks_per_shard_append - 1; + capacity = gaps <= (UINT64_MAX - capacity) / (alignment - 1) + ? capacity + gaps * (alignment - 1) + : 0; + } + if (sink->prepare_shards(sink, (uint8_t)lv, capacity)) { + if (shard_sink_cancel_prepared(sink)) + log_error("sink preparation rollback failed"); + return 1; + } + } + return 0; +} + int shard_sink_init_append(struct shard_sink* sink, const struct dim_info* dims, diff --git a/src/zarr/shard_delivery.h b/src/zarr/shard_delivery.h index 0e1a2605..ad2021e5 100644 --- a/src/zarr/shard_delivery.h +++ b/src/zarr/shard_delivery.h @@ -78,6 +78,13 @@ shard_state_destroy(struct shard_state* ss); size_t shard_state_heap_bytes(const struct level_layout_info* li); +int +shard_sink_prepare_first(struct shard_sink* sink, + const struct shard_state* shards, + int nlod, + size_t alignment, + int padded); + // A new stream advertises no appended data, independently of its configured // capacity. Called once during creation, before any array accepts input. int diff --git a/src/zarr/shard_pool.h b/src/zarr/shard_pool.h index 752f98a3..738dd68d 100644 --- a/src/zarr/shard_pool.h +++ b/src/zarr/shard_pool.h @@ -16,6 +16,18 @@ struct shard_pool uint64_t slot, const char* key); + // These hooks are a single optional capability. One unused resource per + // slot; open adopts it. Cancellation waits, reclaims only owned resources, + // reports failures, and is safe to repeat after the caller stops preparing. + int (*prepare)(struct shard_pool* self, + uint64_t slot, + const char* key, + uint64_t capacity); + int (*wait_prepared)(struct shard_pool* self, uint64_t slot); + int (*cancel_prepared)(struct shard_pool* self, + uint64_t first, + uint64_t count); + // Record a fence capturing the current I/O sequence point. struct io_event (*record_fence)(struct shard_pool* self); diff --git a/src/zarr/shard_pool_fs.c b/src/zarr/shard_pool_fs.c index f0e60cf0..ae0b66b8 100644 --- a/src/zarr/shard_pool_fs.c +++ b/src/zarr/shard_pool_fs.c @@ -16,11 +16,24 @@ struct fs_slot; +struct prepared_file +{ + char* path; + unsigned char* created_dirs; + platform_fd fd; + uint64_t capacity; + struct io_event ready; + int owned; +}; + struct shard_pool_fs { struct shard_pool base; struct io_backend_fs* backend; struct io_scheduler* queue; + struct io_scheduler* prepare_queue; + struct io_backend prepare_backend; + struct prepared_file* prepared; struct fs_slot* slots; uint64_t nslots; int unbuffered; @@ -37,6 +50,7 @@ struct fs_slot struct io_scheduler* queue; size_t alignment; // 0 = normal malloc, >0 = page-aligned allocation int presize; // set the file's size up front + uint64_t prepared_capacity; }; static int @@ -165,7 +179,7 @@ static int fs_slot_presize(struct shard_writer* self, uint64_t nbytes) { struct fs_slot* w = (struct fs_slot*)self; - if (w->token.generation == 0 || !w->presize || nbytes == 0) + if (w->token.generation == 0 || !w->presize || nbytes <= w->prepared_capacity) return 0; return io_scheduler_post(w->queue, @@ -203,7 +217,174 @@ fs_slot_finalize(struct shard_writer* self) return 1; w->token = (struct io_file_token){ 0 }; + w->prepared_capacity = 0; + return 0; +} + +static void +prepared_file_clear(struct prepared_file* f) +{ + free(f->created_dirs); + free(f->path); + *f = (struct prepared_file){ .fd = PLATFORM_FD_INVALID }; +} + +static void +prepare_file(void* ctx, const struct io_request* req) +{ + struct shard_pool_fs* p = ctx; + struct prepared_file* f = (struct prepared_file*)req->payload; + char* directory = strdup(f->path); + CHECK(Fail, directory); + const int flags = + PLATFORM_OPEN_EXCLUSIVE | (p->unbuffered ? PLATFORM_OPEN_UNBUFFERED : 0); + f->fd = platform_open_write(f->path, flags); + if (f->fd == PLATFORM_FD_INVALID) { + const size_t len = strlen(directory); + for (size_t i = 1; i < len; ++i) { + if (directory[i] != '/' && directory[i] != '\\') + continue; + if (i == 2 && directory[1] == ':') + continue; + const char saved = directory[i]; + directory[i] = 0; + const int made = platform_mkdir_new(directory); + directory[i] = saved; + CHECK(Fail, made >= 0); + f->created_dirs[i] = made != 0; + } + f->fd = platform_open_write(f->path, flags); + } + CHECK(Fail, f->fd != PLATFORM_FD_INVALID); + f->owned = 1; + if (f->capacity) + CHECK(Fail, platform_ftruncate(f->fd, f->capacity) == 0); + free(directory); + return; +Fail: + log_error("shard preparation failed: %s", f->path); + atomic_store(&p->io_error, 1); + free(directory); +} + +static int +pool_fs_prepare(struct shard_pool* self, + uint64_t slot, + const char* key, + uint64_t capacity) +{ + struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); + struct strbuf path = { 0 }; + CHECK(Fail, slot < p->nslots && key && !atomic_load(&p->io_error)); + if (!p->prepare_queue) { + CHECK(Fail, p->nslots <= SIZE_MAX / sizeof(*p->prepared)); + p->prepared = calloc((size_t)p->nslots, sizeof(*p->prepared)); + CHECK(Fail, p->prepared); + for (uint64_t i = 0; i < p->nslots; ++i) + p->prepared[i].fd = PLATFORM_FD_INVALID; + p->prepare_queue = io_scheduler_create( + p->prepare_backend, + (struct io_scheduler_limits){ .workers = 1, .max_requests = p->nslots }); + CHECK(Fail, p->prepare_queue); + } + struct prepared_file* f = &p->prepared[slot]; + CHECK(Fail, !f->path); + CHECK(Fail, strbuf_appendf(&path, "%s/%s", strbuf_cstr(&p->root), key) == 0); + f->path = strdup(strbuf_cstr(&path)); + f->created_dirs = calloc(strbuf_len(&path) + 1, 1); + if (!f->path || !f->created_dirs) { + prepared_file_clear(f); + goto Fail; + } + f->capacity = p->slots[slot].presize ? capacity : 0; + if (io_scheduler_post(p->prepare_queue, + (struct io_request){ .op = IO_OP_NOOP, + .payload = f, + .path = f->path, + .logical_size = f->capacity })) { + prepared_file_clear(f); + goto Fail; + } + f->ready = io_scheduler_record(p->prepare_queue); + strbuf_free(&path); return 0; +Fail: + strbuf_free(&path); + atomic_store(&p->io_error, 1); + return 1; +} + +static int +pool_fs_wait_prepared(struct shard_pool* self, uint64_t slot) +{ + struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); + if (slot >= p->nslots) + return 1; + if (p->prepare_queue && p->prepared[slot].path) + io_event_wait(p->prepare_queue, p->prepared[slot].ready); + return atomic_load(&p->io_error); +} + +static int +pool_fs_cancel_prepared(struct shard_pool* self, uint64_t first, uint64_t count) +{ + struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); + if (first > p->nslots || count > p->nslots - first) + return 1; + if (!p->prepared) + return 0; + int failed = 0; + size_t longest = 0; + for (uint64_t i = first; i < first + count; ++i) { + struct prepared_file* f = &p->prepared[i]; + if (!f->path) + continue; + if (p->prepare_queue) + io_event_wait(p->prepare_queue, f->ready); + if (f->fd != PLATFORM_FD_INVALID) { + platform_close(f->fd); + f->fd = PLATFORM_FD_INVALID; + } + if (f->owned) { + if (platform_remove_file(f->path) == 0 || + platform_path_exists(f->path) == 0) + f->owned = 0; + else { + log_error("unused shard cleanup failed: %s", f->path); + failed = 1; + } + } + const size_t len = strlen(f->path); + if (len > longest) + longest = len; + } + if (failed) { + atomic_store(&p->io_error, 1); + return 1; + } + for (size_t pos = longest; pos > 0; --pos) { + for (uint64_t i = first; i < first + count; ++i) { + struct prepared_file* f = &p->prepared[i]; + if (!f->path || strlen(f->path) <= pos || !f->created_dirs[pos]) + continue; + char saved = f->path[pos]; + f->path[pos] = 0; + if (platform_remove_empty_directory(f->path)) { + log_error("unused shard directory cleanup failed: %s", f->path); + failed = 1; + } else { + f->created_dirs[pos] = 0; + } + f->path[pos] = saved; + } + } + if (!failed) { + for (uint64_t i = first; i < first + count; ++i) + prepared_file_clear(&p->prepared[i]); + } else { + atomic_store(&p->io_error, 1); + } + return failed; } static struct shard_writer* @@ -222,6 +403,20 @@ pool_fs_open(struct shard_pool* self, uint64_t slot, const char* key) if (strbuf_appendf(&path, "%s/%s", strbuf_cstr(&p->root), key)) goto Fail; + struct prepared_file* f = p->prepared ? &p->prepared[slot] : NULL; + if (f && f->path) { + CHECK(Fail, strcmp(f->path, strbuf_cstr(&path)) == 0); + CHECK(Fail, pool_fs_wait_prepared(self, slot) == 0); + const struct io_file_token token = + io_backend_fs_adopt_file(p->backend, f->fd); + CHECK(Fail, token.generation != 0); + w->token = token; + w->prepared_capacity = f->capacity; + prepared_file_clear(f); + strbuf_free(&path); + return &w->base; + } + const size_t path_bytes = strbuf_len(&path) + 1; owned_path = (char*)malloc(path_bytes); CHECK(Fail, owned_path); @@ -342,6 +537,15 @@ pool_fs_destroy(struct shard_pool* self) { struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); + if (pool_fs_cancel_prepared(self, 0, p->nslots)) + log_error("shard pool cleanup failed during destroy"); + io_scheduler_destroy(p->prepare_queue); + if (p->prepared) { + for (uint64_t i = 0; i < p->nslots; ++i) + prepared_file_clear(&p->prepared[i]); + free(p->prepared); + } + // Finalize any open slots for (uint64_t i = 0; i < p->nslots; ++i) { if (p->slots[i].token.generation != 0) @@ -404,6 +608,12 @@ shard_pool_fs_create_wrapped(const char* root, CHECK(Fail, p); p->base.open = pool_fs_open; + p->base.prepare = pool_fs_prepare; + p->base.wait_prepared = pool_fs_wait_prepared; + p->base.cancel_prepared = pool_fs_cancel_prepared; + p->prepare_backend = (struct io_backend){ .ctx = p, .execute = prepare_file }; + if (wrapper.wrap_prepare) + p->prepare_backend = wrapper.wrap_prepare(wrapper.ctx, p->prepare_backend); p->base.record_fence = pool_fs_record_fence; p->base.wait_fence = pool_fs_wait_fence; p->base.queue_metadata = pool_fs_queue_metadata; diff --git a/src/zarr/shard_pool_fs.h b/src/zarr/shard_pool_fs.h index d7c10d2f..346fb94a 100644 --- a/src/zarr/shard_pool_fs.h +++ b/src/zarr/shard_pool_fs.h @@ -12,6 +12,7 @@ struct io_scheduler_limits; // Handles belong to current slots or unfinished closes, bounding their count // by nslots plus the scheduler's max_requests (1024 by default). +// Lookahead holds at most nslots additional unused files. struct shard_pool* shard_pool_fs_create(const char* root, uint64_t nslots, int unbuffered); @@ -23,6 +24,7 @@ struct shard_pool_fs_wrapper // The pool's backend is the argument, and its executor calls the result. struct io_backend (*wrap)(void* ctx, struct io_backend inner); struct io_scheduler** queue; // receives the scheduler the pool built + struct io_backend (*wrap_prepare)(void* ctx, struct io_backend inner); }; // Create a filesystem shard pool with the wrapper between its scheduler and its diff --git a/src/zarr/zarr_array.c b/src/zarr/zarr_array.c index 196ee45e..a0806376 100644 --- a/src/zarr/zarr_array.c +++ b/src/zarr/zarr_array.c @@ -8,6 +8,7 @@ #include "zarr/metadata_io.h" #include "zarr/zarr_metadata.h" +#include #include struct zarr_array @@ -23,6 +24,9 @@ struct zarr_array uint64_t chunks_per_shard[MAX_ZARR_RANK]; uint64_t shard_inner_count; uint64_t slot_base; // first pool slot used by this array + uint64_t shard_limit; + uint64_t prepared_capacity; + _Atomic int preparing; // Mutable copy for metadata updates struct dimension dimensions[MAX_ZARR_RANK]; @@ -63,32 +67,88 @@ submit_array_metadata(struct zarr_array* a) // --- shard_sink vtable --- +static int +array_shard_key(struct zarr_array* a, uint64_t shard_index, struct strbuf* key) +{ + char suffix[256]; + if (zarr_shard_key( + suffix, sizeof(suffix), a->rank, a->shard_counts, shard_index)) + return 1; + return strbuf_len(&a->prefix) + ? strbuf_appendf(key, "%s/%s", strbuf_cstr(&a->prefix), suffix) + : strbuf_append_cstr(key, suffix); +} + +static int +prepare_shard(struct zarr_array* a, uint64_t shard_index) +{ + struct strbuf key = { 0 }; + int rc = array_shard_key(a, shard_index, &key); + if (!rc) + rc = a->pool->prepare(a->pool, + a->slot_base + shard_index % a->shard_inner_count, + strbuf_cstr(&key), + a->prepared_capacity); + strbuf_free(&key); + return rc; +} + +static void +zarr_array_stop_preparing(struct shard_sink* self) +{ + struct zarr_array* a = container_of(self, struct zarr_array, base); + atomic_store(&a->preparing, 0); +} + +static int +zarr_array_cancel_prepared(struct shard_sink* self) +{ + struct zarr_array* a = container_of(self, struct zarr_array, base); + zarr_array_stop_preparing(self); + return a->pool->cancel_prepared(a->pool, a->slot_base, a->shard_inner_count); +} + +static int +zarr_array_prepare_shards(struct shard_sink* self, + uint8_t level, + uint64_t capacity) +{ + (void)level; + struct zarr_array* a = container_of(self, struct zarr_array, base); + if (atomic_load(&a->preparing)) + return 1; + a->prepared_capacity = capacity; + for (uint64_t i = 0; i < a->shard_inner_count; ++i) + if (prepare_shard(a, i)) + goto Fail; + for (uint64_t i = 0; i < a->shard_inner_count; ++i) + if (a->pool->wait_prepared(a->pool, a->slot_base + i)) + goto Fail; + atomic_store(&a->preparing, 1); + return 0; +Fail: + if (zarr_array_cancel_prepared(self)) + log_error("zarr_array: preparation rollback failed"); + return 1; +} + static struct shard_writer* zarr_array_open(struct shard_sink* self, uint8_t level, uint64_t shard_index) { (void)level; struct zarr_array* a = container_of(self, struct zarr_array, base); - uint64_t slot = a->slot_base + shard_index % a->shard_inner_count; - - char suffix[256]; - if (zarr_shard_key( - suffix, sizeof(suffix), a->rank, a->shard_counts, shard_index) != 0) { - log_error("zarr_array: shard key too long for shard %llu", - (unsigned long long)shard_index); - return NULL; - } - struct strbuf key = { 0 }; - int ok; - if (strbuf_len(&a->prefix) > 0) - ok = strbuf_appendf(&key, "%s/%s", strbuf_cstr(&a->prefix), suffix) == 0; - else - ok = strbuf_append_cstr(&key, suffix) == 0; - + int rc = array_shard_key(a, shard_index, &key); struct shard_writer* w = - ok ? a->pool->open(a->pool, slot, strbuf_cstr(&key)) : NULL; + rc ? NULL : a->pool->open(a->pool, slot, strbuf_cstr(&key)); strbuf_free(&key); + if (w && atomic_load(&a->preparing) && + shard_index <= UINT64_MAX - a->shard_inner_count) { + const uint64_t next = shard_index + a->shard_inner_count; + if ((!a->shard_limit || next < a->shard_limit) && prepare_shard(a, next)) + return NULL; + } return w; } @@ -230,7 +290,21 @@ zarr_array_init(struct store* store, a->chunks_per_shard[d] = chunks_per_shard[d]; } + a->shard_limit = cfg->dimensions[0].size ? 1 : 0; + for (int d = 0; d < cfg->rank; ++d) { + if (!shard_counts[d] || a->shard_limit > UINT64_MAX / shard_counts[d]) { + a->shard_limit = 0; + break; + } + a->shard_limit *= shard_counts[d]; + } + a->base.open = zarr_array_open; + if (pool->prepare && pool->wait_prepared && pool->cancel_prepared) { + a->base.prepare_shards = zarr_array_prepare_shards; + a->base.stop_preparing = zarr_array_stop_preparing; + a->base.cancel_prepared = zarr_array_cancel_prepared; + } a->base.update_append = zarr_array_update_append; a->base.queue_append = pool->queue_metadata ? zarr_array_queue_append : NULL; a->base.record_fence = zarr_array_record_fence_fn; @@ -319,6 +393,8 @@ zarr_array_destroy(struct zarr_array* a) { if (!a) return; + if (shard_sink_cancel_prepared(&a->base)) + log_error("zarr_array: unused shard cleanup failed during destroy"); // Fallback for callers that didn't flush via the writer/sink — same shape // as the auto-flush log in stream destroys. if (zarr_array_flush_metadata(a)) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c15ad43c..5700c348 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -328,3 +328,14 @@ add_gpu_test_exe(test_gpu_placement test_gpu_placement.c LINKS stream multiarray_gpu test_shard_sink platform threadpool chucky_log CUDA::cuda_driver CUDA::cudart_static ) + +add_test_exe(test_shard_preparation test_shard_preparation.c + shard_pool_fs platform_io test_platform chucky_log) +add_test_exe(test_stream_preparation test_stream_preparation.c + stream_cpu test_zarr_helpers test_shard_verify platform_io test_platform) +add_gpu_test_exe(test_stream_preparation_gpu test_stream_preparation.c + LINKS stream test_zarr_helpers test_shard_verify platform_io test_platform CUDA::cuda_driver) +if(CHUCKY_ENABLE_GPU) + target_compile_definitions(test_stream_preparation_gpu PRIVATE TEST_PREPARATION_GPU=1) +endif() +set_tests_properties(test-test_shard_preparation PROPERTIES TIMEOUT 30) diff --git a/tests/test_shard_preparation.c b/tests/test_shard_preparation.c new file mode 100644 index 00000000..7802b61d --- /dev/null +++ b/tests/test_shard_preparation.c @@ -0,0 +1,228 @@ +#include "platform/platform.h" +#include "platform/platform_io.h" +#include "test_platform.h" +#include "util/prelude.h" +#include "zarr/io_scheduler.h" +#include "zarr/shard_pool_fs.h" + +#include +#include +#include +#include + +struct preparation_gate +{ + struct io_backend inner; + _Atomic int entered; + _Atomic int release; +}; + +static void +blocked_prepare(void* ctx, const struct io_request* request) +{ + struct preparation_gate* g = ctx; + atomic_store(&g->entered, 1); + while (!atomic_load(&g->release)) + platform_sleep_ns(1000000); + g->inner.execute(g->inner.ctx, request); +} + +static struct io_backend +wrap_prepare(void* ctx, struct io_backend inner) +{ + struct preparation_gate* g = ctx; + g->inner = inner; + return (struct io_backend){ .ctx = g, .execute = blocked_prepare }; +} + +struct pool_wait +{ + struct shard_pool* pool; + int cancel; + int result; + _Atomic int entered; + _Atomic int done; +}; + +static void +wait_pool(void* arg) +{ + struct pool_wait* w = arg; + atomic_store(&w->entered, 1); + w->result = w->cancel ? w->pool->cancel_prepared(w->pool, 0, 1) + : w->pool->flush(w->pool); + atomic_store(&w->done, 1); +} + +static int +exists(const char* root, const char* key) +{ + char path[4096]; + snprintf(path, sizeof(path), "%s/%s", root, key); + return platform_path_exists(path); +} + +static int +test_adoption_and_cleanup(void) +{ + char root[256] = { 0 }, path[512]; + struct shard_pool* pool = NULL; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + snprintf(path, sizeof(path), "%s/existing", root); + CHECK(Done, platform_mkdir(path) == 0); + pool = shard_pool_fs_create(root, 3, 0); + CHECK(Done, pool); + CHECK(Done, pool->prepare(pool, 0, "c/0/0", 65536) == 0); + CHECK(Done, pool->prepare(pool, 1, "c/1/0", 65536) == 0); + CHECK(Done, pool->prepare(pool, 2, "existing/future", 65536) == 0); + CHECK(Done, pool->wait_prepared(pool, 2) == 0); + struct shard_writer* w = pool->open(pool, 0, "c/0/0"); + CHECK(Done, w); + const char value[] = "kept"; + CHECK(Done, w->presize(w, 65536) == 0); + CHECK(Done, w->write(w, 0, value, value + sizeof(value)) == 0); + CHECK(Done, w->truncate(w, sizeof(value)) == 0 && w->finalize(w) == 0); + CHECK(Done, pool->prepare(pool, 0, "c/2/0", 65536) == 0); + CHECK(Done, pool->flush(pool) == 0); + CHECK(Done, pool->cancel_prepared(pool, 0, 2) == 0); + CHECK(Done, exists(root, "c/0/0") == 1); + CHECK(Done, exists(root, "c/1") == 0 && exists(root, "c/2") == 0); + CHECK(Done, exists(root, "existing/future") == 1); + CHECK(Done, pool->cancel_prepared(pool, 2, 1) == 0); + CHECK(Done, pool->cancel_prepared(pool, 0, 3) == 0); + CHECK(Done, exists(root, "existing") == 1); + snprintf(path, sizeof(path), "%s/c/0/0", root); + FILE* f = fopen(path, "rb"); + char actual[sizeof(value)]; + CHECK(Done, f); + size_t n = fread(actual, 1, sizeof(actual), f); + int tail = fgetc(f); + fclose(f); + CHECK(Done, + n == sizeof(value) && tail == EOF && + memcmp(actual, value, sizeof(value)) == 0); + error = 0; +Done: + shard_pool_destroy(pool); + test_tmpdir_remove(root); + return error; +} + +static int +test_preparation_does_not_hold_writes(void) +{ + char root[256] = { 0 }; + struct preparation_gate gate = { 0 }; + struct shard_pool* pool = NULL; + test_thread* flush_thread = NULL; + test_thread* cancel_thread = NULL; + struct pool_wait flush = { 0 }, cancel = { .cancel = 1 }; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + pool = shard_pool_fs_create_wrapped( + root, + 1, + 0, + NULL, + (struct shard_pool_fs_wrapper){ .ctx = &gate, + .wrap_prepare = wrap_prepare }); + CHECK(Done, pool); + struct shard_writer* w = pool->open(pool, 0, "c/0/0"); + CHECK(Done, w); + CHECK(Done, pool->prepare(pool, 0, "c/1/0", 4096) == 0); + CHECK(Done, test_wait_flag(&gate.entered, 5000) == 0); + const char value[] = "current"; + CHECK(Done, w->write(w, 0, value, value + sizeof(value)) == 0); + CHECK(Done, w->finalize(w) == 0); + CHECK(Done, pool->queue_metadata(pool, "zarr.json", "{}", 2) == 0); + flush.pool = cancel.pool = pool; + CHECK(Done, test_thread_start(&flush_thread, wait_pool, &flush) == 0); + CHECK(Done, test_wait_flag(&flush.done, 5000) == 0 && flush.result == 0); + CHECK(Done, test_thread_start(&cancel_thread, wait_pool, &cancel) == 0); + CHECK(Done, test_wait_flag(&cancel.entered, 5000) == 0); + CHECK(Done, test_wait_flag(&cancel.done, 50) != 0); + CHECK(Done, exists(root, "c/0/0") == 1 && exists(root, "zarr.json") == 1); + atomic_store(&gate.release, 1); + CHECK(Done, test_wait_flag(&cancel.done, 5000) == 0 && cancel.result == 0); + CHECK(Done, exists(root, "c/1") == 0); + error = 0; +Done: + atomic_store(&gate.release, 1); + test_thread_join(flush_thread); + test_thread_join(cancel_thread); + shard_pool_destroy(pool); + test_tmpdir_remove(root); + return error; +} + +static int +test_existing_file_is_preserved(void) +{ + char root[256] = { 0 }, path[512]; + struct shard_pool* pool = NULL; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + snprintf(path, sizeof(path), "%s/keep", root); + platform_fd fd = platform_open_write(path, 0); + CHECK(Done, fd != PLATFORM_FD_INVALID); + int written = platform_write(fd, "sentinel", 8); + platform_close(fd); + CHECK(Done, written == 0); + pool = shard_pool_fs_create(root, 1, 0); + CHECK(Done, pool && pool->prepare(pool, 0, "keep", 65536) == 0); + CHECK(Done, pool->wait_prepared(pool, 0) != 0); + CHECK(Done, pool->has_error(pool)); + CHECK(Done, pool->cancel_prepared(pool, 0, 1) == 0); + CHECK(Done, pool->cancel_prepared(pool, 0, 1) == 0); + FILE* f = fopen(path, "rb"); + char actual[9] = { 0 }; + CHECK(Done, f); + size_t n = fread(actual, 1, sizeof(actual), f); + fclose(f); + CHECK(Done, n == 8 && strcmp(actual, "sentinel") == 0); + error = 0; +Done: + shard_pool_destroy(pool); + test_tmpdir_remove(root); + return error; +} + +static int +test_cleanup_failure_can_be_retried(void) +{ + char root[256] = { 0 }, path[512], moved[512]; + struct shard_pool* pool = NULL; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + pool = shard_pool_fs_create(root, 1, 0); + CHECK(Done, pool && pool->prepare(pool, 0, "c/3/0", 4096) == 0); + CHECK(Done, pool->wait_prepared(pool, 0) == 0); + snprintf(path, sizeof(path), "%s/c/3/0", root); + snprintf(moved, sizeof(moved), "%s/held", root); +#ifndef _WIN32 + CHECK(Done, platform_rename_replace(path, moved) == 0); + CHECK(Done, platform_mkdir(path) == 0); + CHECK(Done, pool->cancel_prepared(pool, 0, 1) != 0); + CHECK(Done, pool->has_error(pool)); + CHECK(Done, platform_remove_empty_directory(path) == 0); + CHECK(Done, platform_rename_replace(moved, path) == 0); +#endif + CHECK(Done, pool->cancel_prepared(pool, 0, 1) == 0); + CHECK(Done, exists(root, "c") == 0); + error = 0; +Done: + shard_pool_destroy(pool); + test_tmpdir_remove(root); + return error; +} + +int +main(void) +{ + int error = test_adoption_and_cleanup(); + error |= test_preparation_does_not_hold_writes(); + error |= test_existing_file_is_preserved(); + error |= test_cleanup_failure_can_be_retried(); + return error; +} diff --git a/tests/test_stream_preparation.c b/tests/test_stream_preparation.c new file mode 100644 index 00000000..8c586728 --- /dev/null +++ b/tests/test_stream_preparation.c @@ -0,0 +1,301 @@ +#include "platform/platform_io.h" +#include "test_platform.h" +#include "test_shard_verify.h" +#include "test_zarr_helpers.h" +#include "util/prelude.h" +#include "writer.h" +#include "zarr/shard_pool.h" +#include "zarr/shard_pool_fs.h" +#include "zarr/zarr_array.h" + +#ifdef TEST_PREPARATION_GPU +#include "gpu/prelude.cuda.h" +#include "stream.gpu.h" +#define stream_type tile_stream_gpu +#define stream_create tile_stream_gpu_create +#define stream_writer tile_stream_gpu_writer +#define stream_destroy tile_stream_gpu_destroy +#define stream_reset tile_stream_gpu_reset_metrics +#else +#include "stream.cpu.h" +#define stream_type tile_stream_cpu +#define stream_create tile_stream_cpu_create +#define stream_writer tile_stream_cpu_writer +#define stream_destroy tile_stream_cpu_destroy +#define stream_reset tile_stream_cpu_reset_metrics +#endif + +#include +#include +#include +#include + +static int +verify_shard(const char* path, int first, int frames) +{ + FILE* f = fopen(path, "rb"); + if (!f) + return 1; + uint8_t bytes[16384]; + size_t len = fread(bytes, 1, sizeof(bytes), f); + int end = fgetc(f); + fclose(f); + uint64_t offset[2], size[2]; + if (end != EOF || shard_index_parse(bytes, len, 2, offset, size)) + return 1; + for (int t = 0; t < 2; ++t) { + if (first + t >= frames) { + if (size[t] != UINT64_MAX) + return 1; + continue; + } + if (size[t] != 32 || offset[t] > len || size[t] > len - offset[t]) + return 1; + for (int i = 0; i < 16; ++i) { + uint16_t value; + memcpy(&value, bytes + offset[t] + (size_t)i * 2, 2); + if (value != (uint16_t)(100 + first + t)) + return 1; + } + } + return 0; +} + +static int +test_stream(int frames, int bounded, int destroy_only) +{ + char root[256] = { 0 }, path[512]; + struct test_zarr_sink z = { 0 }; + struct stream_type* stream = NULL; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + struct dimension dims[] = { + { .size = bounded ? 4 : 0, + .chunk_size = 1, + .chunks_per_shard = 2, + .name = "t", + .storage_position = 0 }, + { .size = 8, + .chunk_size = 4, + .chunks_per_shard = 1, + .name = "y", + .storage_position = 1 }, + { .size = 8, + .chunk_size = 4, + .chunks_per_shard = 1, + .name = "x", + .storage_position = 2 }, + }; + const struct codec_config codec = { .id = CODEC_NONE }; + CHECK(Done, + test_zarr_sink_open_with_pool( + &z, store_fs_create(root, 1), "0", dims, 3, dtype_u16, codec) == 0); + struct tile_stream_configuration config = { + .buffer_capacity_bytes = 128, + .dtype = dtype_u16, + .rank = 3, + .dimensions = dims, + .codec = codec, + .max_threads = 2, + .max_nlod = 1, + .epochs_per_batch = 1, + .prepare_shards = 1, + }; + stream = stream_create(&config, test_zarr_sink_as_shard_sink(&z)); + CHECK(Done, stream); + for (int y = 0; y < 2; ++y) + for (int x = 0; x < 2; ++x) { + snprintf(path, sizeof(path), "%s/0/c/0/%d/%d", root, y, x); + CHECK(Done, platform_path_exists(path) == 1); + } + struct writer* w = stream_writer(stream); + for (int t = 0; t < frames; ++t) { + uint16_t frame[64]; + for (int i = 0; i < 64; ++i) + frame[i] = (uint16_t)(100 + t); + CHECK(Done, + writer_append_wait(w, (struct slice){ frame, frame + 64 }).error == + 0); + } + if (frames && !destroy_only) { + CHECK(Done, stream_reset(stream) == 0); + for (uint64_t slot = 0; slot < 4; ++slot) + CHECK(Done, z.pool->wait_prepared(z.pool, slot) == 0); + const int next = (frames - 1) / 2 + 1; + for (int y = 0; y < 2; ++y) + for (int x = 0; x < 2; ++x) { + snprintf(path, sizeof(path), "%s/0/c/%d/%d/%d", root, next, y, x); + CHECK(Done, platform_path_exists(path) == (!bounded || next < 2)); + snprintf(path, sizeof(path), "%s/0/c/%d/%d/%d", root, next + 1, y, x); + CHECK(Done, platform_path_exists(path) == 0); + } + } + if (!destroy_only) { + CHECK(Done, writer_flush(w).error == 0); + CHECK(Done, writer_close(w).error == 0); + CHECK(Done, writer_flush(w).error == 0 && writer_close(w).error == 0); + } + stream_destroy(stream); + stream = NULL; + CHECK(Done, test_zarr_sink_has_error(&z) == 0); + for (int g = 0; g < (frames + 1) / 2; ++g) + for (int y = 0; y < 2; ++y) + for (int x = 0; x < 2; ++x) { + snprintf(path, sizeof(path), "%s/0/c/%d/%d/%d", root, g, y, x); + CHECK(Done, verify_shard(path, g * 2, frames) == 0); + } + for (int g = (frames + 1) / 2; g < 5; ++g) { + snprintf(path, sizeof(path), "%s/0/c/%d", root, g); + CHECK(Done, platform_path_exists(path) == 0); + } + if (!frames) { + snprintf(path, sizeof(path), "%s/0/c", root); + CHECK(Done, platform_path_exists(path) == 0); + } + error = 0; +Done: + stream_destroy(stream); + test_zarr_sink_close(&z); + test_tmpdir_remove(root); + return error; +} + +struct resize_count +{ + struct io_backend inner; + _Atomic int resizes; + _Atomic int opens; +}; + +static void +count_resizes(void* ctx, const struct io_request* request) +{ + struct resize_count* c = ctx; + if (request->op == IO_OP_TRUNCATE) + atomic_fetch_add(&c->resizes, 1); + if (request->op == IO_OP_OPEN) + atomic_fetch_add(&c->opens, 1); + c->inner.execute(c->inner.ctx, request); +} + +static struct io_backend +wrap_resizes(void* ctx, struct io_backend inner) +{ + struct resize_count* c = ctx; + c->inner = inner; + return (struct io_backend){ .ctx = c, .execute = count_resizes }; +} + +static int +test_prepared_files_need_no_resize(enum compression_codec codec_id) +{ + char root[256] = { 0 }; + struct test_zarr_sink z = { 0 }; + struct stream_type* stream = NULL; + struct resize_count counts = { 0 }; + uint16_t* frame = NULL; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + struct dimension dims[] = { + { .size = 0, + .chunk_size = 1, + .chunks_per_shard = 2, + .name = "t", + .storage_position = 0 }, + { .size = 128, + .chunk_size = 64, + .chunks_per_shard = 1, + .name = "y", + .storage_position = 1 }, + { .size = 256, + .chunk_size = 128, + .chunks_per_shard = 1, + .name = "x", + .storage_position = 2 }, + }; + struct codec_config codec = { + .id = codec_id, + .level = 3, + .shuffle = CODEC_SHUFFLE_BIT, + .blosc_block_bytes = 16384, + }; + z.store = store_fs_create(root, 1); + CHECK(Done, z.store && z.store->mkdirs(z.store, "0") == 0); + z.pool = shard_pool_fs_create_wrapped( + root, + 4, + 1, + NULL, + (struct shard_pool_fs_wrapper){ .ctx = &counts, .wrap = wrap_resizes }); + CHECK(Done, z.pool); + struct zarr_array_config array_config = { + .data_type = dtype_u16, + .rank = 3, + .dimensions = dims, + .codec = codec, + }; + z.array = zarr_array_create_with_pool(z.store, z.pool, 0, "0", &array_config); + CHECK(Done, z.array); + struct tile_stream_configuration config = { + .buffer_capacity_bytes = 65536, + .dtype = dtype_u16, + .rank = 3, + .dimensions = dims, + .codec = codec, + .max_threads = 2, + .max_nlod = 1, + .epochs_per_batch = 1, + .prepare_shards = 1, + }; + stream = stream_create(&config, test_zarr_sink_as_shard_sink(&z)); + CHECK(Done, stream); + frame = calloc(32768, sizeof(*frame)); + CHECK(Done, frame); + struct writer* w = stream_writer(stream); + for (int t = 0; t < 2; ++t) + CHECK(Done, + writer_append_wait(w, (struct slice){ frame, frame + 32768 }).error == + 0); + CHECK(Done, writer_flush(w).error == 0); + CHECK(Done, writer_close(w).error == 0); + CHECK(Done, atomic_load(&counts.opens) == 0); + CHECK(Done, atomic_load(&counts.resizes) == 4); + error = 0; +Done: + stream_destroy(stream); + test_zarr_sink_close(&z); + free(frame); + test_tmpdir_remove(root); + return error; +} + +int +main(void) +{ +#ifdef TEST_PREPARATION_GPU + CUcontext context = NULL; + CUdevice device; + CU(Fail, cuInit(0)); + CU(Fail, cuDeviceGet(&device, 0)); + CU(Fail, cu_ctx_create(&context, 0, device)); +#endif + int error = test_prepared_files_need_no_resize(CODEC_NONE); + error |= test_prepared_files_need_no_resize(CODEC_BLOSC_LZ4); + for (int destroy_only = 0; destroy_only < 2; ++destroy_only) { + error |= test_stream(0, 0, destroy_only); + error |= test_stream(1, 0, destroy_only); + error |= test_stream(4, 0, destroy_only); + error |= test_stream(5, 0, destroy_only); + error |= test_stream(4, 1, destroy_only); + } +#ifdef TEST_PREPARATION_GPU + cuCtxDestroy(context); +#endif + return error; +#ifdef TEST_PREPARATION_GPU +Fail: + if (context) + cuCtxDestroy(context); + return 1; +#endif +} diff --git a/tests/test_zarr_readback.c b/tests/test_zarr_readback.c index 61b67c15..d4ba7b6e 100644 --- a/tests/test_zarr_readback.c +++ b/tests/test_zarr_readback.c @@ -46,6 +46,7 @@ write_zarr(const char* store_path, struct codec_config codec, int buffered) .dtype = dtype_u16, .rank = 3, .dimensions = dims, + .prepare_shards = 1, .codec = codec, }; diff --git a/tests/test_zarr_readback_gpu.c b/tests/test_zarr_readback_gpu.c index 3a4ee0e1..6c7ff3ef 100644 --- a/tests/test_zarr_readback_gpu.c +++ b/tests/test_zarr_readback_gpu.c @@ -45,6 +45,7 @@ write_zarr(const char* store_path, struct codec_config codec, int buffered) .dtype = dtype_u16, .rank = 3, .dimensions = dims, + .prepare_shards = 1, .codec = codec, .epochs_per_batch = 1, }; From 69f72b0d5ee49c4c8d856b803e978f71237a3df6 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Sun, 20 Sep 2026 19:15:19 -0700 Subject: [PATCH 2/2] Default to shard preparation and centralize resource ownership --- README.md | 5 +- docs/shard-preparation.md | 14 +- src/cpu/stream.body.c | 8 +- src/cpu/stream.c | 13 +- src/gpu/stream.c | 21 +-- src/gpu/stream.init.c | 16 +- src/multiarray/stream.c | 32 ++-- src/multiarray/stream.gpu.c | 17 +- src/platform/platform_io.darwin.c | 5 +- src/platform/platform_io.h | 2 +- src/platform/platform_io.posix.c | 5 +- src/platform/platform_io.win32.c | 7 +- src/types.stream.h | 3 +- src/writer.c | 10 +- src/writer.h | 4 + src/zarr/CMakeLists.txt | 1 + src/zarr/io_request.h | 2 +- src/zarr/shard_delivery.c | 36 ++-- src/zarr/shard_delivery.h | 14 +- src/zarr/shard_pool_fs.c | 194 +++----------------- src/zarr/shard_preparation.c | 284 ++++++++++++++++++++++++++++++ src/zarr/shard_preparation.h | 47 +++++ src/zarr/shard_write_plan.c | 11 +- tests/CMakeLists.txt | 4 +- tests/test_shard_preparation.c | 54 ++++++ tests/test_stream_preparation.c | 179 +++++++++++++++++-- tests/test_zarr_readback.c | 1 - tests/test_zarr_readback_gpu.c | 1 - 28 files changed, 705 insertions(+), 285 deletions(-) create mode 100644 src/zarr/shard_preparation.c create mode 100644 src/zarr/shard_preparation.h diff --git a/README.md b/README.md index 91c8fb48..7637f750 100644 --- a/README.md +++ b/README.md @@ -304,8 +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 can also [prepare one shard generation ahead](docs/shard-preparation.md) -with `config.prepare_shards = 1`, including cleanup of unused files at shutdown. +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]. diff --git a/docs/shard-preparation.md b/docs/shard-preparation.md index 68fa1d81..7010e0a7 100644 --- a/docs/shard-preparation.md +++ b/docs/shard-preparation.md @@ -1,8 +1,10 @@ # Preparing shard files ahead of use -Set `config.prepare_shards = 1` before creating a CPU or GPU stream to -prepare shard files ahead of use. This is opt-in and currently supported by -the filesystem store, including Zarr arrays and NGFF multiscales. +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, @@ -33,7 +35,5 @@ 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 the matching `prepare_shards`, `stop_preparing`, and -`cancel_prepared` hooks. Enabling the option on an unsupported sink, including -the current S3 backend, fails stream creation. With the option left at zero, -existing sinks retain their ordinary open-on-demand behavior. +Custom sinks need all three `prepare_shards`, `stop_preparing`, and +`cancel_prepared` hooks. Missing any hook disables preparation for that sink. diff --git a/src/cpu/stream.body.c b/src/cpu/stream.body.c index 67ae4a24..8849844d 100644 --- a/src/cpu/stream.body.c +++ b/src/cpu/stream.body.c @@ -237,7 +237,7 @@ cpu_stream_flush_batch(struct cpu_stream_view* v) struct writer_result cpu_stream_flush_body(struct cpu_stream_view* v) { - if (v->config->prepare_shards) + 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. @@ -340,7 +340,8 @@ cpu_stream_flush_body(struct cpu_stream_view* v) // they read. if (shard_sink_drain(v->sink)) failed = 1; - if (v->config->prepare_shards && shard_sink_cancel_prepared(v->sink)) + if (!v->config->disable_shard_preparation && + shard_sink_cancel_prepared(v->sink)) failed = 1; return failed ? writer_error() : writer_ok(); @@ -349,7 +350,8 @@ cpu_stream_flush_body(struct cpu_stream_view* v) struct writer_result cpu_stream_close_body(struct cpu_stream_view* v) { - int failed = v->config->prepare_shards && shard_sink_cancel_prepared(v->sink); + 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. diff --git a/src/cpu/stream.c b/src/cpu/stream.c index 37b787da..4230370c 100644 --- a/src/cpu/stream.c +++ b/src/cpu/stream.c @@ -291,14 +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); - if (config->prepare_shards) - CHECK(Fail, - shard_sink_prepare_first( - sink, - s->shard, - s->levels.nlod, - s->shard_alignment, - config->codec.id != CODEC_NONE) == 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; @@ -311,7 +306,7 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, return s; Fail: - if (config->prepare_shards && shard_sink_cancel_prepared(sink)) + 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); diff --git a/src/gpu/stream.c b/src/gpu/stream.c index a7a5873a..05b57eaf 100644 --- a/src/gpu/stream.c +++ b/src/gpu/stream.c @@ -397,7 +397,7 @@ 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.prepare_shards) + 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 @@ -431,6 +431,11 @@ 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)) { @@ -438,14 +443,9 @@ stream_flush_body(struct stream_engine* e, struct stream_context* ctx) } schedule_quiesce_output(e, ctx->sink); - if (ctx->config.prepare_shards) { - for (int fc = 0; fc < 2; ++fc) - if (gpu_delivery_pending(&e->delivery, fc) && - gpu_delivery_join(&e->delivery, fc).error) - r = writer_error(); - if (shard_sink_cancel_prepared(ctx->sink)) - r = writer_error(); - } + if (!ctx->config.disable_shard_preparation && + shard_sink_cancel_prepared(ctx->sink)) + r = writer_error(); collect_ingest_timing(e); return r; @@ -459,7 +459,8 @@ stream_close_body(struct compress_agg_array* ar, struct stream_context* ctx) struct writer_result r = writer_ok(); - if (ctx->config.prepare_shards && shard_sink_cancel_prepared(ctx->sink)) + 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. diff --git a/src/gpu/stream.init.c b/src/gpu/stream.init.c index 0524acac..86a557d3 100644 --- a/src/gpu/stream.init.c +++ b/src/gpu/stream.init.c @@ -350,7 +350,8 @@ 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.prepare_shards && shard_sink_cancel_prepared(s->ctx.sink)) + 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); } @@ -453,13 +454,12 @@ tile_stream_gpu_create(const struct tile_stream_configuration* config, shard_sink_init_append(sink, &out->ctx.dims, out->ctx.levels.nlod) == 0); - if (config->prepare_shards) - CHECK(FailPhase2, - shard_sink_prepare_first(sink, - out->engine.compress_agg.ar.shard, - out->ctx.levels.nlod, - out->ctx.shard_alignment, - config->codec.id != CODEC_NONE) == 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; diff --git a/src/multiarray/stream.c b/src/multiarray/stream.c index cb25285d..eff7294c 100644 --- a/src/multiarray/stream.c +++ b/src/multiarray/stream.c @@ -478,13 +478,12 @@ multiarray_tile_stream_cpu_create( CHECK(Fail, shard_sink_init_append( desc->sink, &desc->cl.dims, desc->levels.nlod) == 0); - if (desc->config.prepare_shards) - CHECK(Fail, - shard_sink_prepare_first(desc->sink, - desc->shard, - desc->levels.nlod, - desc->shard_alignment, - desc->config.codec.id != CODEC_NONE) == 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; @@ -495,7 +494,7 @@ multiarray_tile_stream_cpu_create( Fail: if (ms && ms->arrays) for (int i = 0; i < ms->n_arrays; ++i) - if (ms->arrays[i].config.prepare_shards && + 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. @@ -745,6 +744,8 @@ 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 @@ -752,9 +753,11 @@ flush_impl(struct multiarray_writer* self) if (desc->flushed) continue; if (desc->cursor_elements == 0 && desc->batch_accumulated == 0) { - if (desc->config.prepare_shards && shard_sink_cancel_prepared(desc->sink)) - goto Error; + if (!desc->config.disable_shard_preparation && + shard_sink_cancel_prepared(desc->sink)) + failed = 1; desc->flushed = 1; + desc->closed = 0; continue; } @@ -771,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, }; } diff --git a/src/multiarray/stream.gpu.c b/src/multiarray/stream.gpu.c index 36bdd1d4..fe2c812f 100644 --- a/src/multiarray/stream.gpu.c +++ b/src/multiarray/stream.gpu.c @@ -285,7 +285,7 @@ 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.prepare_shards && + if (!desc->ctx.config.disable_shard_preparation && shard_sink_cancel_prepared(desc->ctx.sink)) failed = 1; desc->flushed = 1; @@ -384,7 +384,7 @@ multiarray_tile_stream_gpu_rollback(struct multiarray_tile_stream_gpu* ms) gpu_streams_sync(&ms->engine.streams); if (ms->arrays) for (int a = 0; a < ms->n_arrays; ++a) - if (ms->arrays[a].ctx.config.prepare_shards && + 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); @@ -504,13 +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); - if (ctx->config.prepare_shards) - CHECK(Fail, - shard_sink_prepare_first(ctx->sink, - ms->arrays[a].st.agg.shard, - ctx->levels.nlod, - ctx->shard_alignment, - ctx->config.codec.id != CODEC_NONE) == 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; diff --git a/src/platform/platform_io.darwin.c b/src/platform/platform_io.darwin.c index b741bc6b..8b95ff67 100644 --- a/src/platform/platform_io.darwin.c +++ b/src/platform/platform_io.darwin.c @@ -30,10 +30,9 @@ platform_mkdir_new(const char* path) int platform_remove_empty_directory(const char* path) { - if (rmdir(path) == 0 || errno == ENOENT || errno == ENOTEMPTY || - errno == EEXIST) + if (rmdir(path) == 0 || errno == ENOENT) return 0; - return -1; + return errno == ENOTEMPTY || errno == EEXIST ? 1 : -1; } int diff --git a/src/platform/platform_io.h b/src/platform/platform_io.h index d4b0524f..d7f1ef9f 100644 --- a/src/platform/platform_io.h +++ b/src/platform/platform_io.h @@ -21,7 +21,7 @@ platform_mkdir(const char* path); int platform_mkdir_new(const char* path); -// Missing and nonempty directories are left alone; other failures return -1. +// Returns 0 if removed or missing, 1 if nonempty, -1 on other failures. int platform_remove_empty_directory(const char* path); diff --git a/src/platform/platform_io.posix.c b/src/platform/platform_io.posix.c index 78ead5be..8b54448d 100644 --- a/src/platform/platform_io.posix.c +++ b/src/platform/platform_io.posix.c @@ -32,10 +32,9 @@ platform_mkdir_new(const char* path) int platform_remove_empty_directory(const char* path) { - if (rmdir(path) == 0 || errno == ENOENT || errno == ENOTEMPTY || - errno == EEXIST) + if (rmdir(path) == 0 || errno == ENOENT) return 0; - return -1; + return errno == ENOTEMPTY || errno == EEXIST ? 1 : -1; } int diff --git a/src/platform/platform_io.win32.c b/src/platform/platform_io.win32.c index ad451110..a7e5dd37 100644 --- a/src/platform/platform_io.win32.c +++ b/src/platform/platform_io.win32.c @@ -32,10 +32,9 @@ platform_remove_empty_directory(const char* path) if (RemoveDirectoryA(path)) return 0; DWORD error = GetLastError(); - return error == ERROR_FILE_NOT_FOUND || error == ERROR_PATH_NOT_FOUND || - error == ERROR_DIR_NOT_EMPTY - ? 0 - : -1; + if (error == ERROR_FILE_NOT_FOUND || error == ERROR_PATH_NOT_FOUND) + return 0; + return error == ERROR_DIR_NOT_EMPTY ? 1 : -1; } int diff --git a/src/types.stream.h b/src/types.stream.h index 210a523f..7a3ff25e 100644 --- a/src/types.stream.h +++ b/src/types.stream.h @@ -154,7 +154,8 @@ struct tile_stream_configuration // reports more pending than this int max_threads; // 0 = OpenMP default int full_memcpy_timing; - int prepare_shards; // opt-in; requires preparation and cleanup from the sink + int disable_shard_preparation; // default: prepare ahead when the sink + // supports it }; struct tile_stream_status diff --git a/src/writer.c b/src/writer.c index 6541259f..f87e3930 100644 --- a/src/writer.c +++ b/src/writer.c @@ -88,10 +88,16 @@ writer_finished_at(const void* beg, const void* end) return r; } +int +shard_sink_supports_preparation(const struct shard_sink* s) +{ + return s && s->prepare_shards && s->stop_preparing && s->cancel_prepared; +} + void shard_sink_stop_preparing(struct shard_sink* s) { - if (s && s->stop_preparing) + if (shard_sink_supports_preparation(s)) s->stop_preparing(s); } @@ -99,7 +105,7 @@ int shard_sink_cancel_prepared(struct shard_sink* s) { shard_sink_stop_preparing(s); - return s && s->cancel_prepared ? s->cancel_prepared(s) : 0; + return shard_sink_supports_preparation(s) ? s->cancel_prepared(s) : 0; } uint64_t diff --git a/src/writer.h b/src/writer.h index 0aa993d6..ba39f9d5 100644 --- a/src/writer.h +++ b/src/writer.h @@ -144,6 +144,10 @@ struct shard_sink void shard_sink_stop_preparing(struct shard_sink* s); +// Preparation requires all three lifecycle hooks; partial support is ignored. +int +shard_sink_supports_preparation(const struct shard_sink* s); + int shard_sink_cancel_prepared(struct shard_sink* s); diff --git a/src/zarr/CMakeLists.txt b/src/zarr/CMakeLists.txt index ad8e75ef..144f4a29 100644 --- a/src/zarr/CMakeLists.txt +++ b/src/zarr/CMakeLists.txt @@ -26,6 +26,7 @@ add_src_lib(io_backend_fs SOURCES io_backend.fs.c io_backend.fs.h ) add_src_lib(shard_pool_fs SOURCES shard_pool_fs.c shard_pool_fs.h shard_pool.h + shard_preparation.c shard_preparation.h LINKS io_backend_fs io_scheduler filesystem_write host_output_pool platform platform_io writer strbuf chucky_log ) diff --git a/src/zarr/io_request.h b/src/zarr/io_request.h index 000e2dd2..7e38bc12 100644 --- a/src/zarr/io_request.h +++ b/src/zarr/io_request.h @@ -6,7 +6,7 @@ enum io_op { - IO_OP_NOOP = 0, // no file or payload; used by fault-injection tests + IO_OP_NOOP = 0, // backend-defined work without a file token IO_OP_OPEN, // barrier: create or truncate a file IO_OP_WRITE, // payload to a file at an offset IO_OP_TRUNCATE, // barrier: set the file's size diff --git a/src/zarr/shard_delivery.c b/src/zarr/shard_delivery.c index 59e1c060..0bcbf8b2 100644 --- a/src/zarr/shard_delivery.c +++ b/src/zarr/shard_delivery.c @@ -243,26 +243,36 @@ record_finalized(struct shard_state* ss, struct shard_sink* sink) } } +uint64_t +shard_state_file_capacity(const struct shard_state* ss, + enum host_batch_storage storage, + size_t alignment) +{ + uint64_t capacity = ss->shard_file_capacity; + if (capacity && storage == HOST_BATCH_PAGE_PADDED && alignment > 1) { + const uint64_t gaps = ss->chunks_per_shard_append - 1; + if (gaps > (UINT64_MAX - capacity) / (alignment - 1)) + return 0; + capacity += gaps * (alignment - 1); + } + return capacity; +} + int shard_sink_prepare_first(struct shard_sink* sink, const struct shard_state* shards, int nlod, size_t alignment, - int padded) + const struct tile_stream_configuration* config) { - if (!sink || !sink->prepare_shards || !sink->stop_preparing || - !sink->cancel_prepared) { - log_error("sink does not support shard preparation with cleanup"); - return 1; - } + if (config->disable_shard_preparation || + !shard_sink_supports_preparation(sink)) + return 0; + const enum host_batch_storage storage = + host_batch_storage_select(config->codec.id == CODEC_NONE, alignment); for (int lv = 0; lv < nlod; ++lv) { - uint64_t capacity = shards[lv].shard_file_capacity; - if (capacity && padded && alignment > 1) { - const uint64_t gaps = shards[lv].chunks_per_shard_append - 1; - capacity = gaps <= (UINT64_MAX - capacity) / (alignment - 1) - ? capacity + gaps * (alignment - 1) - : 0; - } + const uint64_t capacity = + shard_state_file_capacity(&shards[lv], storage, alignment); if (sink->prepare_shards(sink, (uint8_t)lv, capacity)) { if (shard_sink_cancel_prepared(sink)) log_error("sink preparation rollback failed"); diff --git a/src/zarr/shard_delivery.h b/src/zarr/shard_delivery.h index ad2021e5..f6730180 100644 --- a/src/zarr/shard_delivery.h +++ b/src/zarr/shard_delivery.h @@ -62,8 +62,8 @@ struct shard_state size_t footer_capacity; // bytes per shard // Compact upper bound: every chunk at its worst-case compressed size, plus - // the footer. The GPU write plan adds the worst retained padding bound - // when that policy is active. Zero means unknown and turns pre-sizing off. + // the footer. shard_state_file_capacity adds any retained padding bound. + // Zero means unknown and turns pre-sizing off. uint64_t shard_file_capacity; }; @@ -78,12 +78,20 @@ shard_state_destroy(struct shard_state* ss); size_t shard_state_heap_bytes(const struct level_layout_info* li); +// Shared capacity bound for preparation and ordinary presizing. Zero disables +// presizing when the bound cannot be represented. +uint64_t +shard_state_file_capacity(const struct shard_state* ss, + enum host_batch_storage storage, + size_t alignment); + +// Default-on; explicit opt-out and unsupported sinks are successful no-ops. int shard_sink_prepare_first(struct shard_sink* sink, const struct shard_state* shards, int nlod, size_t alignment, - int padded); + const struct tile_stream_configuration* config); // A new stream advertises no appended data, independently of its configured // capacity. Called once during creation, before any array accepts input. diff --git a/src/zarr/shard_pool_fs.c b/src/zarr/shard_pool_fs.c index ae0b66b8..f8c14c3b 100644 --- a/src/zarr/shard_pool_fs.c +++ b/src/zarr/shard_pool_fs.c @@ -7,6 +7,7 @@ #include "zarr/filesystem_write.h" #include "zarr/io_backend.fs.h" #include "zarr/io_scheduler.h" +#include "zarr/shard_preparation.h" #include #include @@ -16,24 +17,13 @@ struct fs_slot; -struct prepared_file -{ - char* path; - unsigned char* created_dirs; - platform_fd fd; - uint64_t capacity; - struct io_event ready; - int owned; -}; - struct shard_pool_fs { struct shard_pool base; struct io_backend_fs* backend; struct io_scheduler* queue; - struct io_scheduler* prepare_queue; - struct io_backend prepare_backend; - struct prepared_file* prepared; + struct shard_preparation* preparation; + struct shard_pool_fs_wrapper wrapper; struct fs_slot* slots; uint64_t nslots; int unbuffered; @@ -221,52 +211,6 @@ fs_slot_finalize(struct shard_writer* self) return 0; } -static void -prepared_file_clear(struct prepared_file* f) -{ - free(f->created_dirs); - free(f->path); - *f = (struct prepared_file){ .fd = PLATFORM_FD_INVALID }; -} - -static void -prepare_file(void* ctx, const struct io_request* req) -{ - struct shard_pool_fs* p = ctx; - struct prepared_file* f = (struct prepared_file*)req->payload; - char* directory = strdup(f->path); - CHECK(Fail, directory); - const int flags = - PLATFORM_OPEN_EXCLUSIVE | (p->unbuffered ? PLATFORM_OPEN_UNBUFFERED : 0); - f->fd = platform_open_write(f->path, flags); - if (f->fd == PLATFORM_FD_INVALID) { - const size_t len = strlen(directory); - for (size_t i = 1; i < len; ++i) { - if (directory[i] != '/' && directory[i] != '\\') - continue; - if (i == 2 && directory[1] == ':') - continue; - const char saved = directory[i]; - directory[i] = 0; - const int made = platform_mkdir_new(directory); - directory[i] = saved; - CHECK(Fail, made >= 0); - f->created_dirs[i] = made != 0; - } - f->fd = platform_open_write(f->path, flags); - } - CHECK(Fail, f->fd != PLATFORM_FD_INVALID); - f->owned = 1; - if (f->capacity) - CHECK(Fail, platform_ftruncate(f->fd, f->capacity) == 0); - free(directory); - return; -Fail: - log_error("shard preparation failed: %s", f->path); - atomic_store(&p->io_error, 1); - free(directory); -} - static int pool_fs_prepare(struct shard_pool* self, uint64_t slot, @@ -276,36 +220,21 @@ pool_fs_prepare(struct shard_pool* self, struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); struct strbuf path = { 0 }; CHECK(Fail, slot < p->nslots && key && !atomic_load(&p->io_error)); - if (!p->prepare_queue) { - CHECK(Fail, p->nslots <= SIZE_MAX / sizeof(*p->prepared)); - p->prepared = calloc((size_t)p->nslots, sizeof(*p->prepared)); - CHECK(Fail, p->prepared); - for (uint64_t i = 0; i < p->nslots; ++i) - p->prepared[i].fd = PLATFORM_FD_INVALID; - p->prepare_queue = io_scheduler_create( - p->prepare_backend, - (struct io_scheduler_limits){ .workers = 1, .max_requests = p->nslots }); - CHECK(Fail, p->prepare_queue); + if (!p->preparation) { + p->preparation = + shard_preparation_create(p->nslots, + p->unbuffered ? PLATFORM_OPEN_UNBUFFERED : 0, + &p->io_error, + p->wrapper.ctx, + p->wrapper.wrap_prepare); + CHECK(Fail, p->preparation); } - struct prepared_file* f = &p->prepared[slot]; - CHECK(Fail, !f->path); CHECK(Fail, strbuf_appendf(&path, "%s/%s", strbuf_cstr(&p->root), key) == 0); - f->path = strdup(strbuf_cstr(&path)); - f->created_dirs = calloc(strbuf_len(&path) + 1, 1); - if (!f->path || !f->created_dirs) { - prepared_file_clear(f); - goto Fail; - } - f->capacity = p->slots[slot].presize ? capacity : 0; - if (io_scheduler_post(p->prepare_queue, - (struct io_request){ .op = IO_OP_NOOP, - .payload = f, - .path = f->path, - .logical_size = f->capacity })) { - prepared_file_clear(f); - goto Fail; - } - f->ready = io_scheduler_record(p->prepare_queue); + CHECK(Fail, + shard_preparation_post(p->preparation, + slot, + strbuf_cstr(&path), + p->slots[slot].presize ? capacity : 0) == 0); strbuf_free(&path); return 0; Fail: @@ -320,9 +249,8 @@ pool_fs_wait_prepared(struct shard_pool* self, uint64_t slot) struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); if (slot >= p->nslots) return 1; - if (p->prepare_queue && p->prepared[slot].path) - io_event_wait(p->prepare_queue, p->prepared[slot].ready); - return atomic_load(&p->io_error); + return shard_preparation_wait(p->preparation, slot) || + atomic_load(&p->io_error); } static int @@ -331,60 +259,7 @@ pool_fs_cancel_prepared(struct shard_pool* self, uint64_t first, uint64_t count) struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); if (first > p->nslots || count > p->nslots - first) return 1; - if (!p->prepared) - return 0; - int failed = 0; - size_t longest = 0; - for (uint64_t i = first; i < first + count; ++i) { - struct prepared_file* f = &p->prepared[i]; - if (!f->path) - continue; - if (p->prepare_queue) - io_event_wait(p->prepare_queue, f->ready); - if (f->fd != PLATFORM_FD_INVALID) { - platform_close(f->fd); - f->fd = PLATFORM_FD_INVALID; - } - if (f->owned) { - if (platform_remove_file(f->path) == 0 || - platform_path_exists(f->path) == 0) - f->owned = 0; - else { - log_error("unused shard cleanup failed: %s", f->path); - failed = 1; - } - } - const size_t len = strlen(f->path); - if (len > longest) - longest = len; - } - if (failed) { - atomic_store(&p->io_error, 1); - return 1; - } - for (size_t pos = longest; pos > 0; --pos) { - for (uint64_t i = first; i < first + count; ++i) { - struct prepared_file* f = &p->prepared[i]; - if (!f->path || strlen(f->path) <= pos || !f->created_dirs[pos]) - continue; - char saved = f->path[pos]; - f->path[pos] = 0; - if (platform_remove_empty_directory(f->path)) { - log_error("unused shard directory cleanup failed: %s", f->path); - failed = 1; - } else { - f->created_dirs[pos] = 0; - } - f->path[pos] = saved; - } - } - if (!failed) { - for (uint64_t i = first; i < first + count; ++i) - prepared_file_clear(&p->prepared[i]); - } else { - atomic_store(&p->io_error, 1); - } - return failed; + return shard_preparation_cancel(p->preparation, first, count); } static struct shard_writer* @@ -403,16 +278,14 @@ pool_fs_open(struct shard_pool* self, uint64_t slot, const char* key) if (strbuf_appendf(&path, "%s/%s", strbuf_cstr(&p->root), key)) goto Fail; - struct prepared_file* f = p->prepared ? &p->prepared[slot] : NULL; - if (f && f->path) { - CHECK(Fail, strcmp(f->path, strbuf_cstr(&path)) == 0); - CHECK(Fail, pool_fs_wait_prepared(self, slot) == 0); - const struct io_file_token token = - io_backend_fs_adopt_file(p->backend, f->fd); - CHECK(Fail, token.generation != 0); - w->token = token; - w->prepared_capacity = f->capacity; - prepared_file_clear(f); + const int adopted = shard_preparation_adopt(p->preparation, + slot, + strbuf_cstr(&path), + p->backend, + &w->token, + &w->prepared_capacity); + CHECK(Fail, adopted >= 0); + if (adopted) { strbuf_free(&path); return &w->base; } @@ -537,14 +410,7 @@ pool_fs_destroy(struct shard_pool* self) { struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); - if (pool_fs_cancel_prepared(self, 0, p->nslots)) - log_error("shard pool cleanup failed during destroy"); - io_scheduler_destroy(p->prepare_queue); - if (p->prepared) { - for (uint64_t i = 0; i < p->nslots; ++i) - prepared_file_clear(&p->prepared[i]); - free(p->prepared); - } + shard_preparation_destroy(p->preparation); // Finalize any open slots for (uint64_t i = 0; i < p->nslots; ++i) { @@ -611,9 +477,7 @@ shard_pool_fs_create_wrapped(const char* root, p->base.prepare = pool_fs_prepare; p->base.wait_prepared = pool_fs_wait_prepared; p->base.cancel_prepared = pool_fs_cancel_prepared; - p->prepare_backend = (struct io_backend){ .ctx = p, .execute = prepare_file }; - if (wrapper.wrap_prepare) - p->prepare_backend = wrapper.wrap_prepare(wrapper.ctx, p->prepare_backend); + p->wrapper = wrapper; p->base.record_fence = pool_fs_record_fence; p->base.wait_fence = pool_fs_wait_fence; p->base.queue_metadata = pool_fs_queue_metadata; diff --git a/src/zarr/shard_preparation.c b/src/zarr/shard_preparation.c new file mode 100644 index 00000000..8d2e92b7 --- /dev/null +++ b/src/zarr/shard_preparation.c @@ -0,0 +1,284 @@ +#include "zarr/shard_preparation.h" + +#include "platform/platform.h" +#include "platform/platform_io.h" +#include "util/prelude.h" +#include "zarr/io_backend.fs.h" +#include "zarr/io_scheduler.h" + +#include +#include + +struct prepared_file +{ + char* path; + platform_fd fd; + uint64_t capacity; + struct io_event ready; + int owned; +}; + +struct prepared_directory +{ + struct prepared_directory* next; + size_t len; + char path[]; +}; + +struct shard_preparation +{ + struct io_scheduler* queue; + struct prepared_file* files; + uint64_t nslots; + int open_flags; + _Atomic int* io_error; + + // Shared by slots; retain nonempty directories across cancellation ranges. + // New children precede their parents, so cleanup visits deepest first. + // The lock also prevents cleanup from removing a directory between mkdir + // and open on the worker. Sizing happens outside this lock. + struct platform_mutex* directories_mutex; + struct prepared_directory* directories; +}; + +static void +prepared_file_clear(struct prepared_file* f) +{ + free(f->path); + *f = (struct prepared_file){ .fd = PLATFORM_FD_INVALID }; +} + +// Called with directories_mutex held. Allocate before mkdir so every new +// directory can be tracked even when a subsequent allocation fails. +static int +prepare_directory(struct shard_preparation* p, const char* path, size_t len) +{ + struct prepared_directory* d = malloc(sizeof(*d) + len + 1); + if (!d) + return 1; + memcpy(d->path, path, len); + d->path[len] = 0; + d->len = len; + const int made = platform_mkdir_new(d->path); + if (made == 1) { + d->next = p->directories; + p->directories = d; + } else { + free(d); + } + return made < 0; +} + +static void +prepare_file(void* ctx, const struct io_request* req) +{ + struct shard_preparation* p = ctx; + struct prepared_file* f = (struct prepared_file*)req->payload; + int failed = 0; + platform_mutex_lock(p->directories_mutex); + f->fd = platform_open_write(f->path, p->open_flags); + if (f->fd == PLATFORM_FD_INVALID) { + const size_t len = strlen(f->path); + for (size_t i = 1; i < len; ++i) { + if (f->path[i] != '/' && f->path[i] != '\\') + continue; + if (i == 2 && f->path[1] == ':') + continue; + if (prepare_directory(p, f->path, i)) { + failed = 1; + break; + } + } + if (!failed) + f->fd = platform_open_write(f->path, p->open_flags); + } + platform_mutex_unlock(p->directories_mutex); + f->owned = f->fd != PLATFORM_FD_INVALID; + if (failed || f->fd == PLATFORM_FD_INVALID || + (f->capacity && platform_ftruncate(f->fd, f->capacity))) { + log_error("shard preparation failed: %s", f->path); + atomic_store(p->io_error, 1); + } +} + +struct shard_preparation* +shard_preparation_create(uint64_t nslots, + int open_flags, + _Atomic int* io_error, + void* wrap_ctx, + struct io_backend (*wrap)(void*, struct io_backend)) +{ + struct shard_preparation* p = calloc(1, sizeof(*p)); + CHECK(Fail, p); + p->io_error = io_error; + p->open_flags = open_flags | PLATFORM_OPEN_EXCLUSIVE; + CHECK(Fail, nslots > 0 && nslots <= SIZE_MAX / sizeof(*p->files)); + p->files = calloc((size_t)nslots, sizeof(*p->files)); + CHECK(Fail, p->files); + p->nslots = nslots; + for (uint64_t i = 0; i < nslots; ++i) + p->files[i].fd = PLATFORM_FD_INVALID; + p->directories_mutex = platform_mutex_new(); + CHECK(Fail, p->directories_mutex); + struct io_backend backend = { .ctx = p, .execute = prepare_file }; + if (wrap) + backend = wrap(wrap_ctx, backend); + p->queue = io_scheduler_create( + backend, + (struct io_scheduler_limits){ .workers = 1, .max_requests = nslots }); + CHECK(Fail, p->queue); + return p; +Fail: + shard_preparation_destroy(p); + return NULL; +} + +int +shard_preparation_post(struct shard_preparation* p, + uint64_t slot, + const char* path, + uint64_t capacity) +{ + CHECK(Fail, slot < p->nslots && path && !atomic_load(p->io_error)); + struct prepared_file* f = &p->files[slot]; + CHECK(Fail, !f->path); + f->path = strdup(path); + CHECK(Fail, f->path); + f->capacity = capacity; + if (io_scheduler_post(p->queue, + (struct io_request){ + .op = IO_OP_NOOP, .payload = f, .path = f->path })) { + prepared_file_clear(f); + goto Fail; + } + f->ready = io_scheduler_record(p->queue); + return 0; +Fail: + atomic_store(p->io_error, 1); + return 1; +} + +int +shard_preparation_wait(struct shard_preparation* p, uint64_t slot) +{ + if (!p) + return 0; + if (slot >= p->nslots) + return 1; + if (p->files[slot].path) + io_event_wait(p->queue, p->files[slot].ready); + return atomic_load(p->io_error); +} + +int +shard_preparation_adopt(struct shard_preparation* p, + uint64_t slot, + const char* path, + struct io_backend_fs* backend, + struct io_file_token* token, + uint64_t* capacity) +{ + if (!p) + return 0; + if (slot >= p->nslots) + return -1; + if (!p->files[slot].path) + return 0; + struct prepared_file* f = &p->files[slot]; + if (strcmp(f->path, path) || shard_preparation_wait(p, slot)) + return -1; + *token = io_backend_fs_adopt_file(backend, f->fd); + if (!token->generation) + return -1; + *capacity = f->capacity; + + // Adopted files make their ancestors permanent, including directories + // created by another slot. Forget them to keep tracking bounded. + platform_mutex_lock(p->directories_mutex); + struct prepared_directory** link = &p->directories; + const size_t len = strlen(path); + while (*link) { + struct prepared_directory* d = *link; + if (len > d->len && !memcmp(path, d->path, d->len) && + (path[d->len] == '/' || path[d->len] == '\\')) { + *link = d->next; + free(d); + } else { + link = &d->next; + } + } + platform_mutex_unlock(p->directories_mutex); + prepared_file_clear(f); + return 1; +} + +int +shard_preparation_cancel(struct shard_preparation* p, + uint64_t first, + uint64_t count) +{ + if (!p) + return 0; + if (first > p->nslots || count > p->nslots - first) + return 1; + int failed = 0; + for (uint64_t i = first; i < first + count; ++i) { + struct prepared_file* f = &p->files[i]; + if (!f->path) + continue; + io_event_wait(p->queue, f->ready); + if (f->fd != PLATFORM_FD_INVALID) { + platform_close(f->fd); + f->fd = PLATFORM_FD_INVALID; + } + if (f->owned && platform_remove_file(f->path) && + platform_path_exists(f->path) != 0) { + log_error("unused shard cleanup failed: %s", f->path); + failed = 1; + continue; + } + prepared_file_clear(f); + } + if (p->directories_mutex) { + platform_mutex_lock(p->directories_mutex); + struct prepared_directory** link = &p->directories; + while (*link) { + struct prepared_directory* d = *link; + const int rc = platform_remove_empty_directory(d->path); + if (rc == 0) { + *link = d->next; + free(d); + } else { + if (rc < 0) { + log_error("unused shard directory cleanup failed: %s", d->path); + failed = 1; + } + link = &d->next; + } + } + platform_mutex_unlock(p->directories_mutex); + } + if (failed) + atomic_store(p->io_error, 1); + return failed; +} + +void +shard_preparation_destroy(struct shard_preparation* p) +{ + if (!p) + return; + if (shard_preparation_cancel(p, 0, p->nslots)) + log_error("shard preparation cleanup failed during destroy"); + io_scheduler_destroy(p->queue); + for (uint64_t i = 0; i < p->nslots; ++i) + prepared_file_clear(&p->files[i]); + while (p->directories) { + struct prepared_directory* next = p->directories->next; + free(p->directories); + p->directories = next; + } + platform_mutex_free(p->directories_mutex); + free(p->files); + free(p); +} diff --git a/src/zarr/shard_preparation.h b/src/zarr/shard_preparation.h new file mode 100644 index 00000000..c1790235 --- /dev/null +++ b/src/zarr/shard_preparation.h @@ -0,0 +1,47 @@ +// Private filesystem preparation worker and ownership of unused resources. +// Callers serialize posting and adoption; cancel only after delivery stops. +#pragma once + +#include "zarr/io_backend.h" + +#include +#include + +struct io_backend_fs; +struct shard_preparation; + +struct shard_preparation* +shard_preparation_create(uint64_t nslots, + int open_flags, + _Atomic int* io_error, + void* wrap_ctx, + struct io_backend (*wrap)(void*, struct io_backend)); + +int +shard_preparation_post(struct shard_preparation* p, + uint64_t slot, + const char* path, + uint64_t capacity); + +// Waiting and cancellation accept NULL when preparation has never started. +int +shard_preparation_wait(struct shard_preparation* p, uint64_t slot); + +// 1 transfers a prepared file to backend, 0 means no preparation, -1 fails. +// Failure leaves ownership here so cancellation can reclaim the file. +int +shard_preparation_adopt(struct shard_preparation* p, + uint64_t slot, + const char* path, + struct io_backend_fs* backend, + struct io_file_token* token, + uint64_t* capacity); + +int +shard_preparation_cancel(struct shard_preparation* p, + uint64_t first, + uint64_t count); + +// Cancels unused files and joins the worker. NULL is harmless. +void +shard_preparation_destroy(struct shard_preparation* p); diff --git a/src/zarr/shard_write_plan.c b/src/zarr/shard_write_plan.c index 95488250..dc4677f6 100644 --- a/src/zarr/shard_write_plan.c +++ b/src/zarr/shard_write_plan.c @@ -621,15 +621,8 @@ ensure_writer(struct shard_sink* sink, if (!sh->writer) return 1; if (sh->writer->presize) { - uint64_t capacity = ss->shard_file_capacity; - if (capacity > 0 && plan->host->storage == HOST_BATCH_PAGE_PADDED) { - const uint64_t gap_count = ss->chunks_per_shard_append - 1; - const uint64_t max_gap = plan->shard_alignment - 1; - if (gap_count > 0 && max_gap > (UINT64_MAX - capacity) / gap_count) - capacity = 0; - else - capacity += gap_count * max_gap; - } + const uint64_t capacity = + shard_state_file_capacity(ss, plan->host->storage, plan->shard_alignment); if (sh->writer->presize(sh->writer, capacity)) return 1; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5700c348..40793a1b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -332,9 +332,9 @@ add_gpu_test_exe(test_gpu_placement test_gpu_placement.c add_test_exe(test_shard_preparation test_shard_preparation.c shard_pool_fs platform_io test_platform chucky_log) add_test_exe(test_stream_preparation test_stream_preparation.c - stream_cpu test_zarr_helpers test_shard_verify platform_io test_platform) + stream_cpu multiarray_cpu test_zarr_helpers test_shard_verify platform_io test_platform) add_gpu_test_exe(test_stream_preparation_gpu test_stream_preparation.c - LINKS stream test_zarr_helpers test_shard_verify platform_io test_platform CUDA::cuda_driver) + LINKS stream multiarray_gpu test_zarr_helpers test_shard_verify platform_io test_platform CUDA::cuda_driver) if(CHUCKY_ENABLE_GPU) target_compile_definitions(test_stream_preparation_gpu PRIVATE TEST_PREPARATION_GPU=1) endif() diff --git a/tests/test_shard_preparation.c b/tests/test_shard_preparation.c index 7802b61d..649c5748 100644 --- a/tests/test_shard_preparation.c +++ b/tests/test_shard_preparation.c @@ -217,6 +217,57 @@ test_cleanup_failure_can_be_retried(void) return error; } +static int +test_shared_directory_cleanup(int first) +{ + char root[256] = { 0 }; + struct shard_pool* pool = NULL; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + pool = shard_pool_fs_create(root, 2, 0); + CHECK(Done, pool); + CHECK(Done, pool->prepare(pool, 0, "c/0/0", 4096) == 0); + CHECK(Done, pool->prepare(pool, 1, "c/0/1", 4096) == 0); + CHECK(Done, pool->wait_prepared(pool, 1) == 0); + CHECK(Done, pool->cancel_prepared(pool, (uint64_t)first, 1) == 0); + CHECK(Done, exists(root, "c/0") == 1); + CHECK(Done, pool->cancel_prepared(pool, (uint64_t)(first ^ 1), 1) == 0); + CHECK(Done, exists(root, "c") == 0); + CHECK(Done, pool->cancel_prepared(pool, 0, 2) == 0); + error = 0; +Done: + shard_pool_destroy(pool); + test_tmpdir_remove(root); + return error; +} + +static int +test_adoption_preserves_shared_directories(void) +{ + char root[256] = { 0 }; + struct shard_pool* pool = NULL; + int error = 1; + CHECK(Done, test_tmpdir_create(root, sizeof(root)) == 0); + pool = shard_pool_fs_create(root, 2, 0); + CHECK(Done, pool); + CHECK(Done, pool->prepare(pool, 0, "c/0/0", 4096) == 0); + CHECK(Done, pool->prepare(pool, 1, "c/0/1", 4096) == 0); + CHECK(Done, pool->wait_prepared(pool, 1) == 0); + // Slot 0 created the shared directories; slot 1 makes them permanent. + struct shard_writer* w = pool->open(pool, 1, "c/0/1"); + CHECK(Done, w && w->finalize(w) == 0); + CHECK(Done, pool->cancel_prepared(pool, 0, 2) == 0); + shard_pool_destroy(pool); + pool = NULL; + CHECK(Done, exists(root, "c/0/0") == 0); + CHECK(Done, exists(root, "c/0/1") == 1); + error = 0; +Done: + shard_pool_destroy(pool); + test_tmpdir_remove(root); + return error; +} + int main(void) { @@ -224,5 +275,8 @@ main(void) error |= test_preparation_does_not_hold_writes(); error |= test_existing_file_is_preserved(); error |= test_cleanup_failure_can_be_retried(); + error |= test_shared_directory_cleanup(0); + error |= test_shared_directory_cleanup(1); + error |= test_adoption_preserves_shared_directories(); return error; } diff --git a/tests/test_stream_preparation.c b/tests/test_stream_preparation.c index 8c586728..612794ac 100644 --- a/tests/test_stream_preparation.c +++ b/tests/test_stream_preparation.c @@ -10,19 +10,29 @@ #ifdef TEST_PREPARATION_GPU #include "gpu/prelude.cuda.h" +#include "multiarray.gpu.h" #include "stream.gpu.h" #define stream_type tile_stream_gpu #define stream_create tile_stream_gpu_create #define stream_writer tile_stream_gpu_writer #define stream_destroy tile_stream_gpu_destroy #define stream_reset tile_stream_gpu_reset_metrics +#define multiarray_type multiarray_tile_stream_gpu +#define multiarray_create multiarray_tile_stream_gpu_create +#define multiarray_writer multiarray_tile_stream_gpu_writer +#define multiarray_destroy multiarray_tile_stream_gpu_destroy #else +#include "multiarray.cpu.h" #include "stream.cpu.h" #define stream_type tile_stream_cpu #define stream_create tile_stream_cpu_create #define stream_writer tile_stream_cpu_writer #define stream_destroy tile_stream_cpu_destroy #define stream_reset tile_stream_cpu_reset_metrics +#define multiarray_type multiarray_tile_stream_cpu +#define multiarray_create multiarray_tile_stream_cpu_create +#define multiarray_writer multiarray_tile_stream_cpu_writer +#define multiarray_destroy multiarray_tile_stream_cpu_destroy #endif #include @@ -62,7 +72,7 @@ verify_shard(const char* path, int first, int frames) } static int -test_stream(int frames, int bounded, int destroy_only) +test_stream(int frames, int bounded, int destroy_only, int disabled) { char root[256] = { 0 }, path[512]; struct test_zarr_sink z = { 0 }; @@ -99,14 +109,14 @@ test_stream(int frames, int bounded, int destroy_only) .max_threads = 2, .max_nlod = 1, .epochs_per_batch = 1, - .prepare_shards = 1, + .disable_shard_preparation = disabled, }; stream = stream_create(&config, test_zarr_sink_as_shard_sink(&z)); CHECK(Done, stream); for (int y = 0; y < 2; ++y) for (int x = 0; x < 2; ++x) { snprintf(path, sizeof(path), "%s/0/c/0/%d/%d", root, y, x); - CHECK(Done, platform_path_exists(path) == 1); + CHECK(Done, platform_path_exists(path) == !disabled); } struct writer* w = stream_writer(stream); for (int t = 0; t < frames; ++t) { @@ -125,7 +135,9 @@ test_stream(int frames, int bounded, int destroy_only) for (int y = 0; y < 2; ++y) for (int x = 0; x < 2; ++x) { snprintf(path, sizeof(path), "%s/0/c/%d/%d/%d", root, next, y, x); - CHECK(Done, platform_path_exists(path) == (!bounded || next < 2)); + CHECK(Done, + platform_path_exists(path) == + (!disabled && (!bounded || next < 2))); snprintf(path, sizeof(path), "%s/0/c/%d/%d/%d", root, next + 1, y, x); CHECK(Done, platform_path_exists(path) == 0); } @@ -245,7 +257,6 @@ test_prepared_files_need_no_resize(enum compression_codec codec_id) .max_threads = 2, .max_nlod = 1, .epochs_per_batch = 1, - .prepare_shards = 1, }; stream = stream_create(&config, test_zarr_sink_as_shard_sink(&z)); CHECK(Done, stream); @@ -269,6 +280,145 @@ test_prepared_files_need_no_resize(enum compression_codec codec_id) return error; } +struct preparation_probe +{ + struct shard_sink base; + int prepares; + int stops; + int cancels; + int fail_prepare; + int fail_cancel; +}; + +static int +probe_prepare(struct shard_sink* self, uint8_t level, uint64_t capacity) +{ + (void)level; + (void)capacity; + struct preparation_probe* p = (struct preparation_probe*)self; + ++p->prepares; + return p->fail_prepare; +} + +static void +probe_stop(struct shard_sink* self) +{ + ++((struct preparation_probe*)self)->stops; +} + +static int +probe_cancel(struct shard_sink* self) +{ + struct preparation_probe* p = (struct preparation_probe*)self; + ++p->cancels; + return p->fail_cancel; +} + +static struct preparation_probe +probe_init(int hooks) +{ + return (struct preparation_probe){ .base = { + .prepare_shards = + hooks & 1 ? probe_prepare : NULL, + .stop_preparing = + hooks & 2 ? probe_stop : NULL, + .cancel_prepared = + hooks & 4 ? probe_cancel : NULL, + } }; +} + +static struct tile_stream_configuration +probe_config(struct dimension dims[2], int disabled) +{ + dims_create(dims, "ty", (uint64_t[]){ 4, 4 }); + dims_set_chunk_sizes(dims, 2, (uint64_t[]){ 1, 4 }); + dims_set_shard_counts(dims, 2, (uint64_t[]){ 2, 1 }); + return (struct tile_stream_configuration){ + .buffer_capacity_bytes = 128, + .dtype = dtype_u16, + .rank = 2, + .dimensions = dims, + .codec = { .id = CODEC_NONE }, + .max_threads = 2, + .max_nlod = 1, + .disable_shard_preparation = disabled, + }; +} + +static int +test_preparation_policy(int hooks, int disabled) +{ + struct dimension dims[2]; + struct tile_stream_configuration config = probe_config(dims, disabled); + struct preparation_probe single = probe_init(hooks); + struct preparation_probe multi = probe_init(hooks); + struct stream_type* stream = NULL; + struct multiarray_type* arrays = NULL; + struct shard_sink* sinks[] = { &multi.base }; + const int enabled = hooks == 7 && !disabled; + int error = 1; + stream = stream_create(&config, &single.base); + CHECK(Done, stream && single.prepares == enabled); + arrays = multiarray_create(1, &config, sinks, 0); + CHECK(Done, arrays && multi.prepares == enabled); + CHECK(Done, writer_flush(stream_writer(stream)).error == 0); + CHECK(Done, + multiarray_writer(arrays)->flush(multiarray_writer(arrays)).error == 0); + stream_destroy(stream); + stream = NULL; + multiarray_destroy(arrays); + arrays = NULL; + CHECK(Done, (single.stops > 0) == enabled && (single.cancels > 0) == enabled); + CHECK(Done, (multi.stops > 0) == enabled && (multi.cancels > 0) == enabled); + error = 0; +Done: + stream_destroy(stream); + multiarray_destroy(arrays); + return error; +} + +static int +test_multiarray_cleanup_failure(int destroy_only) +{ + struct dimension dims[2]; + struct tile_stream_configuration config = probe_config(dims, 0); + struct tile_stream_configuration configs[] = { config, config }; + struct preparation_probe probes[] = { probe_init(7), probe_init(7) }; + probes[0].fail_cancel = 1; + struct shard_sink* sinks[] = { &probes[0].base, &probes[1].base }; + struct multiarray_type* arrays = multiarray_create(2, configs, sinks, 0); + int error = 1; + CHECK(Done, arrays); + if (!destroy_only) { + CHECK(Done, + multiarray_writer(arrays)->flush(multiarray_writer(arrays)).error != + 0); + CHECK(Done, probes[0].cancels > 0 && probes[1].cancels > 0); + } + multiarray_destroy(arrays); + arrays = NULL; + CHECK(Done, probes[0].cancels > 0 && probes[1].cancels > 0); + error = 0; +Done: + multiarray_destroy(arrays); + return error; +} + +static int +test_multiarray_preparation_rollback(void) +{ + struct dimension dims[2]; + struct tile_stream_configuration config = probe_config(dims, 0); + struct tile_stream_configuration configs[] = { config, config }; + struct preparation_probe probes[] = { probe_init(7), probe_init(7) }; + probes[1].fail_prepare = 1; + struct shard_sink* sinks[] = { &probes[0].base, &probes[1].base }; + struct multiarray_type* arrays = multiarray_create(2, configs, sinks, 0); + const int failed = arrays || !probes[0].cancels || !probes[1].cancels; + multiarray_destroy(arrays); + return failed; +} + int main(void) { @@ -281,13 +431,20 @@ main(void) #endif int error = test_prepared_files_need_no_resize(CODEC_NONE); error |= test_prepared_files_need_no_resize(CODEC_BLOSC_LZ4); - for (int destroy_only = 0; destroy_only < 2; ++destroy_only) { - error |= test_stream(0, 0, destroy_only); - error |= test_stream(1, 0, destroy_only); - error |= test_stream(4, 0, destroy_only); - error |= test_stream(5, 0, destroy_only); - error |= test_stream(4, 1, destroy_only); + for (int disabled = 0; disabled < 2; ++disabled) { + for (int hooks = 0; hooks < 8; ++hooks) + error |= test_preparation_policy(hooks, disabled); + for (int destroy_only = 0; destroy_only < 2; ++destroy_only) { + error |= test_stream(0, 0, destroy_only, disabled); + error |= test_stream(1, 0, destroy_only, disabled); + error |= test_stream(4, 0, destroy_only, disabled); + error |= test_stream(5, 0, destroy_only, disabled); + error |= test_stream(4, 1, destroy_only, disabled); + } } + error |= test_multiarray_cleanup_failure(0); + error |= test_multiarray_cleanup_failure(1); + error |= test_multiarray_preparation_rollback(); #ifdef TEST_PREPARATION_GPU cuCtxDestroy(context); #endif diff --git a/tests/test_zarr_readback.c b/tests/test_zarr_readback.c index d4ba7b6e..61b67c15 100644 --- a/tests/test_zarr_readback.c +++ b/tests/test_zarr_readback.c @@ -46,7 +46,6 @@ write_zarr(const char* store_path, struct codec_config codec, int buffered) .dtype = dtype_u16, .rank = 3, .dimensions = dims, - .prepare_shards = 1, .codec = codec, }; diff --git a/tests/test_zarr_readback_gpu.c b/tests/test_zarr_readback_gpu.c index 6c7ff3ef..3a4ee0e1 100644 --- a/tests/test_zarr_readback_gpu.c +++ b/tests/test_zarr_readback_gpu.c @@ -45,7 +45,6 @@ write_zarr(const char* store_path, struct codec_config codec, int buffered) .dtype = dtype_u16, .rank = 3, .dimensions = dims, - .prepare_shards = 1, .codec = codec, .epochs_per_batch = 1, };