diff --git a/README.md b/README.md index 21be57a4..7637f750 100644 --- a/README.md +++ b/README.md @@ -304,6 +304,9 @@ the library handles all tiling, padding, and shard assembly internally. See [docs/guide.md][docs-guide-md] for a quick orientation to the module structure. Paced producers can opt into the [bounded input-buffering writer](docs/buffered-writer.md). +Filesystem streams [prepare one shard generation ahead](docs/shard-preparation.md) +by default, including cleanup of unused files at shutdown. Set +`config.disable_shard_preparation = 1` to open files on demand instead. For writing directly to S3 (or S3-compatible stores), see the [S3 storage guide][s3-storage-guide]. diff --git a/docs/shard-preparation.md b/docs/shard-preparation.md new file mode 100644 index 00000000..7010e0a7 --- /dev/null +++ b/docs/shard-preparation.md @@ -0,0 +1,39 @@ +# Preparing shard files ahead of use + +CPU and GPU streams prepare shard files ahead of use by default when the +sink supports it. The filesystem store supports preparation for Zarr arrays, +multiarray streams, and NGFF multiscales. Set +`config.disable_shard_preparation = 1` to open files on demand instead. +Unsupported sinks, including S3, automatically retain their existing behavior. + +Creation waits until the first shard generation is ready. When a stream +adopts one of those files, a separate worker creates and, where supported, +sets the capacity of the next file for that slot. There is at most one unused +prepared file per slot. A layout with 15 spatial shards therefore holds at +most 15 additional files, independently of the stream's length. A finite +stream does not prepare files beyond its configured capacity. + +Preparation uses one worker separate from the normal write workers. Future +preparation is excluded from the fences used by writes, metadata publication, +and footer reuse. If preparation has not finished when a file is needed, +opening that shard waits for it. + +`writer_flush` stops scheduling future preparation, drains accepted input, +and waits for preparation before removing unused files. `writer_close` +reports cleanup failures, and destruction also attempts cleanup when explicit +flush/close was omitted. Files already adopted follow the ordinary shard +lifecycle. Preparation creates files exclusively: it fails if a target +already exists, without truncating or deleting that file. Cleanup removes +only unused files and empty directories created by preparation; it preserves +pre-existing directories and directories containing used shards. + +Store backends expose preparation through the shard pool. `prepare`, +`wait_prepared`, and `cancel_prepared` form one capability. Supporting +preparation requires supporting cancellation: wait for outstanding work, +reclaim unused owned resources, report failures, and allow cancellation to +be repeated. Opening a prepared shard transfers it to the ordinary writer. +Preparation and open calls are serialized by the stream; cancellation runs +after delivery stops using the prepared resources. + +Custom sinks need all three `prepare_shards`, `stop_preparing`, and +`cancel_prepared` hooks. Missing any hook disables preparation for that sink. diff --git a/src/cpu/stream.body.c b/src/cpu/stream.body.c index da1a84af..8849844d 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->disable_shard_preparation) + shard_sink_stop_preparing(v->sink); // A create that failed before sizing the layout leaves epoch_elements at 0 // and the divisions below would fault. Nothing was sized, so nothing to do. if (v->layout->epoch_elements == 0) @@ -338,6 +340,9 @@ cpu_stream_flush_body(struct cpu_stream_view* v) // they read. if (shard_sink_drain(v->sink)) failed = 1; + if (!v->config->disable_shard_preparation && + shard_sink_cancel_prepared(v->sink)) + failed = 1; return failed ? writer_error() : writer_ok(); } @@ -345,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 = 0; + int failed = !v->config->disable_shard_preparation && + shard_sink_cancel_prepared(v->sink); // A sink IO error is the one case the shape is withheld: which writes landed // is unknowable. diff --git a/src/cpu/stream.c b/src/cpu/stream.c index 7cfa6a9b..4230370c 100644 --- a/src/cpu/stream.c +++ b/src/cpu/stream.c @@ -291,6 +291,9 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, } CHECK(Fail, shard_sink_init_append(sink, &s->cl.dims, s->levels.nlod) == 0); + CHECK(Fail, + shard_sink_prepare_first( + sink, s->shard, s->levels.nlod, s->shard_alignment, config) == 0); s->writer.append = cpu_append; s->writer.flush = cpu_flush_final; @@ -303,6 +306,8 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, return s; Fail: + if (!config->disable_shard_preparation && shard_sink_cancel_prepared(sink)) + log_error("CPU stream preparation cleanup failed during rollback"); // Construction rollback releases resources without finalizing the sink. tile_stream_cpu_release_resources(s); return NULL; diff --git a/src/gpu/stream.c b/src/gpu/stream.c index 4e066ccd..05b57eaf 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.disable_shard_preparation) + shard_sink_stop_preparing(ctx->sink); // A create that fails before sizing the layout leaves epoch_elements at 0; // the divisions below would then fault. Nothing was ever sized, so there is // nothing to flush. @@ -429,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)) { @@ -436,6 +443,9 @@ stream_flush_body(struct stream_engine* e, struct stream_context* ctx) } schedule_quiesce_output(e, ctx->sink); + if (!ctx->config.disable_shard_preparation && + shard_sink_cancel_prepared(ctx->sink)) + r = writer_error(); collect_ingest_timing(e); return r; @@ -449,6 +459,10 @@ stream_close_body(struct compress_agg_array* ar, struct stream_context* ctx) struct writer_result r = writer_ok(); + if (!ctx->config.disable_shard_preparation && + shard_sink_cancel_prepared(ctx->sink)) + r = writer_error(); + // The shape is written after this, so it never names data still queued. const int sink_failed = shard_sink_drain(ctx->sink); if (sink_failed) diff --git a/src/gpu/stream.init.c b/src/gpu/stream.init.c index 0f8286a4..86a557d3 100644 --- a/src/gpu/stream.init.c +++ b/src/gpu/stream.init.c @@ -350,6 +350,9 @@ tile_stream_gpu_rollback(struct tile_stream_gpu* s) const int pushed = cu_ctx_push(s->engine.cuda); gpu_delivery_stop_join(&s->engine.delivery); gpu_streams_sync(&s->engine.streams); + if (!s->ctx.config.disable_shard_preparation && + shard_sink_cancel_prepared(s->ctx.sink)) + log_error("GPU stream preparation cleanup failed during rollback"); tile_stream_gpu_release_resources(s, pushed); } @@ -451,6 +454,13 @@ tile_stream_gpu_create(const struct tile_stream_configuration* config, shard_sink_init_append(sink, &out->ctx.dims, out->ctx.levels.nlod) == 0); + CHECK(FailPhase2, + shard_sink_prepare_first(sink, + out->engine.compress_agg.ar.shard, + out->ctx.levels.nlod, + out->ctx.shard_alignment, + config) == 0); + CHECK(FailPhase2, gpu_placement_leave(&placement_scope) == 0); out->flushed = 0; out->closed = 0; diff --git a/src/multiarray/stream.c b/src/multiarray/stream.c index 2ebd3ff5..eff7294c 100644 --- a/src/multiarray/stream.c +++ b/src/multiarray/stream.c @@ -478,6 +478,12 @@ multiarray_tile_stream_cpu_create( CHECK(Fail, shard_sink_init_append( desc->sink, &desc->cl.dims, desc->levels.nlod) == 0); + CHECK(Fail, + shard_sink_prepare_first(desc->sink, + desc->shard, + desc->levels.nlod, + desc->shard_alignment, + &desc->config) == 0); } for (int i = 0; i < n_arrays; ++i) { ms->arrays[i].flushed = 0; @@ -486,6 +492,11 @@ multiarray_tile_stream_cpu_create( return ms; Fail: + if (ms && ms->arrays) + for (int i = 0; i < ms->n_arrays; ++i) + if (!ms->arrays[i].config.disable_shard_preparation && + shard_sink_cancel_prepared(ms->arrays[i].sink)) + log_error("CPU multiarray preparation cleanup failed during rollback"); // Construction rollback releases resources without finalizing any sink. multiarray_tile_stream_cpu_release_resources(ms); return NULL; @@ -733,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 @@ -740,7 +753,11 @@ flush_impl(struct multiarray_writer* self) if (desc->flushed) continue; if (desc->cursor_elements == 0 && desc->batch_accumulated == 0) { + if (!desc->config.disable_shard_preparation && + shard_sink_cancel_prepared(desc->sink)) + failed = 1; desc->flushed = 1; + desc->closed = 0; continue; } @@ -757,16 +774,11 @@ flush_impl(struct multiarray_writer* self) desc->flushed = 1; desc->closed = 0; if (r.error) - goto Error; + failed = 1; } return (struct multiarray_writer_result){ - .error = multiarray_writer_ok, - }; - -Error: - return (struct multiarray_writer_result){ - .error = multiarray_writer_fail, + .error = failed ? multiarray_writer_fail : multiarray_writer_ok, }; } diff --git a/src/multiarray/stream.gpu.c b/src/multiarray/stream.gpu.c index 2bb93afe..fe2c812f 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.disable_shard_preparation && + shard_sink_cancel_prepared(desc->ctx.sink)) + failed = 1; desc->flushed = 1; continue; } @@ -379,6 +382,11 @@ multiarray_tile_stream_gpu_rollback(struct multiarray_tile_stream_gpu* ms) const int pushed = cu_ctx_push(ms->engine.cuda); gpu_delivery_stop_join(&ms->engine.delivery); gpu_streams_sync(&ms->engine.streams); + if (ms->arrays) + for (int a = 0; a < ms->n_arrays; ++a) + if (!ms->arrays[a].ctx.config.disable_shard_preparation && + shard_sink_cancel_prepared(ms->arrays[a].ctx.sink)) + log_error("GPU multiarray preparation cleanup failed during rollback"); multiarray_tile_stream_gpu_release_resources(ms, pushed); } @@ -496,6 +504,12 @@ multiarray_tile_stream_gpu_create( const struct stream_context* ctx = &ms->arrays[a].ctx; CHECK(Fail, shard_sink_init_append(ctx->sink, &ctx->dims, ctx->levels.nlod) == 0); + CHECK(Fail, + shard_sink_prepare_first(ctx->sink, + ms->arrays[a].st.agg.shard, + ctx->levels.nlod, + ctx->shard_alignment, + &ctx->config) == 0); } for (int a = 0; a < n_arrays; ++a) { ms->arrays[a].flushed = 0; 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..8b95ff67 100644 --- a/src/platform/platform_io.darwin.c +++ b/src/platform/platform_io.darwin.c @@ -17,6 +17,24 @@ platform_mkdir(const char* path) return 0; } +int +platform_mkdir_new(const char* path) +{ + if (mkdir(path, 0755) == 0) + return 1; + struct stat st; + return errno == EEXIST && stat(path, &st) == 0 && S_ISDIR(st.st_mode) ? 0 + : -1; +} + +int +platform_remove_empty_directory(const char* path) +{ + if (rmdir(path) == 0 || errno == ENOENT) + return 0; + return errno == ENOTEMPTY || errno == EEXIST ? 1 : -1; +} + int platform_mkdirp(const char* path) { @@ -40,13 +58,16 @@ platform_mkdirp(const char* path) platform_fd platform_open_write(const char* path, int flags) { - int oflags = O_WRONLY | O_CREAT | O_TRUNC; + int oflags = O_WRONLY | O_CREAT; + oflags |= flags & PLATFORM_OPEN_EXCLUSIVE ? O_EXCL : O_TRUNC; int fd = open(path, oflags, 0644); if (fd < 0) return fd; if (flags & PLATFORM_OPEN_UNBUFFERED) { if (fcntl(fd, F_NOCACHE, 1) != 0) { close(fd); + if (flags & PLATFORM_OPEN_EXCLUSIVE) + unlink(path); return -1; } } diff --git a/src/platform/platform_io.h b/src/platform/platform_io.h index 17fc1a4e..d7f1ef9f 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); + +// Returns 0 if removed or missing, 1 if nonempty, -1 on other failures. +int +platform_remove_empty_directory(const char* path); + // Create a directory and all parent directories. Returns 0 on success. int platform_mkdirp(const char* path); @@ -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..8b54448d 100644 --- a/src/platform/platform_io.posix.c +++ b/src/platform/platform_io.posix.c @@ -19,6 +19,24 @@ platform_mkdir(const char* path) return 0; } +int +platform_mkdir_new(const char* path) +{ + if (mkdir(path, 0755) == 0) + return 1; + struct stat st; + return errno == EEXIST && stat(path, &st) == 0 && S_ISDIR(st.st_mode) ? 0 + : -1; +} + +int +platform_remove_empty_directory(const char* path) +{ + if (rmdir(path) == 0 || errno == ENOENT) + return 0; + return errno == ENOTEMPTY || errno == EEXIST ? 1 : -1; +} + int platform_mkdirp(const char* path) { @@ -42,7 +60,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..a7e5dd37 100644 --- a/src/platform/platform_io.win32.c +++ b/src/platform/platform_io.win32.c @@ -13,6 +13,30 @@ 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(); + if (error == ERROR_FILE_NOT_FOUND || error == ERROR_PATH_NOT_FOUND) + return 0; + return error == ERROR_DIR_NOT_EMPTY ? 1 : -1; +} + int platform_mkdirp(const char* path) { @@ -44,9 +68,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..7a3ff25e 100644 --- a/src/types.stream.h +++ b/src/types.stream.h @@ -154,6 +154,8 @@ struct tile_stream_configuration // reports more pending than this int max_threads; // 0 = OpenMP default int full_memcpy_timing; + 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 dd67494c..f87e3930 100644 --- a/src/writer.c +++ b/src/writer.c @@ -88,6 +88,26 @@ 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 (shard_sink_supports_preparation(s)) + s->stop_preparing(s); +} + +int +shard_sink_cancel_prepared(struct shard_sink* s) +{ + shard_sink_stop_preparing(s); + return shard_sink_supports_preparation(s) ? 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..ba39f9d5 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,16 @@ struct shard_sink size_t (*required_shard_alignment)(const struct shard_sink* self); }; +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); + uint64_t shard_sink_pending_bytes(const 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_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/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 202a1fe8..0bcbf8b2 100644 --- a/src/zarr/shard_delivery.c +++ b/src/zarr/shard_delivery.c @@ -243,6 +243,45 @@ 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, + const struct tile_stream_configuration* config) +{ + 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) { + 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"); + 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..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,6 +78,21 @@ 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, + 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. 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..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 @@ -21,6 +22,8 @@ struct shard_pool_fs struct shard_pool base; struct io_backend_fs* backend; struct io_scheduler* queue; + struct shard_preparation* preparation; + struct shard_pool_fs_wrapper wrapper; struct fs_slot* slots; uint64_t nslots; int unbuffered; @@ -37,6 +40,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 +169,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,9 +207,61 @@ fs_slot_finalize(struct shard_writer* self) return 1; w->token = (struct io_file_token){ 0 }; + w->prepared_capacity = 0; return 0; } +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->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); + } + CHECK(Fail, strbuf_appendf(&path, "%s/%s", strbuf_cstr(&p->root), key) == 0); + CHECK(Fail, + shard_preparation_post(p->preparation, + slot, + strbuf_cstr(&path), + p->slots[slot].presize ? capacity : 0) == 0); + 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; + return shard_preparation_wait(p->preparation, slot) || + 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; + return shard_preparation_cancel(p->preparation, first, count); +} + static struct shard_writer* pool_fs_open(struct shard_pool* self, uint64_t slot, const char* key) { @@ -222,6 +278,18 @@ 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; + 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; + } + const size_t path_bytes = strbuf_len(&path) + 1; owned_path = (char*)malloc(path_bytes); CHECK(Fail, owned_path); @@ -342,6 +410,8 @@ pool_fs_destroy(struct shard_pool* self) { struct shard_pool_fs* p = container_of(self, struct shard_pool_fs, base); + shard_preparation_destroy(p->preparation); + // Finalize any open slots for (uint64_t i = 0; i < p->nslots; ++i) { if (p->slots[i].token.generation != 0) @@ -404,6 +474,10 @@ 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->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_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/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/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..40793a1b 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 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 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() +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..649c5748 --- /dev/null +++ b/tests/test_shard_preparation.c @@ -0,0 +1,282 @@ +#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; +} + +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) +{ + 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(); + 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 new file mode 100644 index 00000000..612794ac --- /dev/null +++ b/tests/test_stream_preparation.c @@ -0,0 +1,458 @@ +#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 "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 +#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, int disabled) +{ + 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, + .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) == !disabled); + } + 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) == + (!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); + } + } + 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, + }; + 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; +} + +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) +{ +#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 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 + return error; +#ifdef TEST_PREPARATION_GPU +Fail: + if (context) + cuCtxDestroy(context); + return 1; +#endif +}