diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index b8d63bb75..b34179302 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -330,6 +330,22 @@ Shutdown exceptions are suppressed, so an integration requiring an authoritative final read must ensure every worker reaches this close path and should separately check native host failures. +Under `storage_backend="persistent"`, the native pack sink stages the pack it +still has open when the stopping ring releases it, so the last records reach +the spool without a flush. With `capture_storage_config` set, `close()` first +drains capture, best effort, with a budget of `close_flush_timeout_s`: it +flushes the sink, stops the ring, waits for the storage service to get the +staged packs into the catalog, and stops the service. What misses the budget is +logged and left for the next start (the spool, or the reconcile); +`flush_and_wait` is the call that raises when captures are not queryable in +time. Past the budget the drain starts no upload and at most one index batch, +so against a catalog or object store that stops answering, `close()` outlasts +the budget by up to about two `clickhouse_request_timeout_s` (the request in +flight, and the lease release), plus up to about 6 s when the budget cuts a +multipart upload, whose abort nothing cuts; `close_flush_timeout_s` documents +the full bound. The same bounds hold for `flush_and_wait(timeout_s)`, less the +lease release. + Closing does not disable or uninstall HookPoints: they retain hook IDs and the old payload tensor. Treat the attached model as terminal too. A later CUDA forward—especially after another engine becomes active—can combine stale hook diff --git a/native/Makefile b/native/Makefile index d7e60e889..2353eeea8 100644 --- a/native/Makefile +++ b/native/Makefile @@ -367,7 +367,7 @@ build/conformance_sign: csrc/store/s3_sign.cpp csrc/store/conformance_sign.cpp c -Icsrc/store -Icsrc/common -lcrypto # A2 store conformance driver + fault-matrix target. Needs libcurl (above). -build/conformance_store: csrc/store/s3_sign.cpp csrc/store/s3_client.cpp csrc/store/spool.cpp csrc/store/uploader.cpp csrc/store/conformance_store.cpp csrc/store/s3_sign.h csrc/store/s3_client.h csrc/store/spool.h csrc/store/uploader.h csrc/common/json.cpp csrc/common/json.h csrc/common/curl_init.cpp csrc/common/curl_init.h | check-libcurl +build/conformance_store: csrc/store/s3_sign.cpp csrc/store/s3_client.cpp csrc/store/spool.cpp csrc/store/uploader.cpp csrc/store/conformance_store.cpp csrc/store/s3_sign.h csrc/store/s3_client.h csrc/store/cancel.h csrc/store/spool.h csrc/store/uploader.h csrc/common/json.cpp csrc/common/json.h csrc/common/curl_init.cpp csrc/common/curl_init.h | check-libcurl mkdir -p $(BUILD_DIR) $(CXX) -std=c++17 -O2 -Wall -Wextra -o $@ \ csrc/store/s3_sign.cpp csrc/store/s3_client.cpp csrc/store/spool.cpp csrc/store/uploader.cpp csrc/store/conformance_store.cpp csrc/common/json.cpp csrc/common/curl_init.cpp \ @@ -375,7 +375,7 @@ build/conformance_store: csrc/store/s3_sign.cpp csrc/store/s3_client.cpp csrc/st # B1 catalog conformance driver: lease coordinator + version allocator over # ClickHouse's HTTP interface. No torch/pybind; needs libcurl like the store. -build/conformance_catalog: csrc/catalog/clickhouse_client.cpp csrc/catalog/lease_coordinator.cpp csrc/catalog/version_allocator.cpp csrc/catalog/catalog_writer.cpp csrc/catalog/pack_index.cpp csrc/catalog/indexer.cpp csrc/catalog/schema.cpp csrc/catalog/reader.cpp csrc/catalog/hydration.cpp csrc/catalog/conformance_catalog.cpp csrc/common/json.cpp csrc/pack/pack_builder.cpp csrc/store/s3_sign.cpp csrc/store/s3_client.cpp csrc/catalog/clickhouse_client.h csrc/catalog/lease_coordinator.h csrc/catalog/version_allocator.h csrc/catalog/catalog_writer.h csrc/catalog/sql_escape.h csrc/catalog/pack_index.h csrc/catalog/indexer.h csrc/catalog/schema.h csrc/catalog/reader.h csrc/catalog/hydration.h csrc/common/json.h csrc/store/s3_client.h csrc/common/curl_init.cpp csrc/common/curl_init.h | check-libcurl +build/conformance_catalog: csrc/catalog/clickhouse_client.cpp csrc/catalog/lease_coordinator.cpp csrc/catalog/version_allocator.cpp csrc/catalog/catalog_writer.cpp csrc/catalog/pack_index.cpp csrc/catalog/indexer.cpp csrc/catalog/schema.cpp csrc/catalog/reader.cpp csrc/catalog/hydration.cpp csrc/catalog/conformance_catalog.cpp csrc/common/json.cpp csrc/pack/pack_builder.cpp csrc/store/s3_sign.cpp csrc/store/s3_client.cpp csrc/catalog/clickhouse_client.h csrc/catalog/lease_coordinator.h csrc/catalog/version_allocator.h csrc/catalog/catalog_writer.h csrc/catalog/sql_escape.h csrc/catalog/pack_index.h csrc/catalog/indexer.h csrc/catalog/schema.h csrc/catalog/reader.h csrc/catalog/hydration.h csrc/common/json.h csrc/store/s3_client.h csrc/store/cancel.h csrc/common/curl_init.cpp csrc/common/curl_init.h | check-libcurl mkdir -p $(BUILD_DIR) $(CXX) -std=c++17 -O2 -Wall -Wextra -o $@ \ csrc/catalog/clickhouse_client.cpp csrc/catalog/lease_coordinator.cpp \ @@ -388,7 +388,7 @@ build/conformance_catalog: csrc/catalog/clickhouse_client.cpp csrc/catalog/lease $(CURL_CPPFLAGS) $(CURL_LDFLAGS) -lcrypto -lcurl -lpthread # A3a spool conformance driver: no network, no torch; runs in any CI job. -build/conformance_spool: csrc/store/spool.cpp csrc/store/conformance_spool.cpp csrc/store/spool.h csrc/common/json.cpp csrc/common/json.h +build/conformance_spool: csrc/store/spool.cpp csrc/store/conformance_spool.cpp csrc/store/spool.h csrc/store/cancel.h csrc/common/json.cpp csrc/common/json.h mkdir -p $(BUILD_DIR) $(CXX) -std=c++17 -O2 -Wall -Wextra -o $@ \ csrc/store/spool.cpp csrc/store/conformance_spool.cpp csrc/common/json.cpp \ diff --git a/native/csrc/catalog/bindings_store.cpp b/native/csrc/catalog/bindings_store.cpp index d07593772..dcaf4cdc6 100644 --- a/native/csrc/catalog/bindings_store.cpp +++ b/native/csrc/catalog/bindings_store.cpp @@ -116,6 +116,7 @@ py::dict snapshot_dict(const dc::StorageServiceSnapshot& s) { out["uploaded_packs"] = s.uploaded_packs; out["uploaded_bytes"] = s.uploaded_bytes; out["upload_failures"] = s.upload_failures; + out["cancelled_uploads"] = s.cancelled_uploads; out["indexed_packs"] = s.indexed_packs; out["indexed_rows"] = s.indexed_rows; out["index_failures"] = s.index_failures; diff --git a/native/csrc/catalog/conformance_catalog.cpp b/native/csrc/catalog/conformance_catalog.cpp index e66c2f19e..be25307e4 100644 --- a/native/csrc/catalog/conformance_catalog.cpp +++ b/native/csrc/catalog/conformance_catalog.cpp @@ -22,6 +22,7 @@ #include #include "../common/json.h" +#include "../store/cancel.h" #include "../store/s3_client.h" #include "catalog_writer.h" #include "clickhouse_client.h" @@ -466,6 +467,66 @@ std::string respond(const std::string& line, Session* session) { out += "]"; return prefix + "true" + out + "}"; } + if (op == "read_pack_rows") { + // Session-less: one pack's descriptor rows, read through the object + // store as the indexer reads them, the client holding a Cancellation + // -- whose deadline passes cancel_after_ms from now, as a flush arms + // one, or which is cancelled once the first exchange has returned + // (cancel_after_exchange), as a cancel that comes in while a failure + // is reported. Pins how a failed read is told apart -- the store's + // answer, the store not answering, and whether a cancel is what cut + // it -- without a catalog. + dmi_store::S3Config s3_config; + s3_config.endpoint = jc::FindString(line, "endpoint"); + s3_config.bucket = jc::FindString(line, "bucket"); + s3_config.region = jc::FindString(line, "region"); + s3_config.access_key = jc::FindString(line, "access"); + s3_config.secret_key = jc::FindString(line, "secret"); + s3_config.allow_insecure_http = jc::FindBool(line, "insecure"); + if (jc::HasKey(line, "read_timeout")) { + s3_config.read_timeout_s = + static_cast(field_int(line, "read_timeout")); + } + if (jc::HasKey(line, "max_attempts")) { + s3_config.max_attempts = + static_cast(field_int(line, "max_attempts")); + } + dmi_store::Cancellation cancel; // outlives the client + dmi_store::S3Client s3(s3_config); + if (jc::HasKey(line, "cancel_after_ms")) { + cancel.set_deadline( + dmi_store::Cancellation::NowNs() + + static_cast(field_int(line, "cancel_after_ms")) * + 1'000'000ull); + s3.set_cancellation(&cancel); + } + if (jc::FindBool(line, "cancel_after_exchange")) { + s3.set_cancellation(&cancel); + s3.SetAfterExchangeHookForTesting([&cancel] { cancel.Cancel(); }); + } + const std::string element = jc::FindObject(line, "ref"); + dmi_catalog::PackRefData ref; + ref.pack_id = jc::FindString(element, "pack_id"); + ref.store_id = jc::FindString(element, "store_id"); + ref.object_key = jc::FindString(element, "object_key"); + ref.object_bytes = + static_cast(field_int(element, "object_bytes")); + ref.checksum = jc::FindString(element, "checksum"); + ref.record_count = + static_cast(field_int(element, "record_count")); + try { + out = ",\"rows\":" + + std::to_string(dmi_catalog::read_pack_descriptor_rows(&s3, ref) + .size()); + } catch (const dmi_catalog::StoreUnavailableError& e) { + std::string message; + escape_into(e.what(), &message); + return prefix + "false,\"error\":\"StoreUnavailable\",\"cancelled\":" + + (e.cancelled() ? "true" : "false") + ",\"message\":" + + message + "}"; + } + return prefix + "true" + out + "}"; + } if (session->writer == nullptr) { return prefix + "false,\"what\":\"call open first\"}"; } diff --git a/native/csrc/catalog/indexer.cpp b/native/csrc/catalog/indexer.cpp index 516eabf61..223c6a9fb 100644 --- a/native/csrc/catalog/indexer.cpp +++ b/native/csrc/catalog/indexer.cpp @@ -239,6 +239,10 @@ void NativeIndexer::read(IndexPlan* planned) { try { rows = read_pack_descriptor_rows(s3_, ref); } catch (const CatalogError& e) { + if (config_.end_read_when_store_unavailable && + dynamic_cast(&e) != nullptr) { + throw; + } std::string message = e.what(); if (message.size() > 512) message.resize(512); result.failures.push_back( diff --git a/native/csrc/catalog/indexer.h b/native/csrc/catalog/indexer.h index d6e7050fa..506c8e2c5 100644 --- a/native/csrc/catalog/indexer.h +++ b/native/csrc/catalog/indexer.h @@ -27,6 +27,14 @@ struct IndexerConfig { int max_rows_per_insert = 10'000; uint64_t max_estimated_bytes = 128ull * 1024 * 1024; int max_publish_attempts = 8; + // read() ends at a pack the object store did not answer for + // (StoreUnavailableError propagates) instead of failing that one pack + // and reading the next. The oracle, CatalogIndexer, fails the pack and + // moves on, and so does the default; the storage service sets it, since + // against a store that stopped answering every further pack costs the + // client's full timeouts and fails the same way, and an outage is not + // the pack's fault to count against it. + bool end_read_when_store_unavailable = false; // Test seam, unset in production, called with each allocated version // just before the publish that carries it. The publish wedges in // CatalogWriter exist for the same reason: a version race needs the @@ -82,7 +90,8 @@ class NativeIndexer { // Deduplicates refs and reads the replay guard (the catalog). IndexPlan plan(const std::vector& refs); // Reads each pending pack's descriptor rows (the object store only). - // Throws kBatchTooLarge past max_estimated_bytes. + // Throws kBatchTooLarge past max_estimated_bytes, and + // StoreUnavailableError under end_read_when_store_unavailable. void read(IndexPlan* plan); // Allocates a version, writes the descriptors, publishes and commits the // inventory (the catalog). Requires read(). diff --git a/native/csrc/catalog/pack_index.cpp b/native/csrc/catalog/pack_index.cpp index 7daa34498..cd2db0d0a 100644 --- a/native/csrc/catalog/pack_index.cpp +++ b/native/csrc/catalog/pack_index.cpp @@ -33,6 +33,17 @@ constexpr uint64_t kMaxLayerNumber = 0x7FFFFFFFull; // 2^31 - 1 throw CatalogError(CatalogError::Kind::kValue, what); } +// A range read that failed: the store's answer about the pack, or no +// answer at all (StoreUnavailableError) -- the cancel's only when the read +// says the Cancellation cut it. One set once the store had failed the read +// (a flush's deadline passing just after a 503) cut nothing: that failure +// is the store's, and is recorded as such. +[[noreturn]] void range_read_failed(const std::string& what, bool unavailable, + bool cancelled) { + if (unavailable) throw StoreUnavailableError(what, cancelled); + throw CatalogError(CatalogError::Kind::kValue, what); +} + // `CaptureMetadata.from_mapping` wraps every ValueError `__post_init__` // raises as `invalid capture metadata: {exc}`, so the refusals that come // from the metadata model's own bounds carry that prefix and the ones that @@ -433,10 +444,12 @@ std::vector read_pack_descriptor_rows( std::vector trailer; const uint64_t trailer_offset = ref.object_bytes - kTrailerSize; if (charge) charge(kTrailerSize); + bool unavailable = false; + bool cancelled = false; if (!s3->GetRange(ref.object_key, trailer_offset, kTrailerSize, &trailer, - &error)) { - throw CatalogError(CatalogError::Kind::kValue, - "pack trailer read failed: " + error); + &error, &unavailable, &cancelled)) { + range_read_failed("pack trailer read failed: " + error, unavailable, + cancelled); } if (trailer.size() != kTrailerSize) format_error("pack trailer is truncated"); const auto read_u64 = [&](size_t at) { @@ -482,9 +495,9 @@ std::vector read_pack_descriptor_rows( // about to be read -- not the object's size, which only bounds it. if (charge) charge(footer_length); if (!s3->GetRange(ref.object_key, footer_offset, footer_length, &footer, - &error)) { - throw CatalogError(CatalogError::Kind::kValue, - "pack footer read failed: " + error); + &error, &unavailable, &cancelled)) { + range_read_failed("pack footer read failed: " + error, unavailable, + cancelled); } if (dmi_pack::Crc32(footer.data(), footer.size()) != footer_crc) { throw CatalogError(CatalogError::Kind::kValue, "footer checksum mismatch"); diff --git a/native/csrc/catalog/pack_index.h b/native/csrc/catalog/pack_index.h index 0d7971e77..dc04ebeac 100644 --- a/native/csrc/catalog/pack_index.h +++ b/native/csrc/catalog/pack_index.h @@ -14,9 +14,28 @@ #include #include "../store/s3_client.h" +#include "lease_coordinator.h" namespace dmi_catalog { +// read_pack_descriptor_rows could not get the object store's answer about +// the pack: a transport error or timeout, a retryable status on every +// attempt, or the S3 client's Cancellation cutting the read (cancelled() +// says which; a cancel that came in once the store had failed it does not +// count, as in S3Client). It says +// nothing about the pack itself, unlike every other refusal the read makes. +// A CatalogError of kind kValue like those, so a caller that treats every +// unreadable pack alike still does. +class StoreUnavailableError : public CatalogError { + public: + StoreUnavailableError(const std::string& what, bool cancelled) + : CatalogError(Kind::kValue, what), cancelled_(cancelled) {} + bool cancelled() const { return cancelled_; } + + private: + bool cancelled_; +}; + struct PackRefData { std::string pack_id; std::string store_id; @@ -29,7 +48,8 @@ struct PackRefData { // Reads the pack the ref names and renders one descriptor VALUES row per // record — the 33 capture_raw columns in schema order, without // index_version (the batch's own version). Throws CatalogError (kValue) -// on any format violation, mirroring PackFormatError / PackIntegrityError. +// on any format violation, mirroring PackFormatError / PackIntegrityError, +// and StoreUnavailableError when the store did not answer a range read. // The bucket is the S3Client's own config; a bucket parameter here would // only invite a caller to believe passing a different one redirects the // read. diff --git a/native/csrc/catalog/storage_service.cpp b/native/csrc/catalog/storage_service.cpp index df57f79ca..c1490eec4 100644 --- a/native/csrc/catalog/storage_service.cpp +++ b/native/csrc/catalog/storage_service.cpp @@ -72,6 +72,40 @@ std::string seconds_text(uint64_t ns) { return out; } +// The service's indexer ends a read pass at a pack the object store did not +// answer for, rather than failing that pack and reading the next. +IndexerConfig service_indexer_config(IndexerConfig config) { + config.end_read_when_store_unavailable = true; + return config; +} + +// How long past a flush's deadline its index reads may still run: the one +// catalog request timeout the flush may overrun by anyway. Capped so the +// nanoseconds cannot overflow. +uint64_t read_grace_ns(const ClickHouseConnection& connection) { + const double seconds = std::min(connection.timeouts.request_s, 1e6); + return seconds > 0 ? static_cast(seconds * 1e9) : 0; +} + +// Arms a Cancellation's deadline for one cycle, at ns (0: none), and +// disarms it at the end. Cancel(), stop()'s, outlives it. +class ArmedDeadline { + public: + ArmedDeadline(dmi_store::Cancellation* cancel, uint64_t ns) + : cancel_(cancel), armed_(ns != 0) { + if (armed_) cancel_->set_deadline(ns); + } + ~ArmedDeadline() { + if (armed_) cancel_->set_deadline(0); + } + ArmedDeadline(const ArmedDeadline&) = delete; + ArmedDeadline& operator=(const ArmedDeadline&) = delete; + + private: + dmi_store::Cancellation* cancel_; + const bool armed_; +}; + } // namespace class CaptureStorageService::LeaseScope { @@ -121,9 +155,10 @@ class CaptureStorageService::LeaseScope { CaptureStorageService::CaptureStorageService(StorageServiceConfig config) : config_(std::move(config)), s3_(config_.s3), + upload_s3_(config_.s3), clickhouse_(std::make_shared(config_.clickhouse)), writer_(clickhouse_, config_.writer), - indexer_(&s3_, &writer_, config_.indexer) { + indexer_(&s3_, &writer_, service_indexer_config(config_.indexer)) { if (config_.spool_root.empty()) { throw std::invalid_argument("storage service: spool_root is required"); } @@ -159,8 +194,11 @@ CaptureStorageService::CaptureStorageService(StorageServiceConfig config) &spool_, &error) != dmi_store::SpoolStatus::kOk) { throw std::runtime_error("storage service: cannot open spool: " + error); } - uploader_ = std::make_unique(&spool_, &s3_, + s3_.set_cancellation(&read_cancel_); + upload_s3_.set_cancellation(&upload_cancel_); + uploader_ = std::make_unique(&spool_, &upload_s3_, config_.uploader); + uploader_->set_cancellation(&upload_cancel_); } CaptureStorageService::~CaptureStorageService() { @@ -173,6 +211,9 @@ CaptureStorageService::~CaptureStorageService() { void CaptureStorageService::start() { std::lock_guard cycle(cycle_mutex_); if (started_) throw std::logic_error("storage service: already started"); + // A stop() before this one cancelled both for good. + upload_cancel_.Reset(); + read_cancel_.Reset(); CatalogSchema(clickhouse_, config_.writer.database, config_.writer.table_prefix) .ensure(&writer_.leases(), config_.schema_retry_sleep_ns); @@ -271,6 +312,14 @@ void CaptureStorageService::stop() { std::lock_guard lock(wake_mutex_); stop_requested_ = true; } + // Before the join: an upload or an index read the store never answers, + // or a retry backoff, would otherwise hold it for the S3 client's + // timeouts on every attempt, with the lease held. A cancelled upload + // leaves its pack in the spool; a cancelled read leaves its pack owed, + // as a read that timed out would, and so in the bucket for the next + // start's reconcile. + upload_cancel_.Cancel(); + read_cancel_.Cancel(); wake_.notify_all(); if (thread_.joinable()) thread_.join(); // Only after the loop: its last cycle may still be indexing, and the @@ -313,6 +362,13 @@ bool CaptureStorageService::flush(double timeout_s) { const auto deadline = std::chrono::steady_clock::now() + std::chrono::duration_cast( std::chrono::duration(timeout_s)); + // The same instant for the cycles' upload cancel, on the steady clock + // steady_ns() reads. Never 0, which would arm none. + const uint64_t deadline_ns = std::max( + 1, static_cast( + std::chrono::duration_cast( + deadline.time_since_epoch()) + .count())); while (true) { { // A cycle in flight -- the loop's, stuck on a slow catalog -- must not @@ -321,7 +377,14 @@ bool CaptureStorageService::flush(double timeout_s) { if (!cycle.try_lock_until(deadline)) return false; if (!started_) throw std::logic_error("storage service: not started"); rethrow_if_failed(); - const bool drained = run_cycle().drained; + // stop() has begun: it cancelled the uploads for good and waits for + // this lock. Cycles now could only index, and would do so until the + // deadline. + if (upload_cancel_.cancelled_for_good()) return false; + // Its own cycle, bounded by the deadline, and without the reconcile: + // the reconcile lists the whole bucket and asks the catalog about + // every page, which no deadline bounds. The loop runs it. + const bool drained = run_cycle(deadline_ns, false).drained; if (!rejected_unreported_.empty()) { std::string message = "storage service: " + std::to_string(rejected_unreported_.size()) + @@ -353,12 +416,15 @@ void CaptureStorageService::rethrow_if_failed() const { void CaptureStorageService::loop() { uint64_t wait_ns = config_.poll_interval_ns; + bool waited_again = false; // the last wake gave way to a flush's cycle while (true) { + bool kicked = false; { std::unique_lock lock(wake_mutex_); wake_.wait_for(lock, std::chrono::nanoseconds(wait_ns), [this] { return stop_requested_ || kick_; }); if (stop_requested_) return; + kicked = kick_; kick_ = false; } { @@ -366,7 +432,27 @@ void CaptureStorageService::loop() { if (failure_) return; // another publisher holds the catalog } std::lock_guard cycle(cycle_mutex_); - run_cycle(); + { + std::lock_guard lock(wake_mutex_); + if (stop_requested_) return; + } + // The wait runs from the end of the last cycle, anyone's. A flush that + // held the lock while this waited for it has just run one; running + // another straight after it would only delay a stop() that follows the + // flush -- close()'s order -- by that cycle's catalog work, which stop() + // cannot cut. So wait out the rest of the interval first, unless a + // fresh lease asked for a cycle now -- once: flushes that keep coming + // do every cycle's work but the reconcile, which only this loop runs, + // so the next wake runs a cycle whatever they did. + const uint64_t since_ns = steady_ns() - last_cycle_end_ns_; + if (!kicked && !waited_again && last_cycle_end_ns_ != 0 && + since_ns < wait_ns) { + wait_ns -= since_ns; + waited_again = true; + continue; + } + waited_again = false; + run_cycle(0, true); // poll_interval * 2^streak, capped: flush() shares the streak, so an // outage it saw also slows the loop, and a success from either resets it. wait_ns = config_.poll_interval_ns; @@ -378,8 +464,28 @@ void CaptureStorageService::loop() { } } -CaptureStorageService::CycleOutcome CaptureStorageService::run_cycle() { +CaptureStorageService::CycleOutcome CaptureStorageService::run_cycle( + uint64_t deadline_ns, bool allow_reconcile) { CycleOutcome outcome; + // stop() has begun: it cancelled the uploads and the reads for good, so a + // cycle could only make catalog requests -- a lease claim, a replay guard + // -- that stop() would have to wait out. Neither drained nor failed. + if (upload_cancel_.cancelled_for_good()) { + outcome.failed = false; + outcome.cut_short = true; + last_cycle_end_ns_ = steady_ns(); + return outcome; + } + // flush()'s deadline cancels this cycle's uploads at that moment, and its + // index reads one catalog request timeout later: an upload cut short + // leaves its pack in the spool, but a read cut short leaves an uploaded + // pack owed, which only this process remembers, so what the cycle + // uploaded gets the time a catalog statement in flight would. The loop's + // cycles are cancelled only by stop(), which cancels both for good. + const ArmedDeadline upload_deadline(&upload_cancel_, deadline_ns); + const ArmedDeadline read_deadline( + &read_cancel_, + deadline_ns == 0 ? 0 : deadline_ns + read_grace_ns(config_.clickhouse)); // The catalog phase needs the lease. Without one -- quarantined after an // unknown outcome, or refused by another holder -- the cycle uploads // nothing either: whatever it uploaded it could only owe, in memory. @@ -390,14 +496,20 @@ CaptureStorageService::CycleOutcome CaptureStorageService::run_cycle() { } // Indexes refs, keeping whatever does not index owed: it is already gone // from the spool, so pending_index_ is the only record of it in-process. - const auto index_or_owe = [this, catalog](std::vector refs) { + // Those a cancel or the deadline left owed are counted in `deferred` + // (index_bounded). + size_t deferred = 0; + const auto index_or_owe = [this, catalog, deadline_ns, + &deferred](std::vector refs) { if (!catalog) { pending_index_.insert(pending_index_.end(), refs.begin(), refs.end()); return; } std::vector unindexed; try { - if (!refs.empty()) index_bounded(std::move(refs), &unindexed); + if (!refs.empty()) { + deferred += index_bounded(std::move(refs), &unindexed, deadline_ns); + } } catch (...) { pending_index_.insert(pending_index_.end(), unindexed.begin(), unindexed.end()); @@ -410,88 +522,141 @@ CaptureStorageService::CycleOutcome CaptureStorageService::run_cycle() { // 1. Retry what earlier cycles uploaded but could not index. While any // of it is still owed, the catalog is down or refusing: upload // nothing new, so new packs stay in the durable spool rather than - // joining a list that only this process remembers. + // joining a list that only this process remembers. Past a flush's + // deadline no batch starts but the first (index_bounded). if (catalog && !pending_index_.empty()) { std::vector owed; owed.swap(pending_index_); index_or_owe(std::move(owed)); } - // 2. Upload everything the sink has staged -- but only with the lease - // and nothing owed. Without the lease an uploaded pack could only be + // 2. Upload what the sink has staged -- but only with the lease and + // nothing owed. Without the lease an uploaded pack could only be // owed, and pending_index_ dies with the process: with // reconcile_on_start off, a crash would leave it in the bucket and // never in the catalog. Left in the spool it survives the crash. - dmi_store::UploadBatchResult batch; - if (catalog && pending_index_.empty()) { - batch = uploader_->UploadPending(-1); - } - std::vector to_index; - uint64_t uploaded_packs = 0; - uint64_t uploaded_bytes = 0; + // The spool is listed once, since listing hashes every staged pack. + // A cancel (the flush deadline, or stop()) stops the listing between + // packs, before the next hash, starts no upload, and cuts those in + // flight: those packs stay staged too. So a flush out of time before + // this point -- flush(0) -- hashes nothing: a listing already + // cancelled stops at the first staged pack, and finds an empty spool + // empty, which then reports drained. + // 3. Index them -- a chunk of indexer.max_packs at a time, each chunk + // indexed before the next is uploaded, so at most one chunk is ever + // out of the spool and not yet in the catalog, remembered by this + // process alone (stop() drops it). A cancel starts no further chunk, + // so what the flush deadline or stop() finds unsent stays in the + // spool, which any later start uploads, not owed in memory; a chunk + // left owed, or a lease lost meanwhile, stops the uploads as in step + // 1. What a chunk uploaded it indexes, a cancel of the uploads or + // not: past a flush's deadline, as its one batch past it + // (index_bounded). Its reads end at stop(), or one request timeout + // past a flush's deadline, leaving what they did not read owed. + // Every staged pack listed and uploaded: drained as far as flush() is + // concerned, since the listing came after the sink's flush. + bool uploaded_all = false; + bool listing_failed = false; + bool lost_lease = false; size_t upload_failures = 0; - for (size_t i = 0; i < batch.refs.size(); ++i) { - const dmi_store::PackRef& ref = batch.refs[i]; - if (!ref.pack_id.empty()) { - to_index.push_back({ref.pack_id, ref.store_id, ref.object_key, - ref.object_bytes, ref.checksum, ref.record_count}); - ++uploaded_packs; - uploaded_bytes += ref.object_bytes; + if (catalog && pending_index_.empty()) { + std::vector staged; + std::string error; + bool cut = false; + if (spool_.ListPending(&staged, &error, &upload_cancel_, &cut) != + dmi_store::SpoolStatus::kOk) { + record_error("spool listing failed: " + error); + listing_failed = true; + } else if (cut) { + outcome.cut_short = true; } else { - // A failed upload stays in the spool, so the next cycle retries it. - ++upload_failures; - if (i < batch.failures.size()) { - record_error("upload failed for " + batch.failures[i].object_key + - ": " + batch.failures[i].error); + const size_t chunk = + static_cast(std::max(1, config_.indexer.max_packs)); + size_t next = 0; + while (next < staged.size()) { + if (next != 0) { + if (!pending_index_.empty()) break; // the chunk before is owed + if (upload_cancel_.cancelled()) { + outcome.cut_short = true; + break; + } + // No request: a lease quarantined or refused meanwhile has + // already been dropped, and a fresh one is the next cycle's. + LeaseScope lease(this); + if (writer_.held_lease() == nullptr) { + lost_lease = true; + break; + } + } + const size_t end = std::min(staged.size(), next + chunk); + const dmi_store::UploadBatchResult batch = uploader_->UploadStaged( + std::vector(staged.begin() + next, + staged.begin() + end)); + next = end; + std::vector to_index; + uint64_t uploaded_packs = 0; + uint64_t uploaded_bytes = 0; + size_t failed_uploads = 0; + size_t cancelled_uploads = 0; + for (size_t i = 0; i < batch.refs.size(); ++i) { + const dmi_store::PackRef& ref = batch.refs[i]; + if (!ref.pack_id.empty()) { + to_index.push_back({ref.pack_id, ref.store_id, ref.object_key, + ref.object_bytes, ref.checksum, + ref.record_count}); + ++uploaded_packs; + uploaded_bytes += ref.object_bytes; + } else if (i < batch.failures.size() && + batch.failures[i].cancelled) { + ++cancelled_uploads; // still staged; not the pack's fault + } else { + // A failed upload stays in the spool, so a later cycle + // retries it. + ++failed_uploads; + if (i < batch.failures.size()) { + record_error("upload failed for " + + batch.failures[i].object_key + ": " + + batch.failures[i].error); + } + } + } + if (cancelled_uploads != 0) outcome.cut_short = true; + upload_failures += failed_uploads; + { + std::lock_guard lock(state_mutex_); + state_.uploaded_packs += uploaded_packs; + state_.uploaded_bytes += uploaded_bytes; + state_.upload_failures += failed_uploads; + state_.cancelled_uploads += cancelled_uploads; + } + index_or_owe(std::move(to_index)); } + uploaded_all = next == staged.size() && upload_failures == 0; } } - { - std::lock_guard lock(state_mutex_); - state_.uploaded_packs += uploaded_packs; - state_.uploaded_bytes += uploaded_bytes; - state_.upload_failures += upload_failures; - } - - // 3. Index them. - index_or_owe(std::move(to_index)); // 4. Reconcile on its interval, or when the pass at start() lost the - // lease before it finished. The lease thread keeps the lease alive. - if (catalog && + // lease before it finished -- in the loop's cycles only, and not + // once a cancel came. The lease thread keeps the lease alive. + if (allow_reconcile && catalog && !upload_cancel_.cancelled() && (reconcile_owed_ || (config_.reconcile_interval_ns > 0 && steady_ns() - last_reconcile_ns_ >= config_.reconcile_interval_ns))) { - reconcile(); - reconcile_owed_ = false; - last_reconcile_ns_ = steady_ns(); - } - - // Drained: nothing failed to upload, every uploaded pack is in the - // catalog, and nothing is pending. A batch that uploaded everything it - // listed is drained as far as flush() is concerned -- its listing came - // after the sink's flush -- so the spool is re-listed only when the batch - // was empty, which UploadPending also returns when its listing FAILED. - // A cycle that already failed is not drained whatever the spool holds, - // so it skips the listing: an empty spool lists for free, but a backlog - // would be re-hashed on every cycle of an outage. - // Without the lease nothing can be confirmed in the catalog, so the - // cycle is not drained, and it counts towards the backoff. - outcome.failed = - !catalog || upload_failures != 0 || !pending_index_.empty(); - bool nothing_pending = !batch.refs.empty(); - if (batch.refs.empty() && !outcome.failed) { - std::vector pending; - std::string error; - const bool listed = - spool_.ListPending(&pending, &error) == dmi_store::SpoolStatus::kOk; - if (!listed) { - record_error("spool listing failed: " + error); - outcome.failed = true; + if (reconcile()) { + reconcile_owed_ = false; + last_reconcile_ns_ = steady_ns(); } - nothing_pending = listed && pending.empty(); } - outcome.drained = nothing_pending && !outcome.failed; + + // Drained: every staged pack uploaded, every uploaded pack in the + // catalog, and nothing pending. Without the lease nothing can be + // confirmed in the catalog, so the cycle is not drained, and it counts + // towards the backoff. Packs a cancel or the deadline left owed make it + // cut short, not failed. + if (deferred != 0) outcome.cut_short = true; + outcome.failed = !catalog || listing_failed || lost_lease || + upload_failures != 0 || pending_index_.size() > deferred; + outcome.drained = uploaded_all && !outcome.failed && !outcome.cut_short; std::lock_guard lock(state_mutex_); ++state_.cycles; } catch (const CatalogError& exc) { @@ -508,48 +673,91 @@ CaptureStorageService::CycleOutcome CaptureStorageService::run_cycle() { std::lock_guard lock(state_mutex_); state_.pending_index = pending_index_.size(); } - failure_streak_ = outcome.failed ? std::min(failure_streak_ + 1, 32) : 0; + // A cycle a cancel cut short proved nothing about the store or the + // catalog: it resets the backoff only by succeeding, and grows it only by + // failing. + if (outcome.failed) { + failure_streak_ = std::min(failure_streak_ + 1, 32); + } else if (!outcome.cut_short) { + failure_streak_ = 0; + } + last_cycle_end_ns_ = steady_ns(); return outcome; } -void CaptureStorageService::index_bounded(std::vector refs, - std::vector* unindexed) { +size_t CaptureStorageService::index_bounded(std::vector refs, + std::vector* unindexed, + uint64_t deadline_ns) { // Batches the indexer can take: at most max_packs, and halved again when the // rendered descriptors exceed max_estimated_bytes. The two bounds are // independent, so packs that each fit can still overflow together. const size_t max_packs = static_cast(std::max(1, config_.indexer.max_packs)); + // A stack, the next batch at the back. The chunks go on it in reverse, so + // the pass runs them front to back and its first batch is a full one: + // past a flush's deadline that batch is the only one, and the remainder + // chunk can be a single pack. A split pushes its halves the same way. std::vector> work; - for (size_t i = 0; i < refs.size(); i += max_packs) { - work.emplace_back(refs.begin() + i, - refs.begin() + std::min(refs.size(), i + max_packs)); + for (size_t chunk = (refs.size() + max_packs - 1) / max_packs; chunk-- > 0;) { + const size_t begin = chunk * max_packs; + work.emplace_back(refs.begin() + begin, + refs.begin() + std::min(refs.size(), begin + max_packs)); } + // What is still queued, in the order it would have run. + const auto drain_queued = [&](uint64_t* count) { + for (auto queued = work.rbegin(); queued != work.rend(); ++queued) { + *count += queued->size(); + unindexed->insert(unindexed->end(), queued->begin(), queued->end()); + } + work.clear(); + }; // A batch that threw indexed nothing, and neither did anything still queued. const auto give_up = [&](std::vector& failed, const std::string& message) { uint64_t count = failed.size(); unindexed->insert(unindexed->end(), failed.begin(), failed.end()); - for (std::vector& queued : work) { - count += queued.size(); - unindexed->insert(unindexed->end(), queued.begin(), queued.end()); - } - work.clear(); + drain_queued(&count); record_error(message); std::lock_guard lock(state_mutex_); state_.index_failures += count; }; + // Left owed and not counted as failed: a cancel or the deadline cut the + // pass short. + size_t deferred = 0; + const auto defer = [&](std::vector& cut) { + uint64_t count = cut.size(); + unindexed->insert(unindexed->end(), cut.begin(), cut.end()); + drain_queued(&count); + deferred += static_cast(count); + }; + bool first = true; while (!work.empty()) { + // Once the reads are cut -- stop(), or a flush one request timeout + // past its deadline -- no batch starts, the first included: it could + // read nothing, and would only send catalog statements (the replay + // guard) whose answer is thrown away, and which stop() waits for. + // Past the deadline no batch starts but the first: the pass overruns + // it by one batch at most, never by the rest of its work, however many + // batches that is. Catalog statements are never cut, so this is where + // the pass can stop. + if (read_cancel_.cancelled() || + (!first && deadline_ns != 0 && steady_ns() >= deadline_ns)) { + std::vector none; + defer(none); + break; + } + first = false; std::vector batch = std::move(work.back()); work.pop_back(); IndexResultData result; try { // The lease lock for the catalog phases only. Between them the pass // reads its packs from the object store, where a stalled GET is bound - // only by the S3 client's own timeouts, and must not keep the lease - // thread from renewing meanwhile. Nothing here stamps a renewal: the - // schedule follows the lease itself (renew_lease_if_due), which only - // a claim that confirmed moves -- a pass that published nothing, over - // packs already committed, renewed nothing. + // by the S3 client's own timeouts and read_cancel_, and must not keep + // the lease thread from renewing meanwhile. Nothing here stamps a + // renewal: the schedule follows the lease itself (renew_lease_if_due), + // which only a claim that confirmed moves -- a pass that published + // nothing, over packs already committed, renewed nothing. IndexPlan plan; { LeaseScope lease(this); @@ -558,6 +766,19 @@ void CaptureStorageService::index_bounded(std::vector refs, indexer_.read(&plan); LeaseScope lease(this); result = indexer_.commit(&plan); + } catch (const StoreUnavailableError& exc) { + // The object store did not answer for a pack. Every later read of + // the pass would cost the client's timeouts and fail alike, so the + // pass ends here, like a batch that threw: this batch and the rest + // stay owed, and nothing counts against the packs themselves, since + // an outage is no fault of theirs. A read a cancel cut -- stop(), or + // a flush out of time -- is not a failure at all. + if (exc.cancelled()) { + defer(batch); + return deferred; + } + give_up(batch, std::string("index failed: ") + exc.what()); + return deferred; } catch (const CatalogError& exc) { if (exc.kind() == CatalogError::Kind::kBatchTooLarge && batch.size() > 1) { const size_t middle = batch.size() / 2; @@ -577,10 +798,10 @@ void CaptureStorageService::index_bounded(std::vector refs, const bool lease_lost = is_lease_refusal(exc); give_up(batch, std::string("index failed: ") + exc.what()); if (lease_lost) throw; - return; + return deferred; } catch (const std::exception& exc) { give_up(batch, std::string("index failed: ") + exc.what()); - return; + return deferred; } // A failure the indexer reports for one pack is that pack's own (it was // read and refused), unlike a batch that threw, which an outage explains. @@ -607,6 +828,7 @@ void CaptureStorageService::index_bounded(std::vector refs, state_.indexed_rows += result.indexed_rows; state_.index_failures += result.failed_packs; } + return deferred; } void CaptureStorageService::reject(const PackRefData& ref, @@ -620,20 +842,23 @@ void CaptureStorageService::reject(const PackRefData& ref, ++state_.index_failures; } -void CaptureStorageService::reconcile() { +bool CaptureStorageService::reconcile() { // List every pack under the prefix, ask the catalog which it already // committed, and HEAD only the rest -- so a steady-state pass over a large // bucket costs listing pages and one catalog query per page, not a HEAD per - // object. + // object. stop() ends it between requests: what it has not reached is + // still in the bucket, for the next pass. std::string token; uint64_t found = 0; uint64_t skipped = 0; uint64_t head_errors = 0; do { + if (upload_cancel_.cancelled()) return false; dmi_store::ListResult page; std::string error; if (!s3_.ListObjects(config_.reconcile_prefix, "", 1000, token, &page, &error)) { + if (read_cancel_.cancelled()) return false; // stop() cut it throw std::runtime_error("reconcile: listing failed: " + error); } token = page.truncated ? page.next_token : ""; @@ -651,6 +876,8 @@ void CaptureStorageService::reconcile() { packs.push_back(&object); } if (packs.empty()) continue; + // A listing answered just as stop() came: nothing after it can be read. + if (upload_cancel_.cancelled()) return false; std::set committed; { LeaseScope lease(this); @@ -660,9 +887,11 @@ void CaptureStorageService::reconcile() { std::vector missing; for (size_t i = 0; i < packs.size(); ++i) { if (committed.count(identities[i]) != 0) continue; + if (upload_cancel_.cancelled()) return false; const dmi_store::ListedObject& object = *packs[i]; std::string head_error; const dmi_store::ObjectHead head = s3_.HeadObject(object.key, &head_error); + if (!head_error.empty() && read_cancel_.cancelled()) return false; if (!head_error.empty()) { // Unread, not foreign: the object may well be a pack, so the pass // reports it rather than counting it as skipped. The next pass @@ -706,6 +935,7 @@ void CaptureStorageService::reconcile() { state_.reconciled_packs += found; state_.reconcile_skipped_objects += skipped; state_.reconcile_head_errors += head_errors; + return true; } void CaptureStorageService::keep_lease() { diff --git a/native/csrc/catalog/storage_service.h b/native/csrc/catalog/storage_service.h index 74d697ee3..54c3fa1f0 100644 --- a/native/csrc/catalog/storage_service.h +++ b/native/csrc/catalog/storage_service.h @@ -4,14 +4,17 @@ // NativePackSink -> spool (the sink, in the capture process) // spool -> SpoolUploader -> object store -> NativeIndexer -> catalog (here) // -// One background thread runs a cycle: upload everything pending, index what -// was uploaded, keep the publisher lease alive, and periodically reconcile the -// bucket against the catalog. The conformance drivers exercise each of these -// pieces; this is what composes them outside a test. +// One background thread runs a cycle: upload what is pending and index it, a +// chunk at a time, keep the publisher lease alive, and periodically reconcile +// the bucket against the catalog. The conformance drivers exercise each of +// these pieces; this is what composes them outside a test. // // SpoolUploader removes a pack from the spool the moment its upload is // verified, before anything indexes it, so the spool alone cannot say what is -// still owed to the catalog. Two things cover the gap: +// still owed to the catalog. A cycle uploads a chunk of indexer.max_packs at +// a time and indexes it before it uploads the next, so at most one chunk is +// in that gap at once, and a cycle cut short leaves the rest in the spool. +// Two things cover the gap: // - in-process, a pack whose indexing fails stays on a retry list, and // flush() does not report drained until that list is empty. Nothing new // is uploaded while it is not, so an outage leaves new packs in the @@ -56,6 +59,27 @@ // (or, once, until the quarantine it left is over), and lease requests that // keep timing out say which knobs bound them (snapshot().lease_timeout_error). The constructor refuses a clock skew // that leaves a renewal too little time to finish. +// +// Cancelled object-store work. The uploads go through an S3 client of their +// own that shares one Cancellation with the uploader; the index reads and +// the reconcile go through the other client, which has a Cancellation of +// its own. stop() cancels both for good. A flush arms a deadline on both +// for the cycle it runs: the uploads' at the flush deadline, the reads' +// one catalog request timeout past it, since a pack cut short on upload +// stays in the durable spool but one uploaded and not yet indexed is owed, +// and only this process remembers it. A cancel starts no further request, +// aborts a transfer in flight (a multipart upload is then aborted, one +// attempt bounded by 5 s) and ends a retry backoff; a cancelled upload +// leaves its pack in the spool, a cancelled read leaves its pack owed. The +// uploads' cancel also stops a spool listing between packs, since listing +// hashes every staged pack and a backlog would hold it for as long. +// Catalog statements are never cut mid-flight: each is bounded by the +// client's request timeout, or under the lease by the lease deadline. +// An index pass ends at its first failure: a batch that threw, or a pack +// the object store did not answer for (a transport error, a timeout, a +// retryable status on every attempt), which leaves the rest of the pass +// owed without counting against the packs. Only a pack the store answered +// for and the indexer refused counts towards max_index_attempts. #pragma once #include @@ -72,6 +96,7 @@ #include "catalog/catalog_writer.h" #include "catalog/clickhouse_client.h" #include "catalog/indexer.h" +#include "store/cancel.h" #include "store/s3_client.h" #include "store/spool.h" #include "store/uploader.h" @@ -131,14 +156,14 @@ struct StorageServiceConfig { // take the lease and sweep while the first is still writing; one on // another (database, table_prefix) never meets the lease at all. Its // Recover() also lists the first's sealed packs, which the first may - // upload too. A cycle checks the lease once, before its upload batch, and - // UploadPending does not stop when the lease is lost: a holder that is - // quarantined, or refused a renewal or publish, while a batch is in - // flight finishes that batch (which can outlast the TTL), and only its - // later cycles upload nothing while it holds no lease. One whose catalog - // requests stall gives the lease up at its deadline, before its row + // upload too. A cycle checks the lease before each chunk it uploads, and + // the uploads in flight do not stop when the lease is lost: a holder that + // is quarantined, or refused a renewal or publish, while a chunk is in + // flight finishes that chunk (which can outlast the TTL), and uploads + // nothing more while it holds no lease. One whose catalog requests + // stall gives the lease up at its deadline, before its row // lapses: requests made under the lease are cut off there, and a cycle's - // check abandons a lease past it (LeaseScope), so no batch starts after + // check abandons a lease past it (LeaseScope), so no chunk starts after // that. Only a holder whose whole process stalls keeps the lease locally // past its row: until it resumes and next checks, or -- after a system // suspend, which the steady clock the deadline runs on does not count -- @@ -155,6 +180,9 @@ struct StorageServiceSnapshot { uint64_t uploaded_packs = 0; uint64_t uploaded_bytes = 0; uint64_t upload_failures = 0; + // Uploads a flush deadline or stop() cancelled, the pack left staged. + // Not failures: nothing was wrong with the pack or the store. + uint64_t cancelled_uploads = 0; uint64_t indexed_packs = 0; uint64_t indexed_rows = 0; uint64_t index_failures = 0; @@ -209,13 +237,44 @@ class CaptureStorageService { // Run cycles until one finds the spool empty with every uploaded pack // indexed, or the timeout passes. Call after the sink's own flush, so // everything it will stage is already staged. Returns false on timeout, - // including while a cycle already in flight outlives the deadline; it can - // overrun only by its own last cycle, whose requests are all bounded. + // including while a cycle already in flight outlives the deadline, and + // once stop() has begun. The cycles it runs honour the deadline: they + // skip the reconcile (the loop runs it), and at the deadline their + // uploads are cancelled, each pack cut short left in the spool, and no + // further chunk starts, so what is not uploaded by then stays in the + // spool. The chunk a cycle has uploaded it still indexes, since until + // then only this process remembers it -- but past the deadline no index + // batch starts after the first: that batch's object-store reads are cut + // one catalog request timeout past the deadline, and its catalog + // statements are never cut mid-flight, each bounded by the client's + // request timeout (under the lease, by the lease deadline). What it + // leaves unindexed stays owed, for the loop or a later flush. A pass + // ends at its first failure, so a flush overruns its deadline by about + // one request timeout against a catalog or an object store that stopped + // answering, and by one batch of statements against a slow catalog that + // still answers. A multipart upload the deadline cuts adds its abort, + // which no request timeout bounds: up to about a second for the stalled + // transfer to see the cancel (libcurl's progress poll), then one abort + // request of at most 5 s that nothing cuts -- past one request timeout + // when that is under about 6 s. At zero it still runs one cycle, which + // indexes one batch of what earlier cycles owe and uploads nothing. // Throws, once, if packs were set aside since the last flush: they are in // the object store but can never reach the catalog. bool flush(double timeout_s); - // Stop the background cycle and release the lease. Does not flush. + // Stop the background cycle and release the lease. Does not flush. Its + // object-store work is cancelled first: a spool listing stops between + // packs; an upload in flight is aborted, its pack left in the spool for + // the next start, with every pack no upload has reached; an index read + // in flight is cut, its pack left owed -- which stop() drops, so it waits + // in the bucket for a start's reconcile (reconcile_on_start); that is at + // most the one chunk a cycle has uploaded and not yet indexed. A retry + // backoff ends, and so does the reconcile. What stop() still waits for + // is the catalog work in flight, never cut mid-flight, the abort of a + // multipart upload it cut (one attempt, 5 s at most), and the lease + // release, each catalog request bounded by the client's request timeout + // (under the lease, by the lease deadline); the lease renews until the + // loop is done. void stop(); StorageServiceSnapshot snapshot() const; @@ -228,6 +287,9 @@ class CaptureStorageService { struct CycleOutcome { bool drained = false; // nothing pending and nothing failed bool failed = true; // an upload or index failed, or the cycle threw + // A cancel left packs in the spool: neither drained nor a failure, so + // it moves the backoff neither way. + bool cut_short = false; }; // Holds lease_mutex_ for a stretch of catalog work, and bounds every @@ -252,12 +314,21 @@ class CaptureStorageService { void sweep_and_reconcile_at_start(); // Stops the lease thread and waits for it. void stop_lease_thread(); - CycleOutcome run_cycle(); // requires cycle_mutex_ + // One cycle. A non-zero deadline_ns (steady ns) cancels its uploads at + // that moment -- flush()'s -- and allow_reconcile false skips the + // periodic and the owed reconcile. Requires cycle_mutex_. + CycleOutcome run_cycle(uint64_t deadline_ns, bool allow_reconcile); // Indexes refs in bounded batches, appending every ref that did not index - // to *unindexed. Only a lost lease propagates; other failures are recorded. - void index_bounded(std::vector refs, - std::vector* unindexed); - void reconcile(); + // to *unindexed. Only a lost lease propagates; other failures are + // recorded. A non-zero deadline_ns (steady ns) starts no batch past it + // but the first, and none starts once read_cancel_ is cancelled. Returns + // how many of *unindexed a cancel or the deadline left there: owed, but + // not failed. + size_t index_bounded(std::vector refs, + std::vector* unindexed, + uint64_t deadline_ns = 0); + // False when stop() cut it short, between two of its requests. + bool reconcile(); void keep_lease(); // the lease thread's body void renew_lease_if_due(); // requires lease_mutex_ // LeaseScope's before_request hook: renew_lease_if_due() before each @@ -292,7 +363,16 @@ class CaptureStorageService { void latch_failure(std::exception_ptr failure, const std::string& message); const StorageServiceConfig config_; + // Cancels the uploads: stop() for good, a flush's cycle at its deadline. + // Before the clients that point at it. + dmi_store::Cancellation upload_cancel_; + // Cancels the index reads and the reconcile's requests: stop() for good, + // a flush's cycle one catalog request timeout past its deadline. + dmi_store::Cancellation read_cancel_; + // Indexing and the reconcile read through s3_, which read_cancel_ cuts; + // the uploader writes through upload_s3_, which upload_cancel_ does. dmi_store::S3Client s3_; + dmi_store::S3Client upload_s3_; std::shared_ptr clickhouse_; CatalogWriter writer_; NativeIndexer indexer_; @@ -328,6 +408,10 @@ class CaptureStorageService { // runs one once it holds a lease again. Guarded by cycle_mutex_. bool reconcile_owed_ = false; int failure_streak_ = 0; // consecutive failed cycles, for the backoff + // steady ns at which the last cycle -- the loop's or a flush's -- ended; + // 0 before the first. The loop waits its interval from it. Guarded by + // cycle_mutex_. + uint64_t last_cycle_end_ns_ = 0; // Uploaded, so gone from the spool, but not yet in the catalog. std::vector pending_index_; std::map index_attempts_; // by pack id diff --git a/native/csrc/sink/bindings_sink.cpp b/native/csrc/sink/bindings_sink.cpp index 696256494..e9bd394aa 100644 --- a/native/csrc/sink/bindings_sink.cpp +++ b/native/csrc/sink/bindings_sink.cpp @@ -13,7 +13,9 @@ #include #include +#include #include +#include #include #include @@ -25,6 +27,28 @@ namespace py = pybind11; namespace { +// Parks every stage of a sink's spool until Open(), or for max_hold at +// most, so a test can wedge the sink's pipeline without a hung filesystem +// -- and still gets it back if what it tests never lets go. +struct StageGate { + std::mutex mutex; + std::condition_variable cv; + bool open = false; + std::chrono::nanoseconds max_hold{0}; + + void Wait() { + std::unique_lock lock(mutex); + cv.wait_for(lock, max_hold, [this] { return open; }); + } + void Open() { + { + std::lock_guard lock(mutex); + open = true; + } + cv.notify_all(); + } +}; + ring::PayloadSlice ParseSlice(const py::dict& row) { ring::PayloadSlice slice; slice.offset_bytes = row["offset"].cast(); @@ -157,7 +181,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { uint64_t max_pack_bytes, uint64_t max_pack_records, uint64_t max_linger_ns, uint64_t spool_max_bytes, const std::string& overload, - std::optional admission_timeout_s) { + std::optional admission_timeout_s, + double release_flush_timeout_s) { dmi_sink::SinkConfig config; config.overload = ParseOverload(overload); config.admission_timeout_s = @@ -170,9 +195,17 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { config.max_pack_bytes = max_pack_bytes; config.max_pack_records = max_pack_records; config.max_linger_ns = max_linger_ns; + if (!std::isfinite(release_flush_timeout_s) || + release_flush_timeout_s < 0.0) { + throw py::value_error( + "release_flush_timeout_s must be a finite, non-negative " + "number"); + } auto sink = std::make_unique(config); return std::make_shared( - std::move(sink), layout); + std::move(sink), layout, + std::chrono::ceil( + std::chrono::duration(release_flush_timeout_s))); }), py::arg("spool_root"), py::arg("layout"), py::arg("num_workers") = 1, @@ -185,7 +218,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // SinkConfig's own defaults: the Python NativeSinkConfig, which // the ring-fed sink is built from, picks block with 2 s. py::arg("overload") = "drop_newest", - py::arg("admission_timeout_s") = py::none()) + py::arg("admission_timeout_s") = py::none(), + // How long releasing the sink from its engine waits for the open + // pack to reach the spool; 0 turns that flush off. + py::arg("release_flush_timeout_s") = + std::chrono::duration( + dmi_sink::NativePackSink::kDefaultReleaseFlushTimeout) + .count()) .def("attach", [](std::shared_ptr self) { // Simulates engine ownership for tests (the real engine takes @@ -215,12 +254,36 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { std::chrono::duration(timeout_s))); }) .def("rethrow_if_failed", &dmi_sink::NativePackSink::rethrow_if_failed) + // Test seam, not API: from now on each stage of the sink's spool + // waits, before it writes anything, until the returned function is + // called or max_hold_s has passed. Wedges the pipeline the way a hung + // filesystem would (tests/test_native_sink_release.py). + .def("_hold_stages_for_testing", + [](dmi_sink::NativePackSink& self, double max_hold_s) { + if (!(max_hold_s > 0.0 && max_hold_s <= 3600.0)) { + throw py::value_error("max_hold_s must be in (0, 3600]"); + } + auto gate = std::make_shared(); + gate->max_hold = + std::chrono::duration_cast( + std::chrono::duration(max_hold_s)); + self.sink_for_testing().SpoolForTesting().SetStageHookForTesting( + [gate] { gate->Wait(); }); + return py::cpp_function([gate] { gate->Open(); }); + }, + py::arg("max_hold_s")) .def_property_readonly("layout", &dmi_sink::NativePackSink::layout) .def_property_readonly( "overload", [](const dmi_sink::NativePackSink& self) { return std::string(OverloadName(self.sink().config().overload)); }) + .def_property_readonly( + "release_flush_timeout_s", + [](const dmi_sink::NativePackSink& self) { + return std::chrono::duration(self.release_flush_timeout()) + .count(); + }) .def_property_readonly( "admission_timeout_s", [](const dmi_sink::NativePackSink& self) -> std::optional { diff --git a/native/csrc/sink/native_pack_sink.cpp b/native/csrc/sink/native_pack_sink.cpp index 9687c73e7..5f8d9bcbf 100644 --- a/native/csrc/sink/native_pack_sink.cpp +++ b/native/csrc/sink/native_pack_sink.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -104,10 +105,16 @@ std::string AtenDtypeName(int32_t scalar_type) { } NativePackSink::NativePackSink(std::unique_ptr sink, - std::string layout) - : sink_(std::move(sink)), layout_(std::move(layout)) { + std::string layout, + Duration release_flush_timeout) + : sink_(std::move(sink)), + layout_(std::move(layout)), + release_flush_timeout_(release_flush_timeout) { if (!sink_) invalid("PackSink is required"); if (layout_.empty()) invalid("layout is required"); + if (release_flush_timeout_ < Duration::zero()) { + invalid("release_flush_timeout must not be negative"); + } std::string error; const std::string start_error = sink_->Start(&error); if (!start_error.empty()) invalid("sink start failed: " + start_error); @@ -250,6 +257,32 @@ bool NativePackSink::flush_and_wait(Duration timeout) { return true; } +void NativePackSink::on_engine_release() noexcept { + if (release_flush_timeout_ == Duration::zero()) return; + try { + std::string error; + const bool flushed = sink_->Flush( + std::chrono::duration(release_flush_timeout_).count(), &error); + if (flushed) return; + // Once released, rethrow_if_failed refuses the sink as not attached, + // and a timeout latches nothing: this line is what says what became of + // the open pack (a failure also counts in the snapshot's failures). + const std::string why = + error.empty() ? "timed out after " + + std::to_string(release_flush_timeout_.count()) + + " ms" + : error; + std::fprintf(stderr, + "NativePackSink: the open pack did not reach the spool when " + "the engine released the sink: %s\n", + why.substr(0, 512).c_str()); + std::fflush(stderr); + } catch (...) { + // noexcept: a release cannot fail. The sink's own Flush does not + // throw; this guards against allocation failure in the message. + } +} + void NativePackSink::rethrow_if_failed() const { if (!engine_owned()) invalid("sink is not attached to a RingEngine"); std::string error = sink_->LastError(); diff --git a/native/csrc/sink/native_pack_sink.h b/native/csrc/sink/native_pack_sink.h index 141450848..97b2dfeaa 100644 --- a/native/csrc/sink/native_pack_sink.h +++ b/native/csrc/sink/native_pack_sink.h @@ -10,6 +10,7 @@ #ifndef DMI_SINK_NATIVE_PACK_SINK_H_ #define DMI_SINK_NATIVE_PACK_SINK_H_ +#include #include #include #include @@ -27,9 +28,16 @@ std::string AtenDtypeName(int32_t scalar_type); class NativePackSink final : public ring::RecordSink { public: + // How long the release backstop (on_engine_release) waits for the open + // pack to reach the spool, unless the constructor is given another bound. + static constexpr std::chrono::milliseconds kDefaultReleaseFlushTimeout{ + 30'000}; + // Takes ownership of a started-or-unstarted PackSink; Start()s it here so // construction failure (bad spool dir) throws instead of wedging attach. - NativePackSink(std::unique_ptr sink, std::string layout); + // release_flush_timeout bounds the release backstop; zero turns it off. + NativePackSink(std::unique_ptr sink, std::string layout, + Duration release_flush_timeout = kDefaultReleaseFlushTimeout); ~NativePackSink() override; NativePackSink(const NativePackSink&) = delete; @@ -53,11 +61,27 @@ class NativePackSink final : public ring::RecordSink { std::optional admission_bound() const override; const PackSink& sink() const { return *sink_; } + // Test seam: the PackSink itself, so a test can park its stages + // (PackSink::SpoolForTesting) and wedge the pipeline. + PackSink& sink_for_testing() { return *sink_; } const std::string& layout() const { return layout_; } + Duration release_flush_timeout() const { return release_flush_timeout_; } protected: void on_engine_acquire() override {} - void on_engine_release() noexcept override {} + // The backstop for a ring that stops without a flush: RingEngine::stop + // drains its record worker into submit() and then releases the sink, and + // until a flush seals it the open pack is only in memory -- for up to + // max_linger_ns, or until this object dies. So the release flushes, + // bounded by release_flush_timeout -- a flush_and_wait in flight on + // another thread included, which it waits for only that long -- and the + // open pack reaches the spool as one .ready file. It cannot throw (it + // runs in RingEngine::stop and its destructor): a flush that fails or + // times out writes one line to stderr, and that line is the only report + // of a timeout. A pipeline failure also counts in snapshot()["failures"]; + // rethrow_if_failed reports it only once the sink is attached again, + // since a released sink refuses the call as not attached. + void on_engine_release() noexcept override; private: // "NativePackSink: pipeline reported ..." for the loss counters that @@ -66,6 +90,7 @@ class NativePackSink final : public ring::RecordSink { std::unique_ptr sink_; const std::string layout_; + const Duration release_flush_timeout_; // Counters at construction. Losses are judged against it, as the // reference adapter judges them against its pipeline's baseline. SinkSnapshot baseline_; diff --git a/native/csrc/sink/pack_sink.cpp b/native/csrc/sink/pack_sink.cpp index adee3592e..f51e129c6 100644 --- a/native/csrc/sink/pack_sink.cpp +++ b/native/csrc/sink/pack_sink.cpp @@ -265,7 +265,19 @@ bool PackSink::Flush(double timeout_s, std::string* error) { return false; } } - std::lock_guard flush_lock(flush_mutex_); + // A flush in flight holds this for as long as its own timeout. This one + // waits for it only until its own deadline, so the release backstop's + // bounded flush does not wait out a concurrent flush_and_wait's. + std::unique_lock flush_lock(flush_mutex_, std::defer_lock); + if (deadline < 0) { + flush_lock.lock(); + } else if (!flush_lock.try_lock_until( + std::chrono::steady_clock::now() + + std::chrono::duration_cast( + std::chrono::duration( + std::max(0.0, deadline - NowS()))))) { + return false; // timed out, as a wait on the barrier would have + } { std::unique_lock lock(mutex_); if (!latched_error_.empty()) { diff --git a/native/csrc/sink/pack_sink.h b/native/csrc/sink/pack_sink.h index daed49088..214f7478c 100644 --- a/native/csrc/sink/pack_sink.h +++ b/native/csrc/sink/pack_sink.h @@ -153,7 +153,8 @@ class PackSink { // Persist everything admitted before this call. False on timeout; the // in-flight barrier is kept for the next call to reuse. timeout_s < 0 - // waits forever. Error string set on pipeline failure. + // waits forever. The timeout covers waiting for a flush already in + // flight on another thread too. Error string set on pipeline failure. bool Flush(double timeout_s, std::string* error); // Terminal close: drains, joins, returns the final snapshot. SinkSnapshot Close(double timeout_s, std::string* error); @@ -212,7 +213,10 @@ class PackSink { std::vector workers_; std::vector stagers_; - std::mutex flush_mutex_; + // Serialises flushes. Timed: a flush waits for another in flight only + // until its own deadline, so a short flush (the release backstop's) is + // not held for a long one's (flush_and_wait's). + std::timed_mutex flush_mutex_; std::shared_ptr pending_; std::string latched_error_; diff --git a/native/csrc/store/cancel.h b/native/csrc/store/cancel.h new file mode 100644 index 000000000..2b887958f --- /dev/null +++ b/native/csrc/store/cancel.h @@ -0,0 +1,112 @@ +// Cancellation for object-store work: a flag set once to stop for good, and +// a deadline armed for one bounded stretch. +// +// The storage service owns two. One cuts its uploads: it goes to the S3 +// client and the uploader that do them, and a flush arms it at its +// deadline. The other cuts its index reads and the reconcile's requests: +// it goes to the S3 client those read through, and a flush arms it one +// catalog request timeout past its deadline. stop() cancels both for good, +// so a stalled PUT or GET or a retry backoff cannot hold it. What a cancel +// cuts short is left where it was: a pack whose upload was cancelled stays +// in the spool, which is where a pack waits for an object store; one +// uploaded but not yet indexed, whose read was cut, stays owed to the +// catalog (in memory, or in the bucket for a start's reconcile once stop() +// drops it). +// +// Header-only, so every target that builds the S3 client gets it without a +// source list to keep in step. + +#ifndef DMI_STORE_CANCEL_H_ +#define DMI_STORE_CANCEL_H_ + +#include +#include +#include +#include +#include +#include + +namespace dmi_store { + +class Cancellation { + public: + Cancellation() = default; + Cancellation(const Cancellation&) = delete; + Cancellation& operator=(const Cancellation&) = delete; + + // Cancels everything from now on, until Reset(), and wakes every sleeper. + void Cancel() { + std::lock_guard lock(mutex_); + cancelled_.store(true, std::memory_order_release); + ++generation_; + cv_.notify_all(); + } + + // Clears Cancel() and the deadline. + void Reset() { + std::lock_guard lock(mutex_); + cancelled_.store(false, std::memory_order_release); + deadline_ns_.store(0, std::memory_order_release); + ++generation_; + cv_.notify_all(); + } + + // Arms a deadline on the steady clock, in ns (NowNs()), or disarms it + // with 0: past it, cancelled() is true. + void set_deadline(uint64_t steady_ns) { + std::lock_guard lock(mutex_); + deadline_ns_.store(steady_ns, std::memory_order_release); + ++generation_; + cv_.notify_all(); + } + + // Cancel()ed, or past the armed deadline. Cheap: libcurl's progress + // callback asks it while a transfer runs. + bool cancelled() const { + if (cancelled_.load(std::memory_order_acquire)) return true; + const uint64_t deadline = deadline_ns_.load(std::memory_order_acquire); + return deadline != 0 && NowNs() >= deadline; + } + + // Cancel()ed, whatever the deadline. + bool cancelled_for_good() const { + return cancelled_.load(std::memory_order_acquire); + } + + // Sleeps for `duration`, or until cancelled, whichever comes first -- a + // deadline armed while it sleeps included. True if it slept it out. + bool SleepFor(std::chrono::nanoseconds duration) const { + const uint64_t end = + NowNs() + static_cast(std::max(0, duration.count())); + std::unique_lock lock(mutex_); + while (true) { + if (cancelled()) return false; + const uint64_t now = NowNs(); + if (now >= end) return true; + uint64_t wake = end; + const uint64_t deadline = deadline_ns_.load(std::memory_order_acquire); + if (deadline != 0 && deadline < wake) wake = deadline; + const uint64_t generation = generation_; + cv_.wait_for(lock, std::chrono::nanoseconds(wake - now), + [&] { return generation_ != generation; }); + } + } + + static uint64_t NowNs() { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); + } + + private: + mutable std::mutex mutex_; + mutable std::condition_variable cv_; + std::atomic cancelled_{false}; + std::atomic deadline_ns_{0}; + uint64_t generation_ = 0; // guarded by mutex_; moves on every change +}; + +} // namespace dmi_store + +#endif // DMI_STORE_CANCEL_H_ diff --git a/native/csrc/store/conformance_store.cpp b/native/csrc/store/conformance_store.cpp index 3456bf043..51fae7ee2 100644 --- a/native/csrc/store/conformance_store.cpp +++ b/native/csrc/store/conformance_store.cpp @@ -16,6 +16,14 @@ // "continuation":"..."} -> {"ok":true,"truncated":bool,"next_token":"...", // "objects":[{"key":"...","size":N,"etag":"..."}...],"attempts":N} // Errors: {"ok":false,"what":"..."}. +// +// Any op may carry "cancel_after_ms":N: the client (and the uploader) get a +// Cancellation, cancelled N ms after the op starts unless it has finished, +// and the response says whether that fired ("cancelled"; for upload_one, +// whether the cancel ended the upload, and per failure for upload_pending). +// With "cancel_uploader_only":true only the uploader gets it, so a request +// the cancel comes in during runs to its own end, as one whose answer +// lands in the gap before libcurl next asks the Cancellation would. #include "s3_client.h" @@ -23,9 +31,13 @@ #include "spool.h" #include "uploader.h" +#include +#include #include #include +#include #include +#include #include @@ -150,11 +162,48 @@ dmi_store::UploaderConfig ReadUploaderConfig(const std::string& line) { // attempt (like botocore), and only its exhaustion surfaces here. const int64_t attempts = Integer(line, "upload_max_attempts"); if (attempts > 0) config.max_attempts = static_cast(attempts); + // The uploader's own backoff between attempts, base * 2^attempt capped + // at max_backoff_s: long enough, a cancel shows whether it wakes it. + const int64_t backoff_ms = Integer(line, "upload_base_backoff_ms"); + if (backoff_ms > 0) config.base_backoff_s = backoff_ms / 1000.0; config.store_id = jc::FindString(line, "store_id"); if (config.store_id.empty()) config.store_id = "s3"; return config; } +// Cancels `cancel` a set time after Arm(), unless Finish() comes first. +struct Canceller { + dmi_store::Cancellation cancel; + std::mutex mutex; + std::condition_variable cv; + bool done = false; + bool fired = false; + std::thread thread; + + void Arm(int64_t after_ms) { + thread = std::thread([this, after_ms] { + std::unique_lock lock(mutex); + if (!cv.wait_for(lock, std::chrono::milliseconds(after_ms), + [this] { return done; })) { + fired = true; + cancel.Cancel(); + } + }); + } + + bool Finish() { + { + std::lock_guard lock(mutex); + done = true; + } + cv.notify_all(); + if (thread.joinable()) thread.join(); + return fired; + } + + ~Canceller() { Finish(); } +}; + int main() { std::string line; std::ios::sync_with_stdio(false); @@ -171,6 +220,8 @@ int main() { g_out_of_range.clear(); const std::string op = jc::FindString(line, "op"); const std::string key = jc::FindString(line, "key"); + const int64_t cancel_after_ms = Integer(line, "cancel_after_ms"); + Canceller canceller; // outlives the client, which points at it dmi_store::S3Client client(ReadConfig(line)); // Before any request goes out: a timeout or attempt count that cannot be // represented must not be replaced by the default. @@ -178,6 +229,15 @@ int main() { refuse_out_of_range(); continue; } + const bool armed = cancel_after_ms > 0; + if (armed) { + if (!jc::FindBool(line, "cancel_uploader_only")) { + client.set_cancellation(&canceller.cancel); + } + canceller.Arm(cancel_after_ms); + } + bool upload_cancelled = false; + bool upload_op = false; std::string error; std::string out = "{\"ok\":"; if (op == "put") { @@ -311,10 +371,13 @@ int main() { continue; } dmi_store::SpoolUploader uploader(&spool, &client, uploader_config); + if (armed) uploader.set_cancellation(&canceller.cancel); dmi_store::PackRef ref; int attempts = 0; std::string error; - const bool ok = uploader.UploadOne(staged, &ref, &attempts, &error); + upload_op = true; + const bool ok = uploader.UploadOne(staged, &ref, &attempts, &error, + &upload_cancelled); out += ok ? "true" : "false"; if (ok) { out += ",\"ref\":"; @@ -352,9 +415,11 @@ int main() { continue; } dmi_store::SpoolUploader uploader(&spool, &client, uploader_config); + if (armed) uploader.set_cancellation(&canceller.cancel); const dmi_store::UploadBatchResult result = uploader.UploadPending(limit < 0 ? -1 : static_cast(limit)); - out += "true,\"refs\":["; + out += std::string("true,\"listing_cancelled\":") + + (result.listing_cancelled ? "true" : "false") + ",\"refs\":["; bool first = true; for (const auto& ref : result.refs) { if (!first) out.push_back(','); @@ -372,6 +437,8 @@ int main() { out += ",\"attempts\":" + std::to_string(failure.attempts); out += ",\"error\":"; jc::EscapeJson(failure.error, &out); + out += std::string(",\"cancelled\":") + + (failure.cancelled ? "true" : "false"); out += "}"; first = false; } @@ -380,7 +447,8 @@ int main() { std::to_string(snap.attempted_packs) + ",\"uploaded_packs\":" + std::to_string(snap.uploaded_packs) + ",\"uploaded_bytes\":" + std::to_string(snap.uploaded_bytes) + ",\"failed_packs\":" + - std::to_string(snap.failed_packs) + ",\"retries\":" + + std::to_string(snap.failed_packs) + ",\"cancelled_packs\":" + + std::to_string(snap.cancelled_packs) + ",\"retries\":" + std::to_string(snap.retries) + ",\"peak_active_uploads\":" + std::to_string(snap.peak_active_uploads) + ",\"peak_in_flight_bytes\":" + @@ -391,6 +459,11 @@ int main() { } else { out += "false,\"what\":\"unknown op\""; } + if (armed) { + const bool fired = canceller.Finish(); + out += std::string(",\"cancelled\":") + + ((upload_op ? upload_cancelled : fired) ? "true" : "false"); + } out += ",\"attempts\":" + std::to_string(client.last_attempts()) + "}\n"; std::cout << out; } diff --git a/native/csrc/store/s3_client.cpp b/native/csrc/store/s3_client.cpp index 370f08562..0f498a976 100644 --- a/native/csrc/store/s3_client.cpp +++ b/native/csrc/store/s3_client.cpp @@ -77,14 +77,29 @@ bool IsRetryableStatus(long status) { status == 504; } -void Backoff(int attempt) { +// False when `cancel` (may be null) woke it before the wait was out. +bool Backoff(int attempt, const Cancellation* cancel) { // 0.2s * 2^attempt, capped at 5s. Deterministic: the fault-matrix tests - // assert attempt counts, not wall time, so no jitter. + // assert attempt counts, not wall time, so no jitter. The shift is capped + // too: max_attempts goes up to 1000, and 200 << 56 overflows int64 (a + // negative wait, so no backoff at all), 200 << 64 is undefined. using namespace std::chrono; - const int64_t ms = std::min(5000, 200LL << attempt); + const int64_t ms = std::min(5000, 200LL << std::min(attempt, 5)); + if (cancel != nullptr) return cancel->SleepFor(milliseconds(ms)); std::this_thread::sleep_for(milliseconds(ms)); + return true; +} + +// CURLOPT_XFERINFOFUNCTION: libcurl calls it throughout a transfer -- +// connecting included, and at least once a second while nothing moves -- +// and aborts the transfer (CURLE_ABORTED_BY_CALLBACK) on a non-zero return. +int AbortWhenCancelled(void* cancel, curl_off_t, curl_off_t, curl_off_t, + curl_off_t) { + return static_cast(cancel)->cancelled() ? 1 : 0; } +constexpr const char* kCancelled = "request cancelled"; + std::string XmlEscape(const std::string& value) { std::string out; for (char c : value) { @@ -186,7 +201,33 @@ S3Response S3Client::Exchange( const std::vector>& query, const std::map& extra_headers, const uint8_t* body, size_t body_len, const std::string& body_hash_hex) { + S3Response response = + ExchangeWith(method, key, query, extra_headers, body, body_len, + body_hash_hex, ExchangeOptions{}); + if (after_exchange_for_testing_) after_exchange_for_testing_(); + return response; +} + +S3Response S3Client::ExchangeWith( + const std::string& method, const std::string& key, + const std::vector>& query, + const std::map& extra_headers, + const uint8_t* body, size_t body_len, const std::string& body_hash_hex, + const ExchangeOptions& options) { S3Response response; + const Cancellation* cancel = options.cancellable ? cancel_ : nullptr; + const int max_attempts = + options.max_attempts > 0 ? options.max_attempts : config_.max_attempts; + const long timeout_s = + options.timeout_s > 0 ? options.timeout_s : config_.read_timeout_s; + const long connect_timeout_s = + std::min(config_.connect_timeout_s, timeout_s); + const auto cancelled = [&response] { + response.ok = false; + response.cancelled = true; + response.error = kCancelled; + return response; + }; if (!config_error_.empty()) { last_attempts_ = 0; response.error = config_error_; @@ -228,7 +269,9 @@ S3Response S3Client::Exchange( (query_text.empty() ? "" : "?" + query_text); last_attempts_ = 0; - for (int attempt = 0; attempt < config_.max_attempts; ++attempt) { + for (int attempt = 0; attempt < max_attempts; ++attempt) { + // Nothing goes out once cancelled, a first attempt included. + if (cancel != nullptr && cancel->cancelled()) return cancelled(); ++last_attempts_; // Signed per attempt, not once before the loop: SigV4 binds the // signature to x-amz-date, and S3 refuses a date more than 15 minutes @@ -259,9 +302,15 @@ S3Response S3Client::Exchange( } curl_easy_setopt(curl, CURLOPT_URL, url.c_str()); curl_easy_setopt(curl, CURLOPT_HTTPHEADER, chunk); - curl_easy_setopt(curl, CURLOPT_CONNECTTIMEOUT, config_.connect_timeout_s); - curl_easy_setopt(curl, CURLOPT_TIMEOUT, config_.read_timeout_s); + curl_easy_setopt(curl, CURLOPT_CONNECTTIMEOUT, connect_timeout_s); + curl_easy_setopt(curl, CURLOPT_TIMEOUT, timeout_s); curl_easy_setopt(curl, CURLOPT_NOSIGNAL, 1L); + if (cancel != nullptr) { + curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L); + curl_easy_setopt(curl, CURLOPT_XFERINFOFUNCTION, AbortWhenCancelled); + curl_easy_setopt(curl, CURLOPT_XFERINFODATA, + const_cast(cancel)); + } if (is_https_) { // https always verifies: peer and host name, stated explicitly rather // than left to libcurl's defaults, and never switched off. A private @@ -310,10 +359,11 @@ S3Response S3Client::Exchange( curl_slist_free_all(chunk); curl_easy_cleanup(curl); + if (code == CURLE_ABORTED_BY_CALLBACK) return cancelled(); if (code != CURLE_OK) { response.error = std::string("curl: ") + curl_easy_strerror(code); - if (IsRetryableCurl(code) && attempt + 1 < config_.max_attempts) { - Backoff(attempt); + if (IsRetryableCurl(code) && attempt + 1 < max_attempts) { + if (!Backoff(attempt, cancel)) return cancelled(); continue; } return response; @@ -322,8 +372,8 @@ S3Response S3Client::Exchange( response.http_status = status; response.headers = std::move(response_headers); response.body = std::move(response_body); - if (IsRetryableStatus(status) && attempt + 1 < config_.max_attempts) { - Backoff(attempt); + if (IsRetryableStatus(status) && attempt + 1 < max_attempts) { + if (!Backoff(attempt, cancel)) return cancelled(); continue; } return response; @@ -335,11 +385,14 @@ S3Response S3Client::Exchange( return response; } -ObjectHead S3Client::HeadObject(const std::string& key, std::string* error) { +ObjectHead S3Client::HeadObject(const std::string& key, std::string* error, + bool* cancelled) { ObjectHead head; + if (cancelled) *cancelled = false; S3Response response = Exchange("HEAD", key, {}, {}, nullptr, 0, Sha256Hex("")); if (!response.ok) { + if (cancelled) *cancelled = response.cancelled; if (error) *error = response.error; return head; } @@ -378,8 +431,11 @@ ObjectHead S3Client::HeadObject(const std::string& key, std::string* error) { bool S3Client::GetRange(const std::string& key, uint64_t offset, uint64_t length, std::vector* out, - std::string* error) { + std::string* error, bool* unavailable, + bool* cancelled) { out->clear(); + if (unavailable) *unavailable = false; + if (cancelled) *cancelled = false; if (length == 0) return true; std::map headers; headers["Range"] = "bytes=" + std::to_string(offset) + "-" + @@ -387,10 +443,14 @@ bool S3Client::GetRange(const std::string& key, uint64_t offset, S3Response response = Exchange("GET", key, {}, headers, nullptr, 0, Sha256Hex("")); if (!response.ok) { + if (unavailable) *unavailable = true; + if (cancelled) *cancelled = response.cancelled; if (error) *error = response.error; return false; } if (response.http_status != 200 && response.http_status != 206) { + // Retried until the attempts ran out: the store, not the object. + if (unavailable) *unavailable = IsRetryableStatus(response.http_status); if (error) { *error = "GetObject returned HTTP " + std::to_string(response.http_status); @@ -415,7 +475,7 @@ bool S3Client::PutSingle( const std::string& key, const uint8_t* data, size_t n, const std::map& metadata, const std::string& content_type, std::string* etag_out, - std::string* error) { + std::string* error, bool* cancelled_out) { std::map headers; headers["Content-Type"] = content_type; for (const auto& [name, value] : metadata) { @@ -424,6 +484,7 @@ bool S3Client::PutSingle( S3Response response = Exchange("PUT", key, {}, headers, data, n, Sha256Hex(data, n)); if (!response.ok) { + if (cancelled_out) *cancelled_out = response.cancelled; if (error) *error = response.error; return false; } @@ -444,7 +505,7 @@ bool S3Client::PutMultipart( const std::string& key, const uint8_t* data, size_t n, const std::map& metadata, const std::string& content_type, std::string* etag_out, - std::string* error) { + std::string* error, bool* cancelled_out) { std::map headers; headers["Content-Type"] = content_type; for (const auto& [name, value] : metadata) { @@ -455,6 +516,7 @@ bool S3Client::PutMultipart( Exchange("POST", key, {{"uploads", ""}}, headers, nullptr, 0, Sha256Hex("")); if (!created.ok || created.http_status != 200) { + if (cancelled_out) *cancelled_out = created.cancelled; if (error) { *error = "CreateMultipartUpload failed: " + (created.ok ? "HTTP " + std::to_string(created.http_status) @@ -482,6 +544,7 @@ bool S3Client::PutMultipart( {"uploadId", upload_id}}, {}, data + offset, len, Sha256Hex(data + offset, len)); if (!part.ok || part.http_status != 200) { + if (cancelled_out) *cancelled_out = part.cancelled; abort_error = "UploadPart failed: " + (part.ok ? "HTTP " + std::to_string(part.http_status) : part.error); @@ -496,8 +559,7 @@ bool S3Client::PutMultipart( offset += len; } if (!abort_error.empty()) { - Exchange("DELETE", key, {{"uploadId", upload_id}}, {}, nullptr, 0, - Sha256Hex("")); + AbortMultipart(key, upload_id, cancelled()); if (error) *error = abort_error; return false; } @@ -518,6 +580,16 @@ bool S3Client::PutMultipart( reinterpret_cast(xml.data()), xml.size(), Sha256Hex(reinterpret_cast(xml.data()), xml.size())); + if (done.cancelled) { + // Cut short with every part sent: the upload may have completed or + // not. Aborting a completed one fails harmlessly (NoSuchUpload), and + // the object it made is the pack's own, which a retry's preflight + // re-reads and blesses. + AbortMultipart(key, upload_id, true); + if (cancelled_out) *cancelled_out = true; + if (error) *error = "CompleteMultipartUpload failed: " + done.error; + return false; + } if (!done.ok || done.http_status != 200) { if (error) { *error = "CompleteMultipartUpload failed: " + @@ -532,14 +604,33 @@ bool S3Client::PutMultipart( return true; } +void S3Client::AbortMultipart(const std::string& key, + const std::string& upload_id, + bool after_cancel) { + // Best effort either way: an upload left behind is invisible, and a + // bucket lifecycle rule for incomplete multipart uploads reaps it. + ExchangeOptions options; + if (after_cancel) { + options.cancellable = false; + options.max_attempts = 1; + options.timeout_s = kAbortAfterCancelTimeoutS; + } + ExchangeWith("DELETE", key, {{"uploadId", upload_id}}, {}, nullptr, 0, + Sha256Hex(""), options); +} + bool S3Client::PutObject(const std::string& key, const uint8_t* data, size_t n, const std::map& metadata, const std::string& content_type, - std::string* etag_out, std::string* error) { + std::string* etag_out, std::string* error, + bool* cancelled) { + if (cancelled) *cancelled = false; if (n >= config_.multipart_threshold_bytes) { - return PutMultipart(key, data, n, metadata, content_type, etag_out, error); + return PutMultipart(key, data, n, metadata, content_type, etag_out, error, + cancelled); } - return PutSingle(key, data, n, metadata, content_type, etag_out, error); + return PutSingle(key, data, n, metadata, content_type, etag_out, error, + cancelled); } bool S3Client::DeleteObject(const std::string& key, std::string* error) { diff --git a/native/csrc/store/s3_client.h b/native/csrc/store/s3_client.h index 84e62c4af..1b9ec2dae 100644 --- a/native/csrc/store/s3_client.h +++ b/native/csrc/store/s3_client.h @@ -20,8 +20,11 @@ #include #include #include +#include #include +#include "cancel.h" + namespace dmi_store { struct S3Config { @@ -56,6 +59,10 @@ struct S3Response { std::map headers; // lowercased names std::string body; std::string error; + // The client's Cancellation cut the exchange short (ok is false): before + // an attempt, during its transfer, or in the backoff before a retry. + // Never retried. + bool cancelled = false; }; // Parsed HEAD metadata the DMI pack layout stores per object. @@ -81,6 +88,10 @@ struct ListResult { // S3's minimum size for every part of a multipart upload but the last. inline constexpr uint64_t kMinMultipartPartBytes = 5ull * 1024 * 1024; +// The whole-request bound, connect included, on the AbortMultipartUpload a +// cancel leads to: one attempt, since whoever cancelled is waiting on it. +inline constexpr int kAbortAfterCancelTimeoutS = 5; + class S3Client { public: // An invalid config (see ValidateConfig) does not throw: the client @@ -96,17 +107,40 @@ class S3Client { S3Client& operator=(const S3Client&) = delete; const S3Config& config() const { return config_; } + + // Every request from now on honours `cancel` (nullptr: none): none is + // sent once it is cancelled, a transfer in flight is aborted (libcurl's + // progress callback asks at least once a second, connecting included), + // and a retry's backoff wakes for it. A cancelled request fails with the + // error "request cancelled" and is not retried; a multipart upload it cut + // short is aborted with a request of its own, which the cancel does not + // cut (bounded by kAbortAfterCancelTimeoutS instead). Set it before the + // client is shared, and keep `cancel` alive as long as the client. + void set_cancellation(const Cancellation* cancel) { cancel_ = cancel; } + bool cancelled() const { return cancel_ != nullptr && cancel_->cancelled(); } // Attempts actually made by the last call (1 + retries), for tests. int last_attempts() const { return last_attempts_.load(std::memory_order_relaxed); } // HEAD /bucket/key. 404 → {found=false}, no error. - ObjectHead HeadObject(const std::string& key, std::string* error); + // + // On failure, *cancelled (when given, here and on GetRange and + // PutObject) says whether the Cancellation cut the call short -- before + // an attempt, in its transfer or in a retry's backoff -- rather than the + // store failing it: a cancel that merely comes in while a failure is + // reported does not count. + ObjectHead HeadObject(const std::string& key, std::string* error, + bool* cancelled = nullptr); // GET /bucket/key, optionally Range: bytes=offset-(offset+length-1). // length==0 returns empty without a request (matches read_range). - // Short/oversized bodies are errors, not truncations. + // Short/oversized bodies are errors, not truncations. On failure, + // *unavailable (when given) says whether the store never answered for + // the object: a transport error or timeout, a retryable status on every + // attempt, or a cancel -- as against an answer about it (a 404, a 403, a + // short body), which says something about the object itself. bool GetRange(const std::string& key, uint64_t offset, uint64_t length, - std::vector* out, std::string* error); + std::vector* out, std::string* error, + bool* unavailable = nullptr, bool* cancelled = nullptr); // PUT /bucket/key with x-amz-content-sha256 over the exact bytes plus the // DMI metadata headers. Over multipart_threshold_bytes the call becomes @@ -114,7 +148,7 @@ class S3Client { bool PutObject(const std::string& key, const uint8_t* data, size_t n, const std::map& metadata, const std::string& content_type, std::string* etag_out, - std::string* error); + std::string* error, bool* cancelled = nullptr); bool DeleteObject(const std::string& key, std::string* error); @@ -124,6 +158,15 @@ class S3Client { int max_keys, const std::string& continuation, ListResult* out, std::string* error); + // Test seam: called each time an exchange has returned (every attempt of + // it done), before the call that made it reads the response. A test can + // cancel there, to stand for a cancel that comes in once a request has + // failed on its own -- the case *cancelled must not report. Set it + // before the client is shared. + void SetAfterExchangeHookForTesting(std::function hook) { + after_exchange_for_testing_ = std::move(hook); + } + // Exposed for the fault-matrix tests: one raw signed exchange. S3Response Exchange(const std::string& method, const std::string& key, const std::vector>& query, @@ -140,16 +183,37 @@ class S3Client { // Atomic because SpoolUploader shares one client across its worker // threads, and every request writes this; a plain int was a data race. std::atomic last_attempts_{0}; + const Cancellation* cancel_ = nullptr; + std::function after_exchange_for_testing_; + + // How one exchange departs from the config: whether the Cancellation + // applies, and (when positive) its own attempt count and whole-request + // timeout in seconds. + struct ExchangeOptions { + bool cancellable = true; + int max_attempts = 0; + int timeout_s = 0; + }; + S3Response ExchangeWith( + const std::string& method, const std::string& key, + const std::vector>& query, + const std::map& extra_headers, + const uint8_t* body, size_t body_len, const std::string& body_hash_hex, + const ExchangeOptions& options); + // Aborts a multipart upload: as configured, or -- after a cancel -- once, + // uncancelled, within kAbortAfterCancelTimeoutS. + void AbortMultipart(const std::string& key, const std::string& upload_id, + bool after_cancel); // Multipart primitives (single PUT when under threshold). bool PutSingle(const std::string& key, const uint8_t* data, size_t n, const std::map& metadata, const std::string& content_type, std::string* etag_out, - std::string* error); + std::string* error, bool* cancelled_out); bool PutMultipart(const std::string& key, const uint8_t* data, size_t n, const std::map& metadata, const std::string& content_type, std::string* etag_out, - std::string* error); + std::string* error, bool* cancelled_out); }; } // namespace dmi_store diff --git a/native/csrc/store/spool.cpp b/native/csrc/store/spool.cpp index d0df3aa80..5be18b7ca 100644 --- a/native/csrc/store/spool.cpp +++ b/native/csrc/store/spool.cpp @@ -1,5 +1,7 @@ #include "spool.h" +#include "cancel.h" + #include #include @@ -627,9 +629,16 @@ SpoolStatus Spool::ListPending(std::vector* out, std::string* error) return Scan(out, false, error); } +SpoolStatus Spool::ListPending(std::vector* out, std::string* error, + const Cancellation* cancel, bool* cut) { + return Scan(out, false, error, cancel, cut); +} + SpoolStatus Spool::Scan(std::vector* out, bool discard_open_files, - std::string* error) { + std::string* error, const Cancellation* cancel, + bool* cut) { out->clear(); + if (cut) *cut = false; std::lock_guard lock(mutex_); std::error_code ec; std::vector readies; @@ -660,6 +669,13 @@ SpoolStatus Spool::Scan(std::vector* out, bool discard_open_files, } std::sort(readies.begin(), readies.end()); for (const std::string& path : readies) { + if (cancel != nullptr && cancel->cancelled()) { + // Before the next pack's hash. The account below is rebuilt from a + // whole listing only; quarantines already made stand. + out->clear(); + if (cut) *cut = true; + return SpoolStatus::kOk; + } const std::string name = fs::path(path).filename().string(); std::string id, sum; uint64_t created = 0, records = 0; diff --git a/native/csrc/store/spool.h b/native/csrc/store/spool.h index ad3558653..60599977f 100644 --- a/native/csrc/store/spool.h +++ b/native/csrc/store/spool.h @@ -25,6 +25,8 @@ namespace dmi_store { +class Cancellation; // cancel.h + struct SpoolConfig { std::string root; uint64_t max_bytes = 0; @@ -91,6 +93,13 @@ class Spool { // Validate and list ready packs without deleting in-progress writes. SpoolStatus ListPending(std::vector* out, std::string* error); + // The same, but stops between packs once `cancel` is cancelled: *cut then + // says so, *out is empty, and the spool's account is left as it was (a + // partial listing must not recount it). Validating hashes every byte of + // every pack, so a backlog makes a listing slow; this bounds the wait a + // cancel has for it by one pack's hash. + SpoolStatus ListPending(std::vector* out, std::string* error, + const Cancellation* cancel, bool* cut); // Remove one staged pack after upload (identity + size verified first). SpoolStatus Remove(const StagedPack& staged, std::string* error); @@ -105,7 +114,8 @@ class Spool { private: SpoolStatus Scan(std::vector* out, bool discard_open_files, - std::string* error); + std::string* error, const Cancellation* cancel = nullptr, + bool* cut = nullptr); // Count/uncount one ready path in the committed account, at most once each // -- Python's _account_ready_locked / _unaccount_ready_locked. `mutex_` // must be held. Both return whether they actually changed the account. diff --git a/native/csrc/store/uploader.cpp b/native/csrc/store/uploader.cpp index 8f8de47b4..bdd756acf 100644 --- a/native/csrc/store/uploader.cpp +++ b/native/csrc/store/uploader.cpp @@ -63,8 +63,9 @@ std::vector ReadFile(const std::string& path, std::string* error) { return data; } -void SleepBackoff(const UploaderConfig& config, int attempt, - std::mt19937_64* rng) { +// False when `cancel` (may be null) woke it before the wait was out. +bool SleepBackoff(const UploaderConfig& config, int attempt, + std::mt19937_64* rng, const Cancellation* cancel) { double wait = config.base_backoff_s * (1 << std::min(attempt, 20)); wait = std::min(wait, config.max_backoff_s); if (config.jitter_ratio > 0) { @@ -72,11 +73,18 @@ void SleepBackoff(const UploaderConfig& config, int attempt, 1.0 + config.jitter_ratio); wait *= jitter(*rng); } - if (wait > 0) { - std::this_thread::sleep_for(std::chrono::duration(wait)); - } + if (wait <= 0) return true; + const auto duration = + std::chrono::duration_cast( + std::chrono::duration(wait)); + if (cancel != nullptr) return cancel->SleepFor(duration); + std::this_thread::sleep_for(duration); + return true; } +constexpr const char* kUploadCancelled = + "upload cancelled; the pack stays staged"; + std::map PackMetadata(const StagedPack& staged) { return { {"dmi-format", "dmi-pack-v1"}, @@ -94,21 +102,39 @@ SpoolUploader::SpoolUploader(Spool* spool, S3Client* client, : spool_(spool), client_(client), config_(std::move(config)) {} bool SpoolUploader::UploadOne(const StagedPack& staged, PackRef* ref, - int* attempts_out, std::string* error) { + int* attempts_out, std::string* error, + bool* cancelled_out) { const std::string& key = staged.object_key; std::mt19937_64 rng( static_cast(std::hash{}(staged.pack_id))); std::string last_error; int attempts = 0; + bool cancelled = false; + // The Cancellation cut this attempt's failed request short. + bool attempt_cut = false; + if (cancelled_out) *cancelled_out = false; for (int attempt = 0; attempt < config_.max_attempts; ++attempt) { + // A cancel ends the retries: the backoff wakes for it, and no attempt + // starts after it. What an attempt cut short left behind is safe: a + // PUT that never completed made no object, and one that did is found + // and re-verified by the next upload's preflight. + if (attempt > 0 && !SleepBackoff(config_, attempt - 1, &rng, cancel_)) { + cancelled = true; + break; + } + if (cancel_ != nullptr && cancel_->cancelled()) { + cancelled = true; + break; + } ++attempts; - if (attempt > 0) SleepBackoff(config_, attempt - 1, &rng); + attempt_cut = false; // 1. Preflight: an object already carrying this pack is re-read and // re-hashed before it is blessed — metadata alone is not proof. { std::string head_error; - const ObjectHead head = client_->HeadObject(key, &head_error); + const ObjectHead head = + client_->HeadObject(key, &head_error, &attempt_cut); if (!head_error.empty()) { last_error = head_error; continue; // transport-level: retryable below @@ -120,7 +146,7 @@ bool SpoolUploader::UploadOne(const StagedPack& staged, PackRef* ref, std::vector existing; std::string get_error; if (!client_->GetRange(key, 0, staged.object_bytes, &existing, - &get_error) || + &get_error, nullptr, &attempt_cut) || Sha256HexBytes(existing.data(), existing.size()) != staged.checksum) { last_error = get_error.empty() @@ -187,14 +213,15 @@ bool SpoolUploader::UploadOne(const StagedPack& staged, PackRef* ref, std::string put_error; if (!client_->PutObject(key, data.data(), data.size(), PackMetadata(staged), config_.content_type, &etag, - &put_error)) { + &put_error, &attempt_cut)) { last_error = put_error; continue; } // 3. Post-upload visibility: the object must be there. { std::string head_error; - const ObjectHead head = client_->HeadObject(key, &head_error); + const ObjectHead head = + client_->HeadObject(key, &head_error, &attempt_cut); if (!head_error.empty() || !head.found || head.size != staged.object_bytes) { last_error = head_error.empty() @@ -217,6 +244,19 @@ bool SpoolUploader::UploadOne(const StagedPack& staged, PackRef* ref, if (attempts_out) *attempts_out = attempts; return true; } + // The last attempt cut short by a cancel ends the loop without the checks + // above seeing it. Only its request's own report says so: attempts that + // ran out on real failures, or the corrupt-bytes break, stay failures + // even when a flush's deadline passed meanwhile, so the error is counted + // and recorded rather than booked as a cancel. + if (!cancelled && attempt_cut) cancelled = true; + if (cancelled) { + last_error = last_error.empty() + ? kUploadCancelled + : std::string(kUploadCancelled) + + " (the attempt before failed: " + last_error + ")"; + if (cancelled_out) *cancelled_out = true; + } if (attempts_out) *attempts_out = attempts; if (error) *error = last_error; return false; @@ -227,14 +267,27 @@ UploadBatchResult SpoolUploader::UploadPending(int limit) { if (limit != -1 && limit <= 0) return result; // invalid limit: empty result std::vector pending; { + // Listing hashes every staged pack; a cancel stops it between packs + // rather than after the whole backlog. std::string error; - if (spool_->ListPending(&pending, &error) != SpoolStatus::kOk) { + bool cut = false; + if (spool_->ListPending(&pending, &error, cancel_, &cut) != + SpoolStatus::kOk) { + return result; + } + if (cut) { + result.listing_cancelled = true; return result; } } if (limit != -1 && static_cast(limit) < pending.size()) { pending.resize(static_cast(limit)); } + return UploadStaged(std::move(pending)); +} + +UploadBatchResult SpoolUploader::UploadStaged(std::vector pending) { + UploadBatchResult result; // Both vectors are positional from the start: sized to the recover() // order up front, oversized refusals written into their own slot, and // workers below fill the rest by index. @@ -303,6 +356,7 @@ UploadBatchResult SpoolUploader::UploadPending(int limit) { // the scan after the wait always finds the pack the predicate saw. cv.wait(lock, [&] { if (remaining.empty()) return true; + if (cancel_ != nullptr && cancel_->cancelled()) return true; for (const auto& candidate : remaining) { if (in_flight + pending[candidate.index].object_bytes <= @@ -313,6 +367,24 @@ UploadBatchResult SpoolUploader::UploadPending(int limit) { return false; }); if (remaining.empty()) return; + if (cancel_ != nullptr && cancel_->cancelled()) { + // No pack starts once cancelled. Each one left is reported at its + // own position, still staged. + for (const Slot& left : remaining) { + UploadFailure& failure = result.failures[left.index]; + failure = {pending[left.index].pack_id, + pending[left.index].object_key, 0, + "upload cancelled before it started; the pack stays " + "staged", + true}; + ++result.snapshot.cancelled_packs; + finished[left.index] = true; + } + remaining.clear(); + lock.unlock(); + cv.notify_all(); + return; + } for (auto it = remaining.begin(); it != remaining.end(); ++it) { if (in_flight + pending[it->index].object_bytes <= config_.max_in_flight_bytes) { @@ -333,7 +405,8 @@ UploadBatchResult SpoolUploader::UploadPending(int limit) { PackRef ref; std::string error; int attempts = 0; - const bool ok = UploadOne(*staged, &ref, &attempts, &error); + bool cancelled = false; + const bool ok = UploadOne(*staged, &ref, &attempts, &error, &cancelled); const int64_t elapsed = NowNs() - started; { std::lock_guard lock(mutex); @@ -355,8 +428,12 @@ UploadBatchResult SpoolUploader::UploadPending(int limit) { } else { result.refs[slot.index] = PackRef{}; result.failures[slot.index] = {staged->pack_id, staged->object_key, - attempts, error}; - ++result.snapshot.failed_packs; + attempts, error, cancelled}; + if (cancelled) { + ++result.snapshot.cancelled_packs; + } else { + ++result.snapshot.failed_packs; + } if (attempts > 1) { result.snapshot.retries += static_cast(attempts - 1); } diff --git a/native/csrc/store/uploader.h b/native/csrc/store/uploader.h index a0d8e6ed3..6219ad050 100644 --- a/native/csrc/store/uploader.h +++ b/native/csrc/store/uploader.h @@ -12,6 +12,11 @@ // 2. The upload stream itself is hashed as curl reads it; a source whose // bytes contradict the staged checksum deletes the upload and fails. // 3. Post-upload HEAD must show the object, else the upload is refused. +// +// Cancellation (set_cancellation): the spool listing stops between packs, +// no pack starts once cancelled, and a pack's retries and their backoff end +// at once. A cancelled pack stays in the spool and is reported as +// cancelled, not failed. #ifndef DMI_STORE_UPLOADER_H_ #define DMI_STORE_UPLOADER_H_ @@ -21,6 +26,7 @@ #include #include +#include "cancel.h" #include "s3_client.h" #include "spool.h" @@ -54,6 +60,11 @@ struct UploadFailure { std::string object_key; int attempts = 0; std::string error; + // Cut short by a cancel: before it started, between its attempts, or in + // a request the cancel cut. The pack is still staged. Not counted in + // failed_packs. Attempts that ran out on failures of their own stay + // failures, even when a cancel came in meanwhile. + bool cancelled = false; }; struct UploadSnapshot { @@ -61,6 +72,7 @@ struct UploadSnapshot { uint64_t uploaded_packs = 0; uint64_t uploaded_bytes = 0; uint64_t failed_packs = 0; + uint64_t cancelled_packs = 0; // left staged by a cancel uint64_t retries = 0; uint64_t peak_active_uploads = 0; uint64_t peak_in_flight_bytes = 0; @@ -74,6 +86,9 @@ struct UploadBatchResult { std::vector refs; std::vector failures; UploadSnapshot snapshot; + // A cancel cut the spool listing short, before any upload started: the + // batch is empty whatever the spool holds. + bool listing_cancelled = false; }; class SpoolUploader { @@ -83,6 +98,14 @@ class SpoolUploader { SpoolUploader(const SpoolUploader&) = delete; SpoolUploader& operator=(const SpoolUploader&) = delete; + // From now on, once `cancel` is cancelled, UploadPending stops listing + // the spool (listing_cancelled, if it had not listed it all) and starts + // no pack, and UploadOne stops retrying (its backoff wakes for it); the + // pack stays staged. A transfer in flight is cut only when the S3 client + // has the same Cancellation (S3Client::set_cancellation). nullptr: none. + // Keep `cancel` alive as long as the uploader. + void set_cancellation(const Cancellation* cancel) { cancel_ = cancel; } + // Recover the spool and upload every entry (or the first `limit`, which // must be positive when set). A pack over max_in_flight_bytes is recorded // as a failure at its own position and the rest of the batch still @@ -91,14 +114,21 @@ class SpoolUploader { // uploads nothing — see the comment at the byte gate in uploader.cpp. UploadBatchResult UploadPending(int limit = -1); - // Upload one staged entry with retry. Public for tests. + // Upload `pending`, packs a ListPending already returned, as UploadPending + // does after its listing: in their order, refs and failures positional. + // A caller that uploads a listing in parts lists it once. + UploadBatchResult UploadStaged(std::vector pending); + + // Upload one staged entry with retry. Public for tests. *cancelled_out + // (when given) says whether a cancel ended it. bool UploadOne(const StagedPack& staged, PackRef* ref, int* attempts_out, - std::string* error); + std::string* error, bool* cancelled_out = nullptr); private: Spool* spool_; S3Client* client_; UploaderConfig config_; + const Cancellation* cancel_ = nullptr; }; } // namespace dmi_store diff --git a/src/dmi/config.py b/src/dmi/config.py index 5e9dca10b..2cef3b591 100644 --- a/src/dmi/config.py +++ b/src/dmi/config.py @@ -161,10 +161,14 @@ class MonitoringConfig: # The rest of the native capture storage path, after the spool: set it # and the engine runs an in-process C++ service that uploads every pack # the sink stages to the object store and indexes it into the ClickHouse - # catalog, and ``flush_and_wait`` returns only once they are queryable. - # Unset, packs stay in the spool for something else to drain. Needs - # ``storage_backend="persistent"`` and ``capture_sink_config``, whose - # spool it drains. + # catalog, and ``flush_and_wait`` returns only once they are queryable + # (raising TimeoutError when they are not by its timeout). ``close()`` + # drains it too, best effort, for the config's ``close_flush_timeout_s`` + # -- whose comment says by how much the drain can outlast it. Unset, + # packs stay in the spool for something else to drain; the sink's open + # pack reaches the spool when the ring releases the sink at close. + # Needs ``storage_backend="persistent"`` and ``capture_sink_config``, + # whose spool it drains. capture_storage_config: Optional["NativeCaptureStorageConfig"] = None storage_backend: StorageBackend = "auto" diff --git a/src/dmi/engine.py b/src/dmi/engine.py index 211ff3ebf..a155b3161 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -613,7 +613,15 @@ def flush_and_wait(self, timeout_s: float = 600.0) -> None: deadline = time.monotonic() + float(timeout_s) transport.flush_records_and_wait(float(timeout_s)) # The sink's boundary is a staged pack; with the storage service it - # is a pack in the catalog. Its flush runs one cycle even at zero. + # is a pack in the catalog. The service's flush returns on time: + # past the deadline it starts no upload and at most one index + # batch, leaving the rest to its background loop -- about one + # clickhouse_request_timeout_s late against a catalog or store that + # stopped answering, and up to about 6 s late when the deadline cuts + # a multipart upload, whose abort nothing cuts + # (NativeCaptureStorageConfig.close_flush_timeout_s has the + # details). At zero it still runs one cycle, so a drained spool + # reports drained. storage = self._capture_storage if storage is not None: storage.flush(max(0.0, deadline - time.monotonic())) @@ -824,11 +832,13 @@ def next_auto_group_id(self) -> int: return gid def _seal_capture_sink(self, deadline: float) -> None: - """Flush the record sink before its ring stops. + """Flush the record sink before its ring stops, within ``deadline``. - Stopping the ring releases the sink WITHOUT flushing it, so the - records of its open pack would still be in memory while the service - drains a spool that does not hold them yet. + Releasing the sink from the stopping ring flushes it too (the native + pack sink's release backstop), but only as a last resort: bounded + by its own timeout, and with nothing but a line on stderr when it + fails. This flush comes first, inside close()'s budget, and logs + why it failed. """ try: self._ring_transport.flush_records_and_wait( @@ -842,9 +852,10 @@ def _retire_capture_storage(self, storage: Any, deadline: float) -> None: return self._capture_storage = None # Best effort: once the sink is sealed, a pack that does not reach - # the catalog here is still in the spool or the bucket, and the next - # start uploads or reconciles it. flush_and_wait is the boundary - # that reports. + # the catalog here is still in the spool, which the next start on it + # uploads, or uploaded but unindexed -- one index batch at most -- + # which only the next start's reconcile finds (reconcile_on_start, + # on by default). flush_and_wait is the boundary that reports. try: storage.flush(max(0.0, deadline - time.monotonic())) except Exception as exc: @@ -853,7 +864,21 @@ def _retire_capture_storage(self, storage: Any, deadline: float) -> None: storage.stop() def close(self) -> None: - """Tear down backend resources.""" + """Tear down backend resources. + + A record ring is stopped, which drains its queued records into the + sink and releases the sink; the native pack sink stages its open + pack in the spool on that release. With a storage service + (``capture_storage_config``), close() drains capture first, best + effort, with a budget of ``close_flush_timeout_s``: it flushes the + sink, stops the ring, waits for the service to get the staged packs + into the catalog, then stops the service. The drain can outlast + the budget -- stopping the service is not bounded by it -- and + ``NativeCaptureStorageConfig.close_flush_timeout_s`` says by how + much. What does not drain in time is logged, not raised, and stays + where the next start recovers it; ``flush_and_wait`` is the call + that raises. + """ storage = self._capture_storage # One budget for the whole capture drain: sealing the sink's open diff --git a/src/dmi/storage/native_capture.py b/src/dmi/storage/native_capture.py index ebd9f76c4..585015ace 100644 --- a/src/dmi/storage/native_capture.py +++ b/src/dmi/storage/native_capture.py @@ -265,8 +265,30 @@ class NativeCaptureStorageConfig: # 0 disables the periodic pass; in-process index failures are retried # regardless. reconcile_interval_s: float = 0.0 - # engine.close()'s total budget for draining capture: sealing the sink's - # open pack, then getting every staged pack into the catalog. + # engine.close()'s budget for draining capture, best effort: it flushes + # the sink (sealing its open pack), then waits for the service to get + # the staged packs into the catalog, and stops the service when the + # budget is spent, whatever is left. That stays in the spool, which the + # next start on it uploads, or -- uploaded but not yet indexed, one index + # batch of packs at most, since the service indexes what it uploads a + # batch at a time -- in the bucket, which only the next start's + # reconcile indexes (reconcile_on_start). Past the budget the drain + # starts no upload -- one in flight is cut, within about a second (the + # transfer's progress poll), and a multipart one then aborted, one + # request of at most 5 s that nothing cuts -- and at most one index + # batch: its object-store reads are cut one clickhouse_request_timeout_s + # past the budget, and its catalog statements are never cut, each + # bounded by that timeout (under the publisher lease, by the lease's + # deadline when that is sooner). Stopping the service then releases the + # lease: one more catalog request, bounded the same way -- or, when a + # lease renewal is in flight, that renewal, which it waits for (a lease + # the renewal loses needs no release). So against a catalog or an + # object store that stopped answering, close() outlasts the budget by + # up to about two request timeouts, and up to about 6 s more when the + # budget cuts a multipart upload; against a slow catalog that still + # answers, by one batch of statements and the release. When the sink + # itself is stuck, the flush its release from the ring makes adds up to + # 30 s. close() logs what did not drain; flush_and_wait is what raises. close_flush_timeout_s: float = 60.0 # Bytes of packs the uploader holds in flight at once. A staged pack # larger than this is never uploaded, so the sink's max_pack_bytes must @@ -298,6 +320,19 @@ class NativeCaptureStorageConfig: # predating these knobs used) outlasts the default wait by 10 s. start_lease_wait_s: Optional[float] = None + # Every object-store request is bounded: each attempt by + # s3_read_timeout_s (whole seconds, connecting included), with up to + # s3_max_attempts attempts for a timeout, a host that does not resolve + # or connect, a connection that broke or answered nothing, a 429 or a + # 500/502/503/504 -- not for another status, nor a TLS failure -- and a + # backoff between them of 0.2 s doubling up to 5 s. That is how long + # one read the store never answers can hold the service's background + # cycle -- 4 x 120 s + 1.4 s on these defaults -- and a reader's + # request. stop() and a flush's deadline cut the service's requests + # short regardless (see close_flush_timeout_s). + s3_read_timeout_s: int = 120 + s3_max_attempts: int = 4 + def __post_init__(self) -> None: for name in ("s3_endpoint", "s3_bucket", "s3_access_key", "s3_secret_key", "store_id", "database", "table_prefix", @@ -332,6 +367,12 @@ def __post_init__(self) -> None: "s3_endpoint") if type(self.clickhouse_port) is not int or not 0 < self.clickhouse_port < 65536: raise ValueError("clickhouse_port must be in 1..65535") + # The native client takes both as C ints; a bool is not a count. + for name, most in (("s3_read_timeout_s", 86_400), + ("s3_max_attempts", 1_000)): + value = getattr(self, name) + if type(value) is not int or not 0 < value <= most: + raise ValueError(f"{name} must be an int in 1..{most}") _positive("poll_interval_s", self.poll_interval_s, float) # The native wait is int(poll_interval_s * 1e9) ns; a zero wait spins. if self.poll_interval_s < 0.001: @@ -473,6 +514,8 @@ def _native_dict(self) -> dict[str, Any]: "s3_allow_insecure_http": self.s3_allow_insecure_http, "s3_ca_file": self.s3_ca_file, "s3_ca_path": self.s3_ca_path, + "s3_read_timeout_s": self.s3_read_timeout_s, + "s3_max_attempts": self.s3_max_attempts, "store_id": self.store_id, "clickhouse_scheme": self.clickhouse_scheme, "clickhouse_host": self.clickhouse_host, @@ -601,6 +644,13 @@ def flush(self, timeout_s: float) -> None: Call after the sink's own flush. Raises TimeoutError, carrying the last upload or index error, if the spool has not drained in time. + Returns on time, to within a bound: past ``timeout_s`` it starts + no upload and at most one index batch, leaving the rest to the + background loop -- about one ``clickhouse_request_timeout_s`` late + against a catalog or store that stopped answering, and up to about + 6 s late when the deadline cuts a multipart upload, whose abort + nothing cuts (see + ``NativeCaptureStorageConfig.close_flush_timeout_s``). """ if not self._service.flush(float(timeout_s)): snapshot = self._service.snapshot() diff --git a/tests/native/test_pack_sink_timeout.cpp b/tests/native/test_pack_sink_timeout.cpp index 7defd3fd8..6842e6cee 100644 --- a/tests/native/test_pack_sink_timeout.cpp +++ b/tests/native/test_pack_sink_timeout.cpp @@ -28,6 +28,10 @@ // timeouts. There the stager is slowed rather than wedged, so each row after // the pipeline fills waits a little under one timeout for room. // +// A third case pins Flush's own timeout while another flush is in flight: +// the wedged stager keeps the first from completing, and the second must +// still return within its own timeout. +// // Built and run by tests/test_native_pack_sink_timeout.py. #include @@ -280,11 +284,84 @@ void TestAnEnvelopeSharesOneAdmissionDeadline() { CHECK(final_snap.persisted_records == static_cast(admitted)); } +// A flush is bounded by its own timeout even while another flush is in +// flight. The second waited for the first to give up before its own wait +// began, since the flush lock was untimed: NativePackSink's release +// backstop, a 30 s flush the ring's stop runs, waited out a concurrent +// flush_and_wait's 600 s. Here the stager is wedged, so neither flush can +// complete: the first is given 4 s, the second 0.5 s, and the second must +// return within about its own 0.5 s. +void TestAFlushIsBoundedWhileAnotherIsInFlight() { + const char* base = std::getenv("SPOOL_TEST_ROOT"); + const std::string root = + std::string(base != nullptr ? base : "/tmp") + "/sink-concurrent-flush"; + fs::remove_all(root); + + dmi_sink::SinkConfig config; + config.spool_root = root; + config.num_workers = 1; + config.max_linger_ns = 3600ull * 1000 * 1000 * 1000; // linger never fires + + dmi_sink::PackSink sink(config); + const std::string start_error = sink.Start(); + CHECK(start_error.empty()); + if (!start_error.empty()) return; + + std::mutex mutex; + std::condition_variable cv; + bool in_hook = false, release = false; + sink.SpoolForTesting().SetStageHookForTesting([&] { + std::unique_lock lock(mutex); + in_hook = true; + cv.notify_all(); + cv.wait(lock, [&] { return release; }); + }); + + CHECK(SubmitRecord(sink, 1) == dmi_sink::Admission::kAccepted); + bool first_flushed = false; + std::thread first([&] { + std::string error; + first_flushed = sink.Flush(4.0, &error); + }); + { + // The first flush sealed the open pack and the stager is parked with + // it: that flush now holds the flush lock, waiting on its barrier. + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return in_hook; }); + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + const auto started = std::chrono::steady_clock::now(); + std::string error; + const bool second_flushed = sink.Flush(0.5, &error); + const double elapsed_s = std::chrono::duration( + std::chrono::steady_clock::now() - started).count(); + CHECK(!second_flushed); + CHECK(elapsed_s < 1.5); + if (elapsed_s >= 1.5) { + std::cerr << "concurrent flush: Flush(0.5) returned after " << elapsed_s + << " s\n"; + } + + { + std::lock_guard lock(mutex); + release = true; + cv.notify_all(); + } + first.join(); + CHECK(first_flushed); // released inside its 4 s, its barrier completes + std::string close_error; + const dmi_sink::SinkSnapshot final_snap = sink.Close(-1.0, &close_error); + CHECK(close_error.empty()); + CHECK(final_snap.persisted_records == 1); +} + } // namespace int main() { TestBlockedPipelineTimesOutAndCountsIt(); TestAnEnvelopeSharesOneAdmissionDeadline(); + TestAFlushIsBoundedWhileAnotherIsInFlight(); if (g_failures != 0) { std::cerr << g_failures << " check(s) failed\n"; return 1; diff --git a/tests/test_native_capture_storage_gpu_e2e.py b/tests/test_native_capture_storage_gpu_e2e.py index da94c41b0..b317ea184 100644 --- a/tests/test_native_capture_storage_gpu_e2e.py +++ b/tests/test_native_capture_storage_gpu_e2e.py @@ -174,6 +174,65 @@ def test_flush_and_wait_means_queryable_and_byte_identical(fake_s3, tmp_path): assert not sorted((tmp_path / "spool").rglob("*.dmi-pack.ready")) +def test_stopping_the_ring_stages_the_open_pack_without_a_service(tmp_path): + """The sink alone, no storage service: nothing in close() flushes the + sink, and stopping the ring releases it. The audit's probe -- ready + packs right after the stop -- reported 0, the tail staying in memory + until the sink object died or its 60 s linger fired. The release now + stages it. The test holds the sink, so its destructor cannot be what + wrote the pack.""" + from dmi.api.v1 import HookPointV1, HookSpecV1, MonitoringEngine, TransportSpec + from dmi.config import MonitoringConfig + from dmi.storage.capture import ( + CaptureRecordFormat, DurablePackSpool, PackReader, + ) + from dmi.storage.capture.native_sink import NativeSinkConfig + + spool = tmp_path / "spool" + config = MonitoringConfig( + storage_backend="persistent", + capture_sink_config=NativeSinkConfig( + spool_root=str(spool), max_pack_records=2, + max_linger_ns=60_000_000_000), + ) + tensors = {f"tail-{i}": torch.arange(8, dtype=torch.float32) + i + for i in range(3)} + engine = MonitoringEngine(config=config, model_id="native-sink-gpu", + ring_config=_ring_config()) + try: + runtime = engine.create_record_runtime(CaptureRecordFormat()) + sink = engine._record_sink + hook = HookPointV1( + HookSpecV1("capture_tensor", (TransportSpec("payload"),))) + hook_runtime = _CaptureHookRuntime(runtime) + runtime.bind_hook(hook, hook_runtime=hook_runtime) + for step, (capture_id, tensor) in enumerate(tensors.items()): + hook_runtime.metadata = _metadata(capture_id, tensor, step=step) + hook(tensor.cuda()) + torch.cuda.synchronize() + + engine.close() # no flush_and_wait, and no service to drain + + # Two packs: the full one of two records, and the tail of one. + assert len(sorted(spool.rglob("*.dmi-pack.ready"))) == 2 + snapshot = sink.snapshot() + assert snapshot["persisted_records"] == len(tensors), snapshot + finally: + engine.close() + + staged = {} + for entry in DurablePackSpool(spool, max_bytes=1 << 40).recover(): + with entry.open() as handle: + reader = PackReader.from_bytes(handle.read()) + for descriptor in reader.descriptors( + store_id="spool", object_key=entry.object_key): + staged[descriptor.metadata.capture_id] = reader.read_payload( + descriptor) + assert sorted(staged) == sorted(tensors) + for capture_id, tensor in tensors.items(): + assert staged[capture_id] == tensor.numpy().tobytes(), capture_id + + def test_close_alone_delivers_the_tail_to_the_catalog(fake_s3, tmp_path): """No flush_and_wait: close() must still seal the sink's open pack and drain it into the catalog. The 60 s linger means nothing but a flush can diff --git a/tests/test_native_capture_storage_live.py b/tests/test_native_capture_storage_live.py index 01218b5ba..7d89fa57b 100644 --- a/tests/test_native_capture_storage_live.py +++ b/tests/test_native_capture_storage_live.py @@ -483,6 +483,16 @@ def restore(self): self._delay_by = None def close(self): + """Refuse new connections and drop the live ones. Closing the + listener does not wake a thread blocked in accept(), which still + takes one more queued connection: the stall settings go first, so + that connection is refused rather than held open for its client's + whole timeout (a stop()'s lease release after a stall() did).""" + self._stalled = False + self._stall_if = None + self._slow_once = None + self._delay_by = None + self._late_by = 0.0 self.cut() self._listener.close() @@ -813,6 +823,956 @@ def _flush(): service.stop() +def test_a_flush_against_a_black_hole_catalog_overruns_by_one_request( + fake_s3, tmp_path): + """A flush's own cycle ran the periodic reconcile when it fell due: past + the index pass that failed on a catalog that accepts connections and + never answers, it listed the bucket and asked the catalog again, one + more request timeout past the deadline. A flush cycle skips the + reconcile now (the loop runs it), and its catalog requests, each + bounded, are never cut mid-flight: flush(1.0) returns within one second + plus the one request in flight at the deadline.""" + from dmi.storage.native_capture import _load_native_store_extension + + request_timeout = 4.0 + spool_root = tmp_path / "spool" + switch = _Switch(CLICKHOUSE_HOST, CLICKHOUSE_HTTP_PORT) + with _catalog() as (_client, catalog): + native = _storage_config( + fake_s3, catalog.table_prefix, + clickhouse_port=switch.port)._native_dict() + native.update( + spool_root=str(spool_root), holder="black-hole-test", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + # Every cycle is due a reconcile. + reconcile_interval_ns=1_000_000, + # No renewal falls due while the test runs. + lease_ttl_ns=30_000_000_000, publish_timeout_ns=5_000_000_000, + clickhouse_request_timeout_s=request_timeout) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + _stage(spool_root, range(2)) + switch.stall() + started = time.monotonic() + drained = service.flush(1.0) + elapsed = time.monotonic() - started + snapshot = service.snapshot() + finally: + switch.close() # releases the stalled connections + service.stop() + + assert drained is False + assert snapshot["reconcile_passes"] == 0, snapshot + assert elapsed < 1.0 + request_timeout + 1.0, elapsed + + +def test_the_stop_after_a_flush_runs_no_further_cycle(fake_s3, tmp_path): + """close() runs flush(budget) and then stop(). A loop that woke while + the flush held the cycle lock took the lock the moment the flush let + go -- before stop() could say anything -- and ran a whole cycle of its + own, which stop() then joined: against a catalog that accepts + connections and never answers, one more request timeout on top of the + flush's and the lease release's. The loop now waits out its interval + from the end of the last cycle, anyone's, so a stop() right after a + flush finds it waiting and it leaves at once.""" + from dmi.storage.native_capture import _load_native_store_extension + + request_timeout = 4.0 + spool_root = tmp_path / "spool" + switch = _Switch(CLICKHOUSE_HOST, CLICKHOUSE_HTTP_PORT) + with _catalog() as (_client, catalog): + native = _storage_config( + fake_s3, catalog.table_prefix, + clickhouse_port=switch.port)._native_dict() + native.update( + spool_root=str(spool_root), holder="stop-after-flush", + # Wakes while the flush below holds the cycle lock. + poll_interval_ns=3_000_000_000, reconcile_on_start=False, + # No renewal falls due while the test runs (a third of the TTL). + lease_ttl_ns=60_000_000_000, publish_timeout_ns=5_000_000_000, + clickhouse_request_timeout_s=request_timeout) + service = _load_native_store_extension().StorageService(native) + service.start() + stopped = False + try: + # Right after one of the loop's cycles, so that it next wakes + # while the flush's cycle waits on the catalog. + _wait_for(lambda: service.snapshot()["cycles"] >= 1, + timeout_s=10.0) + _stage(spool_root, range(2)) + switch.stall() + started = time.monotonic() + drained = service.flush(1.0) + flushed = time.monotonic() - started + cycles = service.snapshot()["cycles"] + started = time.monotonic() + service.stop() + stopped = True + stop_s = time.monotonic() - started + snapshot = service.snapshot() + finally: + switch.close() # releases the stalled connections + if not stopped: + service.stop() + + assert drained is False + assert flushed < 1.0 + request_timeout + 1.0, flushed + # No cycle after the flush's: stop() waited for the lease release + # alone, one request timeout against this catalog, not two. + assert snapshot["cycles"] == cycles, (cycles, stop_s, snapshot) + assert stop_s < request_timeout + 1.5, (stop_s, snapshot) + + +def test_a_flush_out_of_time_does_not_hash_the_spool(fake_s3, tmp_path): + """Past its deadline a flush's cycle starts no upload, but it still + listed the spool through the uploader, and a listing re-hashes every + staged pack: flush(0) -- what flush_and_wait and close() pass once the + sink's flush has spent the budget -- held its caller for as long as + hashing the whole backlog took. A listing already cancelled now stops + before the first pack's hash. One sparse 1 GiB pack stands in for a + backlog.""" + import hashlib + + from dmi.storage.native_capture import _load_native_store_extension + + size = 1 << 30 + digest = hashlib.sha256() + zeros = bytes(1 << 24) + for _ in range(size // len(zeros)): + digest.update(zeros) + spool_root = tmp_path / "spool" + spool_root.mkdir() + ready = spool_root / (f"{uuid.uuid4()}.1.1.{digest.hexdigest()}" + ".dmi-pack.ready") + with open(ready, "wb") as sparse: + sparse.truncate(size) # no blocks on disk + with _catalog() as (_client, catalog): + native = _storage_config(fake_s3, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="flush-out-of-time", + # The loop sleeps through the test, and nothing uploads. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + sweep_spool_on_start=False, + uploader_max_in_flight_bytes=2 * size) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + started = time.monotonic() + drained = service.flush(0.0) + elapsed = time.monotonic() - started + snapshot = service.snapshot() + finally: + service.stop() + + assert drained is False # a pack is staged + assert elapsed < 0.25, (elapsed, snapshot) + assert snapshot["uploaded_packs"] == 0, snapshot + # Not tried, so not cancelled either: it was never listed. + assert snapshot["cancelled_uploads"] == 0, snapshot + assert ready.exists() + + +def _sparse_backlog(spool_root: Path, packs: int, size: int) -> list: + """`packs` ready packs of `size` zero bytes, sparse on disk and named + for their checksum. Listing the spool hashes every byte of them, so + they stand in for a backlog without taking its disk.""" + import hashlib + + digest = hashlib.sha256() + zeros = bytes(1 << 20) + for _ in range(size // len(zeros)): + digest.update(zeros) + spool_root.mkdir(parents=True, exist_ok=True) + paths = [] + for _ in range(packs): + ready = spool_root / (f"{uuid.uuid4()}.1.1.{digest.hexdigest()}" + ".dmi-pack.ready") + with open(ready, "wb") as sparse: + sparse.truncate(size) + paths.append(ready) + return sorted(paths) + + +def test_stop_cuts_a_listing_that_hashes_a_backlog(fake_s3, tmp_path): + """A cycle lists the spool before it uploads, and a listing re-hashes + every staged pack: about 0.8 s a GiB here, over a backlog an outage + can grow to the spool's limit (a TiB by default). stop() waited for + the loop's listing to end before its cancel could turn a single + upload away. The listing stops between packs once cancelled now, so + stop() waits for one pack's hash at most, and nothing was tried. 16 + sparse packs of 256 MiB take the listing about 3 s.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + backlog = _sparse_backlog(spool_root, 16, 256 << 20) + with _catalog() as (_client, catalog): + native = _storage_config(fake_s3, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="listing-stop", + poll_interval_ns=20_000_000, reconcile_on_start=False, + sweep_spool_on_start=False, + uploader_max_in_flight_bytes=1 << 30) + service = _load_native_store_extension().StorageService(native) + service.start() + stopped = False + try: + time.sleep(0.5) # the loop's first cycle is listing the backlog + started = time.monotonic() + service.stop() + stopped = True + elapsed = time.monotonic() - started + snapshot = service.snapshot() + finally: + if not stopped: + service.stop() + + assert elapsed < 1.0, (elapsed, snapshot) + assert snapshot["uploaded_packs"] == 0, snapshot + # Never listed to the end, so never tried: nothing to cancel. + assert snapshot["cancelled_uploads"] == 0, snapshot + assert _ready(spool_root) == backlog + + +def test_a_flush_cuts_the_listing_its_deadline_passes_in(fake_s3, tmp_path): + """A flush whose deadline passes while its cycle lists the spool gave + up only once the listing had hashed the whole backlog, and then turned + every pack away. The listing stops at the deadline now.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + backlog = _sparse_backlog(spool_root, 16, 256 << 20) + with _catalog() as (_client, catalog): + native = _storage_config(fake_s3, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="listing-flush", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + sweep_spool_on_start=False, + uploader_max_in_flight_bytes=1 << 30) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + started = time.monotonic() + drained = service.flush(0.5) + elapsed = time.monotonic() - started + snapshot = service.snapshot() + finally: + service.stop() + + assert drained is False + assert elapsed < 1.2, (elapsed, snapshot) + assert snapshot["uploaded_packs"] == 0, snapshot + assert snapshot["cancelled_uploads"] == 0, snapshot + assert _ready(spool_root) == backlog + + +def _put(request: bytes) -> bool: + return request.startswith(b"PUT ") + + +def test_a_flush_returns_on_time_while_an_upload_stalls(fake_s3, tmp_path): + """The flush's cycle uploads through an object store that accepted the + PUT and never answers. It waited out the S3 read timeout on every + attempt; its deadline now cancels the upload, and the pack stays in the + spool, which is where a pack waits for a store that is not answering.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="stalled-upload-flush", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + s3_read_timeout_s=30) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + s3.stall_requests(_put) + tensors = _stage(spool_root, range(2)) + outcome = {} + + def _flush(): + started = time.monotonic() + outcome["drained"] = service.flush(1.0) + outcome["elapsed"] = time.monotonic() - started + + waiter = threading.Thread(target=_flush, daemon=True) + waiter.start() + waiter.join(timeout=15.0) + assert not waiter.is_alive(), "flush(1.0) still blocked after 15 s" + assert s3.stalled, "the upload never reached the store" + assert outcome["drained"] is False + assert outcome["elapsed"] < 3.0, outcome + assert len(_ready(spool_root)) == 1 + snapshot = service.snapshot() + # Cut short, not failed: nothing for the backoff to count. + assert snapshot["upload_failures"] == 0, snapshot + assert snapshot["cancelled_uploads"] == 1, snapshot + assert snapshot["uploaded_packs"] == 0, snapshot + + s3.restore() + assert service.flush(30.0) + # A flush out of time uploads nothing, but still finds a drained + # spool drained. + assert service.flush(0.0) + finally: + s3.close() + service.stop() + + direct = _storage_config(fake_s3, catalog.table_prefix) + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_a_flush_that_cuts_a_multipart_upload_waits_for_its_abort( + fake_s3, tmp_path): + """The one overrun of a flush's deadline that no catalog request + timeout bounds: a multipart upload the deadline cuts is then aborted, + one request of up to 5 s that nothing cuts, after up to about a second + for the stalled transfer to see the cancel (libcurl's progress poll). + With clickhouse_request_timeout_s at 2 s -- the config allows it -- that + outlasts the request timeout. Here the store holds the part and the + abort both, so the flush pays the whole bound, and no more. The pack is + a sparse file of zeros named for its checksum, over the client's + 64 MiB multipart threshold; it stays staged.""" + import hashlib + + from dmi.storage.native_capture import _load_native_store_extension + + size = 65 << 20 + digest = hashlib.sha256() + zeros = bytes(1 << 20) + for _ in range(size // len(zeros)): + digest.update(zeros) + spool_root = tmp_path / "spool" + spool_root.mkdir() + ready = spool_root / (f"{uuid.uuid4()}.1.1.{digest.hexdigest()}" + ".dmi-pack.ready") + with open(ready, "wb") as sparse: + sparse.truncate(size) + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="multipart-abort-flush", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + clickhouse_request_timeout_s=2.0, + publish_timeout_ns=1_000_000_000, s3_read_timeout_s=30) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + s3.stall_requests(_multipart_part_or_abort) + started = time.monotonic() + drained = service.flush(1.0) + elapsed = time.monotonic() - started + snapshot = service.snapshot() + finally: + s3.close() + service.stop() + + assert drained is False + assert len(s3.stalled) == 2, s3.stalled # the part, then the abort + # The deadline, up to a second for the cut, and the 5 s abort. + assert elapsed < 1.0 + 1.0 + 5.0 + 1.0, (elapsed, snapshot) + assert snapshot["cancelled_uploads"] == 1, snapshot + assert snapshot["upload_failures"] == 0, snapshot + assert _ready(spool_root) == [ready] + + +def test_stop_returns_promptly_while_an_upload_stalls(fake_s3, tmp_path): + """stop() joined a loop whose cycle was inside a PUT the store never + answers, so it waited out the S3 read timeout on every attempt, with the + lease held. It cancels the upload now: the pack stays in the spool, the + lease is released, and the next process uploads the pack.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="stalled-upload-stop", + poll_interval_ns=20_000_000, reconcile_on_start=False, + s3_read_timeout_s=60) + service = _load_native_store_extension().StorageService(native) + service.start() + stopper = None + try: + s3.stall_requests(_put) + tensors = _stage(spool_root, range(2)) + _wait_for(lambda: s3.stalled, timeout_s=10.0) + outcome = {} + + def _stop(): + started = time.monotonic() + service.stop() + outcome["elapsed"] = time.monotonic() - started + + stopper = threading.Thread(target=_stop, daemon=True) + stopper.start() + stopper.join(timeout=15.0) + assert not stopper.is_alive(), "stop() still blocked after 15 s" + assert outcome["elapsed"] < 3.0, outcome + snapshot = service.snapshot() + assert snapshot["lease_state"] == "released", snapshot + assert snapshot["upload_failures"] == 0, snapshot + assert snapshot["cancelled_uploads"] == 1, snapshot + assert len(_ready(spool_root)) == 1 + finally: + s3.close() # releases the stalled PUT + # Never a second stop() beside one still running. + if stopper is None: + service.stop() + else: + stopper.join(timeout=120.0) + + direct = _storage_config(fake_s3, catalog.table_prefix) + successor = _service(direct, spool_root) + successor.start() + try: + successor.flush(30.0) + finally: + successor.stop() + assert _ready(spool_root) == [] + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_a_flush_returns_on_time_while_the_index_reads_stall( + fake_s3, tmp_path): + """The store takes the PUTs, then stops answering GETs. The flush's + cycle indexed what it had uploaded through a client no cancel could + cut, and a pack whose read failed only moved the pass on to the next + pack: 3 packs cost 3 x (4 attempts x s3_read_timeout_s + backoff), 40 s + at a 3 s read timeout and 8 minutes a pack at the defaults. The reads + are cut one request timeout past the deadline now, and the pass ends at + the first read the store does not answer: the packs stay owed, none is + counted against or set aside, and they index once the store answers.""" + from dmi.storage.native_capture import _load_native_store_extension + + request_timeout = 4.0 + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="stalled-read-flush", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + s3_read_timeout_s=3, + clickhouse_request_timeout_s=request_timeout) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + tensors = _stage(spool_root, range(6)) # three packs + s3.stall_requests(_ranged_get) # only the indexer's reads + outcome = {} + + def _flush(): + started = time.monotonic() + outcome["drained"] = service.flush(1.0) + outcome["elapsed"] = time.monotonic() - started + + waiter = threading.Thread(target=_flush, daemon=True) + waiter.start() + waiter.join(timeout=90.0) + assert not waiter.is_alive(), "flush(1.0) still blocked after 90 s" + snapshot = service.snapshot() + assert outcome["drained"] is False + assert outcome["elapsed"] < 1.0 + request_timeout + 1.5, ( + outcome, snapshot) + # One read's attempts, not every pack's. + assert 1 <= len(s3.stalled) <= 2, (s3.stalled, snapshot) + assert snapshot["uploaded_packs"] == 3, snapshot + assert snapshot["pending_index"] == 3, snapshot + assert snapshot["indexed_packs"] == 0, snapshot + # Cut short by the deadline, not the packs' fault. + assert snapshot["index_failures"] == 0, snapshot + assert snapshot["rejected_packs"] == 0, snapshot + assert _ready(spool_root) == [] + + s3.cut() # releases the stalled reads + s3.restore() + assert service.flush(30.0) + finally: + s3.close() + service.stop() + + direct = _storage_config(fake_s3, catalog.table_prefix) + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_stop_returns_promptly_while_the_index_reads_stall(fake_s3, tmp_path): + """stop() cancelled the uploads but joined a cycle whose index pass + read, through the uncancellable client, every pack it had uploaded: a + GET the store never answers cost 4 attempts x s3_read_timeout_s each. + stop() cuts the reads too now. The packs it leaves unindexed are in the + bucket, where the next start's reconcile finds them.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="stalled-read-stop", + poll_interval_ns=20_000_000, reconcile_on_start=False, + s3_read_timeout_s=3) + service = _load_native_store_extension().StorageService(native) + service.start() + stopper = None + try: + s3.stall_requests(_ranged_get) + tensors = _stage(spool_root, range(6)) + _wait_for(lambda: s3.stalled, timeout_s=10.0) + outcome = {} + + def _stop(): + started = time.monotonic() + service.stop() + outcome["elapsed"] = time.monotonic() - started + + stopper = threading.Thread(target=_stop, daemon=True) + stopper.start() + stopper.join(timeout=90.0) + assert not stopper.is_alive(), "stop() still blocked after 90 s" + snapshot = service.snapshot() + assert outcome["elapsed"] < 3.0, (outcome, snapshot) + assert snapshot["lease_state"] == "released", snapshot + assert snapshot["rejected_packs"] == 0, snapshot + finally: + s3.close() # releases the stalled reads + if stopper is None: + service.stop() + else: + stopper.join(timeout=120.0) + + direct = _storage_config(fake_s3, catalog.table_prefix) + successor = _service(direct, spool_root) # reconciles at start + successor.start() + try: + successor.flush(30.0) + finally: + successor.stop() + assert _ready(spool_root) == [] + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_stop_sends_no_catalog_request_for_packs_it_cannot_index( + fake_s3, tmp_path): + """stop() cuts the index reads for good, but a cycle whose uploads it + cut still began the index pass over the packs it had uploaded: the + first batch's replay guard, a catalog SELECT under the lease lock, went + out after stop(), and its answer was thrown away when the read that + followed was refused. stop() waited for that statement too, up to a + request timeout against a slow catalog. No batch starts once the reads + are cut now; the packs stay owed, as the read would have left them. + Two of the four packs upload, the third's PUT is held, and the + catalog answers the replay guard 3 s late.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + switch = _Switch(CLICKHOUSE_HOST, CLICKHOUSE_HTTP_PORT) + puts = [] + guard_reads = [] + + def _third_put_on(request: bytes) -> bool: + if not _put(request): + return False + puts.append(time.monotonic()) + return len(puts) >= 3 + + def _slow_guard(request: bytes) -> float: + if not _replay_guard_read(request): + return 0.0 + guard_reads.append(time.monotonic()) + return 3.0 + + with _catalog() as (_client, catalog): + native = _storage_config( + s3.url, catalog.table_prefix, + clickhouse_port=switch.port)._native_dict() + native.update( + spool_root=str(spool_root), holder="stop-no-replay-guard", + poll_interval_ns=20_000_000, reconcile_on_start=False, + uploader_max_workers=1, s3_read_timeout_s=30) + # Staged first, so the loop's first cycle lists all four. + tensors = _stage(spool_root, range(8)) + s3.stall_requests(_third_put_on) + switch.delay_requests(_slow_guard) + service = _load_native_store_extension().StorageService(native) + service.start() + stopper = None + try: + _wait_for(lambda: s3.stalled, timeout_s=10.0) + outcome = {} + + def _stop(): + outcome["started"] = time.monotonic() + service.stop() + outcome["elapsed"] = time.monotonic() - outcome["started"] + + stopper = threading.Thread(target=_stop, daemon=True) + stopper.start() + stopper.join(timeout=30.0) + assert not stopper.is_alive(), "stop() still blocked after 30 s" + snapshot = service.snapshot() + finally: + s3.close() # releases the held PUT + switch.close() + if stopper is None: + service.stop() + else: + stopper.join(timeout=120.0) + + late = [t for t in guard_reads if t >= outcome["started"]] + assert late == [], (late, outcome, snapshot) + assert outcome["elapsed"] < 2.0, (outcome, snapshot) + assert snapshot["lease_state"] == "released", snapshot + assert snapshot["uploaded_packs"] == 2, snapshot + assert snapshot["pending_index"] == 2, snapshot + assert snapshot["index_failures"] == 0, snapshot + assert len(_ready(spool_root)) == 2 + + direct = _storage_config(fake_s3, catalog.table_prefix) + successor = _service(direct, spool_root) # reconciles at start + successor.start() + try: + successor.flush(30.0) + finally: + successor.stop() + assert _ready(spool_root) == [] + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_an_object_store_read_outage_sets_no_pack_aside(fake_s3, tmp_path): + """A read the store never answered counted against the pack, as if the + pack itself were unreadable: after max_index_attempts cycles of an + outage the pack was set aside for good -- left in the bucket, out of + the flush boundary, and reported by flush() as never indexable. A pass + that meets an unanswered read now ends there and counts nothing + against the pack.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="read-outage", + poll_interval_ns=20_000_000, max_backoff_ns=200_000_000, + reconcile_on_start=False, s3_read_timeout_s=1, + s3_max_attempts=1, max_index_attempts=2) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + s3.stall_requests(_ranged_get) + tensors = _stage(spool_root, range(2)) # one pack + # Three passes over the pack, each ended by an unanswered read. + _wait_for(lambda: service.snapshot()["index_failures"] >= 3, + timeout_s=30.0) + snapshot = service.snapshot() + assert snapshot["rejected_packs"] == 0, snapshot + assert snapshot["pending_index"] == 1, snapshot + + s3.cut() + s3.restore() + assert service.flush(30.0) # raised "set aside" before + snapshot = service.snapshot() + assert snapshot["indexed_packs"] == 1, snapshot + assert snapshot["rejected_packs"] == 0, snapshot + finally: + s3.close() + service.stop() + + direct = _storage_config(fake_s3, catalog.table_prefix) + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_a_flush_against_a_slow_catalog_indexes_one_batch_past_its_deadline( + fake_s3, tmp_path): + """A flush's cycle checked its deadline only while uploading: past it, + it still indexed everything it had uploaded, batch after batch. Against + a catalog that answers slowly but inside the request timeout no request + fails, so the overrun grew with the batches -- 8 one-pack batches at + 0.4 s a statement ran a 1 s flush for 59 s. Past the deadline a flush + then started no further batch, and left the rest of what it had + uploaded owed: gone from the spool, and remembered only by the service + -- so close(), which stops the service after its flush, dropped it, + and only a later start's reconcile could find those packs in the + bucket; with reconcile_on_start off, none did. A cycle now uploads a + chunk of indexer_max_packs at a time and indexes it before the next, so + past the deadline it indexes the one chunk it uploaded and leaves the + rest in the spool, which any later start uploads. (Its index reads are + cut one request timeout past the deadline too; the request timeout + here outlasts the batch, so that is not what ends it.)""" + from dmi.storage.native_capture import _load_native_store_extension + + request_timeout = 60.0 + spool_root = tmp_path / "spool" + switch = _Switch(CLICKHOUSE_HOST, CLICKHOUSE_HTTP_PORT) + with _catalog() as (_client, catalog): + native = _storage_config( + fake_s3, catalog.table_prefix, + clickhouse_port=switch.port)._native_dict() + native.update( + spool_root=str(spool_root), holder="slow-catalog-flush", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + indexer_max_packs=1, + clickhouse_request_timeout_s=request_timeout) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + tensors = _stage(spool_root, range(16)) # eight packs + switch.delay_requests(_slow_catalog(0.4, 0.4)) + # close()'s order: a flush on the budget, then stop(). + started = time.monotonic() + drained = service.flush(1.0) + elapsed = time.monotonic() - started + snapshot = service.snapshot() + finally: + service.stop() + switch.close() + + # One batch past the deadline -- about 7 s of statements at 0.4 s + # each, the lease's included -- not eight (59 s before). + assert elapsed < 1.0 + 12.0, (elapsed, snapshot) + assert drained is False + assert snapshot["uploaded_packs"] == 1, snapshot + assert snapshot["indexed_packs"] == 1, snapshot + # Nothing owed for stop() to drop: the rest never left the spool. + assert snapshot["pending_index"] == 0, snapshot + assert snapshot["index_failures"] == 0, snapshot + assert len(_ready(spool_root)) == 7 + + # A successor that does not reconcile -- a shared bucket's setting + # -- still gets every capture into the catalog. + direct = _storage_config(fake_s3, catalog.table_prefix, + reconcile_on_start=False) + successor = _service(direct, spool_root) + successor.start() + try: + successor.flush(60.0) + finally: + successor.stop() + assert _ready(spool_root) == [] + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_a_close_whose_budget_ends_mid_upload_leaves_nothing_owed( + fake_s3, tmp_path): + """close() flushes the service on what is left of its budget, then + stops it. The flush's cycle uploaded pack after pack until the + deadline, each deleted from the spool once its upload was verified, + and past the deadline indexed one batch of them: the rest were owed, + in memory only, and stop() dropped them -- in the bucket, out of the + spool and out of the catalog, for a later start's reconcile alone to + find. A cycle uploads a chunk of indexer_max_packs at a time now and + indexes it before the next, so at the deadline at most the chunk in + flight is owed, and the flush indexes it; the rest stays in the spool. + Here uploads take about 0.3 s a pack, one at a time, so a 2 s budget + runs out a few packs into the twelve.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="close-mid-upload", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + uploader_max_workers=1, indexer_max_packs=4) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + tensors = _stage(spool_root, range(24)) # twelve packs + s3.delay_requests(lambda request: 0.3 if _put(request) else 0.0) + # close()'s order: a flush on the budget, then stop(). + drained = service.flush(2.0) + snapshot = service.snapshot() + finally: + service.stop() + s3.close() + + assert drained is False + assert 0 < snapshot["uploaded_packs"] < 12, snapshot + assert snapshot["pending_index"] == 0, snapshot + assert snapshot["indexed_packs"] == snapshot["uploaded_packs"], snapshot + assert len(_ready(spool_root)) == 12 - snapshot["uploaded_packs"] + + # A successor that does not reconcile -- a shared bucket's setting + # -- still gets every capture into the catalog. + direct = _storage_config(fake_s3, catalog.table_prefix, + reconcile_on_start=False) + successor = _service(direct, spool_root) + successor.start() + try: + successor.flush(60.0) + finally: + successor.stop() + assert _ready(spool_root) == [] + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_the_one_batch_past_a_flushs_deadline_is_a_full_one(fake_s3, + tmp_path): + """Past its deadline a flush starts no index batch but the first. The + pass cut its packs into batches of indexer_max_packs from the front and + then took them from the back, so that first batch was the remainder: + three packs in batches of two indexed one, and a remainder of one pack + is what a flush out of time made progress by however many it owed. The + first batch is a full one now.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + switch = _Switch(CLICKHOUSE_HOST, CLICKHOUSE_HTTP_PORT) + with _catalog() as (_client, catalog): + native = _storage_config( + fake_s3, catalog.table_prefix, + clickhouse_port=switch.port)._native_dict() + native.update( + spool_root=str(spool_root), holder="full-batch-flush", + # The loop sleeps through the test, so the flush runs the cycle. + poll_interval_ns=60_000_000_000, reconcile_on_start=False, + indexer_max_packs=2, clickhouse_request_timeout_s=60.0) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + tensors = _stage(spool_root, range(6)) # three packs + # Slow enough that one batch outlasts the 1 s deadline. + switch.delay_requests(_slow_catalog(0.4, 0.4)) + drained = service.flush(1.0) + snapshot = service.snapshot() + assert drained is False + assert snapshot["indexed_packs"] == 2, snapshot + assert snapshot["index_failures"] == 0, snapshot + + switch.restore() + assert service.flush(60.0) + assert service.snapshot()["indexed_packs"] == 3 + finally: + switch.close() + service.stop() + + direct = _storage_config(fake_s3, catalog.table_prefix) + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_only_the_loop_reconciles_never_a_flush(fake_s3, tmp_path): + """A flush's cycles skip the periodic reconcile, which lists the whole + bucket and asks the catalog about every page -- work no deadline + bounds; the loop runs it. A reconcile is due on every cycle here, and + the loop's first wake is 2 s off: the flush drains without one, and the + loop's first cycle runs it.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + with _catalog() as (_client, catalog): + native = _storage_config(fake_s3, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="flush-no-reconcile", + poll_interval_ns=2_000_000_000, reconcile_on_start=False, + reconcile_interval_ns=1_000_000) + service = _load_native_store_extension().StorageService(native) + service.start() + try: + tensors = _stage(spool_root, range(2)) + assert service.flush(30.0) + snapshot = service.snapshot() + assert snapshot["indexed_packs"] == 1, snapshot + assert snapshot["reconcile_passes"] == 0, snapshot + _wait_for(lambda: service.snapshot()["reconcile_passes"] >= 1, + timeout_s=10.0) + finally: + service.stop() + + direct = _storage_config(fake_s3, catalog.table_prefix) + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_a_service_started_again_after_stop_uploads_again(fake_s3, tmp_path): + """stop() cancels the service's uploads and reads for good, and start() + clears that: without it, a service object stopped and started again + would cancel every upload from then on, its flushes returning False + with the packs piling up in the spool.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + with _catalog() as (_client, catalog): + native = _storage_config(fake_s3, catalog.table_prefix)._native_dict() + native.update(spool_root=str(spool_root), holder="restarted", + reconcile_on_start=False) + service = _load_native_store_extension().StorageService(native) + service.start() + service.stop() + service.start() + try: + tensors = _stage(spool_root, range(2)) + assert service.flush(10.0), service.snapshot() + snapshot = service.snapshot() + assert snapshot["cancelled_uploads"] == 0, snapshot + assert snapshot["indexed_packs"] == 1, snapshot + finally: + service.stop() + assert service.snapshot()["lease_state"] == "released" + + direct = _storage_config(fake_s3, catalog.table_prefix) + assert sorted(_read_all(direct)) == sorted(tensors) + + +def test_stop_cuts_a_reconcile_whose_listing_stalls(fake_s3, tmp_path): + """The loop's periodic reconcile lists the bucket through the client + the index reads use. A listing the store never answers held stop() for + s3_read_timeout_s on each of the client's attempts; stop() cuts it now, + and the pass ends quietly, to run again on the next start.""" + from dmi.storage.native_capture import _load_native_store_extension + + spool_root = tmp_path / "spool" + s3 = _Switch.to_url(fake_s3) + with _catalog() as (_client, catalog): + native = _storage_config(s3.url, catalog.table_prefix)._native_dict() + native.update( + spool_root=str(spool_root), holder="stalled-reconcile-stop", + poll_interval_ns=20_000_000, reconcile_on_start=False, + reconcile_interval_ns=1_000_000, s3_read_timeout_s=30) + service = _load_native_store_extension().StorageService(native) + service.start() + stopper = None + try: + s3.stall_requests(_listing) + _wait_for(lambda: s3.stalled, timeout_s=10.0) + outcome = {} + + def _stop(): + started = time.monotonic() + service.stop() + outcome["elapsed"] = time.monotonic() - started + + stopper = threading.Thread(target=_stop, daemon=True) + stopper.start() + stopper.join(timeout=60.0) + assert not stopper.is_alive(), "stop() still blocked after 60 s" + snapshot = service.snapshot() + assert outcome["elapsed"] < 3.0, (outcome, snapshot) + assert snapshot["lease_state"] == "released", snapshot + assert snapshot["reconcile_passes"] == 0, snapshot + # Quietly: a listing stop() cut is no failed pass. + assert "reconcile" not in snapshot["last_error"], snapshot + assert len(s3.stalled) == 1, s3.stalled + finally: + s3.close() # releases the stalled listing + if stopper is None: + service.stop() + else: + stopper.join(timeout=120.0) + + def test_dropping_a_running_service_does_not_hold_the_gil(fake_s3, tmp_path): """A service collected without stop() stops itself in its destructor, joining a cycle that may be waiting on the catalog. That ran with the @@ -1770,13 +2730,36 @@ def test_the_lease_renews_while_start_reconciles( assert sorted(captures) == sorted(tensors) +def _multipart_part_or_abort(request: bytes) -> bool: + line = request.partition(b"\r\n")[0] + return ((line.startswith(b"PUT ") and b"partNumber=" in line) + or (line.startswith(b"DELETE ") and b"uploadId=" in line)) + + def test_the_lease_renews_until_the_loops_last_cycle_is_done( fake_s3, tmp_path): """stop() woke the loop and the lease thread together, and the lease - thread left at once while the loop's last cycle still ran: a pass - reading a pack then found its lease abandoned at the commit, and the - snapshot said "held" over a dead row until then. The lease thread now - stops only after the loop has.""" + thread left at once while the loop's last cycle still ran: the snapshot + said "held" over a row that had expired meanwhile, and stop() found its + lease abandoned instead of releasing it. The lease thread now stops only + after the loop has. + + stop() cuts the cycle's object-store requests short, so the cycle needs + something to do that no cancel cuts, outside the lease lock (inside it, + each request renews the lease when due). A multipart upload that stop() + cuts is aborted with a request of its own, one attempt of up to 5 s -- + here held for all of it, past the 3 s lease -- while the lease thread + must go on renewing. (It used to be an index read held past the lease + deadline; stop() cuts reads now.) The pack is a sparse file of zeros + named for its checksum, over the client's 64 MiB multipart threshold: + it is only ever uploaded, and stays staged.""" + import hashlib + + size = 65 << 20 + digest = hashlib.sha256() + zeros = bytes(1 << 20) + for _ in range(size // len(zeros)): + digest.update(zeros) spool_root = tmp_path / "spool" s3 = _Switch.to_url(fake_s3) with _catalog() as (client, catalog): @@ -1785,30 +2768,32 @@ def test_the_lease_renews_until_the_loops_last_cycle_is_done( publish_timeout_s=1) service = _service(config, spool_root) service.start() - reads = [] - - def _first_read(request: bytes) -> bool: - if _ranged_get(request): - reads.append(time.monotonic()) - return True - return False - try: - s3.slow_once(_first_read, 4.0) # past the lease deadline - tensors = _stage(spool_root, range(2)) - _wait_for(lambda: reads, timeout_s=10.0) + s3.stall_requests(_multipart_part_or_abort) + ready = spool_root / (f"{uuid.uuid4()}.1.1.{digest.hexdigest()}" + ".dmi-pack.ready") + with open(ready, "wb") as sparse: + sparse.truncate(size) + _wait_for(lambda: s3.stalled, timeout_s=10.0) # a part, held + started = time.monotonic() log = _sample_during(service.stop, service, client, catalog.table_prefix) + elapsed = time.monotonic() - started snapshot = service.snapshot() finally: service.stop() s3.close() assert not _held_but_dead(log), log - assert snapshot["indexed_packs"] == 1, snapshot - assert snapshot["lease_state"] == "released", snapshot - captures = _read_all(_storage_config(fake_s3, catalog.table_prefix)) - assert sorted(captures) == sorted(tensors) + assert snapshot["lease_state"] == "released", (snapshot, log) + # The part was cut and the abort held to its 5 s bound: the loop's + # last cycle outlived the lease, which kept renewing throughout. + assert len(s3.stalled) == 2, s3.stalled + assert 4.5 < elapsed < 9.0, (elapsed, snapshot) + assert snapshot["lease_renewals"] >= 3, snapshot + assert snapshot["cancelled_uploads"] == 1, snapshot + assert snapshot["upload_failures"] == 0, snapshot + assert _ready(spool_root) == [ready] def _slow_catalog(insert_s: float, read_s: float): diff --git a/tests/test_native_capture_storage_wiring.py b/tests/test_native_capture_storage_wiring.py index 21cead936..f37e686da 100644 --- a/tests/test_native_capture_storage_wiring.py +++ b/tests/test_native_capture_storage_wiring.py @@ -113,6 +113,30 @@ def test_https_with_the_insecure_flag_is_refused(): s3_allow_insecure_http=True) +def test_the_object_store_bounds_default_to_the_native_client_and_reach_it(): + # The read timeout and attempts bound how long one read the store never + # answers holds the service's cycle; they were fixed at the native + # defaults (native/csrc/store/s3_client.h S3Config) with no way to + # lower them. + config = _storage_config() + assert (config.s3_read_timeout_s, config.s3_max_attempts) == (120, 4) + native = _storage_config(s3_read_timeout_s=10, s3_max_attempts=2) + assert native._native_dict()["s3_read_timeout_s"] == 10 + assert native._native_dict()["s3_max_attempts"] == 2 + # The reader reads through the same client configuration. + assert native._native_reader_dict()["s3_read_timeout_s"] == 10 + + +@pytest.mark.parametrize("name, value", [ + ("s3_read_timeout_s", 0), ("s3_read_timeout_s", -1), + ("s3_read_timeout_s", 1.5), ("s3_read_timeout_s", True), + ("s3_read_timeout_s", 86_401), ("s3_max_attempts", 0), + ("s3_max_attempts", 2.0), ("s3_max_attempts", False)]) +def test_the_object_store_bounds_must_be_positive_ints(name, value): + with pytest.raises(ValueError, match=name): + _storage_config(**{name: value}) + + # --- the catalog connection: scheme, credentials, TLS -------------------------- diff --git a/tests/test_native_pack_read_cancel.py b/tests/test_native_pack_read_cancel.py new file mode 100644 index 000000000..879fda6e8 --- /dev/null +++ b/tests/test_native_pack_read_cancel.py @@ -0,0 +1,99 @@ +"""How a pack's failed index read is booked: the store's, or a cancel's. + +The storage service reads a pack's trailer and footer through an S3 client +holding a Cancellation -- stop()'s, and a flush's read deadline. A read the +store did not answer ends the index pass: recorded as the cycle's error, +counted towards the backoff. A read the cancel cut is no failure at all: +the pack is deferred, nothing recorded. read_pack_descriptor_rows told the +two apart by asking the Cancellation, when the read had failed, whether it +was set -- which a flush's deadline makes true from the moment it passes, +whatever ended the read. A store failure answered in the gap between the +deadline and libcurl's next progress poll was then booked as a cancel, and +the flush's TimeoutError named an older error. It asks the read itself now, +as the uploader does since b61db70. + +That gap is under a poll interval wide, so the driver stands in for it: +conformance_catalog's session-less read_pack_rows op cancels its client's +Cancellation once an exchange has returned (cancel_after_exchange), through +the S3 client's test seam, or arms its deadline cancel_after_ms from the +call. No catalog is involved. + +Build: make -C native build/conformance_catalog +""" + +from __future__ import annotations + +import json +import subprocess +import time +import uuid +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +DRIVER = REPO_ROOT / "native" / "build" / "conformance_catalog" + +pytestmark = [ + pytest.mark.cpu, + pytest.mark.skipif( + not DRIVER.exists(), + reason="native/build/conformance_catalog is not built; run " + "`make -C native build/conformance_catalog`", + ), +] + +from tests.test_native_s3_client import ( # noqa: E402 + ACCESS, + BUCKET, + REGION, + SECRET, + fake_s3, +) + + +def _read(endpoint: str, key: str, **fields) -> dict: + request = { + "op": "read_pack_rows", "endpoint": endpoint, "bucket": BUCKET, + "region": REGION, "access": ACCESS, "secret": SECRET, + "insecure": True, "read_timeout": 10, "max_attempts": 1, + "ref": { + "pack_id": str(uuid.uuid4()), "store_id": "s3", + "object_key": key, "object_bytes": 4096, "checksum": "0" * 64, + "record_count": 1, + }, + **fields, + } + proc = subprocess.run([str(DRIVER)], input=json.dumps(request) + "\n", + capture_output=True, text=True, timeout=60) + lines = [line for line in proc.stdout.splitlines() if line.strip()] + assert lines, f"driver produced no output: {proc.stderr}" + return json.loads(lines[0]) + + +@pytest.mark.parametrize("request_cut", [False, True], + ids=["failed-on-its-own", "cut-by-the-cancel"]) +def test_a_read_counts_as_cancelled_only_when_the_cancel_ended_it( + fake_s3, request_cut): + """fault/always-500 answers the trailer read 500 at once, and the + Cancellation is cancelled as that answer comes back -- a cancel that + comes in while the failure is reported: the store failed the read, and + it is booked so. fault/hang holds the read for 5 s, and a deadline + 0.2 s in cuts it: that is a cancel.""" + if request_cut: + key = f"fault/hang/{uuid.uuid4()}.dmi-pack" + armed = {"cancel_after_ms": 200} + else: + key = f"fault/always-500/{uuid.uuid4()}.dmi-pack" + armed = {"cancel_after_exchange": True} + started = time.monotonic() + result = _read(fake_s3, key, **armed) + elapsed = time.monotonic() - started + assert not result["ok"], result + assert result["error"] == "StoreUnavailable", result + assert result["cancelled"] is request_cut, (result, elapsed) + if request_cut: + assert "cancel" in result["message"], result + assert elapsed < 2.5, (elapsed, result) + else: + assert "HTTP 500" in result["message"], result diff --git a/tests/test_native_pack_sink_timeout.py b/tests/test_native_pack_sink_timeout.py index ce78f76ba..b807a95d1 100644 --- a/tests/test_native_pack_sink_timeout.py +++ b/tests/test_native_pack_sink_timeout.py @@ -14,6 +14,10 @@ takes: the whole envelope must time out within one admission_timeout_s, not one per row. +A third wedges the stager under one flush and gives a second, shorter +flush on another thread its own timeout: the second must return within it, +not wait out the first (the release backstop beside a flush_and_wait). + This is the native-tier coverage of the deadline wait loop in PackSink::Submit; the conformance-driver suite (test_native_pack_sink.py) covers the admission bounds around it. diff --git a/tests/test_native_s3_client.py b/tests/test_native_s3_client.py index f47e81589..52178e84c 100644 --- a/tests/test_native_s3_client.py +++ b/tests/test_native_s3_client.py @@ -185,11 +185,27 @@ def under(name: str) -> bool: return 500, b"boom" if under("fault/always-500"): return 500, b"boom" + if under("fault/hang-put") and self.command == "PUT" and \ + "partNumber=" not in self.path: + # Only a single-request PUT: the upload's HEADs are answered. + time.sleep(5) + return None if under("fault/forbidden"): return 403, b"no" if under("fault/hang"): time.sleep(5) return None + if under("fault/hang-parts") and self.command == "PUT" and \ + "partNumber=" in self.path: + # Only a multipart upload's parts: its create and its abort are + # answered at once. + time.sleep(5) + return None + if under("fault/hang-complete") and self.command == "POST" and \ + "uploadId=" in self.path: + # Only CompleteMultipartUpload, with every part already in. + time.sleep(5) + return None return None def _route(self): @@ -868,3 +884,106 @@ def test_a_missing_ca_is_named_before_any_request(fake_s3_tls, tmp_path): key="anything") assert not head["ok"] and str(missing) in head["what"], head assert STATE.calls == [] + + +# --- cancellation -------------------------------------------------------------- +# +# The storage service cancels the object store work of its uploads when it +# stops, and when a flush's deadline passes: an upload stalled on a server +# that accepted it and never answers must not hold stop() or flush() for +# read_timeout x max_attempts. `cancel_after_ms` makes the driver cancel the +# client that long after the request starts. + + +def _timed_call(op: str, **fields) -> tuple[dict, float]: + started = time.monotonic() + result = _call(op, **fields) + return result, time.monotonic() - started + + +def test_a_cancel_aborts_a_put_in_flight(fake_s3): + """fault/hang holds the request for 5 s; the cancel lands at 0.3 s, and + libcurl's progress callback, which runs at least once a second while a + transfer waits, aborts it.""" + put, elapsed = _timed_call( + "put", **_base(fake_s3, read_timeout=30), key="fault/hang/a", + data_b64=base64.b64encode(b"data").decode(), metadata={}, + content_type="application/octet-stream", cancel_after_ms=300) + assert not put["ok"], put + assert put.get("cancelled") is True, put + assert "cancel" in put["what"], put + assert put["attempts"] == 1, put # a cancelled request is not retried + assert elapsed < 3.0, elapsed + + +def test_a_cancel_aborts_a_multipart_upload_and_says_so_to_the_store(fake_s3): + """A cancelled part leaves an upload the store would keep, invisible and + billed, until a lifecycle rule reaps it: the client aborts it, with a + request of its own that the cancel does not cut.""" + payload = bytes((i * 13) & 0xFF for i in range(10 * MIB)) + put, elapsed = _timed_call( + "put", **_base(fake_s3, read_timeout=30), key="fault/hang-parts/b", + data_b64=base64.b64encode(payload).decode(), metadata={}, + content_type="application/vnd.dmi.pack", + multipart_threshold=5 * MIB, multipart_chunk=5 * MIB, + cancel_after_ms=300) + assert not put["ok"], put + assert put.get("cancelled") is True, put + assert elapsed < 4.0, elapsed + aborts = [call for call in STATE.calls + if call["method"] == "DELETE" and "uploadId=" in call["path"]] + assert len(aborts) == 1, STATE.calls + assert STATE.uploads == {} + assert "fault/hang-parts/b" not in STATE.objects + + +def test_a_cancel_that_cuts_the_complete_still_aborts_the_upload(fake_s3): + """Cut short with every part sent, CompleteMultipartUpload may or may + not have taken effect. The client aborts the upload either way: a + completed one refuses the abort harmlessly, and one left open would + otherwise sit in the bucket, invisible and billed.""" + payload = bytes((i * 7) & 0xFF for i in range(10 * MIB)) + put, elapsed = _timed_call( + "put", **_base(fake_s3, read_timeout=30), key="fault/hang-complete/d", + data_b64=base64.b64encode(payload).decode(), metadata={}, + content_type="application/vnd.dmi.pack", + multipart_threshold=5 * MIB, multipart_chunk=5 * MIB, + cancel_after_ms=1000) + assert not put["ok"], put + assert put.get("cancelled") is True, put + assert "CompleteMultipartUpload" in put["what"], put + assert elapsed < 4.0, elapsed + completes = [call for call in STATE.calls + if call["method"] == "POST" and "uploadId=" in call["path"]] + assert len(completes) == 1, STATE.calls + aborts = [call for call in STATE.calls + if call["method"] == "DELETE" and "uploadId=" in call["path"]] + assert len(aborts) == 1, STATE.calls + assert STATE.uploads == {} + + +def test_a_cancel_interrupts_the_retry_backoff(fake_s3): + """Ten attempts against a store that always answers 500 back off for + about 26 s in all; a cancel wakes the backoff instead of sleeping it + out. The attempts start at 0, 0.2, 0.6 and 1.4 s, and the cancel lands + early in the 1.6 s backoff after the fourth: slept out, that backoff + ends at 3.0 s, where the check before the fifth attempt would stop it + anyway -- so only a return well before then shows the backoff woke.""" + head, elapsed = _timed_call( + "head", **_base(fake_s3, max_attempts=10), key="fault/always-500/c", + cancel_after_ms=1500) + assert not head["ok"], head + assert head.get("cancelled") is True, head + assert head["attempts"] == 4, head + assert elapsed < 2.2, elapsed + + +def test_an_uncancelled_request_is_untouched_by_the_cancel_hook(fake_s3): + """A cancel armed for later than the request takes changes nothing.""" + put = _call("put", **_base(fake_s3), key="packs/on-time", + data_b64=base64.b64encode(b"data").decode(), metadata={}, + content_type="application/octet-stream", + cancel_after_ms=30_000) + assert put["ok"], put + assert put.get("cancelled") is False, put + assert STATE.objects["packs/on-time"]["body"] == b"data" diff --git a/tests/test_native_sink_release.py b/tests/test_native_sink_release.py new file mode 100644 index 000000000..84ac52fee --- /dev/null +++ b/tests/test_native_sink_release.py @@ -0,0 +1,245 @@ +"""The native pack sink stages its open pack when a ring releases it. + +A RingEngine stops by draining its record worker into the sink and then +releasing the sink's lease -- without flushing it. The sink seals a pack +only when it fills, when it has lingered ``max_linger_ns``, or on a flush, +so the records of the pack still open at that moment stayed in memory: the +audit's probe counted 0 ready packs right after the ring stopped (1 once the +linger fired), and a process that exited before then lost them. Nothing +drains a spool a pack never reached, whatever runs after the sink. + +``on_engine_release`` now flushes the sink, bounded and without throwing, +so the release leaves the open pack in the spool as one ``.ready`` file. + +Build: make -C native cpu-goals PYTHON=/bin/python +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path + +import pytest + +REPO = Path(__file__).resolve().parents[1] +BUILD = REPO / "native" / "build" + +pytestmark = pytest.mark.cpu + +LAYOUT = "capture_pack_reference_v1" +MiB = 1 << 20 +MINUTE_NS = 60_000_000_000 + + +@pytest.fixture(scope="module") +def native_sink_module(): + if not sorted(BUILD.glob("_dmi_native_sink*.so")): + pytest.fail("native/build/_dmi_native_sink*.so is not built; run " + "`make -C native cpu-goals PYTHON=/bin/python`") + from dmi.storage.capture.native_sink import _load_native_sink_extension + + return _load_native_sink_extension() + + +def _sink(root: Path, **fields): + from dmi.storage.capture.native_sink import create_native_pack_sink + from dmi.storage.native_capture import NativeSinkConfig + + return create_native_pack_sink( + NativeSinkConfig(spool_root=str(root), **fields)).native_sink + + +def _envelope(index: int, nbytes: int = 1024): + import torch + from dmi.storage.capture import CaptureMetadata + + mapping = CaptureMetadata( + capture_id=f"release-{index:04d}", tenant_id="t", experiment_id="e", + run_id="r", session_id="s", request_id=f"q{index}", + sequence_id=f"n{index}", model_id="m", model_revision="mr", + adapter_revision=None, capture_policy_version="v", + hook_name="resid_post", layer_number=0, producer_rank=0, + step_number=index, token_start=index, token_end=index + 1, + batch_position=0, dtype="float32", shape=(nbytes // 4,), + captured_at_ns=1_700_000_000_000_000_000 + index, + ).to_mapping() + rows = [{"metadata_json": json.dumps(mapping), "offset": 0, + "length": nbytes, "dtype": 6, "shape": [nbytes // 4]}] + return rows, torch.full((nbytes // 4,), float(index)).view(torch.uint8) + + +def _ready(root: Path) -> list[Path]: + return sorted(root.rglob("*.dmi-pack.ready")) + + +def _wait_until_admitted(sink, records: int) -> None: + deadline = time.monotonic() + 10.0 + while sink.snapshot()["admitted_records"] < records: + assert time.monotonic() < deadline, sink.snapshot() + time.sleep(0.01) + + +def test_the_open_pack_is_staged_when_the_engine_releases_the_sink( + native_sink_module, tmp_path): + """A 60 s linger: nothing but a flush can seal the open pack before + the release, so the release must be what stages it.""" + sink = _sink(tmp_path, max_linger_ns=MINUTE_NS) + lease = sink.attach() + for index in range(3): + sink.submit_envelope(LAYOUT, *_envelope(index)) + _wait_until_admitted(sink, 3) + time.sleep(0.2) + assert _ready(tmp_path) == [] # still in memory, as the audit found + + del lease # what RingEngine::stop does once its record worker is joined + + assert len(_ready(tmp_path)) == 1 + snapshot = sink.snapshot() + assert snapshot["persisted_records"] == 3, snapshot + assert snapshot["packs_persisted"] == 1, snapshot + # The sink object is still alive: its destructor, which also seals the + # open pack, has not run and so is not what staged it. + assert sink.layout == LAYOUT + + +def test_a_release_with_nothing_open_stages_nothing(native_sink_module, + tmp_path): + sink = _sink(tmp_path, max_linger_ns=MINUTE_NS) + lease = sink.attach() + sink.submit_envelope(LAYOUT, *_envelope(0)) + assert sink.flush_and_wait(30.0) + assert len(_ready(tmp_path)) == 1 + + del lease + + assert len(_ready(tmp_path)) == 1 + assert sink.snapshot()["packs_persisted"] == 1 + + +def _binding_sink(module, root: Path, **fields): + """The binding's own NativePackSink, for the knobs NativeSinkConfig + does not carry (release_flush_timeout_s).""" + fields.setdefault("max_linger_ns", MINUTE_NS) + return module.NativePackSink(str(root), LAYOUT, **fields) + + +RELEASE_LINE = ("NativePackSink: the open pack did not reach the spool when " + "the engine released the sink: ") + + +def test_a_sink_that_cannot_stage_is_released_and_says_why( + native_sink_module, tmp_path, capfd): + """The release runs inside RingEngine::stop and its destructor, which + cannot throw: a sink whose release flush fails -- here the spool's byte + limit refuses the open pack -- still releases, at once, and says why on + stderr, the only report it can make once released.""" + sink = _binding_sink(native_sink_module, tmp_path, spool_max_bytes=4096) + lease = sink.attach() + sink.submit_envelope(LAYOUT, *_envelope(0, 8192)) + _wait_until_admitted(sink, 1) + capfd.readouterr() + + started = time.monotonic() + del lease + assert time.monotonic() - started < 5.0 + + err = capfd.readouterr().err + assert RELEASE_LINE in err, err + assert "spool byte limit exceeded" in err, err + assert _ready(tmp_path) == [] + assert sink.snapshot()["failures"] == 1, sink.snapshot() + # Released: another engine may take the sink, and then it reports the + # failure the release could only print. + lease = sink.attach() + with pytest.raises(RuntimeError, match="spool byte limit exceeded"): + sink.rethrow_if_failed() + del lease + + +def test_a_record_the_pipeline_dropped_does_not_hold_the_release( + native_sink_module, tmp_path): + """A record dropped as oversized is a loss the sink counts, not a + pipeline failure: flush_and_wait reports it, and the release that + follows flushes an empty pipeline, at once.""" + sink = _sink(tmp_path, max_linger_ns=MINUTE_NS, max_pack_bytes=MiB, + max_queue_bytes=4 * MiB) + lease = sink.attach() + # Admitted (the payload alone fits the queue), then dropped by the pack + # worker as oversized: it fits no empty pack with its framing. + sink.submit_envelope(LAYOUT, *_envelope(0, MiB)) + with pytest.raises(RuntimeError, match="oversized_records"): + sink.flush_and_wait(30.0) + + started = time.monotonic() + del lease + assert time.monotonic() - started < 5.0 + + assert _ready(tmp_path) == [] + # Released: another engine may take the sink. + lease = sink.attach() + del lease + + +def test_the_release_flush_is_bounded_by_its_timeout(native_sink_module, + tmp_path, capfd): + """A sink whose pipeline is stuck -- its stager parked, as on a hung + filesystem -- must not hold RingEngine::stop, and so engine.close(), + for longer than release_flush_timeout_s. The release gives up then, + says so on stderr, and leaves the pack to the pipeline: it reaches the + spool once the stager moves again.""" + sink = _binding_sink(native_sink_module, tmp_path, + release_flush_timeout_s=0.5) + assert sink.release_flush_timeout_s == 0.5 + release_stages = sink._hold_stages_for_testing(10.0) + lease = sink.attach() + for index in range(3): + sink.submit_envelope(LAYOUT, *_envelope(index)) + _wait_until_admitted(sink, 3) + capfd.readouterr() + + started = time.monotonic() + del lease + elapsed = time.monotonic() - started + assert 0.4 < elapsed < 2.0, elapsed + + err = capfd.readouterr().err + assert RELEASE_LINE + "timed out after 500 ms" in err, err + assert _ready(tmp_path) == [] # parked + + release_stages() + deadline = time.monotonic() + 10.0 + while sink.snapshot()["persisted_records"] < 3: + assert time.monotonic() < deadline, sink.snapshot() + time.sleep(0.01) + assert len(_ready(tmp_path)) == 1 + + +def test_a_zero_release_timeout_leaves_the_open_pack_to_the_pipeline( + native_sink_module, tmp_path): + """0 turns the backstop off: the release stages nothing, and the open + pack is the pipeline's, for the next flush, the linger or the sink's + own close.""" + sink = _binding_sink(native_sink_module, tmp_path, + release_flush_timeout_s=0.0) + assert sink.release_flush_timeout_s == 0.0 + lease = sink.attach() + sink.submit_envelope(LAYOUT, *_envelope(0)) + _wait_until_admitted(sink, 1) + + del lease + time.sleep(0.2) + assert _ready(tmp_path) == [] + + lease = sink.attach() + assert sink.flush_and_wait(30.0) + assert len(_ready(tmp_path)) == 1 + del lease + + +@pytest.mark.parametrize("timeout_s", [-1.0, float("nan"), float("inf")]) +def test_the_release_timeout_must_be_finite_and_not_negative( + native_sink_module, tmp_path, timeout_s): + with pytest.raises(ValueError, match="release_flush_timeout_s"): + _binding_sink(native_sink_module, tmp_path, + release_flush_timeout_s=timeout_s) diff --git a/tests/test_native_uploader.py b/tests/test_native_uploader.py index 36ccebfb5..c714f6a9f 100644 --- a/tests/test_native_uploader.py +++ b/tests/test_native_uploader.py @@ -16,6 +16,7 @@ import random import subprocess import sys +import time from pathlib import Path import pytest @@ -308,6 +309,208 @@ def test_corrupt_staged_bytes_are_refused(fake_s3, tmp_path): store.close() +def test_a_cancel_ends_the_retries_and_keeps_the_pack_staged(fake_s3, + tmp_path): + """Against a store that answers every request 500, one pack costs four + upload attempts of four transport attempts each, about 7 s of backoff. + The storage service cancels its uploads when it stops or a flush's + deadline passes; the cancel must end the retries and the backoff at + once, and leave the pack in the spool -- the durable place for it -- + rather than count it lost.""" + sink = DriverSession(SINK_DRIVER) + store = DriverSession(STORE_DRIVER) + try: + staged = _stage(sink, tmp_path / "spool", 5) + staged = dict(staged, object_key=( + "fault/always-500/" + staged["object_key"].rsplit("/", 1)[1])) + fields = _store_base(fake_s3) + fields.update( + op="upload_one", root=str(tmp_path / "spool"), + spool_max_bytes=1 << 40, max_workers=1, + max_in_flight_bytes=1 << 30, staged=staged, + cancel_after_ms=500, + ) + started = time.monotonic() + result = store.call(**fields) + elapsed = time.monotonic() - started + assert not result["ok"], result + assert result.get("cancelled") is True, result + assert "cancel" in result["what"], result + assert result["upload_attempts"] <= 2, result + assert elapsed < 3.0, elapsed + assert Path(staged["path"]).exists() + finally: + sink.close() + store.close() + + +def test_a_cancel_wakes_the_uploaders_own_backoff(fake_s3, tmp_path): + """The uploader backs off between its attempts on a pack, for up to + max_backoff_s (10 s), apart from the S3 client's backoff inside each + attempt. Here each attempt is one HEAD answered 500 at once (one + transport attempt), and the uploader's backoff is 2 s (+-20% jitter): + the cancel at 0.3 s lands in the first one. Slept out, it would end + after 1.6 s at the earliest, when the check before the next attempt + stops the upload anyway.""" + sink = DriverSession(SINK_DRIVER) + store = DriverSession(STORE_DRIVER) + try: + staged = _stage(sink, tmp_path / "spool", 7) + staged = dict(staged, object_key=( + "fault/always-500/" + staged["object_key"].rsplit("/", 1)[1])) + fields = _store_base(fake_s3, max_attempts=1) + fields.update( + op="upload_one", root=str(tmp_path / "spool"), + spool_max_bytes=1 << 40, upload_max_attempts=4, + upload_base_backoff_ms=2000, staged=staged, cancel_after_ms=300, + ) + started = time.monotonic() + result = store.call(**fields) + elapsed = time.monotonic() - started + assert not result["ok"], result + assert result["cancelled"] is True, result + assert result["upload_attempts"] == 1, result + assert elapsed < 1.2, (elapsed, result) + assert Path(staged["path"]).exists() + finally: + sink.close() + store.close() + + +def test_a_cancel_stops_the_listing_between_packs(fake_s3, tmp_path): + """UploadPending lists the spool before it uploads anything, and a + listing re-hashes every staged pack: over a backlog, seconds a GiB of + it, all before a single worker looked at the cancel. The listing now + stops between packs once cancelled, and nothing is tried. 16 sparse + packs of 256 MiB, zeros named for their checksum, take it about 3 s.""" + import hashlib + import uuid + + size = 256 << 20 + digest = hashlib.sha256() + zeros = bytes(1 << 20) + for _ in range(size // len(zeros)): + digest.update(zeros) + root = tmp_path / "spool" + root.mkdir() + backlog = [] + for _ in range(16): + ready = root / f"{uuid.uuid4()}.1.1.{digest.hexdigest()}.dmi-pack.ready" + with open(ready, "wb") as sparse: + sparse.truncate(size) + backlog.append(ready) + store = DriverSession(STORE_DRIVER) + try: + started = time.monotonic() + result = _upload_pending(store, fake_s3, root, cancel_after_ms=300) + elapsed = time.monotonic() - started + finally: + store.close() + assert result["ok"], result + assert result["listing_cancelled"] is True, result + assert result["refs"] == [] and result["failures"] == [], result + assert result["snapshot"]["cancelled_packs"] == 0, result + # Between packs, not after the listing: the cut ends within one pack's + # hash of the cancel, about 0.5 s here, where the whole listing takes + # over 3 s. The bound leaves room for a loaded runner, whose hashing + # slows with it (about 1.1 s at 4x oversubscription). + assert elapsed < 2.0, elapsed + assert sorted(root.rglob("*.dmi-pack.ready")) == sorted(backlog) + assert STATE.calls == [] + + +@pytest.mark.parametrize("request_cut", [False, True], + ids=["failed-on-its-own", "cut-by-the-cancel"]) +def test_a_cancel_counts_only_when_it_ended_the_upload(fake_s3, tmp_path, + request_cut): + """UploadOne booked a pack cancelled whenever its Cancellation was set + once its attempts were over -- also when they had run out on real + failures and a flush's deadline merely passed meanwhile. The storage + service then counted the pack in cancelled_uploads, not in + upload_failures, and never recorded its error, so the TimeoutError a + flush raised named no cause. The one attempt's HEAD is answered 500 + after 1.2 s, and the cancel comes at 0.3 s. With only the uploader + holding the Cancellation the HEAD runs to its answer: a failure, and + it says so. With the client holding it too the cancel cuts the HEAD, + and that is a cancel, last attempt or not.""" + sink = DriverSession(SINK_DRIVER) + store = DriverSession(STORE_DRIVER) + try: + staged = _stage(sink, tmp_path / "spool", 6) + staged = dict(staged, object_key=( + "fault/slow-once-500/" + staged["object_key"].rsplit("/", 1)[1])) + fields = _store_base(fake_s3, max_attempts=1) + fields.update( + op="upload_one", root=str(tmp_path / "spool"), + spool_max_bytes=1 << 40, upload_max_attempts=1, + staged=staged, cancel_after_ms=300, + cancel_uploader_only=not request_cut, + ) + result = store.call(**fields) + assert not result["ok"], result + assert result["upload_attempts"] == 1, result + assert result["cancelled"] is request_cut, result + if request_cut: + assert "cancel" in result["what"], result + else: + assert "HTTP 500" in result["what"], result + assert "cancel" not in result["what"], result + assert Path(staged["path"]).exists() + finally: + sink.close() + store.close() + + +@pytest.mark.parametrize("multipart", [False, True], + ids=["put", "multipart-part"]) +def test_a_cancel_that_cuts_the_last_attempts_upload_is_a_cancel( + fake_s3, tmp_path, multipart): + """The upload's one attempt gets its HEAD answered 404 at once, then + its PUT -- or, over the multipart threshold, its first part -- is held + for 5 s, and the cancel at 0.3 s cuts it. No attempt or backoff comes + after it to see the cancel, so only the request's own report says the + cancel ended the upload; a PUT that stopped reporting it would book + the pack as failed -- upload_failures up, the backoff growing, and + "request cancelled" as the last error. The HEAD-cut case above pins + the preflight's report; these pin the PUT's and the part's.""" + sink = DriverSession(SINK_DRIVER) + store = DriverSession(STORE_DRIVER) + try: + root = tmp_path / "spool" + if multipart: + staged = _stage_large_pack(root, 6, MIB) # over 5 MiB: 2 parts + fault = "fault/hang-parts/" + transport = dict(multipart_threshold=5 * MIB, + multipart_chunk=5 * MIB) + else: + staged = _stage(sink, root, 8) + fault = "fault/hang-put/" + transport = {} + staged = dict(staged, object_key=( + fault + staged["object_key"].rsplit("/", 1)[1])) + fields = _store_base(fake_s3, max_attempts=1, **transport) + fields.update( + op="upload_one", root=str(root), spool_max_bytes=1 << 40, + max_in_flight_bytes=1 << 30, upload_max_attempts=1, + staged=staged, cancel_after_ms=300, + ) + started = time.monotonic() + result = store.call(**fields) + elapsed = time.monotonic() - started + assert not result["ok"], result + assert result["upload_attempts"] == 1, result + assert result["cancelled"] is True, result + assert "cancel" in result["what"], result + assert elapsed < 3.0, (elapsed, result) + assert Path(staged["path"]).exists() + assert staged["object_key"] not in STATE.objects + heads = [c for c in STATE.calls if c["method"] == "HEAD"] + assert len(heads) == 1, STATE.calls # the preflight, answered + finally: + sink.close() + store.close() + + def test_pack_over_the_byte_gate_fails_fast(fake_s3, tmp_path): sink = DriverSession(SINK_DRIVER) store = DriverSession(STORE_DRIVER)