diff --git a/Cargo.lock b/Cargo.lock index aaaa28b..8e26b4e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2283,7 +2283,7 @@ dependencies = [ [[package]] name = "rustqueue-bench" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "clap", @@ -2295,7 +2295,7 @@ dependencies = [ [[package]] name = "rustqueue-console" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "axum", @@ -2323,7 +2323,7 @@ dependencies = [ [[package]] name = "rustqueue-discovery" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "axum", @@ -2337,6 +2337,7 @@ dependencies = [ "rustqueue-telemetry", "serde", "serde_json", + "sha2", "tokio", "tower", "tracing", @@ -2345,13 +2346,14 @@ dependencies = [ [[package]] name = "rustqueue-operator" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "axum", "chrono", "crc32c", "futures", + "hex", "k8s-openapi", "kube", "rand 0.8.7", @@ -2361,6 +2363,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", + "sha2", "thiserror", "tokio", "toml", @@ -2370,7 +2373,7 @@ dependencies = [ [[package]] name = "rustqueue-protocol" -version = "0.7.2" +version = "0.8.0" dependencies = [ "bytes", "serde", @@ -2380,7 +2383,7 @@ dependencies = [ [[package]] name = "rustqueue-proxy" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "axum", @@ -2389,6 +2392,7 @@ dependencies = [ "rand 0.8.7", "reqwest", "rustls", + "rustqueue-protocol", "rustqueue-telemetry", "serde", "serde_json", @@ -2400,7 +2404,7 @@ dependencies = [ [[package]] name = "rustqueue-queue" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "bytes", @@ -2422,7 +2426,7 @@ dependencies = [ [[package]] name = "rustqueue-server" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "async-compression", @@ -2457,7 +2461,7 @@ dependencies = [ [[package]] name = "rustqueue-storage" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "crc32c", @@ -2470,14 +2474,14 @@ dependencies = [ [[package]] name = "rustqueue-telemetry" -version = "0.7.2" +version = "0.8.0" dependencies = [ "serde", ] [[package]] name = "rustqueuectl" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "clap", diff --git a/Cargo.toml b/Cargo.toml index d15f782..abe764a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ members = [ exclude = ["fuzz"] [workspace.package] -version = "0.7.2" +version = "0.8.0" edition = "2021" license = "Apache-2.0" rust-version = "1.88" @@ -55,6 +55,7 @@ tower-http = { version = "0.6.2", features = ["fs", "timeout", "trace"] } tracing = "0.1.41" tracing-subscriber = { version = "0.3.19", features = ["env-filter", "json"] } rustqueue-telemetry = { path = "crates/telemetry" } +rustqueue-protocol = { path = "crates/protocol" } x509-parser = "0.17.0" k8s-openapi = { version = "0.28.0", features = ["v1_33", "schemars"] } kube = { version = "4.0.0", features = ["derive", "runtime"] } diff --git a/Makefile b/Makefile index 529aa17..7cc485c 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: test check fmt clippy rustfmt-component clippy-component release-bin image image-from-dist operator-release-bin operator-image console-ui-build console-ui-check \ +.PHONY: test check fmt clippy rustfmt-component clippy-component release-bin image image-from-dist operator-release-bin operator-image console-ui-build console-ui-check kodo-replay kodo-gateway-acceptance \ helm-lint helm-template k8s-acceptance k8s-console-management-acceptance k8s-multi-acceptance up down compat compat-go compat-python \ fuzz-smoke benchmark release-gate @@ -76,6 +76,10 @@ helm-lint: helm lint deploy/helm/rustqueue helm lint deploy/helm/rustqueue --set monitoring.serviceMonitor.enabled=true \ --set monitoring.prometheusRule.enabled=true + helm lint deploy/helm/rustqueue --set queue.kodoCompatibility.enabled=true \ + --set queue.imagePullPolicy=Never + ! helm template rustqueue deploy/helm/rustqueue \ + --set queue.kodoCompatibility.cleanupEnabled=true ! rg -n 'x-kubernetes-preserve-unknown-fields:[[:space:]]*false' \ deploy/helm/rustqueue/crds @@ -84,6 +88,9 @@ helm-template: helm template rustqueue deploy/helm/rustqueue --namespace rustqueue \ --set monitoring.serviceMonitor.enabled=true \ --set monitoring.prometheusRule.enabled=true + helm template rustqueue deploy/helm/rustqueue --namespace rustqueue \ + --set queue.kodoCompatibility.enabled=true \ + --set queue.imagePullPolicy=Never k8s-acceptance: ./scripts/acceptance-k8s.sh @@ -109,6 +116,12 @@ compat-go: compat-python: ./scripts/compat-python.sh +kodo-replay: + bash ./scripts/kodo-replay.sh + +kodo-gateway-acceptance: + bash ./scripts/acceptance-kodo-gateway.sh + fuzz-smoke: ./scripts/fuzz-smoke.sh diff --git a/README.md b/README.md index 7ed29cc..0b37e5a 100644 --- a/README.md +++ b/README.md @@ -59,11 +59,14 @@ operator -> eligible nodes -> StatefulSet ordinal + retained RWO PVC - Delivery is at least once. A restart redelivers messages without durable FIN. - The broker PVC is the only copy; permanent PVC loss loses its messages. - Topics and channels are broker-local. Lookup consumers union all owners. -- Messages are retained for 90 seconds before a channel exists. This covers the - official Go client's default 60-second lookup poll plus its 30% jitter and - lets `SUB` catch a newly selected owner without a normal-path miss. -- The stable v7 wire limit is 32 MiB; the default single-message limit is - 20 MiB and the default MPUB body limit is 64 MiB. +- Messages are retained for 90 seconds before a channel exists. This covers one + official Go client default 60-second lookup poll plus its 30% jitter and lets + `SUB` catch a newly selected owner without a normal-path miss. The Kodo + profile forces 180 seconds so one failed lookup request still gets a second + discovery opportunity before data can age out. +- The stable v7 single-message limit is 100 MiB. The conservative defaults + remain 20 MiB per message and 64 MiB per MPUB body; the opt-in Kodo profile + raises them to 100 MiB and 128 MiB respectively. - Publish bodies are written to the segment with vectored header/metadata/body I/O. The durable path does not build a second batch body or a full record buffer, and admission charges the input plus bounded encoding metadata. @@ -89,6 +92,11 @@ operator -> eligible nodes -> StatefulSet ordinal + retained RWO PVC - `rustqueue-proxy`: bounded producer TCP/HTTP proxy; new TCP connections choose a least-active ready Broker and are rotated after a jittered five-minute default lifetime so a small long-lived producer pool follows fleet changes. + In Kodo mode the same binary runs as three publish Gateway Pods behind one + highly available Service. Each Gateway terminates the NSQ producer protocol, + enforces the 100 MiB single-message limit, and may retry another Broker only + after an explicit pre-commit rejection. A failure after the full body was + sent is returned as ambiguous and is never retried automatically. - `rustqueue-operator`: creates the StatefulSet, retained PVCs, discovery, proxy, RBAC, disruption budgets, PVC expansion and drain-aware one-at-a-time rolling updates. @@ -146,7 +154,7 @@ kubectl label node worker-1 rustqueue.io/eligible=true helm upgrade --install rustqueue deploy/helm/rustqueue \ --namespace rustqueue --create-namespace \ - --set queue.image=registry.example/rustqueue:0.7.2 \ + --set queue.image=registry.example/rustqueue:0.8.0 \ --set queue.storageClassName=ssd-rwo ``` @@ -156,8 +164,145 @@ internal; set `queue.registrySecretName` to use a pre-created Secret containing `admin-token`, `registry-token` and `console-token`. The console token can read broker observations and, only when Console management is explicitly enabled, call the narrow native Topic/Channel management API. It cannot authorize drain, -scrub, upgrade or the NSQ-compatible admin API. Client TLS is optional and always supplied -through an existing Kubernetes Secret; the operator does not run a CA. +scrub, upgrade or the NSQ-compatible admin API. Client TLS is optional and +always supplied through an existing Kubernetes Secret; the operator does not +run a CA. + +Kodo compatibility is a separate, default-off deployment profile: + +```sh +helm upgrade --install rustqueue deploy/helm/rustqueue \ + --namespace rustqueue \ + --set 'queue.image=registry.example/rustqueue@sha256:<64-lowercase-hex-digest>' \ + --set queue.kodoCompatibility.enabled=true +``` + +The chart pins this profile to three Brokers, storage feature level 2, a +180-second bootstrap-retention window, a 100 MiB RustQueue single-message +ceiling, a +one-CPU/2 GiB Broker request, and three stable Gateway Pods with required +hostname anti-affinity. The profile therefore needs three schedulable nodes. Each +Gateway requests one CPU and 768 MiB, is not CPU-limited, and is limited to +1 GiB memory so its fixed 512 MiB in-flight body budget is scheduler-backed. +The reviewed Kodo source enforces an application maximum of 104857500 bytes +(100 MiB minus 100 bytes) and sets the go-nsq producer write deadline to three +seconds. A maximum-size publish therefore needs more than roughly 34 MiB/s of +application goodput from Kodo to the Gateway, with additional headroom for +scheduling and network variance. Validate that exact 104857500-byte path with +the included acceptance before rollout. The profile also requires an immutable +lowercase `@sha256` image; `imagePullPolicy=Never` is accepted only for a +preloaded local test image. + +For an existing deployment whose running binary still advertises the legacy +32 MiB capability, first roll the new immutable image with Kodo compatibility +disabled and wait for `Ready`; enable the Kodo profile in a second change. The +operator will not cross the 100 MiB/feature-level-2 fence while an old Broker is +still running. + +Activation is fail-closed and ordered. The operator first waits for at least +two Ready Brokers, then starts all three Gateways before disrupting an old +Broker. This allows a migration to recover while one of the three target +Brokers is unavailable; the Broker target remains three. It rolls Discovery +Pods to a distinct `kodo` mode label while the Discovery Service still selects +the existing `direct` mode. Only after the complete Kodo-mode set is Ready does +it atomically switch the Service selector. Broker rollout and maintenance +remain blocked until +`KodoCutoverReady=True` and `KodoProducerRestartConfirmed=True`. + +`queue.kodoCompatibility.cutoverGraceSeconds` defaults to 630 seconds because +Kodo's default `nsq.refresh_at` is 300 seconds. Set it to at least two deployed +Kodo `refresh_at` intervals plus 30 seconds. This preserves a second refresh +opportunity if the first request overlaps the Discovery Service cutover. New +installations may point Kodo at the Discovery Service after +`KodoGatewaysAdvertised=True`. + +The reviewed Kodo `refreshNodes()` implementation replaces cached producers +when an advertised address changes. RustQueue still requires an explicit +publisher restart as a cutover fence, so rollout safety does not depend on that +private cache behavior remaining unchanged in a later Kodo build. After +`KodoGatewaysAdvertised=True`, restart every Kodo publisher process and wait +for it to become Ready, then change +`queue.kodoCompatibility.producerRestartNonce` to a new value: + +```sh +helm upgrade rustqueue deploy/helm/rustqueue \ + --namespace rustqueue \ + --reuse-values \ + --set-string queue.kodoCompatibility.producerRestartNonce="$(date +%s)" +``` + +The Operator captures the nonce only after Gateway advertisement and will not +drain, replace, or place a Broker into maintenance until the value changes. +Losing Gateway advertisement resets the fence and requires another restart +confirmation. An existing Kodo deployment is fully migrated only after both +`KodoCutoverReady=True` and `KodoProducerRestartConfirmed=True`. +During an existing-cluster migration, do not publish messages above the +previous Broker limit until the RustQueue resource itself reports `Ready=True`; +Gateway advertisement and explicit producer restart intentionally precede the +Broker configuration rollout. + +`/nodes` then advertises three producer identities on ports `4150`, `4152`, and +`4153` of one ClusterIP Gateway Service. Every port targets every Ready Gateway, +and every Gateway can publish to every publish-ready Broker. A single Gateway +restart therefore does not turn an advertised identity into a dead Pod address. +`/lookup` continues to expose the real Broker owners for consumption. During a +Broker restart it returns the healthy topic shards instead of blocking every +consumer; the separate complete-inventory signal remains false until all three +Stats shards are back. +When ServiceMonitor support is enabled, the dedicated Gateway monitor scrapes +the metrics-only port `4160`; Gateway Stats ports do not accept HTTP publish +requests. Any increase in unknown publish outcomes raises the critical +`RustQueueAmbiguousPublish` alert. + +Gateway ingress accepts only Pods matching +`queue.kodoCompatibility.allowedPodSelector` (by default +`app.kubernetes.io/name=kodo`) in the RustQueue namespace. Label the Kodo +workload accordingly. For a separate namespace, set +`allowedNamespaceSelector` to labels present on that namespace; the chart's +general same-namespace policy explicitly excludes Gateway and Broker Pods. +In Kodo mode, Broker TCP ingress is restricted to RustQueue Pods and that same +Kodo peer selector. Broker HTTP remains reachable for stats and monitoring, but +its publish endpoints require the internal admin token. Kodo consumers and +publishers share the NSQ TCP protocol and port, so a permitted Kodo Pod remains +technically able to publish directly; keep the selector limited to the trusted +Kodo workload and use the advertised Gateway addresses for every producer. +Discovery `/nodes` advertises those Gateways for publishing while `/lookup` +continues to advertise the currently healthy real Broker owners for +consumption. Automatic Kodo channel cleanup is hard-disabled: the CRD accepts only +`cleanupEnabled: false`, every RustQueue binary rejects an attempt to enable +the compatibility cleanup path, and `/channel/delete` returns a 404 whose body +does not contain `CHANNEL_NOT_FOUND`, so unchanged Kodo stops before contacting +Brokers. +Cleanup must remain disabled until RustQueue has a cluster-wide atomic, +authenticated deletion transaction. + +Enabling this profile activates the storage feature-level-2 rollback fence; +records above the legacy 72 MiB boundary require that feature level. Disabling +the profile is an explicit decommission, not a migration back to unsafe direct +publishing. First stop every Kodo workload using this Discovery Service, then +set `queue.kodoCompatibility.enabled=false` and +`queue.kodoCompatibility.decommissionConfirmed=true`. Without that +confirmation the Operator fails closed before changing any Broker, Discovery, +or Gateway resource. After confirmation it completes the target Broker rollout +while the Gateways remain available, switches Discovery to direct mode, and +removes the Gateway resources; no Kodo workload may be restarted against that +direct mode. The profile retains feature level 2 as the monotonic storage floor +and preserves enough delivery memory to consume previously stored 100 MiB +messages. Channel `message_count` remains monotonic across empty and eviction, +while `requeue_count` and `timeout_count` are persisted with the Channel WAL +and checkpoint. Kodo can therefore calculate rate thresholds across Broker +restarts without a counter rollback. Run the real Kodo +`nsqadmin.go` parser and admin-flow replay before releasing this profile. The +replay is fully offline and copies that upstream +file unchanged into an isolated test module; it also rejects an unreviewed +hash for the Kodo producer discovery, producer client, or admin source. Only +unrelated Kodo globals are stubbed: + +```sh +KODO_SOURCE_DIR=/path/to/kodo make kodo-replay +# Or include it in the complete gate: +KODO_ACCEPTANCE=1 KODO_SOURCE_DIR=/path/to/kodo ./scripts/release-gate.sh +``` The chart enables RustQueue Console by default as a ClusterIP-only Service. It does not create an Ingress and has no built-in login. Put the Service behind @@ -261,7 +406,9 @@ The producer proxy defaults to 10,000 simultaneous TCP connections, a 64 MiB request-body limit, a 512 MiB node-wide in-flight HTTP body budget, and a jittered 300-second TCP tunnel lifetime. A zero lifetime disables rotation. Override them with `RUSTQUEUE_PROXY_MAX_CONNECTIONS`, -`RUSTQUEUE_PROXY_MAX_BODY_BYTES`, `RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES`, and +`RUSTQUEUE_PROXY_MAX_MESSAGE_BYTES`, `RUSTQUEUE_PROXY_MAX_BODY_BYTES`, +`RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES`, +`RUSTQUEUE_PROXY_TCP_COMMAND_TIMEOUT_MS`, and `RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS`. Discovery gives every Kubernetes EndpointSlice list an explicit 1.5-second @@ -317,10 +464,11 @@ aggregate gauges/counters. Format v7 is a clean break. A v6 or older directory is refused; there is no in-place migration. Within v7, record tags and existing fields are append-only. -Every binary declares its reader/writer feature range, and every PVC persists -its active writer and minimum reader levels in `/data/COMPATIBILITY`. The -operator probes the target image and every running broker before replacement; -an incompatible target is blocked before a PVC is touched. +Every binary declares its reader/writer feature range and protocol message/body +limits, and every PVC persists its active writer and minimum reader levels in +`/data/COMPATIBILITY`. The operator probes the target image and every running +broker before replacement; an incompatible target is blocked before a PVC is +touched. Rollouts are durable status-driven operations. A replacement must be Ready before another Broker is drained. They may be paused, optionally stop after one diff --git a/console-ui/package.json b/console-ui/package.json index d63178b..f0f63ac 100644 --- a/console-ui/package.json +++ b/console-ui/package.json @@ -1,7 +1,7 @@ { "name": "rustqueue-console-ui", "private": true, - "version": "0.7.2", + "version": "0.8.0", "packageManager": "pnpm@11.9.0", "type": "module", "scripts": { diff --git a/crates/console/src/collector/observe.rs b/crates/console/src/collector/observe.rs index aa884e8..570ab4f 100644 --- a/crates/console/src/collector/observe.rs +++ b/crates/console/src/collector/observe.rs @@ -263,12 +263,12 @@ mod tests { registry_revision: 9, ..Default::default() }; - observation.node.version = "0.7.2".into(); + observation.node.version = "0.8.0".into(); let mut head = BrokerObservationHead { registry_revision: 9, ..Default::default() }; - head.node.version = "0.7.2".into(); + head.node.version = "0.8.0".into(); let cached = CachedObservation { pod_uid: broker.uid.clone(), pod_ip: broker.pod_ip.clone(), diff --git a/crates/console/src/model.rs b/crates/console/src/model.rs index 94d8239..7ca4de9 100644 --- a/crates/console/src/model.rs +++ b/crates/console/src/model.rs @@ -318,6 +318,7 @@ mod tests { .push(rustqueue_queue::TopicStats { name: "events".into(), paused: false, + published_count: 0, message_count: 1, segment_count: 1, segment_bytes: 128, diff --git a/crates/discovery/Cargo.toml b/crates/discovery/Cargo.toml index 2d0e2d3..36bfa38 100644 --- a/crates/discovery/Cargo.toml +++ b/crates/discovery/Cargo.toml @@ -21,6 +21,7 @@ rustls.workspace = true rustqueue-telemetry.workspace = true serde.workspace = true serde_json.workspace = true +sha2.workspace = true tokio.workspace = true tracing.workspace = true tracing-subscriber.workspace = true diff --git a/crates/discovery/src/directory.rs b/crates/discovery/src/directory.rs index 23290c6..06aa7a9 100644 --- a/crates/discovery/src/directory.rs +++ b/crates/discovery/src/directory.rs @@ -1,9 +1,12 @@ use crate::{BrokerEndpoint, BrokerRegistry, BrokerRegistryHead, DiscoveryMetrics, Producer}; use parking_lot::RwLock; +use sha2::{Digest, Sha256}; use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use std::time::{Duration, Instant}; +const KODO_BROKER_COUNT: usize = 3; + #[derive(Clone, Default)] pub struct Directory { inner: Arc>, @@ -18,7 +21,11 @@ struct State { channel_owners: BTreeMap>, publishers: BTreeSet, consumers: BTreeSet, - revision: u64, + kodo_compatibility_enabled: bool, + kodo_gateways: Vec, + kodo_cleanup_enabled: bool, + source_stale_after: Option, + last_source_success: Option, } struct Observation { @@ -31,6 +38,34 @@ impl Directory { &self.metrics } + pub fn configure_kodo(&self, gateways: Vec, cleanup_enabled: bool) { + let mut state = self.inner.write(); + state.kodo_compatibility_enabled = true; + state.kodo_gateways = gateways; + state.kodo_cleanup_enabled = cleanup_enabled; + } + + pub fn configure_source_health(&self, stale_after: Duration) { + self.inner.write().source_stale_after = Some(stale_after); + } + + pub fn mark_source_success(&self) { + self.inner.write().last_source_success = Some(Instant::now()); + } + + pub fn source_ready(&self) -> bool { + let state = self.inner.read(); + source_ready(&state) + } + + pub fn lookup_ready(&self) -> bool { + let state = self.inner.read(); + source_ready(&state) + && (!state.kodo_compatibility_enabled + || (state.kodo_gateways.len() == KODO_BROKER_COUNT + && kodo_broker_inventory_ready(&state))) + } + pub fn replace_endpoints(&self, endpoints: BTreeSet) { let mut state = self.inner.write(); let removed: Vec<_> = state @@ -113,7 +148,7 @@ impl Directory { } pub fn revision(&self) -> u64 { - self.inner.read().revision + routing_revision(&self.inner.read()) } pub fn topics(&self) -> Vec { @@ -138,30 +173,144 @@ impl Directory { let Some(endpoints) = endpoints else { return Vec::new(); }; - endpoints - .iter() - .filter_map(|endpoint| state.brokers.get(endpoint)) - .map(|item| Producer::from_registry(&item.registry)) - .collect() + let mut producers = unique_producers( + endpoints + .iter() + .filter_map(|endpoint| state.brokers.get(endpoint)) + .map(|item| &item.registry), + ); + if state.kodo_cleanup_enabled { + for producer in &mut producers { + producer.http_port = 4152; + } + } + producers } pub fn publishers(&self) -> Vec { let state = self.inner.read(); + publisher_producers(&state) + } + + pub fn brokers(&self) -> Vec { + broker_producers(&self.inner.read()) + } + + pub fn publisher_snapshot(&self) -> (u64, Vec) { + let state = self.inner.read(); + (routing_revision(&state), publisher_producers(&state)) + } + + pub fn broker_snapshot(&self) -> (u64, Vec) { + let state = self.inner.read(); + (routing_revision(&state), broker_producers(&state)) + } + + pub fn publisher_head(&self) -> (u64, usize) { + let state = self.inner.read(); + (routing_revision(&state), publisher_producers(&state).len()) + } + + pub fn node_producers(&self) -> Vec { + let state = self.inner.read(); + if state.kodo_compatibility_enabled { + return state.kodo_gateways.clone(); + } + publisher_producers(&state) + } + + pub fn kodo_nodes_ready(&self) -> bool { + let state = self.inner.read(); + !state.kodo_compatibility_enabled || !state.kodo_gateways.is_empty() + } + + pub fn kodo_cleanup_enabled(&self) -> bool { + self.inner.read().kodo_cleanup_enabled + } + + pub fn broker_count(&self) -> usize { + let state = self.inner.read(); + consumer_node_ids(&state).len() + } + + pub fn publisher_count(&self) -> usize { + let state = self.inner.read(); + publisher_producers(&state).len() + } +} + +fn source_ready(state: &State) -> bool { + state + .last_source_success + .zip(state.source_stale_after) + .is_some_and(|(last_success, stale_after)| last_success.elapsed() <= stale_after) +} + +fn publisher_producers(state: &State) -> Vec { + unique_producers( state .publishers .iter() .filter_map(|endpoint| state.brokers.get(endpoint)) - .map(|item| Producer::from_registry(&item.registry)) - .collect() + .map(|item| &item.registry), + ) +} + +fn broker_producers(state: &State) -> Vec { + unique_producers(state.brokers.values().map(|item| &item.registry)) +} + +fn unique_producers<'a>(registries: impl IntoIterator) -> Vec { + let mut producers = BTreeMap::new(); + for registry in registries { + producers + .entry(registry.node_id) + .or_insert_with(|| Producer::from_registry(registry)); } + producers.into_values().collect() +} - pub fn broker_count(&self) -> usize { - self.inner.read().consumers.len() +fn consumer_node_ids(state: &State) -> BTreeSet { + state + .consumers + .iter() + .filter_map(|endpoint| state.brokers.get(endpoint)) + .map(|item| item.registry.node_id) + .collect() +} + +fn kodo_broker_inventory_ready(state: &State) -> bool { + let node_ids = consumer_node_ids(state); + if node_ids.len() != KODO_BROKER_COUNT || node_ids.contains(&0) { + return false; } + node_ids + .iter() + .map(|node_id| (node_id - 1) % KODO_BROKER_COUNT as u64) + .collect::>() + .len() + == KODO_BROKER_COUNT +} - pub fn publisher_count(&self) -> usize { - self.inner.read().publishers.len() +fn routing_revision(state: &State) -> u64 { + let mut hasher = Sha256::new(); + for producer in publisher_producers(state) { + hash_producer(&mut hasher, b'P', &producer); } + for producer in broker_producers(state) { + hash_producer(&mut hasher, b'B', &producer); + } + let digest = hasher.finalize(); + u64::from_be_bytes(digest[..8].try_into().expect("SHA-256 prefix")).max(1) +} + +fn hash_producer(hasher: &mut Sha256, role: u8, producer: &Producer) { + hasher.update([role]); + hasher.update(producer.node_id.to_be_bytes()); + hasher.update((producer.broadcast_address.len() as u64).to_be_bytes()); + hasher.update(producer.broadcast_address.as_bytes()); + hasher.update(producer.tcp_port.to_be_bytes()); + hasher.update(producer.http_port.to_be_bytes()); } fn add(state: &mut State, endpoint: BrokerEndpoint, registry: BrokerRegistry) { @@ -173,7 +322,6 @@ fn add(state: &mut State, endpoint: BrokerEndpoint, registry: BrokerRegistry) { seen_at: Instant::now(), }, ); - state.revision = state.revision.wrapping_add(1).max(1); } fn remove(state: &mut State, endpoint: &BrokerEndpoint) { @@ -181,7 +329,6 @@ fn remove(state: &mut State, endpoint: &BrokerEndpoint) { return; }; deindex(state, endpoint, &observation.registry); - state.revision = state.revision.wrapping_add(1).max(1); } fn index(state: &mut State, endpoint: &BrokerEndpoint, registry: &BrokerRegistry) { @@ -289,10 +436,55 @@ mod tests { directory.observe(endpoint, registry(index)); } assert_eq!(directory.producers(Some("events")).len(), 500); + assert_eq!(directory.broker_count(), 500); assert_eq!(directory.topics(), vec!["events"]); assert_eq!(directory.channels("events"), vec!["workers"]); } + #[test] + fn duplicate_registry_endpoints_do_not_inflate_broker_inventory() { + let directory = Directory::default(); + let first = endpoint(1); + let duplicate = endpoint(2); + directory.replace_endpoints([first.clone(), duplicate.clone()].into_iter().collect()); + let registry = registry(0); + directory.observe(first, registry.clone()); + directory.observe(duplicate, registry); + + assert_eq!(directory.producers(Some("events")).len(), 1); + assert_eq!(directory.publishers().len(), 1); + assert_eq!(directory.brokers().len(), 1); + assert_eq!(directory.broker_count(), 1); + assert_eq!(directory.publisher_count(), 1); + assert_eq!(directory.publisher_head().1, 1); + } + + #[test] + fn kodo_lookup_requires_one_unique_broker_per_stats_shard() { + let directory = Directory::default(); + directory.configure_source_health(Duration::from_secs(5)); + directory.configure_kodo( + (0..KODO_BROKER_COUNT) + .map(|ordinal| Producer::gateway("gateway".into(), ordinal)) + .collect(), + false, + ); + for (index, node_id) in [1, 2, 4].into_iter().enumerate() { + let endpoint = endpoint(index); + let mut registry = registry(index); + registry.node_id = node_id; + directory.observe(endpoint, registry); + } + directory.mark_source_success(); + assert!(!directory.lookup_ready()); + + let endpoint = endpoint(2); + let mut registry = registry(2); + registry.node_id = 3; + directory.observe(endpoint, registry); + assert!(directory.lookup_ready()); + } + #[test] fn unchanged_head_refreshes_liveness_without_reindexing() { let directory = Directory::default(); @@ -313,4 +505,66 @@ mod tests { )); assert_eq!(directory.revision(), revision); } + + #[test] + fn routing_revision_is_content_based_across_discovery_replicas() { + let first = Directory::default(); + let second = Directory::default(); + let first_endpoint = endpoint(1); + let second_endpoint = endpoint(2); + first.replace_endpoints([first_endpoint.clone()].into_iter().collect()); + second.replace_endpoints([second_endpoint.clone()].into_iter().collect()); + first.observe(first_endpoint, registry(1)); + second.observe(second_endpoint, registry(2)); + + assert_ne!(first.revision(), second.revision()); + + let replica = Directory::default(); + let first_endpoint = endpoint(1); + replica.replace_endpoints([first_endpoint.clone()].into_iter().collect()); + replica.observe(first_endpoint, registry(1)); + assert_eq!(first.revision(), replica.revision()); + } + + #[test] + fn staged_kodo_mode_never_falls_back_to_direct_broker_publishers() { + let directory = Directory::default(); + let endpoint = endpoint(1); + directory.replace_endpoints([endpoint.clone()].into_iter().collect()); + directory.observe(endpoint, registry(1)); + assert_eq!(directory.node_producers().len(), 1); + + directory.configure_kodo(Vec::new(), false); + + assert!(directory.node_producers().is_empty()); + assert_eq!(directory.producers(Some("events")).len(), 1); + } + + #[test] + fn nodes_exclude_a_broker_that_can_consume_but_cannot_publish() { + let directory = Directory::default(); + let endpoint = endpoint(1); + let mut draining = registry(1); + draining.publish_ready = false; + directory.replace_endpoints([endpoint.clone()].into_iter().collect()); + directory.observe(endpoint, draining); + + assert!(directory.node_producers().is_empty()); + assert_eq!(directory.producers(Some("events")).len(), 1); + } + + #[test] + fn cleanup_uses_the_network_isolated_broker_compatibility_port() { + let directory = Directory::default(); + let endpoint = endpoint(1); + directory.replace_endpoints([endpoint.clone()].into_iter().collect()); + directory.observe(endpoint, registry(1)); + directory.configure_kodo(Vec::new(), true); + + let producer = directory.producers(Some("events")).remove(0); + assert_eq!(producer.broadcast_address, "broker-1.rustqueue"); + assert_eq!(producer.tcp_port, 4150); + assert_eq!(producer.http_port, 4152); + assert!(directory.node_producers().is_empty()); + } } diff --git a/crates/discovery/src/kube_source.rs b/crates/discovery/src/kube_source.rs index c9d6a85..d5a77ca 100644 --- a/crates/discovery/src/kube_source.rs +++ b/crates/discovery/src/kube_source.rs @@ -6,6 +6,7 @@ use kube::api::{Api, ListParams}; use kube::Client; use std::collections::BTreeSet; use std::net::IpAddr; +use std::path::{Path, PathBuf}; use std::sync::atomic::Ordering; use std::sync::Arc; use std::time::Duration; @@ -22,11 +23,12 @@ pub struct RefreshConfig { pub poll_interval: Duration, pub endpoint_slice_timeout: Duration, pub stale_after: Duration, - pub registry_token: Option, + pub registry_token_file: Option, pub max_parallel_polls: usize, } pub async fn run_refresh_loop(directory: Directory, config: RefreshConfig) -> anyhow::Result<()> { + directory.configure_source_health(config.stale_after); let client = Client::try_default() .await .context("create Kubernetes client")?; @@ -51,8 +53,21 @@ pub async fn run_refresh_loop(directory: Directory, config: RefreshConfig) -> an { Ok(Ok(slices)) => { let endpoints = endpoints_from_slices(&slices.items, config.fallback_http_port); + let endpoints_empty = endpoints.is_empty(); directory.replace_endpoints(endpoints); - poll_registries(&directory, &http, &config).await; + let token = match load_registry_token(config.registry_token_file.as_deref()) { + Ok(token) => token, + Err(error) => { + tracing::warn!(%error, "registry token reload failed"); + directory.expire(config.stale_after); + continue; + } + }; + let successful_polls = + poll_registries(&directory, &http, &config, token.as_deref()).await; + if endpoints_empty || successful_polls > 0 { + directory.mark_source_success(); + } } Ok(Err(error)) => tracing::warn!(%error, "EndpointSlice refresh failed"), Err(_) => { @@ -109,12 +124,17 @@ fn endpoints_from_slices(slices: &[EndpointSlice], fallback_port: u16) -> BTreeS result } -async fn poll_registries(directory: &Directory, http: &reqwest::Client, config: &RefreshConfig) { +async fn poll_registries( + directory: &Directory, + http: &reqwest::Client, + config: &RefreshConfig, + token: Option<&str>, +) -> usize { let permits = Arc::new(Semaphore::new(config.max_parallel_polls.max(1))); let mut polls = FuturesUnordered::new(); for endpoint in directory.endpoints() { let http = http.clone(); - let token = config.registry_token.clone(); + let token = token.map(str::to_owned); let permits = Arc::clone(&permits); let metrics = directory.metrics().clone(); let directory = directory.clone(); @@ -151,11 +171,27 @@ async fn poll_registries(directory: &Directory, http: &reqwest::Client, config: Some((endpoint, Some(registry))) }); } + let mut successful = 0; while let Some(observation) = polls.next().await { - if let Some((endpoint, Some(registry))) = observation { - directory.observe(endpoint, registry); + if let Some((endpoint, registry)) = observation { + successful += 1; + if let Some(registry) = registry { + directory.observe(endpoint, registry); + } } } + successful +} + +fn load_registry_token(path: Option<&Path>) -> anyhow::Result> { + let Some(path) = path else { + return Ok(None); + }; + let token = std::fs::read_to_string(path) + .with_context(|| format!("read registry token file {}", path.display()))?; + let token = token.trim(); + anyhow::ensure!(!token.is_empty(), "registry token file is empty"); + Ok(Some(token.to_owned())) } async fn read_json_bounded( @@ -215,4 +251,23 @@ mod tests { assert_eq!(endpoints.len(), 2); assert_eq!(endpoints.iter().next().unwrap().http_port, 4151); } + + #[test] + fn registry_token_file_is_reloaded_and_empty_values_fail_closed() { + let path = + std::env::temp_dir().join(format!("rustqueue-discovery-token-{}", std::process::id())); + std::fs::write(&path, "first\n").unwrap(); + assert_eq!( + load_registry_token(Some(&path)).unwrap().as_deref(), + Some("first") + ); + std::fs::write(&path, "second").unwrap(); + assert_eq!( + load_registry_token(Some(&path)).unwrap().as_deref(), + Some("second") + ); + std::fs::write(&path, " \n").unwrap(); + assert!(load_registry_token(Some(&path)).is_err()); + std::fs::remove_file(path).unwrap(); + } } diff --git a/crates/discovery/src/main.rs b/crates/discovery/src/main.rs index 9bc898d..2a439f2 100644 --- a/crates/discovery/src/main.rs +++ b/crates/discovery/src/main.rs @@ -32,6 +32,16 @@ struct Cli { endpoint_slice_timeout_ms: u64, #[arg(long, env = "RUSTQUEUE_REGISTRY_TOKEN_FILE")] registry_token_file: Option, + #[arg( + long, + env = "RUSTQUEUE_KODO_COMPATIBILITY_ENABLED", + default_value_t = false + )] + kodo_compatibility_enabled: bool, + #[arg(long, env = "RUSTQUEUE_KODO_GATEWAY_ADDRESS")] + kodo_gateway_address: Option, + #[arg(long, env = "RUSTQUEUE_KODO_CLEANUP_ENABLED", default_value_t = false)] + kodo_cleanup_enabled: bool, } #[tokio::main] @@ -46,12 +56,38 @@ async fn main() -> anyhow::Result<()> { if cli.endpoint_slice_timeout_ms == 0 { anyhow::bail!("EndpointSlice timeout must be greater than zero"); } - let token = cli - .registry_token_file - .map(std::fs::read_to_string) - .transpose()? - .map(|token| token.trim().to_owned()); + if cli.kodo_cleanup_enabled { + anyhow::bail!( + "Kodo automatic cleanup is disabled until cluster-wide atomic deletion is available" + ); + } + if cli.kodo_compatibility_enabled { + if cli + .kodo_gateway_address + .as_deref() + .is_none_or(|address| address.trim().is_empty()) + { + anyhow::bail!("Kodo Gateway address is required and cannot be empty"); + } + } else if cli.kodo_cleanup_enabled || cli.kodo_gateway_address.is_some() { + anyhow::bail!("Kodo cleanup and Gateway address require Kodo compatibility"); + } let directory = Directory::default(); + if cli.kodo_compatibility_enabled { + directory.configure_kodo( + cli.kodo_gateway_address + .map(|address| { + let address = address.trim().to_owned(); + (0..3) + .map(|ordinal| { + rustqueue_discovery::Producer::gateway(address.clone(), ordinal) + }) + .collect() + }) + .expect("validated Kodo Gateway address"), + cli.kodo_cleanup_enabled, + ); + } let refresh = run_refresh_loop( directory.clone(), RefreshConfig { @@ -61,7 +97,7 @@ async fn main() -> anyhow::Result<()> { poll_interval: Duration::from_secs(2), endpoint_slice_timeout: Duration::from_millis(cli.endpoint_slice_timeout_ms), stale_after: Duration::from_secs(5), - registry_token: token, + registry_token_file: cli.registry_token_file, max_parallel_polls: 128, }, ); diff --git a/crates/discovery/src/model.rs b/crates/discovery/src/model.rs index 8d19d9a..81fb9eb 100644 --- a/crates/discovery/src/model.rs +++ b/crates/discovery/src/model.rs @@ -116,4 +116,18 @@ impl Producer { node_id: registry.node_id, } } + + pub fn gateway(address: String, ordinal: usize) -> Self { + let tcp_port = [4150, 4152, 4153].get(ordinal).copied().unwrap_or(4150); + let http_port = [4151, 4154, 4155].get(ordinal).copied().unwrap_or(4151); + Self { + remote_address: format!("{address}:{tcp_port}"), + hostname: address.clone(), + broadcast_address: address, + tcp_port, + http_port, + version: env!("CARGO_PKG_VERSION").into(), + node_id: 1_000_000u64.saturating_add(ordinal as u64), + } + } } diff --git a/crates/discovery/src/server.rs b/crates/discovery/src/server.rs index 2e6b7aa..0beb498 100644 --- a/crates/discovery/src/server.rs +++ b/crates/discovery/src/server.rs @@ -3,8 +3,8 @@ use axum::body::Body; use axum::extract::{Query, State}; use axum::http::{header, HeaderValue, Request}; use axum::middleware::{self, Next}; -use axum::response::Response; -use axum::routing::get; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; use axum::{Json, Router}; use serde::Deserialize; use serde_json::{json, Value}; @@ -15,6 +15,12 @@ struct TopicQuery { topic: String, } +#[derive(Deserialize)] +struct ChannelQuery { + topic: String, + channel: String, +} + pub fn router(directory: Directory) -> Router { Router::new() .route("/ping", get(|| async { "OK" })) @@ -23,8 +29,10 @@ pub fn router(directory: Directory) -> Router { .route("/topics", get(topics)) .route("/channels", get(channels)) .route("/nodes", get(nodes)) + .route("/channel/delete", post(delete_channel)) .route("/v1/publishers/head", get(publishers_head)) .route("/v1/publishers", get(publishers)) + .route("/v1/brokers", get(brokers)) .route("/v1/health", get(health)) .route("/metrics", get(prometheus)) .layer(middleware::from_fn(nsq_content_negotiation)) @@ -32,7 +40,28 @@ pub fn router(directory: Directory) -> Router { } async fn prometheus(State(directory): State) -> String { - directory.metrics().render() + let mut output = directory.metrics().render(); + output.push_str( + "# HELP rustqueue_discovery_source_ready Whether the Kubernetes and Broker registry source is fresh.\n\ +# TYPE rustqueue_discovery_source_ready gauge\n\ +# HELP rustqueue_discovery_lookup_ready Whether lookup responses have a complete compatible Broker inventory.\n\ +# TYPE rustqueue_discovery_lookup_ready gauge\n\ +# HELP rustqueue_discovery_brokers Number of consume-ready Brokers in the current inventory.\n\ +# TYPE rustqueue_discovery_brokers gauge\n\ +# HELP rustqueue_discovery_publishers Number of publish-ready Brokers in the current inventory.\n\ +# TYPE rustqueue_discovery_publishers gauge\n", + ); + output.push_str(&format!( + "rustqueue_discovery_source_ready {}\n\ +rustqueue_discovery_lookup_ready {}\n\ +rustqueue_discovery_brokers {}\n\ +rustqueue_discovery_publishers {}\n", + u8::from(directory.source_ready()), + u8::from(directory.lookup_ready()), + directory.broker_count(), + directory.publisher_count(), + )); + output } async fn nsq_content_negotiation(request: Request, next: Next) -> Response { @@ -71,59 +100,144 @@ async fn info(State(directory): State) -> Json { })) } -async fn lookup( - State(directory): State, - Query(query): Query, -) -> Json { +async fn lookup(State(directory): State, Query(query): Query) -> Response { + // A Broker outage must not block consumers from discovering the remaining + // topic shards. The complete-inventory signal is still exposed for + // monitoring and gates operations that require all three Brokers. + if !directory.source_ready() { + return source_unavailable(); + } Json(json!({ "channels": directory.channels(&query.topic), "producers": directory.producers(Some(&query.topic)), })) + .into_response() +} + +async fn topics(State(directory): State) -> Response { + if !directory.source_ready() { + return source_unavailable(); + } + Json(json!({"topics": directory.topics()})).into_response() } -async fn topics(State(directory): State) -> Json { - Json(json!({"topics": directory.topics()})) +async fn channels(State(directory): State, Query(query): Query) -> Response { + if !directory.source_ready() { + return source_unavailable(); + } + Json(json!({"channels": directory.channels(&query.topic)})).into_response() } -async fn channels( - State(directory): State, - Query(query): Query, -) -> Json { - Json(json!({"channels": directory.channels(&query.topic)})) +async fn nodes(State(directory): State) -> Response { + if !directory.source_ready() || !directory.kodo_nodes_ready() { + return source_unavailable(); + } + Json(json!({"producers": directory.node_producers()})).into_response() } -async fn nodes(State(directory): State) -> Json { - Json(json!({"producers": directory.producers(None)})) +async fn publishers(State(directory): State) -> Response { + if !directory.source_ready() { + return source_unavailable(); + } + let (revision, producers) = directory.publisher_snapshot(); + Json(json!({ + "revision": revision, + "producers": producers, + })) + .into_response() } -async fn publishers(State(directory): State) -> Json { +async fn publishers_head(State(directory): State) -> Response { + if !directory.source_ready() { + return source_unavailable(); + } + let (revision, broker_count) = directory.publisher_head(); Json(json!({ - "revision": directory.revision(), - "producers": directory.publishers(), + "revision": revision, + "broker_count": broker_count, })) + .into_response() } -async fn publishers_head(State(directory): State) -> Json { +async fn brokers(State(directory): State) -> Response { + if !directory.source_ready() { + return source_unavailable(); + } + let (revision, producers) = directory.broker_snapshot(); Json(json!({ - "revision": directory.revision(), - "broker_count": directory.publisher_count(), + "revision": revision, + "producers": producers, })) + .into_response() +} + +fn source_unavailable() -> Response { + ( + axum::http::StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"message": "discovery source is not ready"})), + ) + .into_response() +} + +async fn delete_channel( + State(directory): State, + Query(query): Query, +) -> Response { + if !directory.kodo_cleanup_enabled() { + return ( + axum::http::StatusCode::NOT_FOUND, + "E_NOT_FOUND Kodo cleanup compatibility is disabled", + ) + .into_response(); + } + if query.topic.is_empty() || query.channel.is_empty() { + return ( + axum::http::StatusCode::BAD_REQUEST, + "E_BAD_REQUEST topic and channel are required", + ) + .into_response(); + } + if directory.broker_count() != 3 { + return ( + axum::http::StatusCode::SERVICE_UNAVAILABLE, + "E_UNAVAILABLE Kodo cleanup requires all three Brokers", + ) + .into_response(); + } + "OK".into_response() } -async fn health(State(directory): State) -> Json { - Json(json!({"status": "ready", "broker_count": directory.broker_count()})) +async fn health(State(directory): State) -> Response { + let ready = directory.source_ready(); + ( + if ready { + axum::http::StatusCode::OK + } else { + axum::http::StatusCode::SERVICE_UNAVAILABLE + }, + Json(json!({ + "status": if ready { "ready" } else { "source-unavailable" }, + "broker_count": directory.broker_count(), + "lookup_ready": directory.lookup_ready() + })), + ) + .into_response() } #[cfg(test)] mod tests { use super::*; + use crate::Producer; use axum::body::Body; use axum::http::Request; use tower::ServiceExt; #[tokio::test] async fn lookup_shape_matches_nsqlookupd() { - let response = router(Directory::default()) + let directory = Directory::default(); + directory.configure_source_health(std::time::Duration::from_secs(5)); + directory.mark_source_success(); + let response = router(directory) .oneshot( Request::builder() .uri("/lookup?topic=events") @@ -139,4 +253,248 @@ mod tests { "nsq; version=1.0" ); } + + #[tokio::test] + async fn nodes_use_stable_kodo_gateways_without_changing_lookup() { + let directory = Directory::default(); + directory.configure_kodo( + (0..3) + .map(|ordinal| Producer::gateway("gateway".into(), ordinal)) + .collect(), + false, + ); + directory.configure_source_health(std::time::Duration::from_secs(5)); + directory.mark_source_success(); + let response = router(directory) + .oneshot( + Request::builder() + .uri("/nodes") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let value: Value = serde_json::from_slice(&body).unwrap(); + let producers = value["producers"].as_array().unwrap(); + assert_eq!(producers.len(), 3); + assert!(producers + .iter() + .all(|producer| producer["broadcast_address"] == "gateway")); + assert_eq!( + producers + .iter() + .map(|producer| producer["tcp_port"].as_u64().unwrap()) + .collect::>(), + vec![4150, 4152, 4153] + ); + assert_eq!( + producers + .iter() + .map(|producer| producer["http_port"].as_u64().unwrap()) + .collect::>(), + vec![4151, 4154, 4155] + ); + } + + #[tokio::test] + async fn kodo_lookup_serves_healthy_shards_while_inventory_is_incomplete() { + let directory = Directory::default(); + directory.configure_source_health(std::time::Duration::from_secs(5)); + directory.configure_kodo( + (0..3) + .map(|ordinal| Producer::gateway("gateway".into(), ordinal)) + .collect(), + false, + ); + for node_id in 1..=2 { + let endpoint = crate::BrokerEndpoint { + address: format!("127.0.0.{node_id}").parse().unwrap(), + http_port: 4151, + }; + directory.observe( + endpoint, + crate::BrokerRegistry { + format: 7, + revision: node_id, + node_id, + ready: true, + publish_ready: true, + consume_ready: true, + broadcast_address: format!("broker-{node_id}"), + tcp_port: 4150, + http_port: 4151, + stored_messages: 0, + depth: 0, + in_flight: 0, + topics: vec![crate::RegistryTopic { + name: "events".into(), + paused: false, + channels: vec!["workers".into()], + stored_messages: 0, + }], + compatibility: None, + }, + ); + } + directory.mark_source_success(); + + let lookup = router(directory.clone()) + .oneshot( + Request::builder() + .uri("/lookup?topic=events") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(lookup.status(), axum::http::StatusCode::OK); + let body = axum::body::to_bytes(lookup.into_body(), usize::MAX) + .await + .unwrap(); + let value: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(value["producers"].as_array().unwrap().len(), 2); + assert!(!directory.lookup_ready()); + let nodes = router(directory.clone()) + .oneshot( + Request::builder() + .uri("/nodes") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(nodes.status(), axum::http::StatusCode::OK); + + let endpoint = crate::BrokerEndpoint { + address: "127.0.0.3".parse().unwrap(), + http_port: 4151, + }; + directory.observe( + endpoint, + crate::BrokerRegistry { + format: 7, + revision: 3, + node_id: 3, + ready: true, + publish_ready: true, + consume_ready: true, + broadcast_address: "broker-3".into(), + tcp_port: 4150, + http_port: 4151, + stored_messages: 0, + depth: 0, + in_flight: 0, + topics: vec![crate::RegistryTopic { + name: "events".into(), + paused: false, + channels: vec!["workers".into()], + stored_messages: 0, + }], + compatibility: None, + }, + ); + let lookup = router(directory.clone()) + .oneshot( + Request::builder() + .uri("/lookup?topic=events") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(lookup.status(), axum::http::StatusCode::OK); + let body = axum::body::to_bytes(lookup.into_body(), usize::MAX) + .await + .unwrap(); + let value: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(value["producers"].as_array().unwrap().len(), 3); + assert!(directory.lookup_ready()); + } + + #[tokio::test] + async fn nodes_and_health_fail_closed_until_the_source_is_ready() { + let directory = Directory::default(); + directory.configure_source_health(std::time::Duration::from_secs(5)); + directory.configure_kodo(Vec::new(), false); + for uri in [ + "/nodes", + "/lookup?topic=events", + "/topics", + "/channels?topic=events", + "/v1/publishers/head", + "/v1/publishers", + "/v1/brokers", + "/v1/health", + ] { + let response = router(directory.clone()) + .oneshot(Request::builder().uri(uri).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!( + response.status(), + axum::http::StatusCode::SERVICE_UNAVAILABLE + ); + } + directory.mark_source_success(); + let response = router(directory) + .oneshot( + Request::builder() + .uri("/v1/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), axum::http::StatusCode::OK); + } + + #[tokio::test] + async fn cleanup_fails_closed_without_the_complete_broker_inventory() { + let directory = Directory::default(); + directory.configure_kodo( + (0..3) + .map(|ordinal| Producer::gateway(format!("gateway-{ordinal}"), ordinal)) + .collect(), + true, + ); + let response = router(directory) + .oneshot( + Request::builder() + .method("POST") + .uri("/channel/delete?topic=events&channel=workers") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + response.status(), + axum::http::StatusCode::SERVICE_UNAVAILABLE + ); + } + + #[tokio::test] + async fn metrics_expose_source_and_lookup_readiness() { + let directory = Directory::default(); + directory.configure_source_health(std::time::Duration::from_secs(5)); + directory.mark_source_success(); + let response = router(directory) + .oneshot( + Request::builder() + .uri("/metrics") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert!(body.contains("rustqueue_discovery_source_ready 1\n")); + assert!(body.contains("rustqueue_discovery_lookup_ready 1\n")); + } } diff --git a/crates/operator/Cargo.toml b/crates/operator/Cargo.toml index 728a9f4..c7d8059 100644 --- a/crates/operator/Cargo.toml +++ b/crates/operator/Cargo.toml @@ -11,6 +11,7 @@ axum.workspace = true chrono.workspace = true crc32c.workspace = true futures.workspace = true +hex.workspace = true k8s-openapi.workspace = true kube.workspace = true rand.workspace = true @@ -20,6 +21,7 @@ schemars.workspace = true serde.workspace = true serde_json.workspace = true serde_yaml.workspace = true +sha2.workspace = true thiserror.workspace = true tokio.workspace = true toml.workspace = true diff --git a/crates/operator/src/controller/apply.rs b/crates/operator/src/controller/apply.rs index b3d973d..2ef3860 100644 --- a/crates/operator/src/controller/apply.rs +++ b/crates/operator/src/controller/apply.rs @@ -6,7 +6,7 @@ use k8s_openapi::api::core::v1::{ConfigMap, Service, ServiceAccount}; use k8s_openapi::api::networking::v1::NetworkPolicy; use k8s_openapi::api::policy::v1::PodDisruptionBudget; use k8s_openapi::api::rbac::v1::{Role, RoleBinding}; -use kube::api::{Api, Patch, PatchParams}; +use kube::api::{Api, DeleteParams, ObjectMeta, Patch, PatchParams}; use kube::ResourceExt; use serde_json::json; @@ -16,31 +16,23 @@ pub(super) async fn resources( set: ResourceSet, ) -> anyhow::Result<()> { let params = PatchParams::apply(MANAGER).force(); + let cluster_name = set.brokers.name_any(); + let owner_uid = set + .brokers + .metadata + .owner_references + .as_deref() + .and_then(|owners| owners.iter().find(|owner| owner.controller == Some(true))) + .map(|owner| owner.uid.clone()); + let network_policies = Api::::namespaced(client.clone(), namespace); Api::::namespaced(client.clone(), namespace) .patch(&set.config.name_any(), ¶ms, &Patch::Apply(&set.config)) .await?; - Api::::namespaced(client.clone(), namespace) + network_policies .patch( - &set.service_account.name_any(), + &set.broker_network_policy.name_any(), ¶ms, - &Patch::Apply(&set.service_account), - ) - .await?; - Api::::namespaced(client.clone(), namespace) - .patch(&set.role.name_any(), ¶ms, &Patch::Apply(&set.role)) - .await?; - Api::::namespaced(client.clone(), namespace) - .patch( - &set.role_binding.name_any(), - ¶ms, - &Patch::Apply(&set.role_binding), - ) - .await?; - Api::::namespaced(client.clone(), namespace) - .patch( - &set.broker_service.name_any(), - ¶ms, - &Patch::Apply(&set.broker_service), + &Patch::Apply(&set.broker_network_policy), ) .await?; Api::::namespaced(client.clone(), namespace) @@ -57,47 +49,172 @@ pub(super) async fn resources( &Patch::Apply(&set.broker_pdb), ) .await?; - Api::::namespaced(client.clone(), namespace) + runtime_resources(client, namespace, &set).await?; + + let services = Api::::namespaced(client.clone(), namespace); + let gateway_publish_service_name = format!("{cluster_name}-kodo-publish"); + let statefulsets = Api::::namespaced(client.clone(), namespace); + let gateway_name = format!("{cluster_name}-kodo-gateway"); + if set.kodo_gateway.is_none() + && !set.retain_existing_kodo_resources + && statefulsets + .get_opt(&gateway_name) + .await? + .as_ref() + .is_some_and(|resource| controlled_by(&resource.metadata, owner_uid.as_deref())) + { + statefulsets + .delete(&gateway_name, &DeleteParams::default()) + .await?; + } + if !set.retain_existing_kodo_resources + && set.kodo_gateway_service.is_none() + && services + .get_opt(&gateway_publish_service_name) + .await? + .as_ref() + .is_some_and(|resource| controlled_by(&resource.metadata, owner_uid.as_deref())) + { + services + .delete(&gateway_publish_service_name, &DeleteParams::default()) + .await?; + } + let gateway_headless_service_name = format!("{cluster_name}-kodo-gateways"); + if !set.retain_existing_kodo_resources + && set.kodo_gateway_headless_service.is_none() + && services + .get_opt(&gateway_headless_service_name) + .await? + .as_ref() + .is_some_and(|resource| controlled_by(&resource.metadata, owner_uid.as_deref())) + { + services + .delete(&gateway_headless_service_name, &DeleteParams::default()) + .await?; + } + let pdbs = Api::::namespaced(client.clone(), namespace); + let gateway_pdb_name = format!("{gateway_name}-pdb"); + if set.kodo_gateway_pdb.is_none() + && !set.retain_existing_kodo_resources + && pdbs + .get_opt(&gateway_pdb_name) + .await? + .as_ref() + .is_some_and(|resource| controlled_by(&resource.metadata, owner_uid.as_deref())) + { + pdbs.delete(&gateway_pdb_name, &DeleteParams::default()) + .await?; + } + let gateway_policy_name = format!("{cluster_name}-kodo-gateway-ingress"); + if !set.retain_existing_kodo_resources + && set.kodo_gateway_network_policy.is_none() + && network_policies + .get_opt(&gateway_policy_name) + .await? + .as_ref() + .is_some_and(|resource| controlled_by(&resource.metadata, owner_uid.as_deref())) + { + network_policies + .delete(&gateway_policy_name, &DeleteParams::default()) + .await?; + } + Ok(()) +} + +pub(super) async fn runtime_resources( + client: &kube::Client, + namespace: &str, + set: &ResourceSet, +) -> anyhow::Result<()> { + let params = PatchParams::apply(MANAGER).force(); + Api::::namespaced(client.clone(), namespace) .patch( - &set.discovery_service.name_any(), + &set.service_account.name_any(), ¶ms, - &Patch::Apply(&set.discovery_service), + &Patch::Apply(&set.service_account), ) .await?; - Api::::namespaced(client.clone(), namespace) + Api::::namespaced(client.clone(), namespace) + .patch(&set.role.name_any(), ¶ms, &Patch::Apply(&set.role)) + .await?; + Api::::namespaced(client.clone(), namespace) .patch( - &set.discovery.name_any(), + &set.role_binding.name_any(), ¶ms, - &Patch::Apply(&set.discovery), + &Patch::Apply(&set.role_binding), ) .await?; - Api::::namespaced(client.clone(), namespace) + + let network_policies = Api::::namespaced(client.clone(), namespace); + network_policies .patch( - &set.discovery_pdb.name_any(), + &set.network_policy.name_any(), ¶ms, - &Patch::Apply(&set.discovery_pdb), + &Patch::Apply(&set.network_policy), ) .await?; - Api::::namespaced(client.clone(), namespace) + if let Some(policy) = &set.kodo_gateway_network_policy { + network_policies + .patch(&policy.name_any(), ¶ms, &Patch::Apply(policy)) + .await?; + } + + let services = Api::::namespaced(client.clone(), namespace); + for service in [ + Some(&set.broker_service), + Some(&set.discovery_service), + Some(&set.proxy_service), + set.kodo_gateway_service.as_ref(), + set.kodo_gateway_headless_service.as_ref(), + ] + .into_iter() + .flatten() + { + services + .patch(&service.name_any(), ¶ms, &Patch::Apply(service)) + .await?; + } + + Api::::namespaced(client.clone(), namespace) .patch( - &set.proxy_service.name_any(), + &set.discovery.name_any(), ¶ms, - &Patch::Apply(&set.proxy_service), + &Patch::Apply(&set.discovery), ) .await?; Api::::namespaced(client.clone(), namespace) .patch(&set.proxy.name_any(), ¶ms, &Patch::Apply(&set.proxy)) .await?; - Api::::namespaced(client.clone(), namespace) - .patch( - &set.network_policy.name_any(), - ¶ms, - &Patch::Apply(&set.network_policy), - ) - .await?; + if let Some(gateway) = &set.kodo_gateway { + Api::::namespaced(client.clone(), namespace) + .patch(&gateway.name_any(), ¶ms, &Patch::Apply(gateway)) + .await?; + } + + let pdbs = Api::::namespaced(client.clone(), namespace); + pdbs.patch( + &set.discovery_pdb.name_any(), + ¶ms, + &Patch::Apply(&set.discovery_pdb), + ) + .await?; + if let Some(pdb) = &set.kodo_gateway_pdb { + pdbs.patch(&pdb.name_any(), ¶ms, &Patch::Apply(pdb)) + .await?; + } Ok(()) } +fn controlled_by(metadata: &ObjectMeta, owner_uid: Option<&str>) -> bool { + owner_uid.is_some_and(|owner_uid| { + metadata.owner_references.as_deref().is_some_and(|owners| { + owners + .iter() + .any(|owner| owner.controller == Some(true) && owner.uid == owner_uid) + }) + }) +} + pub(super) async fn status( client: &kube::Client, cluster: &RustQueue, @@ -113,3 +230,27 @@ pub(super) async fn status( .await?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; + + #[test] + fn optional_resource_cleanup_requires_the_cluster_owner() { + let metadata = ObjectMeta { + owner_references: Some(vec![OwnerReference { + api_version: "rustqueue.io/v1alpha1".into(), + kind: "RustQueue".into(), + name: "queue".into(), + uid: "queue-uid".into(), + controller: Some(true), + block_owner_deletion: Some(true), + }]), + ..Default::default() + }; + assert!(controlled_by(&metadata, Some("queue-uid"))); + assert!(!controlled_by(&metadata, Some("other-uid"))); + assert!(!controlled_by(&metadata, None)); + } +} diff --git a/crates/operator/src/controller/auth.rs b/crates/operator/src/controller/auth.rs index de22bf5..ebc85f7 100644 --- a/crates/operator/src/controller/auth.rs +++ b/crates/operator/src/controller/auth.rs @@ -45,6 +45,12 @@ pub(super) async fn ensure( "console-token".into(), Alphanumeric.sample_string(&mut rand::thread_rng(), 48), ); + if cluster.spec.kodo_compatibility.effective_cleanup_enabled() { + string_data.insert( + "kodo-cleanup-token".into(), + Alphanumeric.sample_string(&mut rand::thread_rng(), 48), + ); + } api.create( &PostParams::default(), &Secret { @@ -71,17 +77,30 @@ pub(super) async fn ensure( .data .as_ref() .is_some_and(|data| data.contains_key("observer-token")); - if generated && (missing_console || has_observer) { - let patch = if missing_console { - json!({ - "stringData": { - "console-token": Alphanumeric.sample_string(&mut rand::thread_rng(), 48), - }, - "data": {"observer-token": null}, - }) - } else { - json!({"data": {"observer-token": null}}) - }; + let missing_cleanup = secret + .data + .as_ref() + .is_none_or(|data| !data.contains_key("kodo-cleanup-token")); + let needs_cleanup = + cluster.spec.kodo_compatibility.effective_cleanup_enabled() && missing_cleanup; + if generated && (missing_console || needs_cleanup || has_observer) { + let mut additions = BTreeMap::new(); + if missing_console { + additions.insert( + "console-token", + Alphanumeric.sample_string(&mut rand::thread_rng(), 48), + ); + } + if needs_cleanup { + additions.insert( + "kodo-cleanup-token", + Alphanumeric.sample_string(&mut rand::thread_rng(), 48), + ); + } + let patch = json!({ + "stringData": additions, + "data": {"observer-token": null}, + }); secret = api .patch(&name, &PatchParams::default(), &Patch::Merge(patch)) .await?; @@ -98,15 +117,40 @@ pub(super) async fn ensure( } Ok(value.trim().into()) }; - let _ = read("console-token")?; + let admin_token = read("admin-token")?; + let registry_token = read("registry-token")?; + let console_token = read("console-token")?; + let mut tokens = vec![ + ("admin-token", admin_token.as_str()), + ("registry-token", registry_token.as_str()), + ("console-token", console_token.as_str()), + ]; + let cleanup_token = if cluster.spec.kodo_compatibility.effective_cleanup_enabled() { + Some(read("kodo-cleanup-token")?) + } else { + None + }; + if let Some(cleanup_token) = cleanup_token.as_deref() { + tokens.push(("kodo-cleanup-token", cleanup_token)); + } + validate_distinct_tokens(&tokens)?; Ok(AuthSecret { name, - admin_token: read("admin-token")?, - registry_token: read("registry-token")?, + admin_token, + registry_token, revision: secret.resource_version().unwrap_or_default(), }) } +fn validate_distinct_tokens(tokens: &[(&str, &str)]) -> anyhow::Result<()> { + for (index, (left_name, left)) in tokens.iter().enumerate() { + if let Some((right_name, _)) = tokens[index + 1..].iter().find(|(_, right)| left == right) { + bail!("auth Secret {left_name} and {right_name} must be distinct"); + } + } + Ok(()) +} + pub(super) async fn mounted_secret_revision( client: &kube::Client, cluster: &RustQueue, @@ -135,3 +179,28 @@ pub(super) async fn mounted_secret_revision( secret.resource_version().unwrap_or_default() )) } + +#[cfg(test)] +mod tests { + use super::validate_distinct_tokens; + + #[test] + fn auth_roles_require_pairwise_distinct_tokens() { + assert!(validate_distinct_tokens(&[ + ("admin-token", "admin"), + ("registry-token", "registry"), + ("console-token", "console"), + ("kodo-cleanup-token", "cleanup"), + ]) + .is_ok()); + let error = validate_distinct_tokens(&[ + ("admin-token", "shared"), + ("registry-token", "registry"), + ("console-token", "shared"), + ]) + .unwrap_err(); + assert!(error + .to_string() + .contains("admin-token and console-token must be distinct")); + } +} diff --git a/crates/operator/src/controller/drain.rs b/crates/operator/src/controller/drain.rs index ecbe7dd..0de9e4d 100644 --- a/crates/operator/src/controller/drain.rs +++ b/crates/operator/src/controller/drain.rs @@ -1,6 +1,7 @@ use super::auth::AuthSecret; use super::ContextData; use crate::RustQueue; +use futures::{stream, StreamExt}; use k8s_openapi::api::apps::v1::StatefulSet; use k8s_openapi::api::core::v1::Pod; use kube::api::{Api, DeleteParams, ListParams, Patch, PatchParams}; @@ -82,6 +83,48 @@ pub(super) async fn resume_current( Ok(()) } +pub(super) async fn resume_all( + context: &ContextData, + cluster: &RustQueue, + namespace: &str, + auth: &AuthSecret, +) -> anyhow::Result> { + let api: Api = Api::namespaced(context.client.clone(), namespace); + let checks = broker_pods(&api, cluster) + .await? + .into_iter() + .filter_map(|pod| { + let name = pod.name_any(); + let ip = pod.status?.pod_ip?; + Some((name, ip)) + }); + let mut resumed: Vec<_> = stream::iter(checks) + .map(|(name, ip)| async move { + match drain_status(context, &ip, auth).await { + Ok(status) if status.draining => { + match set_drain(context, &ip, auth, false, false).await { + Ok(()) => Some(name), + Err(error) => { + tracing::warn!(broker = %name, %error, "could not resume drained Broker"); + None + } + } + } + Ok(_) => None, + Err(error) => { + tracing::warn!(broker = %name, %error, "could not inspect Broker drain state"); + None + } + } + }) + .buffer_unordered(32) + .filter_map(async move |name| name) + .collect() + .await; + resumed.sort(); + Ok(resumed) +} + pub(super) async fn scale_down_one( context: &ContextData, cluster: &RustQueue, @@ -91,6 +134,21 @@ pub(super) async fn scale_down_one( ) -> anyhow::Result { let next = current - 1; let pod_name = format!("{}-{next}", cluster.name_any()); + let pods = broker_pods( + &Api::::namespaced(context.client.clone(), namespace), + cluster, + ) + .await?; + if !other_brokers_ready(&pods, &pod_name, current) { + resume_one(context, namespace, &pod_name, auth).await?; + return Ok(Progress { + target: Some(pod_name.clone()), + phase: "WaitingForReady", + message: format!( + "resumed scale-down target {pod_name}; waiting for every other Broker to become Ready" + ), + }); + } if drain_pod(context, namespace, &pod_name, DrainGoal::Empty, auth).await? { Api::::namespaced(context.client.clone(), namespace) .patch( @@ -122,7 +180,9 @@ pub(super) async fn rollout_one( ) -> anyhow::Result { let api: Api = Api::namespaced(context.client.clone(), namespace); let pods = broker_pods(&api, cluster).await?; - let decision = rollout_decision(&pods, &options); + let draining = + intentionally_draining_pods(context, &pods, options.desired_revision, auth).await; + let decision = rollout_decision(&pods, &options, &draining); match decision { Decision::Completed => Ok(Progress { target: None, @@ -152,6 +212,16 @@ pub(super) async fn rollout_one( phase: "Blocked", message, }), + Decision::Resume(names, message) => { + for name in names { + resume_one(context, namespace, &name, auth).await?; + } + Ok(Progress { + target: None, + phase: "WaitingForReady", + message, + }) + } Decision::Replace(name) => { if drain_pod(context, namespace, &name, DrainGoal::Quiesced, auth).await? { api.delete(&name, &DeleteParams::default()).await?; @@ -171,6 +241,28 @@ pub(super) async fn rollout_one( } } +pub(super) async fn resume_one( + context: &ContextData, + namespace: &str, + broker: &str, + auth: &AuthSecret, +) -> anyhow::Result<()> { + let Some(pod) = Api::::namespaced(context.client.clone(), namespace) + .get_opt(broker) + .await? + else { + return Ok(()); + }; + let Some(ip) = pod + .status + .as_ref() + .and_then(|status| status.pod_ip.as_deref()) + else { + return Ok(()); + }; + set_drain(context, ip, auth, false, false).await +} + pub(super) async fn maintenance( context: &ContextData, cluster: &RustQueue, @@ -186,10 +278,29 @@ pub(super) async fn maintenance( cluster.name_any() ); let api: Api = Api::namespaced(context.client.clone(), namespace); - let pod = api - .get_opt(broker) - .await? + let pods = broker_pods(&api, cluster).await?; + let pod = pods + .iter() + .find(|pod| pod.name_any() == broker) .ok_or_else(|| anyhow::anyhow!("maintenance target {broker} does not exist"))?; + if enabled { + let replicas = Api::::namespaced(context.client.clone(), namespace) + .get(&cluster.name_any()) + .await? + .spec + .and_then(|spec| spec.replicas) + .unwrap_or_default(); + if !other_brokers_ready(&pods, broker, replicas) { + resume_one(context, namespace, broker, auth).await?; + return Ok(Progress { + target: Some(broker.into()), + phase: "Blocked", + message: format!( + "resumed maintenance target {broker}; maintenance is blocked until every other Broker is Ready" + ), + }); + } + } let ip = pod .status .as_ref() @@ -232,6 +343,21 @@ pub(super) async fn previous_image( .map(ToOwned::to_owned)) } +pub(super) async fn brokers_current_and_ready( + context: &ContextData, + cluster: &RustQueue, + namespace: &str, + desired_revision: &str, + replicas: i32, +) -> anyhow::Result { + let api = Api::::namespaced(context.client.clone(), namespace); + Ok(broker_set_current_and_ready( + &broker_pods(&api, cluster).await?, + desired_revision, + replicas, + )) +} + async fn broker_pods(api: &Api, cluster: &RustQueue) -> anyhow::Result> { let selector = format!( "app.kubernetes.io/instance={},app.kubernetes.io/component=broker", @@ -243,41 +369,92 @@ async fn broker_pods(api: &Api, cluster: &RustQueue) -> anyhow::Result bool { + if !exact_broker_pod_set(pods, replicas) { + return false; + } + let ordinals: std::collections::BTreeSet<_> = pods + .iter() + .filter(|pod| pod_revision(pod) == Some(desired_revision) && pod_ready(pod)) + .filter_map(|pod| pod_ordinal(&pod.name_any())) + .collect(); + ordinals == (0..replicas as u32).collect() +} + enum Decision { Completed, Wait(String), Paused, AwaitingApproval, Blocked(String), + Resume(Vec, String), Replace(String), } -fn rollout_decision(pods: &[Pod], options: &RolloutOptions<'_>) -> Decision { - if pods.len() < options.replicas as usize { +fn rollout_decision( + pods: &[Pod], + options: &RolloutOptions<'_>, + intentionally_draining: &std::collections::BTreeSet, +) -> Decision { + if !exact_broker_pod_set(pods, options.replicas) { return Decision::Wait(format!( - "waiting for {} of {} broker Pods to exist", + "waiting for the exact {}-Pod Broker ordinal set; observed {} Pods", + options.replicas, pods.len(), - options.replicas )); } let current: Vec<_> = pods .iter() .filter(|pod| pod_revision(pod) == Some(options.desired_revision)) .collect(); + let mut outdated: Vec<_> = pods + .iter() + .filter(|pod| pod_revision(pod) != Some(options.desired_revision)) + .collect(); + if outdated.is_empty() { + return current + .iter() + .find(|pod| !pod_ready(pod)) + .map_or(Decision::Completed, |pod| { + Decision::Wait(format!( + "waiting for current broker {} to become Ready", + pod.name_any() + )) + }); + } + let draining: Vec<_> = outdated + .iter() + .filter(|pod| !pod_ready(pod) && intentionally_draining.contains(&pod.name_any())) + .map(|pod| pod.name_any()) + .collect(); if let Some(pod) = current.iter().find(|pod| !pod_ready(pod)) { + if !draining.is_empty() { + return Decision::Resume( + draining, + format!( + "resumed an in-progress drain because current broker {} is unavailable", + pod.name_any() + ), + ); + } return Decision::Wait(format!( "waiting for replacement broker {} to become Ready", pod.name_any() )); } - let mut outdated: Vec<_> = pods + if let Some(pod) = outdated .iter() - .filter(|pod| pod_revision(pod) != Some(options.desired_revision)) - .collect(); - if outdated.is_empty() { - return Decision::Completed; - } - if let Some(pod) = outdated.iter().find(|pod| !pod_ready(pod)) { + .find(|pod| !pod_ready(pod) && !intentionally_draining.contains(&pod.name_any())) + { + if !draining.is_empty() { + return Decision::Resume( + draining, + format!( + "resumed an in-progress drain because outdated broker {} is also unavailable", + pod.name_any() + ), + ); + } return Decision::Wait(format!( "outdated broker {} is not Ready; refusing to make another broker unavailable", pod.name_any() @@ -286,6 +463,15 @@ fn rollout_decision(pods: &[Pod], options: &RolloutOptions<'_>) -> Decision { if options.replicas < 2 { return Decision::Blocked("rolling replacement needs at least two brokers".into()); } + if draining.len() > 1 { + return Decision::Resume( + draining, + "resumed multiple in-progress drains; only one Broker may be disrupted".into(), + ); + } + if let Some(name) = draining.into_iter().next() { + return Decision::Replace(name); + } if options.paused { return Decision::Paused; } @@ -299,6 +485,41 @@ fn rollout_decision(pods: &[Pod], options: &RolloutOptions<'_>) -> Decision { Decision::Replace(outdated.pop().expect("non-empty outdated set").name_any()) } +async fn intentionally_draining_pods( + context: &ContextData, + pods: &[Pod], + desired_revision: &str, + auth: &AuthSecret, +) -> std::collections::BTreeSet { + let mut draining = std::collections::BTreeSet::new(); + for pod in pods + .iter() + .filter(|pod| !pod_ready(pod) && pod_revision(pod) != Some(desired_revision)) + { + let Some(ip) = pod + .status + .as_ref() + .and_then(|status| status.pod_ip.as_deref()) + else { + continue; + }; + match drain_status(context, ip, auth).await { + Ok(status) if status.draining => { + draining.insert(pod.name_any()); + } + Ok(_) => {} + Err(error) => { + tracing::debug!( + broker = %pod.name_any(), + %error, + "could not verify whether unavailable Broker is intentionally draining" + ); + } + } + } + draining +} + async fn drain_pod( context: &ContextData, namespace: &str, @@ -389,6 +610,41 @@ fn pod_ready(pod: &Pod) -> bool { }) } +fn other_brokers_ready(pods: &[Pod], target: &str, replicas: i32) -> bool { + if !exact_broker_pod_set(pods, replicas) { + return false; + } + let Some(replicas) = u32::try_from(replicas).ok() else { + return false; + }; + let Some(target_ordinal) = pod_ordinal(target).filter(|ordinal| *ordinal < replicas) else { + return false; + }; + let ready: std::collections::BTreeSet<_> = pods + .iter() + .filter(|pod| pod.name_any() != target && pod_ready(pod)) + .filter_map(|pod| pod_ordinal(&pod.name_any())) + .collect(); + ready + == (0..replicas) + .filter(|ordinal| *ordinal != target_ordinal) + .collect() +} + +fn exact_broker_pod_set(pods: &[Pod], replicas: i32) -> bool { + let Ok(replicas) = u32::try_from(replicas) else { + return false; + }; + if replicas == 0 || pods.len() != replicas as usize { + return false; + } + let ordinals: std::collections::BTreeSet<_> = pods + .iter() + .filter_map(|pod| pod_ordinal(&pod.name_any())) + .collect(); + ordinals.len() == pods.len() && ordinals == (0..replicas).collect() +} + fn origin(ip: &str) -> String { if ip.contains(':') { format!("http://[{ip}]:4151") @@ -446,7 +702,7 @@ mod tests { pod("queue-2", "new", false), ]; assert!(matches!( - rollout_decision(&pods, &options("new")), + rollout_decision(&pods, &options("new"), &Default::default()), Decision::Wait(_) )); } @@ -459,11 +715,105 @@ mod tests { pod("queue-2", "old", true), ]; assert!(matches!( - rollout_decision(&pods, &options("new")), + rollout_decision(&pods, &options("new"), &Default::default()), Decision::Wait(message) if message.contains("refusing") )); } + #[test] + fn rollout_continues_the_one_intentionally_draining_broker() { + let pods = vec![ + pod("queue-0", "old", true), + pod("queue-1", "old", true), + pod("queue-2", "old", false), + ]; + assert!(matches!( + rollout_decision( + &pods, + &options("new"), + &["queue-2".into()].into_iter().collect() + ), + Decision::Replace(name) if name == "queue-2" + )); + } + + #[test] + fn rollout_resumes_its_drain_if_another_broker_becomes_unavailable() { + let pods = vec![ + pod("queue-0", "old", false), + pod("queue-1", "old", true), + pod("queue-2", "old", false), + ]; + assert!(matches!( + rollout_decision( + &pods, + &options("new"), + &["queue-2".into()].into_iter().collect() + ), + Decision::Resume(names, _) if names == ["queue-2"] + )); + } + + #[test] + fn rollout_resumes_its_drain_if_a_current_broker_becomes_unavailable() { + let pods = vec![ + pod("queue-0", "new", false), + pod("queue-1", "old", true), + pod("queue-2", "old", false), + ]; + assert!(matches!( + rollout_decision( + &pods, + &options("new"), + &["queue-2".into()].into_iter().collect() + ), + Decision::Resume(names, _) if names == ["queue-2"] + )); + } + + #[test] + fn maintenance_and_scale_down_require_every_other_broker_to_be_ready() { + let pods = vec![ + pod("queue-0", "current", true), + pod("queue-1", "current", true), + pod("queue-2", "current", false), + ]; + assert!(other_brokers_ready(&pods, "queue-2", 3)); + assert!(!other_brokers_ready(&pods, "queue-1", 3)); + assert!(!other_brokers_ready( + &[pods[0].clone(), pods[2].clone()], + "queue-2", + 3 + )); + let mut with_stale = pods.clone(); + with_stale.push(pod("queue-3", "old", true)); + assert!(!other_brokers_ready(&with_stale, "queue-2", 3)); + } + + #[test] + fn rollout_waits_for_stale_or_malformed_broker_pods_to_disappear() { + let mut extra = vec![ + pod("queue-0", "old", true), + pod("queue-1", "old", true), + pod("queue-2", "old", true), + ]; + extra.push(pod("queue-3", "old", true)); + assert!(matches!( + rollout_decision(&extra, &options("new"), &Default::default()), + Decision::Wait(message) if message.contains("exact") + )); + + let duplicate = vec![ + pod("queue-0", "old", true), + pod("another-0", "old", true), + pod("queue-2", "old", true), + ]; + assert!(matches!( + rollout_decision(&duplicate, &options("new"), &Default::default()), + Decision::Wait(message) if message.contains("exact") + )); + } + #[test] fn canary_requires_explicit_revision_approval() { let pods = vec![ @@ -474,12 +824,12 @@ mod tests { let mut options = options("new"); options.require_canary_approval = true; assert!(matches!( - rollout_decision(&pods, &options), + rollout_decision(&pods, &options, &Default::default()), Decision::AwaitingApproval )); options.approved_revision = Some("new"); assert!(matches!( - rollout_decision(&pods, &options), + rollout_decision(&pods, &options, &Default::default()), Decision::Replace(_) )); } @@ -490,11 +840,24 @@ mod tests { let mut options = options("current"); options.replicas = 1; assert!(matches!( - rollout_decision(&pods, &options), + rollout_decision(&pods, &options, &Default::default()), Decision::Completed )); } + #[test] + fn current_revision_is_not_complete_until_every_broker_is_ready() { + let pods = vec![ + pod("queue-0", "current", true), + pod("queue-1", "current", false), + pod("queue-2", "current", true), + ]; + assert!(matches!( + rollout_decision(&pods, &options("current"), &Default::default()), + Decision::Wait(message) if message.contains("queue-1") + )); + } + #[test] fn numeric_ordinal_sorting_selects_the_highest_broker() { let mut names = vec!["queue-9", "queue-10", "queue-499", "queue-100"]; @@ -502,6 +865,25 @@ mod tests { assert_eq!(names.pop(), Some("queue-499")); } + #[test] + fn target_broker_readiness_requires_the_complete_current_revision() { + let ready = vec![ + pod("queue-0", "current", true), + pod("queue-1", "current", true), + pod("queue-2", "current", true), + ]; + assert!(broker_set_current_and_ready(&ready, "current", 3)); + + let mut outdated = ready.clone(); + outdated[1] = pod("queue-1", "old", true); + assert!(!broker_set_current_and_ready(&outdated, "current", 3)); + + let mut unready = ready.clone(); + unready[2] = pod("queue-2", "current", false); + assert!(!broker_set_current_and_ready(&unready, "current", 3)); + assert!(!broker_set_current_and_ready(&ready[..2], "current", 3)); + } + #[test] fn formats_ipv4_and_ipv6_origins() { assert_eq!(origin("10.0.0.1"), "http://10.0.0.1:4151"); diff --git a/crates/operator/src/controller/kodo_cutover.rs b/crates/operator/src/controller/kodo_cutover.rs new file mode 100644 index 0000000..691dcc5 --- /dev/null +++ b/crates/operator/src/controller/kodo_cutover.rs @@ -0,0 +1,66 @@ +#[derive(Debug, Eq, PartialEq)] +pub(super) struct ProducerRestartFence { + pub baseline_nonce: Option, + pub confirmed: bool, +} + +pub(super) fn producer_restart_fence( + enabled: bool, + gateways_advertised: bool, + current_nonce: &str, + previous_baseline: Option<&str>, +) -> ProducerRestartFence { + if !enabled || !gateways_advertised { + return ProducerRestartFence { + baseline_nonce: None, + confirmed: false, + }; + } + let Some(baseline) = previous_baseline else { + return ProducerRestartFence { + baseline_nonce: Some(current_nonce.into()), + confirmed: false, + }; + }; + ProducerRestartFence { + baseline_nonce: Some(baseline.into()), + confirmed: current_nonce != baseline, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn confirmation_requires_a_nonce_change_after_advertisement() { + let first = producer_restart_fence(true, true, "before", None); + assert_eq!( + first, + ProducerRestartFence { + baseline_nonce: Some("before".into()), + confirmed: false, + } + ); + assert!(!producer_restart_fence(true, true, "before", Some("before")).confirmed); + assert!(producer_restart_fence(true, true, "after", Some("before")).confirmed); + } + + #[test] + fn losing_advertisement_resets_the_confirmation_fence() { + assert_eq!( + producer_restart_fence(true, false, "after", Some("before")), + ProducerRestartFence { + baseline_nonce: None, + confirmed: false, + } + ); + assert_eq!( + producer_restart_fence(false, true, "after", Some("before")), + ProducerRestartFence { + baseline_nonce: None, + confirmed: false, + } + ); + } +} diff --git a/crates/operator/src/controller/leadership.rs b/crates/operator/src/controller/leadership.rs index f97169b..d1fb028 100644 --- a/crates/operator/src/controller/leadership.rs +++ b/crates/operator/src/controller/leadership.rs @@ -1,53 +1,116 @@ use k8s_openapi::api::coordination::v1::{Lease, LeaseSpec}; +use k8s_openapi::apimachinery::pkg::apis::meta::v1::MicroTime; use kube::api::{Api, PostParams}; -use std::collections::BTreeMap; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::time::{Duration, Instant}; +use tokio::sync::watch; const LEASE_NAME: &str = "rustqueue-operator-leader"; const LEASE_SECONDS: u64 = 20; const RENEW_INTERVAL: Duration = Duration::from_secs(5); const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); -const RENEWED_AT: &str = "rustqueue.io/renewed-at-unix"; pub(super) fn start( client: kube::Client, namespace: String, leader: Arc, -) -> tokio::task::JoinHandle<()> { +) -> (tokio::task::JoinHandle<()>, watch::Receiver) { let identity = std::env::var("POD_NAME") .or_else(|_| std::env::var("HOSTNAME")) .unwrap_or_else(|_| format!("operator-{}", std::process::id())); - tokio::spawn(async move { - let _reset = LeaderReset(Arc::clone(&leader)); + let (updates, receiver) = watch::channel(false); + let task = tokio::spawn(async move { + let _reset = LeaderReset { + leader: Arc::clone(&leader), + updates: updates.clone(), + }; + let mut observer = LeaseObserver::default(); loop { - match tokio::time::timeout( + let acquired = match tokio::time::timeout( REQUEST_TIMEOUT, - acquire_or_renew(&client, &namespace, &identity), + acquire_or_renew(&client, &namespace, &identity, &mut observer), ) .await { - Ok(Ok(acquired)) => leader.store(acquired, Ordering::Release), + Ok(Ok(acquired)) => acquired, Ok(Err(error)) => { - leader.store(false, Ordering::Release); tracing::warn!(%error, "operator leader lease renewal failed"); + false } Err(_) => { - leader.store(false, Ordering::Release); tracing::warn!("operator leader lease renewal timed out"); + false } - } + }; + set_leader(&leader, &updates, acquired); tokio::time::sleep(RENEW_INTERVAL).await; } - }) + }); + (task, receiver) } -struct LeaderReset(Arc); +fn set_leader(leader: &AtomicBool, updates: &watch::Sender, value: bool) { + leader.store(value, Ordering::Release); + updates.send_replace(value); +} + +struct LeaderReset { + leader: Arc, + updates: watch::Sender, +} impl Drop for LeaderReset { fn drop(&mut self) { - self.0.store(false, Ordering::Release); + set_leader(&self.leader, &self.updates, false); + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct LeaseFingerprint { + holder: Option, + resource_version: Option, + renew_time: Option, +} + +#[derive(Default)] +struct LeaseObserver { + observed: Option<(LeaseFingerprint, Instant)>, +} + +impl LeaseObserver { + fn takeover_ready(&mut self, lease: &Lease, now: Instant) -> bool { + let holder = lease + .spec + .as_ref() + .and_then(|spec| spec.holder_identity.as_deref()); + if holder.is_none_or(str::is_empty) { + return true; + } + let fingerprint = LeaseFingerprint { + holder: holder.map(str::to_owned), + resource_version: lease.metadata.resource_version.clone(), + renew_time: lease + .spec + .as_ref() + .and_then(|spec| spec.renew_time.as_ref()) + .and_then(|time| serde_json::to_string(time).ok()), + }; + let duration = lease + .spec + .as_ref() + .and_then(|spec| spec.lease_duration_seconds) + .unwrap_or(LEASE_SECONDS as i32) + .max(1) as u64; + match self.observed.as_mut() { + Some((observed, since)) if observed == &fingerprint => { + now.saturating_duration_since(*since) >= Duration::from_secs(duration) + } + _ => { + self.observed = Some((fingerprint, now)); + false + } + } } } @@ -55,13 +118,13 @@ async fn acquire_or_renew( client: &kube::Client, namespace: &str, identity: &str, + observer: &mut LeaseObserver, ) -> anyhow::Result { let api = Api::::namespaced(client.clone(), namespace); - let now = unix_seconds(); let mut lease = match api.get_opt(LEASE_NAME).await? { Some(lease) => lease, None => { - let lease = desired_lease(namespace, identity, now); + let lease = desired_lease(namespace, identity); return match api.create(&PostParams::default(), &lease).await { Ok(_) => Ok(true), Err(kube::Error::Api(error)) if error.code == 409 => Ok(false), @@ -73,30 +136,20 @@ async fn acquire_or_renew( .spec .as_ref() .and_then(|spec| spec.holder_identity.as_deref()); - let renewed = lease - .metadata - .annotations - .as_ref() - .and_then(|annotations| annotations.get(RENEWED_AT)) - .and_then(|value| value.parse::().ok()) - .unwrap_or(0); - if holder != Some(identity) && !expired(renewed, now) { + let renewing = holder == Some(identity); + if !renewing && !observer.takeover_ready(&lease, Instant::now()) { return Ok(false); } - lease - .spec - .get_or_insert_with(LeaseSpec::default) - .holder_identity = Some(identity.into()); - lease - .spec - .as_mut() - .expect("inserted lease spec") - .lease_duration_seconds = Some(LEASE_SECONDS as i32); - lease - .metadata - .annotations - .get_or_insert_with(BTreeMap::new) - .insert(RENEWED_AT.into(), now.to_string()); + + let now = current_micro_time(); + let spec = lease.spec.get_or_insert_with(LeaseSpec::default); + if !renewing { + spec.acquire_time = Some(now.clone()); + spec.lease_transitions = Some(spec.lease_transitions.unwrap_or_default().saturating_add(1)); + } + spec.holder_identity = Some(identity.into()); + spec.lease_duration_seconds = Some(LEASE_SECONDS as i32); + spec.renew_time = Some(now); match api .replace(LEASE_NAME, &PostParams::default(), &lease) .await @@ -107,31 +160,27 @@ async fn acquire_or_renew( } } -fn desired_lease(namespace: &str, identity: &str, now: u64) -> Lease { +fn desired_lease(namespace: &str, identity: &str) -> Lease { + let now = current_micro_time(); Lease { metadata: kube::api::ObjectMeta { name: Some(LEASE_NAME.into()), namespace: Some(namespace.into()), - annotations: Some(BTreeMap::from([(RENEWED_AT.into(), now.to_string())])), ..Default::default() }, spec: Some(LeaseSpec { + acquire_time: Some(now.clone()), holder_identity: Some(identity.into()), lease_duration_seconds: Some(LEASE_SECONDS as i32), + lease_transitions: Some(0), + renew_time: Some(now), ..Default::default() }), } } -fn expired(renewed: u64, now: u64) -> bool { - now.saturating_sub(renewed) >= LEASE_SECONDS -} - -fn unix_seconds() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() +fn current_micro_time() -> MicroTime { + MicroTime(k8s_openapi::jiff::Timestamp::now()) } #[cfg(test)] @@ -139,17 +188,28 @@ mod tests { use super::*; #[test] - fn lease_expiry_is_monotonic_under_clock_underflow() { - assert!(!expired(100, 90)); - assert!(!expired(100, 119)); - assert!(expired(100, 120)); + fn takeover_uses_local_monotonic_observation_time() { + let mut lease = desired_lease("queue", "operator-1"); + lease.metadata.resource_version = Some("7".into()); + let start = Instant::now(); + let mut observer = LeaseObserver::default(); + assert!(!observer.takeover_ready(&lease, start)); + assert!(!observer.takeover_ready(&lease, start + Duration::from_secs(19))); + assert!(observer.takeover_ready(&lease, start + Duration::from_secs(20))); + + lease.metadata.resource_version = Some("8".into()); + assert!(!observer.takeover_ready(&lease, start + Duration::from_secs(21))); } #[test] - fn desired_lease_has_a_holder_and_duration() { - let lease = desired_lease("queue", "operator-1", 42); + fn desired_lease_uses_standard_holder_duration_and_timestamps() { + let lease = desired_lease("queue", "operator-1"); + assert!(lease.metadata.annotations.is_none()); let spec = lease.spec.unwrap(); assert_eq!(spec.holder_identity.as_deref(), Some("operator-1")); assert_eq!(spec.lease_duration_seconds, Some(20)); + assert_eq!(spec.lease_transitions, Some(0)); + assert!(spec.acquire_time.is_some()); + assert!(spec.renew_time.is_some()); } } diff --git a/crates/operator/src/controller/mod.rs b/crates/operator/src/controller/mod.rs index 2187076..213b0db 100644 --- a/crates/operator/src/controller/mod.rs +++ b/crates/operator/src/controller/mod.rs @@ -1,6 +1,7 @@ mod apply; mod auth; mod drain; +mod kodo_cutover; mod leadership; mod nodes; mod operations; @@ -12,7 +13,8 @@ use crate::resources::{self, BuildInput}; use crate::RustQueue; use anyhow::{bail, Context as _}; use futures::StreamExt; -use k8s_openapi::api::apps::v1::StatefulSet; +use k8s_openapi::api::apps::v1::{Deployment, StatefulSet}; +use k8s_openapi::api::core::v1::Service; use kube::api::Api; use kube::runtime::controller::{Action, Controller}; use kube::runtime::watcher; @@ -21,11 +23,15 @@ use status::{OperationUpdate, StatusBuilder}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Duration; +use tokio::sync::watch; + +const KODO_BOOTSTRAP_RETENTION_SECONDS: u64 = 180; pub(super) struct ContextData { pub client: Client, pub http: reqwest::Client, pub leader: Arc, + pub leadership: watch::Receiver, } #[derive(Debug, thiserror::Error)] @@ -35,7 +41,8 @@ pub struct ReconcileError(#[from] anyhow::Error); pub async fn run(leader: Arc) -> anyhow::Result<()> { let client = Client::try_default().await?; let namespace = watch_namespace(); - let mut leadership = leadership::start(client.clone(), namespace.clone(), Arc::clone(&leader)); + let (mut leadership, leadership_updates) = + leadership::start(client.clone(), namespace.clone(), Arc::clone(&leader)); let context = Arc::new(ContextData { client: client.clone(), http: reqwest::Client::builder() @@ -44,12 +51,17 @@ pub async fn run(leader: Arc) -> anyhow::Result<()> { .redirect(reqwest::redirect::Policy::none()) .build()?, leader, + leadership: leadership_updates, }); let clusters = Api::::namespaced(client.clone(), &namespace); - let stateful_sets = Api::::namespaced(client, &namespace); + let stateful_sets = Api::::namespaced(client.clone(), &namespace); + let deployments = Api::::namespaced(client.clone(), &namespace); + let services = Api::::namespaced(client, &namespace); tracing::info!(%namespace, "share-nothing RustQueue Operator started"); let controller = Controller::new(clusters, watcher::Config::default()) .owns(stateful_sets, watcher::Config::default()) + .owns(deployments, watcher::Config::default()) + .owns(services, watcher::Config::default()) .run(reconcile, error_policy, context) .for_each(|result| async move { match result { @@ -72,7 +84,19 @@ async fn reconcile( cluster: Arc, context: Arc, ) -> Result { - match reconcile_inner(Arc::clone(&cluster), Arc::clone(&context)).await { + if !context.leader.load(Ordering::Acquire) { + return Ok(Action::requeue(Duration::from_secs(2))); + } + let result = tokio::select! { + result = reconcile_inner(Arc::clone(&cluster), Arc::clone(&context)) => result, + _ = wait_for_leadership_loss(context.leadership.clone()) => { + return Ok(Action::requeue(Duration::from_secs(2))); + } + }; + if !context.leader.load(Ordering::Acquire) { + return Ok(Action::requeue(Duration::from_secs(2))); + } + match result { Ok(action) => Ok(action), Err(error) => { if context.leader.load(Ordering::Acquire) { @@ -83,14 +107,24 @@ async fn reconcile( } } +async fn wait_for_leadership_loss(mut leadership: watch::Receiver) { + if !*leadership.borrow() { + return; + } + while leadership.changed().await.is_ok() { + if !*leadership.borrow() { + return; + } + } +} + async fn reconcile_inner( cluster: Arc, context: Arc, ) -> anyhow::Result { - if !context.leader.load(Ordering::Acquire) { - return Ok(Action::requeue(Duration::from_secs(2))); - } - validate(&cluster)?; + let active_feature_floor = previous_feature_level(&cluster); + validate(&cluster, active_feature_floor)?; + let mut cluster = with_storage_feature_floor(cluster, active_feature_floor); let namespace = cluster .namespace() .context("RustQueue must be namespaced")?; @@ -100,9 +134,7 @@ async fn reconcile_inner( .rollback_to_image .as_deref() .unwrap_or(&cluster.spec.image); - let auth = auth::ensure(&context.client, &cluster, &namespace).await?; - let eligible = nodes::eligible(&context.client, &cluster.spec.eligible_node_selector).await?; - let desired = (eligible as i32).min(cluster.spec.max_brokers); + let effective_image = effective_image.to_owned(); let statefulsets: Api = Api::namespaced(context.client.clone(), &namespace); let current_set = statefulsets.get_opt(&cluster.name_any()).await?; let current = current_set @@ -110,21 +142,109 @@ async fn reconcile_inner( .and_then(|set| set.spec.as_ref()) .and_then(|spec| spec.replicas) .unwrap_or(0); + let current_kodo_gateway = statefulsets + .get_opt(&format!("{}-kodo-gateway", cluster.name_any())) + .await?; + let current_gateways_active = statefulset_has_replicas(current_kodo_gateway.as_ref()); + let decommission_blocked = kodo_decommission_blocked( + cluster.spec.kodo_compatibility.enabled, + cluster.spec.kodo_compatibility.decommission_confirmed, + current_gateways_active, + ); + let auth = auth::ensure(&context.client, &cluster, &namespace).await?; + let eligible = nodes::eligible(&context.client, &cluster.spec.eligible_node_selector).await?; + let ordinary_desired = + ordinary_desired_brokers(eligible, cluster.spec.min_brokers, cluster.spec.max_brokers); + let desired = if cluster.spec.kodo_compatibility.enabled { + 3 + } else if current_gateways_active && !cluster.spec.kodo_compatibility.decommission_confirmed { + ordinary_desired.max(3) + } else { + ordinary_desired + }; + let current_discovery = Api::::namespaced(context.client.clone(), &namespace) + .get_opt(&format!("{}-discovery", cluster.name_any())) + .await?; + let current_discovery_service = Api::::namespaced(context.client.clone(), &namespace) + .get_opt(&format!("{}-discovery", cluster.name_any())) + .await?; + if decommission_blocked { + let active_image = statefulset_broker_image(current_set.as_ref()) + .context("current Broker StatefulSet has no broker image")?; + recover_runtime_resources( + &context, + &cluster, + &namespace, + &auth, + active_image, + current_set.as_ref(), + current_discovery.as_ref(), + current_discovery_service.as_ref(), + current_gateways_active, + ) + .await?; + let message = "Kodo Gateway decommission is blocked: stop every Kodo workload using this \ + Discovery service, then set \ + spec.kodoCompatibility.decommissionConfirmed=true"; + let ready = nodes::ready_brokers(&context.client, &namespace, &cluster.name_any()).await?; + let status = StatusBuilder::new( + &cluster, + desired, + ready.min(current), + previous_feature_level(&cluster), + ) + .summary("KodoDecommissionBlocked", message) + .condition("Ready", false, "KodoDecommissionBlocked", message) + .condition("Progressing", false, "KodoDecommissionBlocked", message) + .condition("Degraded", true, "KodoDecommissionBlocked", message) + .condition( + "KodoGatewaysActive", + true, + "RetainedForSafety", + "Kodo Gateways and their existing Discovery route remain active", + ) + .condition( + "KodoDecommissionConfirmed", + false, + "ConfirmationRequired", + message, + ) + .build(); + apply::status(&context.client, &cluster, status).await?; + return Ok(Action::requeue(Duration::from_secs(5))); + } let target_preflight = - preflight::target_image(&context, &cluster, &namespace, effective_image).await?; + preflight::target_image(&context, &cluster, &namespace, &effective_image).await?; + if !matches!(target_preflight, preflight::Outcome::Ready { .. }) && current > 0 { + let active_image = statefulset_broker_image(current_set.as_ref()) + .context("current Broker StatefulSet has no broker image")?; + recover_runtime_resources( + &context, + &cluster, + &namespace, + &auth, + active_image, + current_set.as_ref(), + current_discovery.as_ref(), + current_discovery_service.as_ref(), + current_gateways_active, + ) + .await?; + } if preflight_status( &context, &cluster, desired, current, - effective_image, + &effective_image, &target_preflight, ) .await? { return Ok(Action::requeue(Duration::from_secs(5))); } + let mut active_feature_level = previous_feature_level(&cluster); if current > 0 { let broker_preflight = @@ -134,23 +254,57 @@ async fn reconcile_inner( } = &broker_preflight { active_feature_level = *active; - } else if preflight_status( - &context, - &cluster, - desired, - current, - effective_image, - &broker_preflight, - ) - .await? - { - return Ok(Action::requeue(Duration::from_secs(5))); + cluster = with_storage_feature_floor(cluster, active_feature_level); + validate(&cluster, active_feature_level)?; + } else { + let active_image = statefulset_broker_image(current_set.as_ref()) + .context("current Broker StatefulSet has no broker image")?; + recover_runtime_resources( + &context, + &cluster, + &namespace, + &auth, + active_image, + current_set.as_ref(), + current_discovery.as_ref(), + current_discovery_service.as_ref(), + current_gateways_active, + ) + .await?; + if preflight_status( + &context, + &cluster, + desired, + current, + &effective_image, + &broker_preflight, + ) + .await? + { + return Ok(Action::requeue(Duration::from_secs(5))); + } } } - preflight::cleanup_old_probes(&context, &cluster, &namespace, effective_image).await?; + preflight::cleanup_old_probes(&context, &cluster, &namespace, &effective_image).await?; let storage = storage::reconcile(&context.client, &cluster, &namespace, desired).await?; if storage.state != storage::StorageState::Ready { + if current > 0 { + let active_image = statefulset_broker_image(current_set.as_ref()) + .context("current Broker StatefulSet has no broker image")?; + recover_runtime_resources( + &context, + &cluster, + &namespace, + &auth, + active_image, + current_set.as_ref(), + current_discovery.as_ref(), + current_discovery_service.as_ref(), + current_gateways_active, + ) + .await?; + } let ready = nodes::ready_brokers(&context.client, &namespace, &cluster.name_any()).await?; let blocked = storage.state == storage::StorageState::Blocked; let phase = if blocked { @@ -186,10 +340,146 @@ async fn reconcile_inner( auth::mounted_secret_revision(&context.client, &cluster, &namespace, &auth).await?; let claim_template_size = storage::claim_template_size(current_set.as_ref(), &cluster.spec.storage_size); + let staged = resources::build(BuildInput { + cluster: &cluster, + replicas: applied_replicas, + kodo_gateway_replicas: 0, + advertise_kodo_gateways: false, + discovery_service_kodo: Some(false), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: !cluster.spec.kodo_compatibility.enabled + && current_gateways_active, + image: &effective_image, + claim_template_size: &claim_template_size, + secret_name: &auth.name, + mounted_secret_revision: &mounted_secret_revision, + })?; + let target_brokers_ready = + drain::brokers_current_and_ready(&context, &cluster, &namespace, &staged.revision, desired) + .await?; + let ready_brokers = + nodes::ready_brokers(&context.client, &namespace, &cluster.name_any()).await?; + let brokers_available_for_kodo = cluster.spec.kodo_compatibility.enabled + && kodo_gateway_activation_ready(current, ready_brokers); + let kodo_gateways_active = cluster.spec.kodo_compatibility.enabled + && (current_gateways_active || brokers_available_for_kodo); + let current_gateways_ready = if current_gateways_active { + let updated = nodes::ready_component( + &context.client, + &namespace, + &cluster.name_any(), + "kodo-gateway", + ) + .await?; + let configured = nodes::ready_component_revision( + &context.client, + &namespace, + &cluster.name_any(), + "kodo-gateway", + &staged.revision, + ) + .await?; + complete_kodo_gateway_set(updated, configured) + } else { + false + }; + let current_discovery_kodo = deployment_requests_kodo_mode(current_discovery.as_ref()); + let current_discovery_route = discovery_service_kodo_route(current_discovery_service.as_ref()); + let direct_inventory_safe = + cluster.spec.kodo_compatibility.decommission_confirmed || ordinary_desired >= 3; + let hold_kodo_for_disable = should_hold_kodo_for_disable( + cluster.spec.kodo_compatibility.enabled, + current_gateways_active, + target_brokers_ready, + direct_inventory_safe, + ); + let desired_discovery_kodo = cluster.spec.kodo_compatibility.enabled || hold_kodo_for_disable; + let (advertise_kodo_gateways, discovery_label_adoption) = discovery_target_mode( + desired_discovery_kodo, + current_gateways_ready, + current_discovery.is_some(), + current_discovery_kodo, + current_discovery_route, + deployment_template_has_discovery_mode(current_discovery.as_ref(), current_discovery_kodo), + deployment_is_fully_ready(current_discovery.as_ref()), + ); + let ready_target_discovery = nodes::ready_discovery_mode( + &context.client, + &namespace, + &cluster.name_any(), + advertise_kodo_gateways, + ) + .await?; + let ready_fallback_discovery = nodes::ready_discovery_mode( + &context.client, + &namespace, + &cluster.name_any(), + !advertise_kodo_gateways, + ) + .await?; + let discovery_replicas = cluster.spec.discovery_replicas.max(2); + let discovery_service_kodo = discovery_target_route( + current_discovery.is_some(), + discovery_label_adoption, + current_discovery_route, + current_discovery_kodo, + advertise_kodo_gateways, + (ready_target_discovery, ready_fallback_discovery), + discovery_replicas, + ); + let kodo_gateways_advertised = cluster.spec.kodo_compatibility.enabled + && advertise_kodo_gateways + && discovery_service_kodo == Some(true) + && ready_target_discovery >= discovery_replicas; + let direct_brokers_advertised = !cluster.spec.kodo_compatibility.enabled + && !advertise_kodo_gateways + && current_discovery_route == Some(false) + && deployment_runs_without_kodo(current_discovery.as_ref()) + && ready_target_discovery >= discovery_replicas; + let kodo_cutover_elapsed = condition_true_for( + &cluster, + "KodoGatewaysAdvertised", + cluster.spec.kodo_compatibility.cutover_grace_seconds, + ); + let producer_restart = kodo_cutover::producer_restart_fence( + cluster.spec.kodo_compatibility.enabled, + kodo_gateways_advertised, + &cluster.spec.kodo_compatibility.producer_restart_nonce, + cluster + .status + .as_ref() + .and_then(|status| status.kodo_producer_restart_baseline_nonce.as_deref()), + ); + let direct_cutover_elapsed = condition_true_for( + &cluster, + "KodoDirectBrokersAdvertised", + cluster.spec.kodo_compatibility.cutover_grace_seconds, + ); + let retain_existing_kodo_resources = !cluster.spec.kodo_compatibility.enabled + && current_gateways_active + && (!direct_inventory_safe + || !direct_brokers_advertised + || (!cluster.spec.kodo_compatibility.decommission_confirmed + && !direct_cutover_elapsed)); + let activate_kodo_cleanup = cluster.spec.kodo_compatibility.enabled + && cluster.spec.kodo_compatibility.effective_cleanup_enabled() + && target_brokers_ready; + let kodo_cleanup_advertised = cluster.spec.kodo_compatibility.enabled + && cluster.spec.kodo_compatibility.effective_cleanup_enabled() + && deployment_advertises_kodo_cleanup(current_discovery.as_ref(), &staged.revision); let set = resources::build(BuildInput { cluster: &cluster, replicas: applied_replicas, - image: effective_image, + kodo_gateway_replicas: if kodo_gateways_active || retain_existing_kodo_resources { + 3 + } else { + 0 + }, + advertise_kodo_gateways, + discovery_service_kodo, + activate_kodo_cleanup, + retain_existing_kodo_resources, + image: &effective_image, claim_template_size: &claim_template_size, secret_name: &auth.name, mounted_secret_revision: &mounted_secret_revision, @@ -197,7 +487,17 @@ async fn reconcile_inner( let revision = set.revision.clone(); apply::resources(&context.client, &namespace, set).await?; - let (phase, message, operation) = operations::reconcile( + let disruptions_allowed = if cluster.spec.kodo_compatibility.enabled { + current_gateways_ready + && kodo_gateways_advertised + && kodo_cutover_elapsed + && producer_restart.confirmed + } else if current_gateways_active && direct_brokers_advertised { + cluster.spec.kodo_compatibility.decommission_confirmed || direct_cutover_elapsed + } else { + true + }; + let (mut phase, mut message, operation) = operations::reconcile( &context, &cluster, &namespace, @@ -206,16 +506,92 @@ async fn reconcile_inner( desired, current, &revision, - effective_image, + &effective_image, + disruptions_allowed, ) .await?; let broker_health = nodes::broker_health(&context.client, &namespace, &cluster.name_any()).await?; let ready = broker_health.ready; - let ready_condition = phase == "Ready" && ready == desired; + let ready_kodo_gateways = if kodo_gateways_active { + nodes::ready_component( + &context.client, + &namespace, + &cluster.name_any(), + "kodo-gateway", + ) + .await? + } else { + 0 + }; + let kodo_gateways_ready = !cluster.spec.kodo_compatibility.enabled + || (kodo_gateways_active && ready_kodo_gateways == 3); + if phase == "Ready" && !kodo_gateways_ready { + phase = "WaitingForKodoGateways".into(); + message = format!("{ready_kodo_gateways} of 3 Kodo publish gateways are Ready"); + } + let kodo_advertisement_ready = + !cluster.spec.kodo_compatibility.enabled || kodo_gateways_advertised; + if phase == "Ready" && !kodo_advertisement_ready { + phase = "WaitingForKodoAdvertisement".into(); + message = "waiting for the replacement Discovery set to become Ready".into(); + } + let kodo_cutover_ready = !cluster.spec.kodo_compatibility.enabled + || (current_gateways_ready && kodo_gateways_advertised && kodo_cutover_elapsed); + if phase == "Ready" && !kodo_cutover_ready { + phase = "WaitingForKodoCutover".into(); + message = format!( + "waiting {} seconds for the Gateway Discovery cutover to remain stable", + cluster.spec.kodo_compatibility.cutover_grace_seconds + ); + } + let kodo_producer_restart_ready = + !cluster.spec.kodo_compatibility.enabled || producer_restart.confirmed; + let producer_restart_is_only_blocker = cluster.spec.kodo_compatibility.enabled + && current_gateways_ready + && kodo_gateways_advertised + && kodo_cutover_elapsed + && !producer_restart.confirmed; + if !kodo_producer_restart_ready + && (phase == "Ready" + || (phase == "WaitingForKodoCutover" && producer_restart_is_only_blocker)) + { + phase = "WaitingForKodoProducerRestart".into(); + message = "restart every Kodo publisher after Gateway advertisement, then change \ + spec.kodoCompatibility.producerRestartNonce" + .into(); + } + let kodo_cleanup_ready = !cluster.spec.kodo_compatibility.enabled + || !cluster.spec.kodo_compatibility.effective_cleanup_enabled() + || kodo_cleanup_advertised; + if phase == "Ready" && !kodo_cleanup_ready { + phase = "WaitingForKodoCleanup".into(); + message = "waiting for Discovery to advertise cleanup-capable Broker endpoints".into(); + } + if phase == "Ready" && retain_existing_kodo_resources { + phase = "DisablingKodo".into(); + message = if !direct_inventory_safe { + "at least three direct Brokers are required before Kodo Gateways can be removed".into() + } else if !direct_brokers_advertised { + "waiting for Discovery to advertise direct Broker addresses".into() + } else { + format!( + "waiting {} seconds before removing the previous Kodo Gateways", + cluster.spec.kodo_compatibility.cutover_grace_seconds + ) + }; + } + let ready_condition = phase == "Ready" + && ready == desired + && kodo_gateways_ready + && kodo_advertisement_ready + && kodo_cutover_ready + && kodo_producer_restart_ready + && kodo_cleanup_ready + && !retain_existing_kodo_resources; let degraded = matches!( phase.as_str(), - "InsufficientNodes" | "RolloutBlocked" | "RolloutFailed" + "InsufficientNodes" | "MaintenanceBlocked" | "RolloutBlocked" | "RolloutFailed" ); let progressing = !ready_condition && !degraded && phase != "Maintenance"; let maintenance_enabled = cluster @@ -277,7 +653,166 @@ async fn reconcile_inner( broker_health.unavailable.join("; ") }, ) - .orphaned_pvcs(storage.orphaned_pvcs); + .condition( + "KodoGatewaysActive", + kodo_gateways_active, + if !cluster.spec.kodo_compatibility.enabled { + "Disabled" + } else if kodo_gateways_active { + "Active" + } else { + "WaitingForBrokers" + }, + if !cluster.spec.kodo_compatibility.enabled { + "Kodo compatibility is disabled" + } else if kodo_gateways_active { + "Kodo publish gateway replicas are activated" + } else { + "waiting for two Ready Brokers before starting Kodo Gateways" + }, + ) + .condition( + "KodoGatewaysReady", + cluster.spec.kodo_compatibility.enabled && kodo_gateways_ready, + if !cluster.spec.kodo_compatibility.enabled { + "Disabled" + } else if kodo_gateways_ready { + "AllReady" + } else { + "PodsUnavailable" + }, + if !cluster.spec.kodo_compatibility.enabled { + "Kodo compatibility is disabled".into() + } else { + format!("{ready_kodo_gateways} of 3 Kodo publish gateways are Ready") + }, + ) + .condition( + "KodoGatewaysAdvertised", + kodo_gateways_advertised, + if !cluster.spec.kodo_compatibility.enabled { + "Disabled" + } else if kodo_gateways_advertised { + "Advertised" + } else { + "WaitingForGateways" + }, + if !cluster.spec.kodo_compatibility.enabled { + "Kodo compatibility is disabled" + } else if kodo_gateways_advertised { + "Discovery advertises the stable Kodo Gateway addresses" + } else { + "Discovery Service has not switched to the stable Kodo Gateway addresses" + }, + ) + .condition( + "KodoCutoverReady", + cluster.spec.kodo_compatibility.enabled && kodo_cutover_ready, + if !cluster.spec.kodo_compatibility.enabled { + "Disabled" + } else if kodo_cutover_ready { + "DiscoveryGraceElapsed" + } else { + "WaitingForDiscoveryGrace" + }, + if !cluster.spec.kodo_compatibility.enabled { + "Kodo compatibility is disabled".into() + } else if kodo_cutover_ready { + "Kodo Gateway Discovery stability grace has elapsed".into() + } else { + format!( + "waiting {} seconds after Gateway advertisement", + cluster.spec.kodo_compatibility.cutover_grace_seconds + ) + }, + ) + .condition( + "KodoProducerRestartConfirmed", + cluster.spec.kodo_compatibility.enabled && producer_restart.confirmed, + if !cluster.spec.kodo_compatibility.enabled { + "Disabled" + } else if producer_restart.confirmed { + "NonceChangedAfterAdvertisement" + } else if kodo_gateways_advertised { + "WaitingForRestart" + } else { + "WaitingForAdvertisement" + }, + if !cluster.spec.kodo_compatibility.enabled { + "Kodo compatibility is disabled" + } else if producer_restart.confirmed { + "Kodo producer restart was explicitly confirmed after Gateway advertisement" + } else if kodo_gateways_advertised { + "restart every Kodo publisher, wait for it to become Ready, then change producerRestartNonce" + } else { + "producer restart confirmation is accepted only after Gateway advertisement" + }, + ) + .condition( + "KodoDirectBrokersAdvertised", + direct_brokers_advertised, + if cluster.spec.kodo_compatibility.enabled { + "GatewayMode" + } else if direct_brokers_advertised { + "Advertised" + } else { + "WaitingForSafeInventory" + }, + if cluster.spec.kodo_compatibility.enabled { + "Discovery advertises Kodo Gateway addresses" + } else if direct_brokers_advertised { + "Discovery advertises direct Broker addresses" + } else if !direct_inventory_safe { + "fewer than three direct Brokers would be visible to Kodo" + } else { + "Discovery has not completed its direct Broker cutover" + }, + ) + .condition( + "KodoCleanupActive", + cluster.spec.kodo_compatibility.enabled + && cluster.spec.kodo_compatibility.effective_cleanup_enabled() + && kodo_cleanup_ready, + if !cluster.spec.kodo_compatibility.enabled { + "Disabled" + } else if !cluster.spec.kodo_compatibility.effective_cleanup_enabled() { + "NotRequested" + } else if kodo_cleanup_ready { + "Advertised" + } else { + "WaitingForBrokers" + }, + if !cluster.spec.kodo_compatibility.enabled { + "Kodo compatibility is disabled" + } else if !cluster.spec.kodo_compatibility.effective_cleanup_enabled() { + "Kodo automatic cleanup is disabled" + } else if kodo_cleanup_ready { + "Discovery advertises authenticated cleanup-capable Broker endpoints" + } else { + "cleanup remains inactive until all target Brokers and Discovery are Ready" + }, + ) + .condition( + "KodoDecommissionConfirmed", + !cluster.spec.kodo_compatibility.enabled + && cluster.spec.kodo_compatibility.decommission_confirmed, + if cluster.spec.kodo_compatibility.enabled { + "CompatibilityActive" + } else if cluster.spec.kodo_compatibility.decommission_confirmed { + "Confirmed" + } else { + "NotRequired" + }, + if cluster.spec.kodo_compatibility.enabled { + "Kodo compatibility is active" + } else if cluster.spec.kodo_compatibility.decommission_confirmed { + "all Kodo workloads were explicitly confirmed stopped" + } else { + "no Kodo Gateway decommission is in progress" + }, + ) + .orphaned_pvcs(storage.orphaned_pvcs) + .kodo_producer_restart_baseline_nonce(producer_restart.baseline_nonce); if let Some(operation) = &operation { operation.audit(&cluster); builder = operation.apply(builder); @@ -286,6 +821,114 @@ async fn reconcile_inner( Ok(Action::requeue(Duration::from_secs(5))) } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct KodoRuntimeRecovery { + gateway_replicas: i32, + advertise_gateways: bool, + discovery_service_kodo: bool, + retain_gateways: bool, +} + +#[derive(Clone, Copy, Debug, Default)] +struct KodoRuntimeSignals { + gateway_active: bool, + discovery_mode: Option, + discovery_route: Option, + previously_active: bool, + previously_advertised: bool, +} + +fn kodo_runtime_recovery( + enabled: bool, + current_brokers: i32, + ready_brokers: i32, + signals: KodoRuntimeSignals, +) -> KodoRuntimeRecovery { + let gateway_was_active = + signals.gateway_active || signals.previously_active || signals.previously_advertised; + let advertise_gateways = signals + .discovery_mode + .or(signals.discovery_route) + .unwrap_or(gateway_was_active); + let discovery_service_kodo = signals.discovery_route.unwrap_or(advertise_gateways); + let gateway_was_active = gateway_was_active || advertise_gateways || discovery_service_kodo; + let gateway_replicas = if gateway_was_active + || (enabled && kodo_gateway_activation_ready(current_brokers, ready_brokers)) + { + 3 + } else { + 0 + }; + KodoRuntimeRecovery { + gateway_replicas, + advertise_gateways, + discovery_service_kodo, + retain_gateways: !enabled && gateway_was_active, + } +} + +#[allow(clippy::too_many_arguments)] +async fn recover_runtime_resources( + context: &ContextData, + cluster: &RustQueue, + namespace: &str, + auth: &auth::AuthSecret, + image: &str, + current_set: Option<&StatefulSet>, + current_discovery: Option<&Deployment>, + current_discovery_service: Option<&Service>, + current_gateway_active: bool, +) -> anyhow::Result<()> { + let resumed = drain::resume_all(context, cluster, namespace, auth).await?; + if !resumed.is_empty() { + tracing::warn!( + brokers = %resumed.join(","), + "resumed Broker drains while reconciliation is blocked" + ); + } + let current_brokers = current_set + .and_then(|set| set.spec.as_ref()) + .and_then(|spec| spec.replicas) + .unwrap_or_default(); + let ready_brokers = + nodes::ready_brokers(&context.client, namespace, &cluster.name_any()).await?; + let recovery = kodo_runtime_recovery( + cluster.spec.kodo_compatibility.enabled, + current_brokers, + ready_brokers, + KodoRuntimeSignals { + gateway_active: current_gateway_active, + discovery_mode: current_discovery + .map(|deployment| deployment_requests_kodo_mode(Some(deployment))), + discovery_route: discovery_service_kodo_route(current_discovery_service), + previously_active: condition_true(cluster, "KodoGatewaysActive"), + previously_advertised: condition_true(cluster, "KodoGatewaysAdvertised"), + }, + ); + let mounted_secret_revision = + auth::mounted_secret_revision(&context.client, cluster, namespace, auth).await?; + let claim_template_size = storage::claim_template_size(current_set, &cluster.spec.storage_size); + let set = resources::build(BuildInput { + cluster, + replicas: current_brokers, + kodo_gateway_replicas: recovery.gateway_replicas, + advertise_kodo_gateways: recovery.advertise_gateways, + discovery_service_kodo: Some(recovery.discovery_service_kodo), + activate_kodo_cleanup: deployment_requests_kodo_cleanup(current_discovery), + retain_existing_kodo_resources: recovery.retain_gateways, + image, + claim_template_size: &claim_template_size, + secret_name: &auth.name, + mounted_secret_revision: &mounted_secret_revision, + })?; + apply::runtime_resources(&context.client, namespace, &set).await?; + tracing::warn!( + ready_brokers, + "reconciled runtime entrypoints without mutating Brokers" + ); + Ok(()) +} + async fn preflight_status( context: &ContextData, cluster: &RustQueue, @@ -369,7 +1012,291 @@ fn previous_feature_level(cluster: &RustQueue) -> u32 { .map_or(1, |status| status.active_storage_feature_level.max(1)) } +fn statefulset_has_replicas(statefulset: Option<&StatefulSet>) -> bool { + statefulset + .and_then(|set| set.spec.as_ref()) + .and_then(|spec| spec.replicas) + .is_some_and(|replicas| replicas > 0) +} + +fn statefulset_broker_image(statefulset: Option<&StatefulSet>) -> Option<&str> { + statefulset? + .spec + .as_ref()? + .template + .spec + .as_ref()? + .containers + .iter() + .find(|container| container.name == "broker")? + .image + .as_deref() +} + +fn kodo_decommission_blocked(enabled: bool, confirmed: bool, gateways_active: bool) -> bool { + gateways_active && !enabled && !confirmed +} + +fn kodo_gateway_activation_ready(current_brokers: i32, ready_brokers: i32) -> bool { + current_brokers >= 2 && ready_brokers >= 2 +} + +fn complete_kodo_gateway_set(updated: i32, configured: i32) -> bool { + updated == 3 && configured == 3 +} + +fn condition_true(cluster: &RustQueue, type_: &str) -> bool { + cluster.status.as_ref().is_some_and(|status| { + status + .conditions + .iter() + .any(|condition| condition.type_ == type_ && condition.status == "True") + }) +} + +fn condition_true_for(cluster: &RustQueue, type_: &str, seconds: u64) -> bool { + cluster + .status + .as_ref() + .and_then(|status| { + status + .conditions + .iter() + .find(|condition| condition.type_ == type_ && condition.status == "True") + }) + .and_then(|condition| status::elapsed_seconds(&condition.last_transition_time)) + .is_some_and(|elapsed| elapsed >= seconds) +} + +fn should_hold_kodo_for_disable( + enabled: bool, + gateways_active: bool, + target_brokers_ready: bool, + direct_inventory_safe: bool, +) -> bool { + !enabled && gateways_active && (!target_brokers_ready || !direct_inventory_safe) +} + +fn deployment_env_value<'a>(deployment: Option<&'a Deployment>, name: &str) -> Option<&'a str> { + deployment + .and_then(|deployment| deployment.spec.as_ref()) + .and_then(|spec| spec.template.spec.as_ref()) + .and_then(|spec| { + spec.containers.iter().find_map(|container| { + container.env.as_deref().and_then(|environment| { + environment + .iter() + .find(|variable| variable.name == name) + .and_then(|variable| variable.value.as_deref()) + }) + }) + }) +} + +fn deployment_requests_kodo_mode(deployment: Option<&Deployment>) -> bool { + deployment_env_value(deployment, "RUSTQUEUE_KODO_COMPATIBILITY_ENABLED") == Some("true") +} + +fn deployment_requests_kodo_cleanup(deployment: Option<&Deployment>) -> bool { + deployment_env_value(deployment, "RUSTQUEUE_KODO_CLEANUP_ENABLED") == Some("true") +} + +fn deployment_template_has_discovery_mode(deployment: Option<&Deployment>, kodo: bool) -> bool { + deployment + .and_then(|deployment| deployment.spec.as_ref()) + .and_then(|spec| spec.template.metadata.as_ref()) + .and_then(|metadata| metadata.labels.as_ref()) + .and_then(|labels| labels.get(resources::DISCOVERY_MODE_LABEL)) + .is_some_and(|mode| mode == if kodo { "kodo" } else { "direct" }) +} + +fn discovery_service_kodo_route(service: Option<&Service>) -> Option { + match service + .and_then(|service| service.spec.as_ref()) + .and_then(|spec| spec.selector.as_ref()) + .and_then(|selector| selector.get(resources::DISCOVERY_MODE_LABEL)) + .map(String::as_str) + { + Some("kodo") => Some(true), + Some("direct") => Some(false), + _ => None, + } +} + +fn discovery_target_mode( + desired_kodo: bool, + gateways_ready: bool, + has_deployment: bool, + current_kodo: bool, + current_route: Option, + template_mode_matches: bool, + deployment_fully_ready: bool, +) -> (bool, bool) { + let adopting_labels = has_deployment + && current_route.is_none() + && (!template_mode_matches || !deployment_fully_ready); + let target_kodo = if adopting_labels { + current_kodo + } else { + desired_kodo && (current_kodo || gateways_ready) + }; + (target_kodo, adopting_labels) +} + +fn discovery_target_route( + has_deployment: bool, + adopting_labels: bool, + current_route: Option, + current_kodo: bool, + target_kodo: bool, + ready: (i32, i32), + replicas: i32, +) -> Option { + let (target_ready, fallback_ready) = ready; + if !has_deployment { + return Some(target_kodo); + } + if adopting_labels { + return None; + } + if target_ready >= replicas { + return Some(target_kodo); + } + if current_route == Some(target_kodo) && fallback_ready >= replicas { + return Some(!target_kodo); + } + current_route.or(Some(current_kodo)) +} + +fn deployment_is_fully_ready(deployment: Option<&Deployment>) -> bool { + let Some(deployment) = deployment else { + return false; + }; + let generation = deployment.metadata.generation.unwrap_or_default(); + let desired = deployment + .spec + .as_ref() + .and_then(|spec| spec.replicas) + .unwrap_or_default(); + let ready = deployment + .status + .as_ref() + .and_then(|status| status.ready_replicas) + .unwrap_or_default(); + let updated = deployment + .status + .as_ref() + .and_then(|status| status.updated_replicas) + .unwrap_or_default(); + let available = deployment + .status + .as_ref() + .and_then(|status| status.available_replicas) + .unwrap_or_default(); + let observed = deployment + .status + .as_ref() + .and_then(|status| status.observed_generation) + .unwrap_or_default(); + desired > 0 + && ready >= desired + && updated >= desired + && available >= desired + && observed >= generation +} + +fn deployment_matches_revision(deployment: &Deployment, target_revision: &str) -> bool { + deployment + .spec + .as_ref() + .and_then(|spec| spec.template.metadata.as_ref()) + .and_then(|metadata| metadata.annotations.as_ref()) + .and_then(|annotations| annotations.get("rustqueue.io/revision")) + .is_some_and(|revision| revision == target_revision) +} + +fn deployment_advertises_kodo_cleanup( + deployment: Option<&Deployment>, + target_revision: &str, +) -> bool { + let Some(deployment) = deployment else { + return false; + }; + deployment_requests_kodo_cleanup(Some(deployment)) + && deployment_matches_revision(deployment, target_revision) + && deployment_is_fully_ready(Some(deployment)) +} + +fn deployment_runs_without_kodo(deployment: Option<&Deployment>) -> bool { + deployment.is_some() + && !deployment_requests_kodo_mode(deployment) + && deployment_is_fully_ready(deployment) +} + +fn with_storage_feature_floor( + cluster: Arc, + active_feature_floor: u32, +) -> Arc { + let effective = + effective_storage_feature_level(cluster.spec.storage_feature_level, active_feature_floor); + let (connection_delivery, node_delivery) = effective_delivery_limits( + effective, + cluster.spec.max_message_bytes, + cluster.spec.connection_delivery_inflight_bytes, + cluster.spec.node_delivery_inflight_bytes, + ); + if effective == cluster.spec.storage_feature_level + && connection_delivery == cluster.spec.connection_delivery_inflight_bytes + && node_delivery == cluster.spec.node_delivery_inflight_bytes + { + return cluster; + } + let mut adjusted = cluster.as_ref().clone(); + adjusted.spec.storage_feature_level = effective; + adjusted.spec.connection_delivery_inflight_bytes = connection_delivery; + adjusted.spec.node_delivery_inflight_bytes = node_delivery; + Arc::new(adjusted) +} + +fn effective_storage_feature_level(requested: u32, active: u32) -> u32 { + requested.max(active).max(1) +} + +fn effective_delivery_limits( + storage_feature_level: u32, + max_message_bytes: usize, + connection: usize, + node: usize, +) -> (usize, usize) { + let retained_message_bound = if storage_feature_level >= 2 { + 100 * 1024 * 1024 + } else { + max_message_bytes + }; + let connection = connection.max(retained_message_bound); + (connection, node.max(connection.saturating_mul(2))) +} + async fn record_reconcile_error(context: &ContextData, cluster: &RustQueue, error: &anyhow::Error) { + if let Some(namespace) = cluster.namespace() { + match auth::ensure(&context.client, cluster, &namespace).await { + Ok(auth) => match drain::resume_all(context, cluster, &namespace, &auth).await { + Ok(resumed) if !resumed.is_empty() => { + tracing::warn!( + brokers = %resumed.join(","), + "resumed Broker drains after reconciliation failure" + ); + } + Ok(_) => {} + Err(resume_error) => { + tracing::warn!(%resume_error, "failed to inspect Broker drains after reconciliation failure"); + } + }, + Err(auth_error) => { + tracing::warn!(%auth_error, "could not load Broker credentials for failure recovery"); + } + } + } let previous = cluster.status.as_ref(); let desired = previous.map_or(cluster.spec.min_brokers, |status| status.desired_brokers); let ready = previous.map_or(0, |status| status.ready_brokers); @@ -385,7 +1312,7 @@ async fn record_reconcile_error(context: &ContextData, cluster: &RustQueue, erro } } -fn validate(cluster: &RustQueue) -> anyhow::Result<()> { +fn validate(cluster: &RustQueue, active_feature_floor: u32) -> anyhow::Result<()> { if cluster.spec.image.trim().is_empty() { bail!("spec.image is required"); } @@ -407,7 +1334,7 @@ fn validate(cluster: &RustQueue) -> anyhow::Result<()> { } if cluster.spec.bootstrap_retention_seconds == 0 || cluster.spec.max_message_bytes == 0 - || cluster.spec.max_message_bytes > 32 * 1024 * 1024 + || cluster.spec.max_message_bytes > 100 * 1024 * 1024 || cluster.spec.message_index_cache_bytes == 0 || cluster.spec.connection_delivery_inflight_bytes < cluster.spec.max_message_bytes || cluster @@ -423,6 +1350,93 @@ fn validate(cluster: &RustQueue) -> anyhow::Result<()> { { bail!("queue limits are outside the stable v7 contract"); } + validate_message_storage_contract( + cluster.spec.max_message_bytes, + effective_storage_feature_level(cluster.spec.storage_feature_level, active_feature_floor), + )?; + if cluster.spec.kodo_compatibility.enabled + && (cluster.spec.min_brokers != 3 + || cluster.spec.max_brokers != 3 + || cluster.spec.storage_feature_level != 2 + || cluster.spec.bootstrap_retention_seconds < KODO_BOOTSTRAP_RETENTION_SECONDS + || cluster.spec.max_message_bytes != 100 * 1024 * 1024 + || cluster.spec.connection_delivery_inflight_bytes < 128 * 1024 * 1024 + || cluster.spec.node_delivery_inflight_bytes < 512 * 1024 * 1024) + { + bail!( + "Kodo compatibility requires exactly 3 brokers, storageFeatureLevel 2, \ + bootstrapRetentionSeconds >= 180, \ + maxMessageBytes 104857600, connectionDeliveryInflightBytes >= 134217728, \ + and nodeDeliveryInflightBytes >= 536870912" + ); + } + if !(630..=86_400).contains(&cluster.spec.kodo_compatibility.cutover_grace_seconds) { + bail!("Kodo compatibility cutoverGraceSeconds must be between 630 and 86400"); + } + if cluster.spec.kodo_compatibility.enabled { + if cluster.spec.kodo_compatibility.decommission_confirmed { + bail!("decommissionConfirmed must be false while Kodo compatibility is enabled"); + } + if cluster.spec.kodo_compatibility.cleanup_enabled { + bail!( + "Kodo automatic cleanup is disabled until cluster-wide atomic deletion is available" + ); + } + let target_image = cluster + .spec + .rollout + .rollback_to_image + .as_deref() + .unwrap_or(&cluster.spec.image); + if cluster.spec.image_pull_policy != "Never" && !has_sha256_digest(target_image) { + bail!( + "Kodo compatibility requires an immutable @sha256 image or imagePullPolicy Never" + ); + } + if cluster + .spec + .kodo_compatibility + .allowed_pod_selector + .is_empty() + || cluster + .spec + .kodo_compatibility + .allowed_pod_selector + .iter() + .any(|(key, value)| key.trim().is_empty() || value.trim().is_empty()) + || cluster + .spec + .kodo_compatibility + .allowed_namespace_selector + .iter() + .any(|(key, value)| key.trim().is_empty() || value.trim().is_empty()) + { + bail!("Kodo compatibility requires a non-empty allowedPodSelector and valid selector labels"); + } + let memory_request = storage::parse_quantity(&cluster.spec.broker_resources.memory_request) + .context("parse broker memory request")?; + if memory_request < 2_u128 << 30 { + bail!("Kodo compatibility requires brokerResources.memoryRequest >= 2Gi"); + } + let cpu_request = parse_cpu_millis(&cluster.spec.broker_resources.cpu_request) + .context("parse broker CPU request")?; + if cpu_request < 1_000.0 { + bail!("Kodo compatibility requires brokerResources.cpuRequest >= 1 CPU"); + } + if let Some(cpu_limit) = cluster.spec.broker_resources.cpu_limit.as_deref() { + let cpu_limit = parse_cpu_millis(cpu_limit).context("parse broker CPU limit")?; + if cpu_limit < cpu_request { + bail!("brokerResources.cpuLimit must be greater than or equal to cpuRequest"); + } + } + if let Some(memory_limit) = cluster.spec.broker_resources.memory_limit.as_deref() { + let memory_limit = + storage::parse_quantity(memory_limit).context("parse broker memory limit")?; + if memory_limit < memory_request { + bail!("brokerResources.memoryLimit must be greater than or equal to memoryRequest"); + } + } + } if cluster.spec.rollout.timeout_seconds == 0 || cluster.spec.rollout.timeout_seconds > 86_400 { bail!("rollout timeoutSeconds must be between 1 and 86400"); } @@ -454,6 +1468,58 @@ fn validate(cluster: &RustQueue) -> anyhow::Result<()> { Ok(()) } +fn ordinary_desired_brokers(eligible: usize, minimum: i32, maximum: i32) -> i32 { + i32::try_from(eligible) + .unwrap_or(i32::MAX) + .clamp(minimum, maximum) +} + +fn parse_cpu_millis(value: &str) -> anyhow::Result { + let value = value.trim(); + let (number, multiplier) = if let Some(number) = value.strip_suffix('m') { + (number, 1.0) + } else if let Some(number) = value.strip_suffix('u') { + (number, 0.001) + } else if let Some(number) = value.strip_suffix('n') { + (number, 0.000_001) + } else { + (value, 1_000.0) + }; + let number: f64 = number.parse()?; + anyhow::ensure!( + number.is_finite() && number >= 0.0, + "invalid CPU quantity {value}" + ); + Ok(number * multiplier) +} + +fn validate_message_storage_contract( + max_message_bytes: usize, + storage_feature_level: u32, +) -> anyhow::Result<()> { + const LEGACY_MAX_RECORD_BYTES: usize = 72 * 1024 * 1024; + const SINGLE_MESSAGE_ENVELOPE_BYTES: usize = 24; + const MPUB_ENTRY_BYTES: usize = 16; + const MAX_MPUB_MESSAGES: usize = 65_536; + let max_body_bytes = (64 * 1024 * 1024).max(max_message_bytes); + let maximum_record = max_message_bytes + .saturating_add(SINGLE_MESSAGE_ENVELOPE_BYTES) + .max(max_body_bytes.saturating_add(MPUB_ENTRY_BYTES.saturating_mul(MAX_MPUB_MESSAGES))); + if maximum_record > LEGACY_MAX_RECORD_BYTES && storage_feature_level < 2 { + bail!("messages above the v7 legacy record bound require storageFeatureLevel 2"); + } + Ok(()) +} + +fn has_sha256_digest(image: &str) -> bool { + image.rsplit_once("@sha256:").is_some_and(|(_, digest)| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + }) +} + fn error_policy( cluster: Arc, error: &ReconcileError, @@ -473,3 +1539,455 @@ fn watch_namespace() -> String { .map(|value| value.trim().to_owned()) .unwrap_or_else(|| "default".into()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn existing_gateway_replicas_preserve_activation_without_status() { + let gateway: StatefulSet = serde_json::from_value(serde_json::json!({ + "apiVersion": "apps/v1", + "kind": "StatefulSet", + "metadata": {"name": "queue-kodo-gateway"}, + "spec": { + "serviceName": "queue-kodo-gateways", + "replicas": 3, + "selector": {"matchLabels": {"app": "gateway"}}, + "template": { + "metadata": {"labels": {"app": "gateway"}}, + "spec": {"containers": []} + } + } + })) + .unwrap(); + assert!(statefulset_has_replicas(Some(&gateway))); + assert!(!statefulset_has_replicas(None)); + } + + #[test] + fn blocked_target_preflight_reuses_the_active_broker_image() { + let brokers: StatefulSet = serde_json::from_value(serde_json::json!({ + "apiVersion": "apps/v1", + "kind": "StatefulSet", + "metadata": {"name": "queue"}, + "spec": { + "serviceName": "queue-brokers", + "selector": {"matchLabels": {"app": "broker"}}, + "template": { + "metadata": {"labels": {"app": "broker"}}, + "spec": {"containers": [ + {"name": "sidecar", "image": "sidecar:v1"}, + {"name": "broker", "image": "rustqueue:active"} + ]} + } + } + })) + .unwrap(); + assert_eq!( + statefulset_broker_image(Some(&brokers)), + Some("rustqueue:active") + ); + assert_eq!(statefulset_broker_image(None), None); + } + + #[test] + fn kodo_decommission_requires_confirmation_only_for_a_live_gateway_set() { + assert!(kodo_decommission_blocked(false, false, true)); + assert!(!kodo_decommission_blocked(false, true, true)); + assert!(!kodo_decommission_blocked(true, false, true)); + assert!(!kodo_decommission_blocked(false, false, false)); + } + + #[test] + fn kodo_gateways_can_bootstrap_during_one_broker_outage() { + assert!(kodo_gateway_activation_ready(3, 2)); + assert!(kodo_gateway_activation_ready(2, 2)); + assert!(!kodo_gateway_activation_ready(3, 1)); + assert!(!kodo_gateway_activation_ready(1, 1)); + } + + #[test] + fn runtime_recovery_recreates_gateways_without_advancing_discovery() { + assert_eq!( + kodo_runtime_recovery( + true, + 3, + 2, + KodoRuntimeSignals { + discovery_mode: Some(false), + discovery_route: Some(false), + ..Default::default() + }, + ), + KodoRuntimeRecovery { + gateway_replicas: 3, + advertise_gateways: false, + discovery_service_kodo: false, + retain_gateways: false, + } + ); + } + + #[test] + fn runtime_recovery_preserves_an_in_progress_kodo_disable() { + assert_eq!( + kodo_runtime_recovery( + false, + 3, + 2, + KodoRuntimeSignals { + discovery_mode: Some(false), + discovery_route: Some(true), + previously_active: true, + ..Default::default() + }, + ), + KodoRuntimeRecovery { + gateway_replicas: 3, + advertise_gateways: false, + discovery_service_kodo: true, + retain_gateways: true, + } + ); + } + + #[test] + fn runtime_recovery_keeps_live_gateways_advertised_when_discovery_is_missing() { + assert_eq!( + kodo_runtime_recovery( + false, + 3, + 2, + KodoRuntimeSignals { + gateway_active: true, + ..Default::default() + }, + ), + KodoRuntimeRecovery { + gateway_replicas: 3, + advertise_gateways: true, + discovery_service_kodo: true, + retain_gateways: true, + } + ); + } + + #[test] + fn runtime_recovery_does_not_resurrect_a_completed_decommission() { + assert_eq!( + kodo_runtime_recovery(false, 3, 2, KodoRuntimeSignals::default()), + KodoRuntimeRecovery { + gateway_replicas: 0, + advertise_gateways: false, + discovery_service_kodo: false, + retain_gateways: false, + } + ); + } + + #[test] + fn kodo_gateway_readiness_requires_both_revision_fences() { + assert!(complete_kodo_gateway_set(3, 3)); + assert!(!complete_kodo_gateway_set(2, 3)); + assert!(!complete_kodo_gateway_set(3, 2)); + } + + #[test] + fn discovery_advertisement_is_recovered_from_the_live_deployment() { + let deployment: Deployment = serde_json::from_value(serde_json::json!({ + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": {"name": "queue-discovery", "generation": 4}, + "spec": { + "replicas": 2, + "selector": {"matchLabels": {"app": "discovery"}}, + "template": { + "metadata": { + "labels": { + "app": "discovery", + "rustqueue.io/discovery-publisher-mode": "kodo" + }, + "annotations": {"rustqueue.io/revision": "current"} + }, + "spec": { + "containers": [{ + "name": "discovery", + "image": "rustqueue:test", + "env": [{ + "name": "RUSTQUEUE_KODO_COMPATIBILITY_ENABLED", + "value": "true" + }, { + "name": "RUSTQUEUE_KODO_GATEWAY_ADDRESS", + "value": "queue-kodo-publish.test.svc" + }, { + "name": "RUSTQUEUE_KODO_CLEANUP_ENABLED", + "value": "true" + }] + }] + } + } + }, + "status": { + "observedGeneration": 4, + "readyReplicas": 2, + "updatedReplicas": 2, + "availableReplicas": 2 + } + })) + .unwrap(); + assert!(deployment_requests_kodo_mode(Some(&deployment))); + assert!(deployment_requests_kodo_cleanup(Some(&deployment))); + assert!(deployment_template_has_discovery_mode( + Some(&deployment), + true + )); + assert!(deployment_advertises_kodo_cleanup( + Some(&deployment), + "current" + )); + assert!(!deployment_runs_without_kodo(Some(&deployment))); + let direct_service: Service = serde_json::from_value(serde_json::json!({ + "apiVersion": "v1", + "kind": "Service", + "metadata": {"name": "queue-discovery"}, + "spec": { + "selector": { + "rustqueue.io/discovery-publisher-mode": "direct" + }, + "ports": [{"port": 4161}] + } + })) + .unwrap(); + assert_eq!( + discovery_service_kodo_route(Some(&direct_service)), + Some(false) + ); + let mut kodo_service = direct_service; + kodo_service + .spec + .as_mut() + .unwrap() + .selector + .as_mut() + .unwrap() + .insert(resources::DISCOVERY_MODE_LABEL.into(), "kodo".into()); + assert_eq!( + discovery_service_kodo_route(Some(&kodo_service)), + Some(true) + ); + assert_eq!(discovery_service_kodo_route(None), None); + + let mut unready = deployment.clone(); + unready.status.as_mut().unwrap().ready_replicas = Some(1); + assert!(!deployment_advertises_kodo_cleanup( + Some(&unready), + "current" + )); + + let mut disabled = deployment.clone(); + disabled + .spec + .as_mut() + .unwrap() + .template + .spec + .as_mut() + .unwrap() + .containers[0] + .env + .as_mut() + .unwrap() + .retain(|variable| !variable.name.starts_with("RUSTQUEUE_KODO_")); + disabled + .spec + .as_mut() + .unwrap() + .template + .metadata + .as_mut() + .unwrap() + .labels + .as_mut() + .unwrap() + .insert(resources::DISCOVERY_MODE_LABEL.into(), "direct".into()); + assert!(deployment_runs_without_kodo(Some(&disabled))); + disabled.status.as_mut().unwrap().updated_replicas = Some(1); + assert!(!deployment_runs_without_kodo(Some(&disabled))); + } + + #[test] + fn discovery_cutover_never_selects_direct_and_kodo_pods_together() { + assert_eq!( + discovery_target_mode(true, true, true, false, None, false, true), + (false, true), + "legacy unlabeled Pods are first relabeled without changing topology" + ); + assert_eq!( + discovery_target_route(true, true, None, false, false, (2, 0), 2), + None, + "label adoption keeps the compatibility selector broad" + ); + + let (target, adopting) = discovery_target_mode(true, true, true, false, None, true, true); + assert_eq!((target, adopting), (true, false)); + assert_eq!( + discovery_target_route(true, adopting, None, false, target, (0, 2), 2), + Some(false), + "Gateway-mode Pods roll out behind the direct-only Service" + ); + assert_eq!( + discovery_target_route(true, adopting, Some(false), false, target, (2, 2), 2), + Some(true), + "the Service switches only after every Gateway-mode Pod is Ready" + ); + assert_eq!( + discovery_target_route(true, adopting, Some(true), true, target, (0, 2), 2), + Some(false), + "the Service rolls back when Gateway-mode Pods fail during the soak" + ); + + let (target, adopting) = + discovery_target_mode(false, false, true, true, Some(true), true, true); + assert_eq!((target, adopting), (false, false)); + assert_eq!( + discovery_target_route(true, adopting, Some(true), true, target, (0, 2), 2), + Some(true) + ); + assert_eq!( + discovery_target_route(true, adopting, Some(true), true, target, (2, 2), 2), + Some(false), + "disable switches back only after every direct-mode Pod is Ready" + ); + assert_eq!( + discovery_target_route(true, adopting, Some(false), false, target, (0, 2), 2), + Some(true), + "the Service restores Gateway mode when direct Pods fail during the soak" + ); + } + + #[test] + fn kodo_disable_holds_gateways_until_target_brokers_are_safe() { + assert!(should_hold_kodo_for_disable(false, true, false, true)); + assert!(should_hold_kodo_for_disable(false, true, true, false)); + assert!(!should_hold_kodo_for_disable(false, true, true, true)); + assert!(!should_hold_kodo_for_disable(true, true, false, false)); + } + + #[test] + fn kodo_cutover_grace_uses_the_condition_transition_time() { + let mut cluster: RustQueue = serde_json::from_value(serde_json::json!({ + "apiVersion": "rustqueue.io/v1alpha1", + "kind": "RustQueue", + "metadata": {"name": "queue", "namespace": "test"}, + "spec": {"image": "rustqueue:test"}, + "status": { + "desiredBrokers": 3, + "readyBrokers": 3, + "phase": "Ready", + "message": "ready", + "activeStorageFeatureLevel": 1, + "conditions": [{ + "type": "KodoGatewaysAdvertised", + "status": "True", + "reason": "Advertised", + "message": "ready", + "lastTransitionTime": "2020-01-01T00:00:00Z" + }] + } + })) + .unwrap(); + assert!(condition_true_for(&cluster, "KodoGatewaysAdvertised", 630)); + cluster.status.as_mut().unwrap().conditions[0].status = "False".into(); + assert!(!condition_true_for(&cluster, "KodoGatewaysAdvertised", 630)); + } + + #[test] + fn active_storage_feature_level_is_a_monotonic_floor() { + assert_eq!(effective_storage_feature_level(1, 2), 2); + assert_eq!(effective_storage_feature_level(2, 1), 2); + } + + #[test] + fn ordinary_scaling_honors_the_configured_floor_and_ceiling() { + assert_eq!(ordinary_desired_brokers(0, 3, 10), 3); + assert_eq!(ordinary_desired_brokers(7, 3, 10), 7); + assert_eq!(ordinary_desired_brokers(12, 3, 10), 10); + assert_eq!(ordinary_desired_brokers(usize::MAX, 3, 10), 10); + } + + #[test] + fn cpu_quantities_compare_in_millicores() { + assert_eq!(parse_cpu_millis("1").unwrap(), 1_000.0); + assert_eq!(parse_cpu_millis("1000m").unwrap(), 1_000.0); + assert_eq!(parse_cpu_millis("500m").unwrap(), 500.0); + assert_eq!(parse_cpu_millis("500000u").unwrap(), 500.0); + assert!(parse_cpu_millis("-1").is_err()); + } + + #[test] + fn feature_two_retains_large_message_delivery_capacity() { + assert_eq!( + effective_delivery_limits(2, 20 * 1024 * 1024, 32 * 1024 * 1024, 64 * 1024 * 1024), + (100 * 1024 * 1024, 200 * 1024 * 1024) + ); + } + + #[test] + fn legacy_feature_rejects_messages_that_cross_the_record_bound() { + let maximum = 71 * 1024 * 1024; + assert!(validate_message_storage_contract(maximum, 1).is_ok()); + assert!(validate_message_storage_contract(maximum + 1, 1).is_err()); + assert!(validate_message_storage_contract(100 * 1024 * 1024, 2).is_ok()); + } + + #[test] + fn immutable_image_detection_requires_a_complete_sha256_digest() { + let digest = "a".repeat(64); + assert!(has_sha256_digest(&format!( + "registry/rustqueue@sha256:{digest}" + ))); + assert!(!has_sha256_digest(&format!( + "registry/rustqueue@sha256:{}", + "A".repeat(64) + ))); + assert!(!has_sha256_digest("registry/rustqueue:latest")); + assert!(!has_sha256_digest("registry/rustqueue@sha256:abcd")); + } + + #[test] + fn kodo_contract_requires_a_second_lookup_poll_retention_window() { + let mut cluster: RustQueue = serde_json::from_value(serde_json::json!({ + "apiVersion": "rustqueue.io/v1alpha1", + "kind": "RustQueue", + "metadata": {"name": "queue", "namespace": "test"}, + "spec": { + "image": "rustqueue:test", + "imagePullPolicy": "Never", + "minBrokers": 3, + "maxBrokers": 3, + "storageFeatureLevel": 2, + "bootstrapRetentionSeconds": 180, + "maxMessageBytes": 104857600, + "connectionDeliveryInflightBytes": 134217728, + "nodeDeliveryInflightBytes": 536870912, + "brokerResources": {"cpuRequest": "1", "memoryRequest": "2Gi"}, + "kodoCompatibility": {"enabled": true} + } + })) + .unwrap(); + + assert!(validate(&cluster, 2).is_ok()); + cluster.spec.broker_resources.cpu_request = "999m".into(); + assert!(validate(&cluster, 2) + .unwrap_err() + .to_string() + .contains("cpuRequest >= 1 CPU")); + cluster.spec.broker_resources.cpu_request = "1".into(); + cluster.spec.bootstrap_retention_seconds = KODO_BOOTSTRAP_RETENTION_SECONDS - 1; + assert!(validate(&cluster, 2) + .unwrap_err() + .to_string() + .contains("bootstrapRetentionSeconds >= 180")); + } +} diff --git a/crates/operator/src/controller/nodes.rs b/crates/operator/src/controller/nodes.rs index 757fa6a..1de9cc1 100644 --- a/crates/operator/src/controller/nodes.rs +++ b/crates/operator/src/controller/nodes.rs @@ -1,3 +1,4 @@ +use k8s_openapi::api::apps::v1::StatefulSet; use k8s_openapi::api::core::v1::{Node, Pod}; use kube::api::{Api, ListParams}; @@ -10,24 +11,10 @@ pub(super) async fn eligible(client: &kube::Client, selector: &str) -> anyhow::R let nodes = Api::::all(client.clone()) .list(&ListParams::default().labels(selector)) .await?; - Ok(nodes - .items - .iter() - .filter(|node| { - node.spec - .as_ref() - .is_none_or(|spec| spec.unschedulable != Some(true)) - && node - .status - .as_ref() - .and_then(|status| status.conditions.as_ref()) - .is_some_and(|conditions| { - conditions.iter().any(|condition| { - condition.type_ == "Ready" && condition.status == "True" - }) - }) - }) - .count()) + // Label membership is the scaling intent. Node readiness and cordons are + // transient health signals and must not silently trigger a second Broker + // disruption while the cluster is already degraded. + Ok(nodes.items.len()) } pub(super) async fn ready_brokers( @@ -38,6 +25,66 @@ pub(super) async fn ready_brokers( Ok(broker_health(client, namespace, name).await?.ready) } +pub(super) async fn ready_component( + client: &kube::Client, + namespace: &str, + name: &str, + component: &str, +) -> anyhow::Result { + let statefulset = Api::::namespaced(client.clone(), namespace) + .get(&format!("{name}-{component}")) + .await?; + let Some(revision) = observed_update_revision(&statefulset) else { + return Ok(0); + }; + let selector = + format!("app.kubernetes.io/instance={name},app.kubernetes.io/component={component}"); + let pods = Api::::namespaced(client.clone(), namespace) + .list(&ListParams::default().labels(&selector)) + .await?; + Ok(pods + .items + .iter() + .filter(|pod| pod_current_and_ready(pod, revision)) + .count() as i32) +} + +pub(super) async fn ready_component_revision( + client: &kube::Client, + namespace: &str, + name: &str, + component: &str, + revision: &str, +) -> anyhow::Result { + let selector = + format!("app.kubernetes.io/instance={name},app.kubernetes.io/component={component}"); + let pods = Api::::namespaced(client.clone(), namespace) + .list(&ListParams::default().labels(&selector)) + .await?; + Ok(pods + .items + .iter() + .filter(|pod| pod_ready_on_revision(pod, revision)) + .count() as i32) +} + +pub(super) async fn ready_discovery_mode( + client: &kube::Client, + namespace: &str, + name: &str, + kodo: bool, +) -> anyhow::Result { + let mode = if kodo { "kodo" } else { "direct" }; + let selector = format!( + "app.kubernetes.io/instance={name},app.kubernetes.io/component=discovery,{}={mode}", + crate::resources::DISCOVERY_MODE_LABEL + ); + let pods = Api::::namespaced(client.clone(), namespace) + .list(&ListParams::default().labels(&selector)) + .await?; + Ok(pods.items.iter().filter(|pod| pod_ready(pod)).count() as i32) +} + pub(super) async fn broker_health( client: &kube::Client, namespace: &str, @@ -50,15 +97,7 @@ pub(super) async fn broker_health( let mut ready = 0; let mut unavailable = Vec::new(); for pod in &pods.items { - let is_ready = pod - .status - .as_ref() - .and_then(|status| status.conditions.as_ref()) - .is_some_and(|conditions| { - conditions - .iter() - .any(|condition| condition.type_ == "Ready" && condition.status == "True") - }); + let is_ready = pod_ready(pod); if is_ready { ready += 1; } else { @@ -73,6 +112,47 @@ pub(super) async fn broker_health( Ok(BrokerHealth { ready, unavailable }) } +fn pod_ready(pod: &Pod) -> bool { + pod.metadata.deletion_timestamp.is_none() + && pod + .status + .as_ref() + .and_then(|status| status.conditions.as_ref()) + .is_some_and(|conditions| { + conditions + .iter() + .any(|condition| condition.type_ == "Ready" && condition.status == "True") + }) +} + +fn pod_current_and_ready(pod: &Pod, revision: &str) -> bool { + pod_ready(pod) + && pod + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get("controller-revision-hash")) + .is_some_and(|value| value == revision) +} + +fn pod_ready_on_revision(pod: &Pod, revision: &str) -> bool { + pod_ready(pod) + && pod + .metadata + .annotations + .as_ref() + .and_then(|annotations| annotations.get("rustqueue.io/revision")) + .is_some_and(|value| value == revision) +} + +fn observed_update_revision(statefulset: &StatefulSet) -> Option<&str> { + let generation = statefulset.metadata.generation?; + let status = statefulset.status.as_ref()?; + (status.observed_generation.unwrap_or_default() >= generation) + .then_some(status.update_revision.as_deref()) + .flatten() +} + fn pod_unavailable_reason(pod: &Pod) -> String { if let Some(waiting) = pod .status @@ -110,3 +190,66 @@ fn pod_unavailable_reason(pod: &Pod) -> String { }) .unwrap_or_else(|| "Pod has not reported Ready".into()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn pod(revision: &str, ready: bool) -> Pod { + serde_json::from_value(serde_json::json!({ + "apiVersion": "v1", + "kind": "Pod", + "metadata": { + "name": "queue-kodo-gateway-0", + "labels": {"controller-revision-hash": revision} + }, + "status": { + "conditions": [{ + "type": "Ready", + "status": if ready { "True" } else { "False" } + }] + } + })) + .unwrap() + } + + #[test] + fn component_readiness_requires_the_target_revision() { + assert!(pod_current_and_ready(&pod("current", true), "current")); + assert!(!pod_current_and_ready(&pod("old", true), "current")); + assert!(!pod_current_and_ready(&pod("current", false), "current")); + } + + #[test] + fn statefulset_revision_must_observe_its_latest_generation() { + let statefulset: StatefulSet = serde_json::from_value(serde_json::json!({ + "apiVersion": "apps/v1", + "kind": "StatefulSet", + "metadata": {"name": "queue-kodo-gateway", "generation": 4}, + "status": {"observedGeneration": 3, "updateRevision": "new"} + })) + .unwrap(); + assert_eq!(observed_update_revision(&statefulset), None); + let mut observed = statefulset; + observed.status.as_mut().unwrap().observed_generation = Some(4); + assert_eq!(observed_update_revision(&observed), Some("new")); + } + + #[test] + fn component_readiness_requires_the_requested_rustqueue_revision() { + let pod: Pod = serde_json::from_value(serde_json::json!({ + "apiVersion": "v1", + "kind": "Pod", + "metadata": { + "name": "queue-kodo-gateway-0", + "annotations": {"rustqueue.io/revision": "current"} + }, + "status": { + "conditions": [{"type": "Ready", "status": "True"}] + } + })) + .unwrap(); + assert!(pod_ready_on_revision(&pod, "current")); + assert!(!pod_ready_on_revision(&pod, "old")); + } +} diff --git a/crates/operator/src/controller/operations.rs b/crates/operator/src/controller/operations.rs index 18b060f..03bc88e 100644 --- a/crates/operator/src/controller/operations.rs +++ b/crates/operator/src/controller/operations.rs @@ -59,21 +59,42 @@ pub(super) async fn reconcile( current: i32, revision: &str, effective_image: &str, + disruptions_allowed: bool, ) -> anyhow::Result<(String, String, Option)> { - if desired < cluster.spec.min_brokers { + if insufficient_eligible_nodes(eligible, desired, cluster.spec.min_brokers) { + let resumed = drain::resume_all(context, cluster, namespace, auth).await?; + let resumed = if resumed.is_empty() { + String::new() + } else { + format!("; resumed disrupted Brokers {}", resumed.join(", ")) + }; return Ok(( "InsufficientNodes".into(), format!( - "{eligible} eligible nodes; minimum is {}", + "{eligible} eligible nodes; desired is {desired} and minimum is {}{resumed}", cluster.spec.min_brokers ), None, )); } if let Some(request) = &cluster.spec.maintenance { - return reconcile_maintenance(context, cluster, namespace, auth, revision, request).await; + if request.enabled && !disruptions_allowed { + drain::resume_all(context, cluster, namespace, auth).await?; + return Ok(waiting_for_kodo_cutover()); + } + if request.enabled { + return reconcile_maintenance(context, cluster, namespace, auth, revision, request) + .await; + } + // A completed maintenance request may remain in spec. Resume its + // target, then continue with scale and rollout reconciliation. + drain::resume_one(context, namespace, &request.broker, auth).await?; } if current > desired { + if !disruptions_allowed { + drain::resume_all(context, cluster, namespace, auth).await?; + return Ok(waiting_for_kodo_cutover()); + } return reconcile_scale_down( context, cluster, namespace, auth, current, desired, revision, ) @@ -97,6 +118,17 @@ pub(super) async fn reconcile( if desired == 0 { return Ok(("Ready".into(), "no eligible Broker nodes".into(), None)); } + if !disruptions_allowed { + drain::resume_all(context, cluster, namespace, auth).await?; + if drain::brokers_current_and_ready(context, cluster, namespace, revision, current).await? { + return Ok(( + "Ready".into(), + "all Brokers already match the target revision".into(), + None, + )); + } + return Ok(waiting_for_kodo_cutover()); + } reconcile_rollout( context, cluster, @@ -109,6 +141,21 @@ pub(super) async fn reconcile( .await } +fn insufficient_eligible_nodes(eligible: usize, desired: i32, minimum: i32) -> bool { + desired < minimum + || usize::try_from(desired) + .map(|desired| eligible < desired) + .unwrap_or(true) +} + +fn waiting_for_kodo_cutover() -> (String, String, Option) { + ( + "WaitingForKodoCutover".into(), + "waiting for Kodo Gateway cutover safety gates before disrupting a Broker".into(), + None, + ) +} + async fn reconcile_maintenance( context: &ContextData, cluster: &RustQueue, @@ -146,7 +193,9 @@ async fn reconcile_maintenance( previous_image: None, current_broker: Some(request.broker.clone()), }; - let phase = if request.enabled { + let phase = if progress.phase == "Blocked" { + "MaintenanceBlocked" + } else if request.enabled { "Maintenance" } else { "Ready" @@ -214,6 +263,11 @@ async fn reconcile_rollout( Some(image) => Some(image), None => drain::previous_image(context, cluster, namespace, revision).await?, }; + if timed_out || previous.is_some_and(|operation| operation.phase == "Failed") { + if let Some(target) = previous.and_then(|operation| operation.current_broker.as_deref()) { + drain::resume_one(context, namespace, target, auth).await?; + } + } let progress = if let Some(failed) = previous.filter(|operation| operation.phase == "Failed") { drain::Progress { target: failed.current_broker.clone(), @@ -280,6 +334,13 @@ fn rollout_phase_can_timeout(phase: &str) -> bool { mod tests { use super::*; + #[test] + fn fixed_three_broker_profile_reports_insufficient_eligible_nodes() { + assert!(insufficient_eligible_nodes(2, 3, 3)); + assert!(!insufficient_eligible_nodes(3, 3, 3)); + assert!(insufficient_eligible_nodes(0, 0, 1)); + } + #[test] fn terminal_and_human_gated_rollouts_never_time_out() { for phase in [ diff --git a/crates/operator/src/controller/preflight.rs b/crates/operator/src/controller/preflight.rs index dd81d5c..f6689ad 100644 --- a/crates/operator/src/controller/preflight.rs +++ b/crates/operator/src/controller/preflight.rs @@ -7,6 +7,7 @@ use kube::api::{Api, DeleteParams, ListParams, PostParams}; use kube::{Resource, ResourceExt}; use serde::Deserialize; use serde_json::json; +use sha2::{Digest, Sha256}; const DATA_FORMAT: u32 = 7; @@ -23,6 +24,10 @@ struct BinaryCapabilities { minimum_reader_feature_level: u32, maximum_reader_feature_level: u32, maximum_writer_feature_level: u32, + #[serde(default = "legacy_maximum_message_bytes")] + maximum_message_bytes: usize, + #[serde(default = "legacy_maximum_batch_bytes")] + maximum_batch_bytes: usize, } #[derive(Clone, Debug, Deserialize)] @@ -45,8 +50,29 @@ pub(super) async fn target_image( target_image: &str, ) -> anyhow::Result { let api = Api::::namespaced(context.client.clone(), namespace); - let name = probe_name(&cluster.name_any(), target_image); + let contract = probe_contract(cluster, target_image); + let fingerprint = probe_fingerprint(&contract); + let name = probe_name(&cluster.name_any(), &contract); + let owner_uid = cluster.metadata.uid.as_deref(); let pod = match api.get_opt(&name).await? { + Some(pod) if !controlled_by(&pod, owner_uid) => { + return Ok(Outcome::Blocked(format!( + "target image capability probe {name} is not owned by this RustQueue" + ))) + } + Some(pod) + if !probe_matches( + &pod, + target_image, + &cluster.spec.image_pull_policy, + &fingerprint, + ) => + { + api.delete(&name, &DeleteParams::default()).await?; + return Ok(Outcome::Pending(format!( + "replacing stale target image capability probe {name}" + ))); + } Some(pod) => pod, None => { let owner = cluster @@ -56,6 +82,9 @@ pub(super) async fn target_image( "apiVersion": "v1", "kind": "Pod", "metadata": { "name": name, "namespace": namespace, + "annotations": { + "rustqueue.io/preflight-contract": fingerprint + }, "labels": { "app.kubernetes.io/instance": cluster.name_any(), "app.kubernetes.io/component": "capability-preflight", @@ -101,11 +130,20 @@ pub(super) async fn target_image( .and_then(|status| status.phase.as_deref()) { Some("Succeeded") => { + if let Err(message) = validate_probe_image(&pod) { + return Ok(Outcome::Blocked(message)); + } let message = terminated_message(&pod).ok_or_else(|| { anyhow::anyhow!("target image probe succeeded without capabilities") })?; let capabilities: BinaryCapabilities = serde_json::from_str(message)?; - validate_binary(&capabilities, cluster.spec.storage_feature_level).map_or_else( + validate_binary( + &capabilities, + cluster.spec.storage_feature_level, + cluster.spec.max_message_bytes, + required_batch_bytes(cluster), + ) + .map_or_else( |message| Ok(Outcome::Blocked(message)), |_| { Ok(Outcome::Ready { @@ -183,7 +221,12 @@ pub(super) async fn current_brokers( "broker capability preflight is waiting for {name}" ))); }; - if let Err(message) = validate_report(&report, desired) { + if let Err(message) = validate_report( + &report, + desired, + cluster.spec.max_message_bytes, + required_batch_bytes(cluster), + ) { return Ok(Outcome::Blocked(format!("broker {name}: {message}"))); } maximum_active = maximum_active.max(report.storage.active_writer_feature_level); @@ -200,17 +243,18 @@ pub(super) async fn cleanup_old_probes( target_image: &str, ) -> anyhow::Result<()> { let api = Api::::namespaced(context.client.clone(), namespace); - let current = probe_name(&cluster.name_any(), target_image); + let current = probe_name(&cluster.name_any(), &probe_contract(cluster, target_image)); let selector = format!( "app.kubernetes.io/instance={},app.kubernetes.io/component=capability-preflight", cluster.name_any() ); + let owner_uid = cluster.metadata.uid.as_deref(); for pod in api .list(&ListParams::default().labels(&selector)) .await? .items { - if pod.name_any() != current { + if pod.name_any() != current && controlled_by(&pod, owner_uid) { api.delete(&pod.name_any(), &DeleteParams::default()) .await?; } @@ -218,7 +262,25 @@ pub(super) async fn cleanup_old_probes( Ok(()) } -fn validate_binary(capabilities: &BinaryCapabilities, desired: u32) -> Result<(), String> { +fn controlled_by(pod: &Pod, owner_uid: Option<&str>) -> bool { + owner_uid.is_some_and(|owner_uid| { + pod.metadata + .owner_references + .as_deref() + .is_some_and(|owners| { + owners + .iter() + .any(|owner| owner.controller == Some(true) && owner.uid == owner_uid) + }) + }) +} + +fn validate_binary( + capabilities: &BinaryCapabilities, + desired: u32, + message_bytes: usize, + batch_bytes: usize, +) -> Result<(), String> { if capabilities.binary_version.trim().is_empty() || capabilities.data_format != DATA_FORMAT { return Err("target image does not advertise RustQueue format v7".into()); } @@ -230,32 +292,125 @@ fn validate_binary(capabilities: &BinaryCapabilities, desired: u32) -> Result<() "target image cannot read and write requested storage feature level {desired}" )); } + if message_bytes > capabilities.maximum_message_bytes + || batch_bytes > capabilities.maximum_batch_bytes + { + return Err(format!( + "target image supports at most {} message bytes and {} batch bytes", + capabilities.maximum_message_bytes, capabilities.maximum_batch_bytes + )); + } Ok(()) } -fn validate_report(report: &CompatibilityReport, desired: u32) -> Result<(), String> { - validate_binary(&report.binary, desired)?; +fn validate_report( + report: &CompatibilityReport, + desired: u32, + message_bytes: usize, + batch_bytes: usize, +) -> Result<(), String> { + validate_binary( + &report.binary, + desired.max(report.storage.active_writer_feature_level), + message_bytes, + batch_bytes, + )?; if report.storage.data_format != DATA_FORMAT { return Err("PVC compatibility state is not format v7".into()); } if report.storage.minimum_reader_feature_level > report.binary.maximum_reader_feature_level { return Err("binary is behind the PVC rollback fence".into()); } - if desired < report.storage.active_writer_feature_level { - return Err(format!( - "requested feature level {desired} is below PVC rollback fence {}", - report.storage.active_writer_feature_level - )); - } Ok(()) } -fn probe_name(instance: &str, image: &str) -> String { - let suffix = format!("-preflight-{:08x}", crc32c::crc32c(image.as_bytes())); +fn required_batch_bytes(cluster: &RustQueue) -> usize { + if cluster.spec.kodo_compatibility.enabled { + 128 * 1024 * 1024 + } else { + (64 * 1024 * 1024).max(cluster.spec.max_message_bytes) + } +} + +fn legacy_maximum_message_bytes() -> usize { + 32 * 1024 * 1024 +} + +fn legacy_maximum_batch_bytes() -> usize { + 64 * 1024 * 1024 +} + +fn probe_contract(cluster: &RustQueue, image: &str) -> String { + format!( + "{}\0{}\0{}\0{}\0{}\0{}\0{}", + cluster.metadata.uid.as_deref().unwrap_or_default(), + image, + cluster.spec.image_pull_policy, + cluster.spec.storage_feature_level, + cluster.spec.max_message_bytes, + required_batch_bytes(cluster), + cluster.spec.rollout.retry_nonce, + ) +} + +fn probe_fingerprint(contract: &str) -> String { + let digest = Sha256::digest(contract.as_bytes()); + hex::encode(&digest[..16]) +} + +fn probe_name(instance: &str, contract: &str) -> String { + let suffix = format!("-preflight-{}", probe_fingerprint(contract)); let maximum = 63usize.saturating_sub(suffix.len()); format!("{}{}", &instance[..instance.len().min(maximum)], suffix) } +fn probe_matches( + pod: &Pod, + target_image: &str, + image_pull_policy: &str, + fingerprint: &str, +) -> bool { + let annotation_matches = pod + .metadata + .annotations + .as_ref() + .and_then(|annotations| annotations.get("rustqueue.io/preflight-contract")) + .is_some_and(|value| value == fingerprint); + let container_matches = pod + .spec + .as_ref() + .and_then(|spec| spec.containers.first()) + .is_some_and(|container| { + container.name == "probe" + && container.image.as_deref() == Some(target_image) + && container.image_pull_policy.as_deref() == Some(image_pull_policy) + && container + .command + .as_ref() + .is_some_and(|values| values.iter().map(String::as_str).eq(["rustqueued"])) + && container.args.as_ref().is_some_and(|values| { + values + .iter() + .map(String::as_str) + .eq(["--capabilities-output", "/dev/termination-log"]) + }) + }); + annotation_matches && container_matches +} + +fn validate_probe_image(pod: &Pod) -> Result<(), String> { + let status = pod + .status + .as_ref() + .and_then(|status| status.container_statuses.as_ref()) + .and_then(|statuses| statuses.first()) + .ok_or_else(|| "target image probe has no container image identity".to_owned())?; + if status.image_id.trim().is_empty() { + return Err("target image probe has an empty container image identity".into()); + } + Ok(()) +} + fn terminated_message(pod: &Pod) -> Option<&str> { pod.status .as_ref()? @@ -301,22 +456,39 @@ mod tests { fn binary(maximum: u32) -> BinaryCapabilities { BinaryCapabilities { - binary_version: "0.7.2".into(), + binary_version: "0.8.0".into(), data_format: 7, minimum_reader_feature_level: 1, maximum_reader_feature_level: maximum, maximum_writer_feature_level: maximum, + maximum_message_bytes: 100 * 1024 * 1024, + maximum_batch_bytes: 128 * 1024 * 1024, } } #[test] fn target_image_must_support_the_requested_feature() { - assert!(validate_binary(&binary(2), 2).is_ok()); - assert!(validate_binary(&binary(1), 2).is_err()); + assert!(validate_binary(&binary(2), 2, 100 * 1024 * 1024, 128 * 1024 * 1024).is_ok()); + assert!(validate_binary(&binary(1), 2, 20 * 1024 * 1024, 64 * 1024 * 1024).is_err()); } #[test] - fn broker_preflight_rejects_a_feature_downgrade() { + fn target_image_must_support_the_requested_protocol_limits() { + let legacy: BinaryCapabilities = serde_json::from_value(json!({ + "binary_version": "0.8.0", + "data_format": 7, + "minimum_reader_feature_level": 1, + "maximum_reader_feature_level": 2, + "maximum_writer_feature_level": 2 + })) + .unwrap(); + assert_eq!(legacy.maximum_message_bytes, legacy_maximum_message_bytes()); + assert_eq!(legacy.maximum_batch_bytes, legacy_maximum_batch_bytes()); + assert!(validate_binary(&legacy, 2, 100 * 1024 * 1024, 128 * 1024 * 1024).is_err()); + } + + #[test] + fn broker_preflight_recovers_the_live_feature_floor() { let report = CompatibilityReport { binary: binary(2), storage: CompatibilityState { @@ -325,11 +497,69 @@ mod tests { minimum_reader_feature_level: 2, }, }; - assert!(validate_report(&report, 1).is_err()); + assert!(validate_report(&report, 1, 20 * 1024 * 1024, 64 * 1024 * 1024).is_ok()); + assert_eq!(report.storage.active_writer_feature_level, 2); } #[test] fn probe_names_are_dns_bounded() { assert!(probe_name(&"q".repeat(80), "registry/rustqueue:v7").len() <= 63); } + + #[test] + fn probe_cleanup_requires_the_cluster_owner() { + let pod: Pod = serde_json::from_value(json!({ + "apiVersion": "v1", + "kind": "Pod", + "metadata": { + "name": "queue-preflight-old", + "ownerReferences": [{ + "apiVersion": "rustqueue.io/v1alpha1", + "kind": "RustQueue", + "name": "queue", + "uid": "queue-uid", + "controller": true + }] + } + })) + .unwrap(); + assert!(controlled_by(&pod, Some("queue-uid"))); + assert!(!controlled_by(&pod, Some("other-uid"))); + } + + #[test] + fn cached_probe_must_match_its_contract_and_have_an_image_identity() { + let digest = "a".repeat(64); + let image = format!("registry/rustqueue@sha256:{digest}"); + let pod: Pod = serde_json::from_value(json!({ + "apiVersion": "v1", + "kind": "Pod", + "metadata": { + "name": "queue-preflight-current", + "annotations": {"rustqueue.io/preflight-contract": "deadbeef"} + }, + "spec": { + "containers": [{ + "name": "probe", + "image": image, + "imagePullPolicy": "IfNotPresent", + "command": ["rustqueued"], + "args": ["--capabilities-output", "/dev/termination-log"] + }] + }, + "status": { + "containerStatuses": [{ + "name": "probe", + "image": image, + "imageID": format!("docker-pullable://registry/rustqueue@sha256:{digest}"), + "ready": false, + "restartCount": 0 + }] + } + })) + .unwrap(); + assert!(probe_matches(&pod, &image, "IfNotPresent", "deadbeef")); + assert!(validate_probe_image(&pod).is_ok()); + assert!(!probe_matches(&pod, &image, "IfNotPresent", "stale")); + } } diff --git a/crates/operator/src/controller/status.rs b/crates/operator/src/controller/status.rs index 3bcda46..0f8222c 100644 --- a/crates/operator/src/controller/status.rs +++ b/crates/operator/src/controller/status.rs @@ -37,6 +37,7 @@ impl<'a> StatusBuilder<'a> { operation_history: previous.operation_history, orphaned_pvcs: previous.orphaned_pvcs, desired_storage_size: cluster.spec.storage_size.clone(), + kodo_producer_restart_baseline_nonce: previous.kodo_producer_restart_baseline_nonce, }, } } @@ -94,6 +95,7 @@ impl<'a> StatusBuilder<'a> { let started_at = current .as_ref() .filter(|operation| operation.id == update.id) + .filter(|operation| !resume_resets_operation_clock(&operation.phase, update.phase)) .map(|operation| operation.started_at.clone()) .unwrap_or_else(|| timestamp.clone()); if let Some(previous) = current.filter(|operation| operation.id != update.id) { @@ -135,11 +137,24 @@ impl<'a> StatusBuilder<'a> { self } + pub fn kodo_producer_restart_baseline_nonce(mut self, nonce: Option) -> Self { + self.status.kodo_producer_restart_baseline_nonce = nonce; + self + } + pub fn build(self) -> RustQueueStatus { self.status } } +fn resume_resets_operation_clock(previous: &str, next: &str) -> bool { + matches!(previous, "Paused" | "AwaitingCanaryApproval" | "Blocked") + && !matches!( + next, + "Paused" | "AwaitingCanaryApproval" | "Blocked" | "Completed" | "Failed" + ) +} + pub(super) fn operation_id(kind: &str, target: &str, revision: &str) -> String { format!( "{kind}-{:08x}", @@ -206,6 +221,7 @@ mod tests { proxy_node_selector: BTreeMap::new(), proxy_tcp_max_connection_age_seconds: 300, discovery_replicas: 2, + kodo_compatibility: crate::crd::KodoCompatibility::default(), maintenance: None, rollout: RolloutPolicy::default(), broker_scheduling: BrokerScheduling::default(), @@ -287,4 +303,48 @@ mod tests { .build(); assert_eq!(second.current_operation.unwrap().completed_at, first); } + + #[test] + fn resuming_a_human_gate_restarts_the_operation_timeout_clock() { + let mut resource = cluster(); + resource.status = Some( + StatusBuilder::new(&resource, 3, 3, 1) + .operation(OperationUpdate { + id: "rollout-1", + kind: "Rollout", + phase: "Paused", + target: "queue:v2", + revision: "r2", + message: "paused", + previous_image: Some("queue:v1".into()), + current_broker: None, + }) + .build(), + ); + resource + .status + .as_mut() + .unwrap() + .current_operation + .as_mut() + .unwrap() + .started_at = "2000-01-01T00:00:00Z".into(); + + let resumed = StatusBuilder::new(&resource, 3, 3, 1) + .operation(OperationUpdate { + id: "rollout-1", + kind: "Rollout", + phase: "Draining", + target: "queue:v2", + revision: "r2", + message: "running", + previous_image: Some("queue:v1".into()), + current_broker: Some("queue-2".into()), + }) + .build() + .current_operation + .unwrap(); + assert_ne!(resumed.started_at, "2000-01-01T00:00:00Z"); + assert!(!resume_resets_operation_clock("Draining", "Replacing")); + } } diff --git a/crates/operator/src/controller/storage.rs b/crates/operator/src/controller/storage.rs index 3a036b6..adc4d0a 100644 --- a/crates/operator/src/controller/storage.rs +++ b/crates/operator/src/controller/storage.rs @@ -169,16 +169,16 @@ fn claim_ordinal(name: &str, cluster: &str) -> Option { name.strip_prefix(&format!("data-{cluster}-"))?.parse().ok() } -fn parse_quantity(value: &str) -> anyhow::Result { +pub(super) fn parse_quantity(value: &str) -> anyhow::Result { let split = value .find(|character: char| !character.is_ascii_digit() && character != '.') .unwrap_or(value.len()); let (number, suffix) = value.split_at(split); - anyhow::ensure!(!number.is_empty(), "invalid storage quantity {value}"); + anyhow::ensure!(!number.is_empty(), "invalid resource quantity {value}"); let number: f64 = number.parse()?; anyhow::ensure!( number.is_finite() && number >= 0.0, - "invalid storage quantity {value}" + "invalid resource quantity {value}" ); let multiplier = match suffix { "" => 1_u128, @@ -194,7 +194,7 @@ fn parse_quantity(value: &str) -> anyhow::Result { "Ti" => 1_u128 << 40, "Pi" => 1_u128 << 50, "Ei" => 1_u128 << 60, - _ => anyhow::bail!("unsupported storage quantity suffix in {value}"), + _ => anyhow::bail!("unsupported resource quantity suffix in {value}"), }; Ok((number * multiplier as f64) as u128) } diff --git a/crates/operator/src/crd.rs b/crates/operator/src/crd.rs index 2be79f1..6993d80 100644 --- a/crates/operator/src/crd.rs +++ b/crates/operator/src/crd.rs @@ -43,7 +43,7 @@ pub struct RustQueueSpec { pub disk_high_watermark_percent: u8, #[serde(default = "default_disk_low_watermark")] pub disk_low_watermark_percent: u8, - #[serde(default = "enabled")] + #[serde(default)] pub protective_eviction_enabled: bool, #[serde(default = "default_disk_pressure_grace")] pub disk_pressure_grace_seconds: u64, @@ -74,6 +74,8 @@ pub struct RustQueueSpec { #[serde(default = "default_discovery_replicas")] pub discovery_replicas: i32, #[serde(default)] + pub kodo_compatibility: KodoCompatibility, + #[serde(default)] pub maintenance: Option, #[serde(default)] pub rollout: RolloutPolicy, @@ -83,6 +85,57 @@ pub struct RustQueueSpec { pub broker_resources: WorkloadResources, } +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +pub struct KodoCompatibility { + #[serde(default)] + pub enabled: bool, + #[serde(default)] + pub decommission_confirmed: bool, + #[serde(default)] + pub producer_restart_nonce: String, + #[serde(default)] + pub cleanup_enabled: bool, + #[serde(default = "default_kodo_cutover_grace")] + pub cutover_grace_seconds: u64, + #[serde(default = "default_kodo_allowed_pod_selector")] + pub allowed_pod_selector: BTreeMap, + #[serde(default)] + pub allowed_namespace_selector: BTreeMap, +} + +impl Default for KodoCompatibility { + fn default() -> Self { + Self { + enabled: false, + decommission_confirmed: false, + producer_restart_nonce: String::new(), + cleanup_enabled: false, + cutover_grace_seconds: default_kodo_cutover_grace(), + allowed_pod_selector: default_kodo_allowed_pod_selector(), + allowed_namespace_selector: BTreeMap::new(), + } + } +} + +impl KodoCompatibility { + pub(crate) fn effective_cleanup_enabled(&self) -> bool { + false + } + + pub(crate) fn effective_allowed_pod_selector(&self) -> BTreeMap { + if self.allowed_pod_selector.is_empty() { + default_kodo_allowed_pod_selector() + } else { + self.allowed_pod_selector.clone() + } + } +} + +fn default_kodo_allowed_pod_selector() -> BTreeMap { + BTreeMap::from([("app.kubernetes.io/name".into(), "kodo".into())]) +} + #[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)] #[serde(rename_all = "camelCase")] pub struct BrokerMaintenance { @@ -195,6 +248,8 @@ pub struct RustQueueStatus { pub orphaned_pvcs: Vec, #[serde(default)] pub desired_storage_size: String, + #[serde(default)] + pub kodo_producer_restart_baseline_nonce: Option, } #[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] @@ -267,6 +322,9 @@ fn default_bootstrap_retention() -> u64 { fn default_proxy_tcp_connection_age() -> u64 { 300 } +fn default_kodo_cutover_grace() -> u64 { + 630 +} fn default_max_message_bytes() -> usize { 20 * 1024 * 1024 } @@ -323,7 +381,31 @@ mod tests { assert!(schema .to_string() .contains("proxyTcpMaxConnectionAgeSeconds")); + assert!(schema.to_string().contains("kodoCompatibility")); + assert!(schema.to_string().contains("decommissionConfirmed")); + assert!(schema.to_string().contains("producerRestartNonce")); assert!(!schema.to_string().contains("replicationFactor")); assert!(!schema.to_string().contains("cell")); } + + #[test] + fn disabling_kodo_also_disables_a_stale_cleanup_request() { + let mut compatibility = KodoCompatibility { + cleanup_enabled: true, + ..KodoCompatibility::default() + }; + assert!(!compatibility.effective_cleanup_enabled()); + compatibility.enabled = true; + assert!(!compatibility.effective_cleanup_enabled()); + } + + #[test] + fn empty_kodo_selector_falls_back_to_the_fail_closed_default() { + let mut compatibility = KodoCompatibility::default(); + compatibility.allowed_pod_selector.clear(); + assert_eq!( + compatibility.effective_allowed_pod_selector(), + BTreeMap::from([("app.kubernetes.io/name".into(), "kodo".into())]) + ); + } } diff --git a/crates/operator/src/lib.rs b/crates/operator/src/lib.rs index bee2cb9..e0dc0bb 100644 --- a/crates/operator/src/lib.rs +++ b/crates/operator/src/lib.rs @@ -4,8 +4,9 @@ pub mod management_crd; pub mod resources; pub use crd::{ - BrokerMaintenance, BrokerScheduling, BrokerToleration, OperationStatus, RolloutPolicy, - RustQueue, RustQueueCondition, RustQueueSpec, RustQueueStatus, WorkloadResources, + BrokerMaintenance, BrokerScheduling, BrokerToleration, KodoCompatibility, OperationStatus, + RolloutPolicy, RustQueue, RustQueueCondition, RustQueueSpec, RustQueueStatus, + WorkloadResources, }; pub use management_crd::{ ManagedResourceAction, ManagedResourceOperation, ManagedResourcePhase, RustQueueChannel, diff --git a/crates/operator/src/resources.rs b/crates/operator/src/resources.rs index 08bd61a..83928c4 100644 --- a/crates/operator/src/resources.rs +++ b/crates/operator/src/resources.rs @@ -1,3 +1,6 @@ +#[path = "resources/kodo.rs"] +mod kodo; + use crate::RustQueue; use anyhow::{bail, Context}; use k8s_openapi::api::apps::v1::{DaemonSet, Deployment, StatefulSet}; @@ -8,12 +11,15 @@ use k8s_openapi::api::rbac::v1::{Role, RoleBinding}; use kube::{Resource, ResourceExt}; use serde::de::DeserializeOwned; use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; use std::collections::BTreeMap; pub const MANAGER: &str = "rustqueue-operator"; +pub const DISCOVERY_MODE_LABEL: &str = "rustqueue.io/discovery-publisher-mode"; pub struct ResourceSet { pub revision: String, + pub retain_existing_kodo_resources: bool, pub config: ConfigMap, pub service_account: ServiceAccount, pub role: Role, @@ -21,17 +27,28 @@ pub struct ResourceSet { pub broker_service: Service, pub brokers: StatefulSet, pub broker_pdb: PodDisruptionBudget, + pub broker_network_policy: NetworkPolicy, pub discovery_service: Service, pub discovery: Deployment, pub discovery_pdb: PodDisruptionBudget, pub proxy_service: Service, pub proxy: DaemonSet, + pub kodo_gateway_service: Option, + pub kodo_gateway_headless_service: Option, + pub kodo_gateway: Option, + pub kodo_gateway_pdb: Option, + pub kodo_gateway_network_policy: Option, pub network_policy: NetworkPolicy, } pub struct BuildInput<'a> { pub cluster: &'a RustQueue, pub replicas: i32, + pub kodo_gateway_replicas: i32, + pub advertise_kodo_gateways: bool, + pub discovery_service_kodo: Option, + pub activate_kodo_cleanup: bool, + pub retain_existing_kodo_resources: bool, pub image: &'a str, pub claim_template_size: &'a str, pub secret_name: &'a str, @@ -58,22 +75,24 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { let broker_service_name = format!("{name}-brokers"); let discovery_name = format!("{name}-discovery"); let proxy_name = format!("{name}-proxy"); + let kodo_gateway_publish_service_name = format!("{name}-kodo-publish"); let service_account_name = format!("{name}-runtime"); let config_name = format!("{name}-config"); let config_text = broker_config(cluster, input.secret_name); - let revision = format!( - "{:08x}", - crc32c::crc32c( - format!( - "{}\0{}\0{}\0{}", - input.image, - config_text, - input.mounted_secret_revision, - cluster.spec.message_index_cache_bytes - ) - .as_bytes(), - ) - ); + let revision_contract = serde_json::to_vec(&json!({ + "image": input.image, + "imagePullPolicy": cluster.spec.image_pull_policy, + "config": &config_text, + "secretName": input.secret_name, + "mountedSecretRevision": input.mounted_secret_revision, + "clientTlsSecretName": cluster.spec.client_tls_secret_name, + "messageIndexCacheBytes": cluster.spec.message_index_cache_bytes, + "eligibleNodeSelector": cluster.spec.eligible_node_selector, + "brokerScheduling": &cluster.spec.broker_scheduling, + "brokerResources": &cluster.spec.broker_resources, + }))?; + let revision_digest = Sha256::digest(revision_contract); + let revision = hex::encode(&revision_digest[..16]); let node_selector = pod_node_selector(&cluster.spec.eligible_node_selector)?; let broker_resources = workload_resources(&cluster.spec.broker_resources); let tolerations: Vec<_> = cluster @@ -169,7 +188,11 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { "name": "broker", "image": input.image, "imagePullPolicy": cluster.spec.image_pull_policy, "command": ["rustqueued"], "args": ["--config", "/etc/rustqueue/rustqueue.toml"], - "ports": [{"name": "tcp", "containerPort": 4150}, {"name": "http", "containerPort": 4151}], + "ports": [ + {"name": "tcp", "containerPort": 4150}, + {"name": "http", "containerPort": 4151}, + {"name": "kodo-http", "containerPort": 4152} + ], "env": [ {"name": "POD_NAME", "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}}}, {"name": "POD_NAMESPACE", "valueFrom": {"fieldRef": {"fieldPath": "metadata.namespace"}}}, @@ -208,19 +231,70 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { } }))?; + let discovery_mode = if input.advertise_kodo_gateways { + "kodo" + } else { + "direct" + }; + let mut discovery_pod_labels = labels_for(&labels, "discovery"); + discovery_pod_labels.insert(DISCOVERY_MODE_LABEL.into(), discovery_mode.into()); + let mut discovery_service_selector = labels_for(&labels, "discovery"); + if let Some(kodo) = input.discovery_service_kodo { + discovery_service_selector.insert( + DISCOVERY_MODE_LABEL.into(), + if kodo { "kodo" } else { "direct" }.into(), + ); + } let discovery_service = typed(service_json( metadata(&discovery_name, "discovery"), - labels_for(&labels, "discovery"), + discovery_service_selector, vec![json!({"name": "http", "port": 4161, "targetPort": "http"})], ))?; + let mut discovery_env = vec![ + json!({"name": "POD_NAMESPACE", "valueFrom": {"fieldRef": {"fieldPath": "metadata.namespace"}}}), + json!({"name": "RUSTQUEUE_BROKER_SERVICE", "value": broker_service_name}), + json!({"name": "RUSTQUEUE_REGISTRY_TOKEN_FILE", "value": "/run/secrets/rustqueue/registry-token"}), + ]; + if input.advertise_kodo_gateways { + let cleanup_enabled = input.activate_kodo_cleanup + && cluster.spec.kodo_compatibility.effective_cleanup_enabled(); + discovery_env.extend([ + json!({"name": "RUSTQUEUE_KODO_COMPATIBILITY_ENABLED", "value": "true"}), + json!({"name": "RUSTQUEUE_KODO_CLEANUP_ENABLED", "value": cleanup_enabled.to_string()}), + ]); + let address = format!("{kodo_gateway_publish_service_name}.{namespace}.svc"); + discovery_env.push(json!({"name": "RUSTQUEUE_KODO_GATEWAY_ADDRESS", "value": address})); + } + let discovery_strategy = json!({ + "type": "RollingUpdate", + "rollingUpdate": {"maxUnavailable": 0, "maxSurge": "100%"} + }); + let discovery_min_ready_seconds = + if cluster.spec.kodo_compatibility.enabled || input.retain_existing_kodo_resources { + i32::try_from( + cluster + .spec + .kodo_compatibility + .cutover_grace_seconds + .saturating_add(30), + ) + .context("Kodo cutover grace exceeds the Deployment minReadySeconds range")? + } else { + 30 + }; let discovery = typed(json!({ "apiVersion": "apps/v1", "kind": "Deployment", "metadata": metadata(&discovery_name, "discovery"), "spec": { "replicas": cluster.spec.discovery_replicas.max(2), + "minReadySeconds": discovery_min_ready_seconds, + "strategy": discovery_strategy, "selector": {"matchLabels": labels_for(&labels, "discovery")}, "template": { - "metadata": {"labels": labels_for(&labels, "discovery"), "annotations": {"rustqueue.io/revision": revision}}, + "metadata": { + "labels": discovery_pod_labels, + "annotations": {"rustqueue.io/revision": revision} + }, "spec": { "serviceAccountName": service_account_name, "securityContext": { @@ -232,11 +306,7 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { "imagePullPolicy": cluster.spec.image_pull_policy, "command": ["rustqueue-discovery"], "ports": [{"name": "http", "containerPort": 4161}], - "env": [ - {"name": "POD_NAMESPACE", "valueFrom": {"fieldRef": {"fieldPath": "metadata.namespace"}}}, - {"name": "RUSTQUEUE_BROKER_SERVICE", "value": broker_service_name}, - {"name": "RUSTQUEUE_REGISTRY_TOKEN_FILE", "value": "/run/secrets/rustqueue/registry-token"} - ], + "env": discovery_env, "volumeMounts": [{"name": "auth", "mountPath": "/run/secrets/rustqueue", "readOnly": true}], "securityContext": { "allowPrivilegeEscalation": false, @@ -244,9 +314,16 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { "capabilities": {"drop": ["ALL"]} }, "readinessProbe": {"httpGet": {"path": "/v1/health", "port": "http"}, "periodSeconds": 2}, + "livenessProbe": {"httpGet": {"path": "/ping", "port": "http"}, "periodSeconds": 10, "failureThreshold": 3}, "resources": {"requests": {"cpu": "50m", "memory": "64Mi"}} }], - "volumes": [{"name": "auth", "secret": {"secretName": input.secret_name}}] + "volumes": [{ + "name": "auth", + "secret": { + "secretName": input.secret_name, + "items": [{"key": "registry-token", "path": "registry-token"}] + } + }] } } } @@ -277,6 +354,7 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { "metadata": {"labels": labels_for(&labels, "proxy"), "annotations": {"rustqueue.io/revision": revision}}, "spec": { "nodeSelector": cluster.spec.proxy_node_selector, + "terminationGracePeriodSeconds": 45, "securityContext": { "runAsNonRoot": true, "runAsUser": 65532, "runAsGroup": 65532, "seccompProfile": {"type": "RuntimeDefault"} @@ -288,7 +366,8 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { "ports": [{"name": "tcp", "containerPort": 4150}, {"name": "http", "containerPort": 4151}], "env": [ {"name": "RUSTQUEUE_DISCOVERY_URLS", "value": format!("http://{discovery_name}:4161")}, - {"name": "RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS", "value": cluster.spec.proxy_tcp_max_connection_age_seconds.to_string()} + {"name": "RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS", "value": cluster.spec.proxy_tcp_max_connection_age_seconds.to_string()}, + {"name": "RUSTQUEUE_PROXY_SHUTDOWN_GRACE_SECONDS", "value": "30"} ], "securityContext": { "allowPrivilegeEscalation": false, @@ -296,24 +375,103 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { "capabilities": {"drop": ["ALL"]} }, "readinessProbe": {"httpGet": {"path": "/v1/health", "port": "http"}, "periodSeconds": 2}, + "livenessProbe": {"httpGet": {"path": "/ping", "port": "http"}, "periodSeconds": 10, "failureThreshold": 3}, "resources": {"requests": {"cpu": "50m", "memory": "64Mi"}} }] } } } }))?; + let kodo_resources = (cluster.spec.kodo_compatibility.enabled + || input.retain_existing_kodo_resources) + .then(|| { + kodo::build(kodo::BuildInput { + cluster, + replicas: input.kodo_gateway_replicas, + image: input.image, + secret_name: input.secret_name, + revision: &revision, + name: &name, + namespace: &namespace, + owner: owner.clone(), + labels: &labels, + discovery_name: &discovery_name, + service_account_name: &service_account_name, + }) + }) + .transpose()?; + let ( + kodo_gateway_service, + kodo_gateway_headless_service, + kodo_gateway, + kodo_gateway_pdb, + kodo_gateway_network_policy, + ) = kodo_resources.map_or((None, None, None, None, None), |resources| { + ( + Some(resources.service), + Some(resources.headless_service), + Some(resources.gateway), + Some(resources.pdb), + Some(resources.network_policy), + ) + }); + let mut broker_ingress = if cluster.spec.kodo_compatibility.enabled { + vec![ + json!({ + "from": [ + {"podSelector": {"matchLabels": {"app.kubernetes.io/instance": name}}}, + kodo::allowed_peer(cluster) + ], + "ports": [{"protocol": "TCP", "port": 4150}] + }), + json!({ + "from": [{"namespaceSelector": {}}], + "ports": [{"protocol": "TCP", "port": 4151}] + }), + ] + } else { + vec![json!({ + "from": [{"namespaceSelector": {}}], + "ports": [ + {"protocol": "TCP", "port": 4150}, + {"protocol": "TCP", "port": 4151} + ] + })] + }; + if cluster.spec.kodo_compatibility.effective_cleanup_enabled() { + broker_ingress.push(json!({ + "from": [kodo::allowed_peer(cluster)], + "ports": [{"protocol": "TCP", "port": 4152}] + })); + } + let broker_network_policy = typed(json!({ + "apiVersion": "networking.k8s.io/v1", "kind": "NetworkPolicy", + "metadata": metadata(&format!("{name}-broker-ingress"), "broker"), + "spec": { + "podSelector": {"matchLabels": labels_for(&labels, "broker")}, + "policyTypes": ["Ingress"], + "ingress": broker_ingress + } + }))?; let network_policy = typed(json!({ "apiVersion": "networking.k8s.io/v1", "kind": "NetworkPolicy", "metadata": metadata(&format!("{name}-internal"), "runtime"), "spec": { - "podSelector": {"matchLabels": {"app.kubernetes.io/instance": name}}, + "podSelector": { + "matchLabels": {"app.kubernetes.io/instance": name}, + "matchExpressions": [{ + "key": "app.kubernetes.io/component", + "operator": "NotIn", + "values": ["broker", "kodo-gateway"] + }] + }, "policyTypes": ["Ingress"], "ingress": [{"from": [{"namespaceSelector": {}}]}] } }))?; - Ok(ResourceSet { revision, + retain_existing_kodo_resources: input.retain_existing_kodo_resources, config, service_account, role, @@ -321,18 +479,62 @@ pub fn build(input: BuildInput<'_>) -> anyhow::Result { broker_service, brokers, broker_pdb, + broker_network_policy, discovery_service, discovery, discovery_pdb, proxy_service, proxy, + kodo_gateway_service, + kodo_gateway_headless_service, + kodo_gateway, + kodo_gateway_pdb, + kodo_gateway_network_policy, network_policy, }) } fn broker_config(cluster: &RustQueue, secret_name: &str) -> String { + let kodo = cluster.spec.kodo_compatibility.enabled; + let large_messages = cluster.spec.max_message_bytes > 64 * 1024 * 1024; + let max_segment_bytes = if kodo || large_messages { + 256 * 1024 * 1024 + } else { + 100 * 1024 * 1024 + }; + let max_body_bytes = if kodo { + 128 * 1024 * 1024 + } else { + (64 * 1024 * 1024).max(cluster.spec.max_message_bytes) + }; + let connection_publish_inflight_bytes = if kodo || large_messages { + 160 * 1024 * 1024 + } else { + 80 * 1024 * 1024 + }; + let node_publish_inflight_bytes = if kodo || large_messages { + 1024 * 1024 * 1024 + } else { + 512 * 1024 * 1024 + }; + let cleanup_enabled = cluster.spec.kodo_compatibility.effective_cleanup_enabled(); + let kodo_cleanup_token = if cleanup_enabled { + "kodo_cleanup_token_file = \"/run/secrets/rustqueue/kodo-cleanup-token\"\n" + } else { + "" + }; + let publish_token = if kodo { + "publish_token_file = \"/run/secrets/rustqueue/admin-token\"\n" + } else { + "" + }; + let kodo_network = if cleanup_enabled { + "[network]\nkodo_http_address = \"0.0.0.0:4152\"\n\n" + } else { + "" + }; let mut output = format!( - "[storage]\ndata_path = \"/data\"\nfeature_level = {}\nmin_free_bytes = {}\ndisk_high_watermark_percent = {}\ndisk_low_watermark_percent = {}\nprotective_eviction_enabled = {}\ndisk_pressure_grace_seconds = {}\nmaintenance_startup_delay_seconds = {}\n\n[queue]\nbootstrap_retention_seconds = {}\nmax_message_bytes = {}\nmax_topics = {}\nmax_publish_workers = {}\npublish_worker_idle_seconds = {}\n\n[limits]\nnode_delivery_inflight_bytes = {}\nconnection_delivery_inflight_bytes = {}\n\n[metrics]\ndetailed_queue_metrics = {}\nmax_detailed_series = {}\n\n[security]\nadmin_token_file = \"/run/secrets/rustqueue/admin-token\"\nregistry_token_file = \"/run/secrets/rustqueue/registry-token\"\nconsole_token_file = \"/run/secrets/rustqueue/console-token\"\nconsole_management_enabled = {}\n# secret: {secret_name}\n", + "{kodo_network}[storage]\ndata_path = \"/data\"\nfeature_level = {}\nmax_segment_bytes = {max_segment_bytes}\nmin_free_bytes = {}\ndisk_high_watermark_percent = {}\ndisk_low_watermark_percent = {}\nprotective_eviction_enabled = {}\ndisk_pressure_grace_seconds = {}\nmaintenance_startup_delay_seconds = {}\n\n[queue]\nbootstrap_retention_seconds = {}\nmax_message_bytes = {}\nmax_topics = {}\nmax_publish_workers = {}\npublish_worker_idle_seconds = {}\n\n[limits]\nmax_body_bytes = {max_body_bytes}\nnode_publish_inflight_bytes = {node_publish_inflight_bytes}\nconnection_publish_inflight_bytes = {connection_publish_inflight_bytes}\nnode_delivery_inflight_bytes = {}\nconnection_delivery_inflight_bytes = {}\ndisconnect_on_retriable_publish_error = false\n\n[metrics]\ndetailed_queue_metrics = {}\nmax_detailed_series = {}\n\n[security]\nadmin_token_file = \"/run/secrets/rustqueue/admin-token\"\n{publish_token}registry_token_file = \"/run/secrets/rustqueue/registry-token\"\nconsole_token_file = \"/run/secrets/rustqueue/console-token\"\n{kodo_cleanup_token}console_management_enabled = {}\nkodo_cleanup_enabled = {}\n# secret: {secret_name}\n", cluster.spec.storage_feature_level, cluster.spec.min_free_bytes, cluster.spec.disk_high_watermark_percent, @@ -350,6 +552,7 @@ fn broker_config(cluster: &RustQueue, secret_name: &str) -> String { cluster.spec.detailed_queue_metrics, cluster.spec.max_detailed_metric_series, cluster.spec.console_management_enabled, + cleanup_enabled, ); if cluster.spec.client_tls_secret_name.is_some() { output.push_str("\n[security.tls]\ncertificate_file = \"/run/tls/rustqueue/tls.crt\"\nprivate_key_file = \"/run/tls/rustqueue/tls.key\"\nclient_ca_file = \"/run/tls/rustqueue/ca.crt\"\nrequire_client_certificate = false\nrequired = false\n"); @@ -440,7 +643,7 @@ mod tests { min_free_bytes: 1024, disk_high_watermark_percent: 85, disk_low_watermark_percent: 75, - protective_eviction_enabled: true, + protective_eviction_enabled: false, disk_pressure_grace_seconds: 60, bootstrap_retention_seconds: 90, max_message_bytes: 20 * 1024 * 1024, @@ -455,6 +658,7 @@ mod tests { proxy_node_selector: BTreeMap::new(), proxy_tcp_max_connection_age_seconds: 300, discovery_replicas: 2, + kodo_compatibility: crate::crd::KodoCompatibility::default(), maintenance: None, rollout: crate::crd::RolloutPolicy::default(), broker_scheduling: crate::crd::BrokerScheduling::default(), @@ -464,11 +668,33 @@ mod tests { } } + fn enable_kodo(cluster: &mut RustQueue, cleanup_enabled: bool) { + cluster.spec.min_brokers = 3; + cluster.spec.max_brokers = 3; + cluster.spec.storage_feature_level = 2; + cluster.spec.bootstrap_retention_seconds = 180; + cluster.spec.max_message_bytes = 100 * 1024 * 1024; + cluster.spec.connection_delivery_inflight_bytes = 128 * 1024 * 1024; + cluster.spec.node_delivery_inflight_bytes = 512 * 1024 * 1024; + cluster.spec.kodo_compatibility.enabled = true; + cluster.spec.kodo_compatibility.cleanup_enabled = cleanup_enabled; + cluster + .spec + .kodo_compatibility + .allowed_pod_selector + .insert("app.kubernetes.io/name".into(), "kodo".into()); + } + #[test] fn renders_one_share_nothing_statefulset_and_retained_pvcs() { let resources = build(BuildInput { cluster: &cluster(), replicas: 3, + kodo_gateway_replicas: 0, + advertise_kodo_gateways: false, + discovery_service_kodo: Some(false), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: false, image: "rustqueue:test", claim_template_size: "100Gi", secret_name: "queue-auth", @@ -489,6 +715,24 @@ mod tests { && config.contains("connection_delivery_inflight_bytes = 33554432") && !config.contains("message_index_cache_bytes") })); + assert!(resources.kodo_gateway.is_none()); + assert!(resources.kodo_gateway_service.is_none()); + assert!(resources.kodo_gateway_headless_service.is_none()); + assert!(resources.kodo_gateway_network_policy.is_none()); + let policy = serde_json::to_value(&resources.network_policy).unwrap(); + assert_eq!( + policy["spec"]["podSelector"]["matchExpressions"][0]["values"], + json!(["broker", "kodo-gateway"]) + ); + let broker_policy = serde_json::to_value(&resources.broker_network_policy).unwrap(); + assert_eq!( + broker_policy["spec"]["ingress"].as_array().unwrap().len(), + 1 + ); + assert_eq!( + broker_policy["spec"]["ingress"][0]["from"], + json!([{"namespaceSelector": {}}]) + ); let spec = resources.brokers.spec.unwrap(); assert_eq!(spec.replicas, Some(3)); assert_eq!(spec.service_name.as_deref(), Some("queue-brokers")); @@ -498,6 +742,39 @@ mod tests { resources.broker_pdb.spec.unwrap().min_available, Some(k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::Int(2)) ); + assert_eq!( + resources + .discovery + .spec + .as_ref() + .unwrap() + .strategy + .as_ref() + .unwrap() + .type_ + .as_deref(), + Some("RollingUpdate") + ); + let discovery_value = serde_json::to_value(&resources.discovery).unwrap(); + assert_eq!(discovery_value["spec"]["minReadySeconds"], 30); + assert_eq!( + discovery_value["spec"]["template"]["spec"]["containers"][0]["livenessProbe"] + ["httpGet"]["path"], + "/ping" + ); + assert_eq!( + discovery_value["spec"]["strategy"]["rollingUpdate"]["maxSurge"], + "100%" + ); + assert_eq!( + discovery_value["spec"]["template"]["metadata"]["labels"][DISCOVERY_MODE_LABEL], + "direct" + ); + let discovery_service = serde_json::to_value(&resources.discovery_service).unwrap(); + assert_eq!( + discovery_service["spec"]["selector"][DISCOVERY_MODE_LABEL], + "direct" + ); assert_eq!( spec.template .spec @@ -524,9 +801,11 @@ mod tests { assert!(resources .proxy .spec + .as_ref() .unwrap() .template .spec + .as_ref() .unwrap() .containers .first() @@ -535,5 +814,614 @@ mod tests { variable.name == "RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS" && variable.value.as_deref() == Some("300") }))); + let proxy_value = serde_json::to_value(&resources.proxy).unwrap(); + assert_eq!( + proxy_value["spec"]["template"]["spec"]["containers"][0]["livenessProbe"]["httpGet"] + ["path"], + "/ping" + ); + assert_eq!( + resources + .proxy + .spec + .as_ref() + .unwrap() + .template + .spec + .as_ref() + .unwrap() + .termination_grace_period_seconds, + Some(45) + ); + } + + #[test] + fn kodo_profile_renders_three_stable_gateways_and_large_message_limits() { + let mut cluster = cluster(); + enable_kodo(&mut cluster, false); + let mut resources = build(BuildInput { + cluster: &cluster, + replicas: 3, + kodo_gateway_replicas: 3, + advertise_kodo_gateways: true, + discovery_service_kodo: Some(true), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: false, + image: "rustqueue:test", + claim_template_size: "100Gi", + secret_name: "queue-auth", + mounted_secret_revision: "1", + }) + .unwrap(); + let config = resources + .config + .data + .as_ref() + .unwrap() + .get("rustqueue.toml") + .unwrap(); + assert!(config.contains("feature_level = 2")); + assert!(config.contains("bootstrap_retention_seconds = 180")); + assert!(config.contains("max_message_bytes = 104857600")); + assert!(config.contains("max_segment_bytes = 268435456")); + assert!(config.contains("connection_publish_inflight_bytes = 167772160")); + assert!(config.contains("node_publish_inflight_bytes = 1073741824")); + assert!(config.contains("disconnect_on_retriable_publish_error = false")); + assert!(config.contains("publish_token_file = \"/run/secrets/rustqueue/admin-token\"")); + assert!(config.contains("kodo_cleanup_enabled = false")); + assert!(!config.contains("kodo_cleanup_token_file")); + assert!(!config.contains("kodo_http_address")); + assert_eq!( + resources + .discovery + .spec + .as_ref() + .unwrap() + .strategy + .as_ref() + .unwrap() + .type_ + .as_deref(), + Some("RollingUpdate") + ); + assert_eq!( + resources.discovery.spec.as_ref().unwrap().min_ready_seconds, + Some(660) + ); + assert_eq!( + resources + .discovery + .spec + .as_ref() + .unwrap() + .template + .metadata + .as_ref() + .unwrap() + .labels + .as_ref() + .unwrap() + .get(DISCOVERY_MODE_LABEL) + .map(String::as_str), + Some("kodo") + ); + assert_eq!( + resources + .discovery_service + .spec + .as_ref() + .unwrap() + .selector + .as_ref() + .unwrap() + .get(DISCOVERY_MODE_LABEL) + .map(String::as_str), + Some("kodo") + ); + + let gateway = resources.kodo_gateway.take().unwrap(); + let gateway_policy = resources.kodo_gateway_network_policy.take().unwrap(); + let gateway_policy = serde_json::to_value(gateway_policy).unwrap(); + assert_eq!( + gateway_policy["spec"]["ingress"][0]["from"][0], + json!({"podSelector": { + "matchLabels": {"app.kubernetes.io/name": "kodo"} + }}) + ); + assert_eq!( + gateway_policy["spec"]["ingress"][0]["ports"] + .as_array() + .unwrap() + .iter() + .map(|port| port["port"].as_u64().unwrap()) + .collect::>(), + vec![4150, 4152, 4153, 4151, 4154, 4155] + ); + assert_eq!( + gateway_policy["spec"]["ingress"][1], + json!({ + "from": [{"namespaceSelector": {}}], + "ports": [{"protocol": "TCP", "port": 4160}] + }) + ); + let broker_policy = serde_json::to_value(&resources.broker_network_policy).unwrap(); + assert_eq!( + broker_policy["spec"]["ingress"].as_array().unwrap().len(), + 2 + ); + assert_eq!( + broker_policy["spec"]["ingress"][0]["from"], + json!([ + {"podSelector": { + "matchLabels": {"app.kubernetes.io/instance": "queue"} + }}, + {"podSelector": { + "matchLabels": {"app.kubernetes.io/name": "kodo"} + }} + ]) + ); + assert_eq!( + broker_policy["spec"]["ingress"][0]["ports"], + json!([{"protocol": "TCP", "port": 4150}]) + ); + assert_eq!( + broker_policy["spec"]["ingress"][1], + json!({ + "from": [{"namespaceSelector": {}}], + "ports": [{"protocol": "TCP", "port": 4151}] + }) + ); + assert_eq!(gateway.spec.as_ref().unwrap().replicas, Some(3)); + assert_eq!( + gateway.spec.as_ref().unwrap().service_name.as_deref(), + Some("queue-kodo-gateways") + ); + let gateway_service = resources.kodo_gateway_service.take().unwrap(); + assert_eq!( + gateway_service.metadata.name.as_deref(), + Some("queue-kodo-publish") + ); + assert_eq!( + gateway_service + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get("rustqueue.io/metrics")) + .map(String::as_str), + Some("true") + ); + let gateway_service = gateway_service.spec.unwrap(); + assert_ne!(gateway_service.cluster_ip.as_deref(), Some("None")); + assert_eq!( + gateway_service + .ports + .unwrap() + .into_iter() + .map(|port| (port.name.unwrap(), port.port, port.target_port.unwrap())) + .collect::>(), + vec![ + ( + "tcp-0".into(), + 4150, + k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::String("tcp".into()) + ), + ( + "tcp-1".into(), + 4152, + k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::String("tcp".into()) + ), + ( + "tcp-2".into(), + 4153, + k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::String("tcp".into()) + ), + ( + "http-0".into(), + 4151, + k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::String( + "http-0".into() + ) + ), + ( + "http-1".into(), + 4154, + k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::String( + "http-1".into() + ) + ), + ( + "http-2".into(), + 4155, + k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::String( + "http-2".into() + ) + ), + ( + "metrics".into(), + 4160, + k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::String( + "metrics".into() + ) + ), + ] + ); + let headless_service = resources.kodo_gateway_headless_service.take().unwrap(); + assert_eq!( + headless_service.metadata.name.as_deref(), + Some("queue-kodo-gateways") + ); + let headless_service = headless_service.spec.unwrap(); + assert_eq!(headless_service.cluster_ip.as_deref(), Some("None")); + assert_eq!(headless_service.publish_not_ready_addresses, Some(true)); + let gateway_value = serde_json::to_value(&gateway).unwrap(); + assert_eq!( + gateway_value["spec"]["template"]["spec"]["terminationGracePeriodSeconds"], + 75 + ); + assert_eq!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["livenessProbe"]["httpGet"] + ["path"], + "/ping" + ); + assert_eq!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["resources"]["limits"] + ["memory"], + "1Gi" + ); + assert_eq!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["resources"]["requests"] + ["memory"], + "768Mi" + ); + assert_eq!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["resources"]["requests"] + ["cpu"], + "1" + ); + assert_eq!( + gateway_value["spec"]["template"]["spec"]["affinity"]["podAntiAffinity"] + ["requiredDuringSchedulingIgnoredDuringExecution"][0]["topologyKey"], + "kubernetes.io/hostname" + ); + assert!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["env"] + .as_array() + .unwrap() + .iter() + .any(|variable| { + variable["name"] == "RUSTQUEUE_PROXY_MAX_MESSAGE_BYTES" + && variable["value"] == "104857600" + }) + ); + assert!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["env"] + .as_array() + .unwrap() + .iter() + .any(|variable| { + variable["name"] == "RUSTQUEUE_PROXY_SHUTDOWN_GRACE_SECONDS" + && variable["value"] == "60" + }) + ); + assert!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["env"] + .as_array() + .unwrap() + .iter() + .any(|variable| { + variable["name"] == "RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES" + && variable["value"] == "536870912" + }) + ); + assert!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["env"] + .as_array() + .unwrap() + .iter() + .any(|variable| { + variable["name"] == "RUSTQUEUE_PROXY_TCP_COMMAND_TIMEOUT_MS" + && variable["value"] == "120000" + }) + ); + assert!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["env"] + .as_array() + .unwrap() + .iter() + .any(|variable| { + variable["name"] == "RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS" + && variable["value"] == "0" + }) + ); + assert!( + gateway_value["spec"]["template"]["spec"]["containers"][0]["env"] + .as_array() + .unwrap() + .iter() + .all(|variable| { + variable["name"] != "RUSTQUEUE_KODO_CLEANUP_TOKEN_FILE" + && variable["name"] != "RUSTQUEUE_REGISTRY_TOKEN_FILE" + }) + ); + assert_eq!( + resources + .kodo_gateway_pdb + .take() + .unwrap() + .spec + .unwrap() + .min_available, + Some(k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::Int(2)) + ); + let discovery_env = resources + .discovery + .spec + .take() + .unwrap() + .template + .spec + .unwrap() + .containers + .remove(0) + .env + .unwrap(); + assert!(discovery_env.iter().any(|item| { + item.name == "RUSTQUEUE_KODO_GATEWAY_ADDRESS" + && item.value.as_deref() == Some("queue-kodo-publish.test.svc") + })); + assert!(discovery_env.iter().any(|item| { + item.name == "RUSTQUEUE_KODO_CLEANUP_ENABLED" && item.value.as_deref() == Some("false") + })); + } + + #[test] + fn staged_kodo_gateways_fail_closed_without_mounting_cleanup_credentials() { + let mut cluster = cluster(); + enable_kodo(&mut cluster, false); + let mut resources = build(BuildInput { + cluster: &cluster, + replicas: 3, + kodo_gateway_replicas: 0, + advertise_kodo_gateways: false, + discovery_service_kodo: Some(false), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: false, + image: "rustqueue:test", + claim_template_size: "100Gi", + secret_name: "queue-auth", + mounted_secret_revision: "1", + }) + .unwrap(); + let discovery = resources.discovery.spec.as_ref().unwrap(); + assert_eq!( + discovery + .template + .metadata + .as_ref() + .unwrap() + .labels + .as_ref() + .unwrap() + .get(DISCOVERY_MODE_LABEL) + .map(String::as_str), + Some("direct") + ); + assert_eq!( + resources + .discovery_service + .spec + .as_ref() + .unwrap() + .selector + .as_ref() + .unwrap() + .get(DISCOVERY_MODE_LABEL) + .map(String::as_str), + Some("direct") + ); + let discovery_env = discovery + .template + .spec + .as_ref() + .unwrap() + .containers + .first() + .unwrap() + .env + .as_ref() + .unwrap(); + assert!(discovery_env + .iter() + .all(|variable| variable.name != "RUSTQUEUE_KODO_COMPATIBILITY_ENABLED")); + assert!(discovery_env + .iter() + .all(|variable| variable.name != "RUSTQUEUE_KODO_GATEWAY_ADDRESS")); + let gateway = resources.kodo_gateway.take().unwrap(); + assert_eq!(gateway.spec.as_ref().unwrap().replicas, Some(0)); + let pod = gateway.spec.unwrap().template.spec.unwrap(); + assert!(pod.volumes.unwrap().is_empty()); + let container = pod.containers.first().unwrap(); + assert!(container.volume_mounts.as_ref().unwrap().is_empty()); + assert!(container.env.as_ref().unwrap().iter().all(|variable| { + variable.name != "RUSTQUEUE_KODO_CLEANUP_TOKEN_FILE" + && variable.name != "RUSTQUEUE_REGISTRY_TOKEN_FILE" + })); + assert_eq!( + resources + .kodo_gateway_pdb + .take() + .unwrap() + .spec + .unwrap() + .min_available, + Some(k8s_openapi::apimachinery::pkg::util::intstr::IntOrString::Int(0)) + ); + } + + #[test] + fn requested_kodo_cleanup_remains_disabled_in_rendered_resources() { + let mut cluster = cluster(); + enable_kodo(&mut cluster, true); + let resources = build(BuildInput { + cluster: &cluster, + replicas: 3, + kodo_gateway_replicas: 3, + advertise_kodo_gateways: true, + discovery_service_kodo: Some(true), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: false, + image: "rustqueue:test", + claim_template_size: "100Gi", + secret_name: "queue-auth", + mounted_secret_revision: "1", + }) + .unwrap(); + let config = resources + .config + .data + .as_ref() + .unwrap() + .get("rustqueue.toml") + .unwrap(); + assert!(config.contains("kodo_cleanup_enabled = false")); + assert!(!config.contains("kodo_http_address")); + let discovery_env = resources + .discovery + .spec + .as_ref() + .unwrap() + .template + .spec + .as_ref() + .unwrap() + .containers[0] + .env + .as_ref() + .unwrap(); + assert!(discovery_env.iter().any(|variable| { + variable.name == "RUSTQUEUE_KODO_CLEANUP_ENABLED" + && variable.value.as_deref() == Some("false") + })); + } + + #[test] + fn disabling_kodo_keeps_the_gateway_recoverable_without_cleanup() { + let mut cluster = cluster(); + cluster.spec.kodo_compatibility.cleanup_enabled = true; + let resources = build(BuildInput { + cluster: &cluster, + replicas: 3, + kodo_gateway_replicas: 3, + advertise_kodo_gateways: false, + discovery_service_kodo: Some(false), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: true, + image: "rustqueue:test", + claim_template_size: "100Gi", + secret_name: "queue-auth", + mounted_secret_revision: "1", + }) + .unwrap(); + let config = resources + .config + .data + .as_ref() + .unwrap() + .get("rustqueue.toml") + .unwrap(); + assert!(config.contains("kodo_cleanup_enabled = false")); + assert!(!config.contains("kodo_cleanup_token_file")); + assert!(!config.contains("kodo_http_address")); + assert_eq!( + resources + .kodo_gateway + .as_ref() + .and_then(|gateway| gateway.spec.as_ref()) + .and_then(|spec| spec.replicas), + Some(3) + ); + let gateway_env = resources + .kodo_gateway + .as_ref() + .unwrap() + .spec + .as_ref() + .unwrap() + .template + .spec + .as_ref() + .unwrap() + .containers[0] + .env + .as_ref() + .unwrap(); + assert!(gateway_env.iter().all(|variable| { + variable.name != "RUSTQUEUE_KODO_CLEANUP_TOKEN_FILE" + && variable.name != "RUSTQUEUE_REGISTRY_TOKEN_FILE" + })); + assert!(resources.retain_existing_kodo_resources); + assert_eq!( + resources.discovery.spec.as_ref().unwrap().min_ready_seconds, + Some(660) + ); + } + + #[test] + fn standalone_large_message_profile_has_matching_publish_limits() { + let mut cluster = cluster(); + cluster.spec.storage_feature_level = 2; + cluster.spec.max_message_bytes = 100 * 1024 * 1024; + cluster.spec.connection_delivery_inflight_bytes = 100 * 1024 * 1024; + cluster.spec.node_delivery_inflight_bytes = 200 * 1024 * 1024; + let resources = build(BuildInput { + cluster: &cluster, + replicas: 3, + kodo_gateway_replicas: 0, + advertise_kodo_gateways: false, + discovery_service_kodo: Some(false), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: false, + image: "rustqueue:test", + claim_template_size: "100Gi", + secret_name: "queue-auth", + mounted_secret_revision: "1", + }) + .unwrap(); + let mut data = resources.config.data.unwrap(); + let config = data.remove("rustqueue.toml").unwrap(); + assert!(config.contains("max_body_bytes = 104857600")); + assert!(config.contains("connection_publish_inflight_bytes = 167772160")); + assert!(config.contains("max_segment_bytes = 268435456")); + assert!(resources.kodo_gateway.is_none()); + assert!(resources.kodo_gateway_service.is_none()); + assert!(resources.kodo_gateway_headless_service.is_none()); + } + + #[test] + fn broker_pod_template_changes_advance_the_rollout_revision() { + let render = |cluster: &RustQueue, secret_name: &str| { + build(BuildInput { + cluster, + replicas: 3, + kodo_gateway_replicas: 0, + advertise_kodo_gateways: false, + discovery_service_kodo: Some(false), + activate_kodo_cleanup: false, + retain_existing_kodo_resources: false, + image: "rustqueue:test", + claim_template_size: "100Gi", + secret_name, + mounted_secret_revision: "same-resource-version", + }) + .unwrap() + .revision + }; + let mut cluster = cluster(); + let original = render(&cluster, "queue-auth"); + cluster.spec.broker_resources.cpu_request = "2".into(); + let resources_changed = render(&cluster, "queue-auth"); + assert_ne!(original, resources_changed); + assert_ne!(resources_changed, render(&cluster, "replacement-auth")); + assert_eq!(original.len(), 32); } } diff --git a/crates/operator/src/resources/kodo.rs b/crates/operator/src/resources/kodo.rs new file mode 100644 index 0000000..4397374 --- /dev/null +++ b/crates/operator/src/resources/kodo.rs @@ -0,0 +1,246 @@ +use super::{labels_for, typed}; +use crate::RustQueue; +use k8s_openapi::api::apps::v1::StatefulSet; +use k8s_openapi::api::core::v1::Service; +use k8s_openapi::api::networking::v1::NetworkPolicy; +use k8s_openapi::api::policy::v1::PodDisruptionBudget; +use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; +use serde_json::{json, Value}; +use std::collections::BTreeMap; + +pub(super) struct KodoResources { + pub service: Service, + pub headless_service: Service, + pub gateway: StatefulSet, + pub pdb: PodDisruptionBudget, + pub network_policy: NetworkPolicy, +} + +pub(super) struct BuildInput<'a> { + pub cluster: &'a RustQueue, + pub replicas: i32, + pub image: &'a str, + pub secret_name: &'a str, + pub revision: &'a str, + pub name: &'a str, + pub namespace: &'a str, + pub owner: OwnerReference, + pub labels: &'a BTreeMap, + pub discovery_name: &'a str, + pub service_account_name: &'a str, +} + +pub(super) fn build(input: BuildInput<'_>) -> anyhow::Result { + let gateway_name = format!("{}-kodo-gateway", input.name); + let publish_service_name = format!("{}-kodo-publish", input.name); + let headless_service_name = format!("{}-kodo-gateways", input.name); + let gateway_labels = labels_for(input.labels, "kodo-gateway"); + let allowed_peer = allowed_peer(input.cluster); + let metadata = |name: &str| { + json!({ + "name": name, + "namespace": input.namespace, + "labels": gateway_labels, + "ownerReferences": [input.owner], + }) + }; + let mut publish_metadata = metadata(&publish_service_name); + publish_metadata["labels"]["rustqueue.io/metrics"] = json!("true"); + let service = typed(json!({ + "apiVersion": "v1", "kind": "Service", + "metadata": publish_metadata, + "spec": { + "selector": gateway_labels, + "ports": [ + {"name": "tcp-0", "port": 4150, "targetPort": "tcp"}, + {"name": "tcp-1", "port": 4152, "targetPort": "tcp"}, + {"name": "tcp-2", "port": 4153, "targetPort": "tcp"}, + {"name": "http-0", "port": 4151, "targetPort": "http-0"}, + {"name": "http-1", "port": 4154, "targetPort": "http-1"}, + {"name": "http-2", "port": 4155, "targetPort": "http-2"}, + {"name": "metrics", "port": 4160, "targetPort": "metrics"} + ] + } + }))?; + let headless_service = typed(json!({ + "apiVersion": "v1", "kind": "Service", + "metadata": metadata(&headless_service_name), + "spec": { + "clusterIP": "None", + "publishNotReadyAddresses": true, + "selector": gateway_labels, + "ports": [ + {"name": "tcp", "port": 4150, "targetPort": "tcp"}, + {"name": "http-0", "port": 4151, "targetPort": "http-0"} + ] + } + }))?; + let mut env = vec![ + json!({"name": "POD_NAME", "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}}}), + json!({"name": "RUSTQUEUE_DISCOVERY_URLS", "value": format!("http://{}:4161", input.discovery_name)}), + json!({"name": "RUSTQUEUE_KODO_COMPATIBILITY_ENABLED", "value": "true"}), + json!({"name": "RUSTQUEUE_KODO_CLEANUP_ENABLED", "value": input.cluster.spec.kodo_compatibility.effective_cleanup_enabled().to_string()}), + json!({"name": "RUSTQUEUE_PROXY_MAX_MESSAGE_BYTES", "value": (100 * 1024 * 1024).to_string()}), + json!({"name": "RUSTQUEUE_PROXY_MAX_BODY_BYTES", "value": (128 * 1024 * 1024).to_string()}), + json!({"name": "RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES", "value": (512 * 1024 * 1024).to_string()}), + json!({"name": "RUSTQUEUE_PROXY_TCP_COMMAND_TIMEOUT_MS", "value": "120000"}), + json!({"name": "RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS", "value": "0"}), + json!({"name": "RUSTQUEUE_PROXY_SHUTDOWN_GRACE_SECONDS", "value": "60"}), + ]; + let (volume_mounts, volumes) = if input + .cluster + .spec + .kodo_compatibility + .effective_cleanup_enabled() + { + env.extend([ + json!({"name": "RUSTQUEUE_KODO_CLEANUP_TOKEN_FILE", "value": "/run/secrets/rustqueue/kodo-cleanup-token"}), + json!({"name": "RUSTQUEUE_REGISTRY_TOKEN_FILE", "value": "/run/secrets/rustqueue/registry-token"}), + ]); + ( + vec![json!({ + "name": "auth", + "mountPath": "/run/secrets/rustqueue", + "readOnly": true + })], + vec![json!({ + "name": "auth", + "secret": { + "secretName": input.secret_name, + "items": [ + {"key": "kodo-cleanup-token", "path": "kodo-cleanup-token"}, + {"key": "registry-token", "path": "registry-token"} + ] + } + })], + ) + } else { + (Vec::new(), Vec::new()) + }; + let gateway = typed(json!({ + "apiVersion": "apps/v1", "kind": "StatefulSet", + "metadata": metadata(&gateway_name), + "spec": { + "serviceName": headless_service_name, + "replicas": input.replicas, + "podManagementPolicy": "Parallel", + "selector": {"matchLabels": gateway_labels}, + "template": { + "metadata": { + "labels": gateway_labels, + "annotations": {"rustqueue.io/revision": input.revision} + }, + "spec": { + "serviceAccountName": input.service_account_name, + "terminationGracePeriodSeconds": 75, + "nodeSelector": input.cluster.spec.proxy_node_selector, + "affinity": {"podAntiAffinity": {"requiredDuringSchedulingIgnoredDuringExecution": [{ + "labelSelector": {"matchLabels": gateway_labels}, + "topologyKey": "kubernetes.io/hostname" + }]}}, + "securityContext": { + "runAsNonRoot": true, "runAsUser": 65532, "runAsGroup": 65532, + "seccompProfile": {"type": "RuntimeDefault"} + }, + "containers": [{ + "name": "gateway", "image": input.image, + "imagePullPolicy": input.cluster.spec.image_pull_policy, + "command": ["rustqueue-proxy"], + "ports": [ + {"name": "tcp", "containerPort": 4150}, + {"name": "http-0", "containerPort": 4151}, + {"name": "http-1", "containerPort": 4154}, + {"name": "http-2", "containerPort": 4155}, + {"name": "metrics", "containerPort": 4160} + ], + "env": env, + "volumeMounts": volume_mounts, + "securityContext": { + "allowPrivilegeEscalation": false, + "readOnlyRootFilesystem": true, + "capabilities": {"drop": ["ALL"]} + }, + "readinessProbe": { + "httpGet": {"path": "/v1/health", "port": "http-0"}, + "periodSeconds": 2, + "failureThreshold": 2 + }, + "livenessProbe": { + "httpGet": {"path": "/ping", "port": "http-0"}, + "periodSeconds": 10, + "failureThreshold": 3 + }, + "resources": { + "requests": {"cpu": "1", "memory": "768Mi"}, + "limits": {"memory": "1Gi"} + } + }], + "volumes": volumes + } + } + } + }))?; + let pdb = typed(json!({ + "apiVersion": "policy/v1", "kind": "PodDisruptionBudget", + "metadata": metadata(&format!("{gateway_name}-pdb")), + "spec": { + "minAvailable": if input.replicas == 0 { 0 } else { 2 }, + "selector": {"matchLabels": gateway_labels} + } + }))?; + let network_policy = typed(json!({ + "apiVersion": "networking.k8s.io/v1", "kind": "NetworkPolicy", + "metadata": metadata(&format!("{}-kodo-gateway-ingress", input.name)), + "spec": { + "podSelector": {"matchLabels": gateway_labels}, + "policyTypes": ["Ingress"], + "ingress": [ + { + "from": [allowed_peer], + "ports": [ + {"protocol": "TCP", "port": 4150}, + {"protocol": "TCP", "port": 4152}, + {"protocol": "TCP", "port": 4153}, + {"protocol": "TCP", "port": 4151}, + {"protocol": "TCP", "port": 4154}, + {"protocol": "TCP", "port": 4155} + ] + }, + { + "from": [{"namespaceSelector": {}}], + "ports": [{"protocol": "TCP", "port": 4160}] + } + ] + } + }))?; + Ok(KodoResources { + service, + headless_service, + gateway, + pdb, + network_policy, + }) +} + +pub(super) fn allowed_peer(cluster: &RustQueue) -> Value { + let allowed_pod_selector = cluster + .spec + .kodo_compatibility + .effective_allowed_pod_selector(); + let mut peer = json!({ + "podSelector": { + "matchLabels": allowed_pod_selector + } + }); + if !cluster + .spec + .kodo_compatibility + .allowed_namespace_selector + .is_empty() + { + peer["namespaceSelector"] = json!({ + "matchLabels": cluster.spec.kodo_compatibility.allowed_namespace_selector + }); + } + peer +} diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index d2dc273..0c1b11e 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -15,3 +15,5 @@ pub const HEARTBEAT: &[u8] = b"_heartbeat_"; pub const OK: &[u8] = b"OK"; pub const CLOSE_WAIT: &[u8] = b"CLOSE_WAIT"; pub const MESSAGE_ID_LEN: usize = 16; +pub const MAX_MESSAGE_BYTES: usize = 100 * 1024 * 1024; +pub const MAX_BATCH_BYTES: usize = 128 * 1024 * 1024; diff --git a/crates/proxy/Cargo.toml b/crates/proxy/Cargo.toml index f00fada..ba1a129 100644 --- a/crates/proxy/Cargo.toml +++ b/crates/proxy/Cargo.toml @@ -16,6 +16,7 @@ parking_lot.workspace = true rand.workspace = true reqwest.workspace = true rustls.workspace = true +rustqueue-protocol.workspace = true rustqueue-telemetry.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/crates/proxy/src/backend.rs b/crates/proxy/src/backend.rs index eb781db..fc63c71 100644 --- a/crates/proxy/src/backend.rs +++ b/crates/proxy/src/backend.rs @@ -4,7 +4,7 @@ use serde::Deserialize; use std::collections::BTreeMap; use std::ops::Deref; use std::sync::Arc; -use std::time::{Duration, Instant}; +use tokio::sync::watch; #[derive(Clone, Debug, Deserialize, PartialEq, Eq)] pub struct Backend { @@ -16,7 +16,11 @@ pub struct Backend { impl Backend { pub fn tcp_address(&self) -> String { - format!("{}:{}", self.broadcast_address, self.tcp_port) + if self.broadcast_address.contains(':') && !self.broadcast_address.starts_with('[') { + format!("[{}]:{}", self.broadcast_address, self.tcp_port) + } else { + format!("{}:{}", self.broadcast_address, self.tcp_port) + } } pub fn http_origin(&self) -> String { if self.broadcast_address.contains(':') && !self.broadcast_address.starts_with('[') { @@ -36,13 +40,14 @@ pub struct BackendPool { struct PoolState { backends: Vec, active_connections: BTreeMap, + invalidation: BTreeMap>, next_tie: usize, - updated_at: Option, } pub struct BackendLease { backend: Backend, pool: BackendPool, + invalidation: watch::Receiver, } impl Deref for BackendLease { @@ -77,6 +82,18 @@ impl BackendPool { unique.insert(backend.node_id, backend); } let mut state = self.inner.write(); + let next: BTreeMap<_, _> = unique.iter().map(|(id, backend)| (*id, backend)).collect(); + let changed: Vec<_> = state + .backends + .iter() + .filter(|backend| next.get(&backend.node_id).copied() != Some(*backend)) + .map(|backend| backend.node_id) + .collect(); + for node_id in changed { + if let Some(sender) = state.invalidation.get(&node_id) { + sender.send_modify(|generation| *generation = generation.wrapping_add(1)); + } + } state.backends = unique.into_values().collect(); let node_ids: std::collections::BTreeSet<_> = state .backends @@ -86,16 +103,14 @@ impl BackendPool { state .active_connections .retain(|node_id, active| *active > 0 || node_ids.contains(node_id)); - state.updated_at = Some(Instant::now()); - } - - pub fn clear_if_stale(&self, maximum_age: Duration) { - let mut state = self.inner.write(); - if state - .updated_at - .is_none_or(|updated| updated.elapsed() > maximum_age) - { - state.backends.clear(); + state + .invalidation + .retain(|node_id, _| node_ids.contains(node_id)); + for node_id in node_ids { + state + .invalidation + .entry(node_id) + .or_insert_with(|| watch::channel(0).0); } } @@ -135,21 +150,49 @@ impl BackendPool { let backend = candidates[state.next_tie % candidates.len()].clone(); state.next_tie = state.next_tie.wrapping_add(1); *state.active_connections.entry(backend.node_id).or_default() += 1; + let invalidation = state + .invalidation + .entry(backend.node_id) + .or_insert_with(|| watch::channel(0).0) + .subscribe(); Some(BackendLease { backend, pool: self.clone(), + invalidation, }) } + pub fn all(&self) -> Vec { + self.inner.read().backends.clone() + } + pub fn len(&self) -> usize { self.inner.read().backends.len() } } +impl BackendLease { + pub fn invalidation(&self) -> watch::Receiver { + self.invalidation.clone() + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn backend_addresses_bracket_ipv6_literals() { + let backend = Backend { + broadcast_address: "2001:db8::1".into(), + tcp_port: 4150, + http_port: 4151, + node_id: 1, + }; + assert_eq!(backend.tcp_address(), "[2001:db8::1]:4150"); + assert_eq!(backend.http_origin(), "http://[2001:db8::1]:4151"); + } + #[test] fn pool_deduplicates_five_hundred_discovered_brokers() { let pool = BackendPool::default(); @@ -205,4 +248,19 @@ mod tests { let second = pool.lease().unwrap().node_id; assert_ne!(first, second); } + + #[tokio::test] + async fn removing_a_backend_invalidates_its_active_lease() { + let pool = BackendPool::default(); + pool.replace(vec![Backend { + broadcast_address: "broker-1".into(), + tcp_port: 4150, + http_port: 4151, + node_id: 1, + }]); + let lease = pool.lease().unwrap(); + let mut invalidation = lease.invalidation(); + pool.replace(Vec::new()); + assert!(invalidation.changed().await.is_ok() || invalidation.has_changed().is_err()); + } } diff --git a/crates/proxy/src/discovery.rs b/crates/proxy/src/discovery.rs index 427901e..75af806 100644 --- a/crates/proxy/src/discovery.rs +++ b/crates/proxy/src/discovery.rs @@ -7,6 +7,7 @@ use std::time::{Duration, Instant}; const MAX_HEAD_BYTES: usize = 64 * 1024; const MAX_NODES_BYTES: usize = 4 * 1024 * 1024; const FULL_REFRESH_INTERVAL: Duration = Duration::from_secs(3); +const ROUTING_STALE_AFTER: Duration = Duration::from_secs(5); #[derive(Deserialize)] struct NodesResponse { @@ -23,12 +24,14 @@ struct HeadResponse { struct Source { revision: u64, producers: Vec, + brokers: Vec, seen_at: Instant, full_at: Instant, } pub async fn run( - pool: BackendPool, + publish_pool: BackendPool, + broker_pool: BackendPool, addresses: Vec, metrics: ProxyMetrics, ) -> anyhow::Result<()> { @@ -38,6 +41,7 @@ pub async fn run( .redirect(reqwest::redirect::Policy::none()) .build()?; let mut sources = BTreeMap::::new(); + let mut last_coherent_at = None; let mut interval = tokio::time::interval(Duration::from_secs(2)); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { @@ -46,20 +50,58 @@ pub async fn run( for address in &addresses { refresh_source(&client, address, &mut sources).await; } - sources.retain(|_, source| source.seen_at.elapsed() <= Duration::from_secs(5)); - if sources.is_empty() { - pool.clear_if_stale(Duration::from_secs(5)); - } else { - pool.replace( - sources - .values() - .flat_map(|source| source.producers.iter().cloned()) - .collect(), - ); + sources.retain(|_, source| source.seen_at.elapsed() <= ROUTING_STALE_AFTER); + match coherent_sources(&sources) { + Some((producers, brokers, observed_at)) => { + publish_pool.replace(producers); + broker_pool.replace(brokers); + last_coherent_at = Some(observed_at); + } + None if last_coherent_is_fresh(last_coherent_at, Instant::now()) => { + tracing::debug!( + "retaining the last coherent routing snapshot during Discovery revision skew" + ); + } + None => { + if !sources.is_empty() { + tracing::debug!( + "active discovery replicas returned different routing revisions" + ); + } + publish_pool.replace(Vec::new()); + broker_pool.replace(Vec::new()); + } } } } +fn coherent_sources( + sources: &BTreeMap, +) -> Option<(Vec, Vec, Instant)> { + let revision = sources.values().next()?.revision; + if sources.values().any(|source| source.revision != revision) { + return None; + } + let observed_at = sources.values().map(|source| source.seen_at).min()?; + Some(( + sources + .values() + .flat_map(|source| source.producers.iter().cloned()) + .collect(), + sources + .values() + .flat_map(|source| source.brokers.iter().cloned()) + .collect(), + observed_at, + )) +} + +fn last_coherent_is_fresh(last_coherent_at: Option, now: Instant) -> bool { + last_coherent_at.is_some_and(|observed_at| { + now.saturating_duration_since(observed_at) <= ROUTING_STALE_AFTER + }) +} + async fn refresh_source( client: &reqwest::Client, address: &str, @@ -95,44 +137,72 @@ async fn refresh_source( return; } } - let url = format!("{address}/v1/publishers"); - let nodes = match client.get(url).send().await { - Ok(response) => match response.error_for_status() { - Ok(response) => { - match read_json_bounded::(response, MAX_NODES_BYTES).await { - Ok(nodes) => nodes, - Err(error) => { - tracing::debug!(%error, "discovery response was invalid"); - return; - } - } - } - Err(error) => { - tracing::debug!(%error, "discovery returned an error"); - return; - } - }, + let nodes = match fetch_nodes(client, &format!("{address}/v1/publishers")).await { + Ok(nodes) => nodes, Err(error) => { - tracing::debug!(%error, "discovery request failed"); + tracing::debug!(%error, "discovery publishers response was invalid"); return; } }; + let broker_nodes = match fetch_optional_nodes(client, &format!("{address}/v1/brokers")).await { + Ok(nodes) => nodes, + Err(error) => { + tracing::debug!(%error, "discovery brokers response was invalid"); + return; + } + }; + let Some(revision) = coherent_revision(head.as_ref(), &nodes, broker_nodes.as_ref()) else { + tracing::debug!("discovery routing snapshot revisions did not match"); + return; + }; + let brokers = broker_nodes + .map(|nodes| nodes.producers) + .unwrap_or_else(|| nodes.producers.clone()); sources.insert( address.to_owned(), Source { - revision: head.map_or(nodes.revision, |head| head.revision), + revision, producers: nodes.producers, + brokers, seen_at: now, full_at: now, }, ); } +async fn fetch_nodes(client: &reqwest::Client, url: &str) -> anyhow::Result { + let response = client.get(url).send().await?.error_for_status()?; + read_json_bounded(response, MAX_NODES_BYTES).await +} + +async fn fetch_optional_nodes( + client: &reqwest::Client, + url: &str, +) -> anyhow::Result> { + let response = client.get(url).send().await?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + Ok(Some( + read_json_bounded(response.error_for_status()?, MAX_NODES_BYTES).await?, + )) +} + fn source_is_current(source: &Source, revision: u64, now: Instant) -> bool { source.revision == revision && now.saturating_duration_since(source.full_at) < FULL_REFRESH_INTERVAL } +fn coherent_revision( + head: Option<&HeadResponse>, + publishers: &NodesResponse, + brokers: Option<&NodesResponse>, +) -> Option { + let revision = head.map_or(publishers.revision, |head| head.revision); + (publishers.revision == revision && brokers.is_none_or(|brokers| brokers.revision == revision)) + .then_some(revision) +} + async fn read_json_bounded( mut response: reqwest::Response, maximum: usize, @@ -161,6 +231,7 @@ mod tests { Source { revision: 7, producers: Vec::new(), + brokers: Vec::new(), seen_at: now, full_at: now, } @@ -179,4 +250,55 @@ mod tests { assert!(!source_is_current(&source, 7, now + FULL_REFRESH_INTERVAL)); assert!(!source_is_current(&source, 8, now)); } + + #[test] + fn mixed_discovery_revisions_are_rejected() { + let publishers = NodesResponse { + revision: 7, + producers: Vec::new(), + }; + let brokers = NodesResponse { + revision: 8, + producers: Vec::new(), + }; + assert_eq!( + coherent_revision( + Some(&HeadResponse { revision: 7 }), + &publishers, + Some(&brokers) + ), + None + ); + assert_eq!( + coherent_revision(Some(&HeadResponse { revision: 7 }), &publishers, None), + Some(7) + ); + } + + #[test] + fn active_discovery_replicas_must_agree_before_routing() { + let now = Instant::now(); + let mut sources = BTreeMap::from([ + ("first".into(), source(now)), + ("second".into(), source(now)), + ]); + assert!(coherent_sources(&sources).is_some()); + + sources.get_mut("second").unwrap().revision = 8; + assert!(coherent_sources(&sources).is_none()); + + sources.remove("second"); + assert!(coherent_sources(&sources).is_some()); + } + + #[test] + fn last_coherent_routing_only_bridges_bounded_revision_skew() { + let now = Instant::now(); + assert!(last_coherent_is_fresh(Some(now), now + ROUTING_STALE_AFTER)); + assert!(!last_coherent_is_fresh( + Some(now), + now + ROUTING_STALE_AFTER + Duration::from_millis(1) + )); + assert!(!last_coherent_is_fresh(None, now)); + } } diff --git a/crates/proxy/src/http.rs b/crates/proxy/src/http.rs index d59aec8..d3e3b94 100644 --- a/crates/proxy/src/http.rs +++ b/crates/proxy/src/http.rs @@ -1,8 +1,9 @@ use crate::backend::BackendPool; +use crate::kodo::{self, KodoConfig}; use crate::metrics::ProxyMetrics; use axum::body::{Body, Bytes}; use axum::extract::State; -use axum::http::{header, Request, StatusCode}; +use axum::http::{header, Method, Request, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing::{any, get}; use axum::{Json, Router}; @@ -10,27 +11,38 @@ use rustqueue_proxy::{parse_forward_metadata, ForwardMetadataError}; use serde_json::json; use std::net::SocketAddr; use std::sync::Arc; -use tokio::sync::Semaphore; +use tokio::sync::{watch, Semaphore}; const MAX_BACKEND_RESPONSE_BYTES: usize = 16 * 1024 * 1024; +const KODO_STATS_HTTP_PORTS: [u16; 3] = [4151, 4154, 4155]; +const KODO_METRICS_HTTP_PORT: u16 = 4160; #[derive(Clone)] struct ProxyState { pool: BackendPool, + broker_pool: BackendPool, client: reqwest::Client, max_body_bytes: usize, body_timeout: std::time::Duration, inflight_bytes: Arc, metrics: ProxyMetrics, + kodo: Option, +} + +pub struct Limits { + pub max_body_bytes: usize, + pub inflight_bytes: Arc, + pub body_timeout: std::time::Duration, } pub async fn serve( address: SocketAddr, pool: BackendPool, - max_body_bytes: usize, - max_inflight_bytes: usize, - body_timeout: std::time::Duration, + broker_pool: BackendPool, + limits: Limits, metrics: ProxyMetrics, + kodo: Option, + shutdown: watch::Receiver, ) -> anyhow::Result<()> { let client = reqwest::Client::builder() .connect_timeout(std::time::Duration::from_millis(500)) @@ -39,25 +51,159 @@ pub async fn serve( .build()?; let state = ProxyState { pool, + broker_pool, client, - max_body_bytes, - body_timeout, - inflight_bytes: Arc::new(Semaphore::new(max_inflight_bytes)), + max_body_bytes: limits.max_body_bytes, + body_timeout: limits.body_timeout, + inflight_bytes: limits.inflight_bytes, metrics, + kodo, }; + if state.kodo.is_some() { + if address.port() != KODO_STATS_HTTP_PORTS[0] { + anyhow::bail!("Kodo Gateway HTTP address must use port 4151"); + } + let first = state_for_stats_shard(&state, 0); + let second = state_for_stats_shard(&state, 1); + let third = state_for_stats_shard(&state, 2); + let mut second_address = address; + second_address.set_port(KODO_STATS_HTTP_PORTS[1]); + let mut third_address = address; + third_address.set_port(KODO_STATS_HTTP_PORTS[2]); + let mut metrics_address = address; + metrics_address.set_port(KODO_METRICS_HTTP_PORT); + let metrics = state.clone(); + tokio::try_join!( + serve_one(address, first, shutdown.clone()), + serve_one(second_address, second, shutdown.clone()), + serve_one(third_address, third, shutdown.clone()), + serve_metrics(metrics_address, metrics, shutdown), + )?; + return Ok(()); + } + serve_one(address, state, shutdown).await +} + +fn state_for_stats_shard(state: &ProxyState, ordinal: usize) -> ProxyState { + let mut state = state.clone(); + state.kodo.as_mut().expect("Kodo state was checked").ordinal = ordinal; + state +} + +async fn serve_one( + address: SocketAddr, + state: ProxyState, + shutdown: watch::Receiver, +) -> anyhow::Result<()> { let router = Router::new() + .route("/ping", get(ping)) .route("/v1/health", get(health)) .route("/metrics", get(prometheus)) + .route("/stats", get(kodo_stats)) + .route("/channel/delete", axum::routing::post(kodo_delete_channel)) .fallback(any(forward)) .with_state(state); let listener = tokio::net::TcpListener::bind(address).await?; tracing::info!(%address, "producer HTTP proxy listening"); - axum::serve(listener, router).await?; + axum::serve(listener, router) + .with_graceful_shutdown(wait_for_shutdown(shutdown)) + .await?; + Ok(()) +} + +async fn serve_metrics( + address: SocketAddr, + state: ProxyState, + shutdown: watch::Receiver, +) -> anyhow::Result<()> { + let router = Router::new() + .route("/ping", get(ping)) + .route("/metrics", get(prometheus)) + .with_state(state); + let listener = tokio::net::TcpListener::bind(address).await?; + tracing::info!(%address, "Kodo Gateway metrics listening"); + axum::serve(listener, router) + .with_graceful_shutdown(wait_for_shutdown(shutdown)) + .await?; Ok(()) } +async fn wait_for_shutdown(mut shutdown: watch::Receiver) { + if *shutdown.borrow() { + return; + } + while shutdown.changed().await.is_ok() { + if *shutdown.borrow() { + return; + } + } +} + +async fn kodo_stats(State(state): State, request: Request) -> Response { + let Some(config) = &state.kodo else { + return forward(State(state), request).await; + }; + kodo::stats( + config, + &state.broker_pool, + &state.client, + &request.uri().to_string(), + ) + .await +} + +#[derive(serde::Deserialize)] +struct ChannelQuery { + topic: String, + channel: String, +} + +async fn kodo_delete_channel(State(state): State, request: Request) -> Response { + let Some(config) = &state.kodo else { + return forward(State(state), request).await; + }; + let Ok(axum::extract::Query(query)) = + axum::extract::Query::::try_from_uri(request.uri()) + else { + return ( + StatusCode::BAD_REQUEST, + "E_BAD_REQUEST topic and channel are required", + ) + .into_response(); + }; + kodo::delete_channel( + config, + &state.broker_pool, + &state.client, + &query.topic, + &query.channel, + ) + .await +} + async fn prometheus(State(state): State) -> String { - state.metrics.render() + let mut output = state.metrics.render(); + output.push_str( + "# HELP rustqueue_proxy_publish_backends Number of publish-ready Brokers known to the proxy.\n\ +# TYPE rustqueue_proxy_publish_backends gauge\n\ +# HELP rustqueue_proxy_broker_backends Number of Brokers known to the proxy for compatibility reads.\n\ +# TYPE rustqueue_proxy_broker_backends gauge\n\ +# HELP rustqueue_proxy_kodo_gateway Whether this process terminates the Kodo producer protocol.\n\ +# TYPE rustqueue_proxy_kodo_gateway gauge\n", + ); + output.push_str(&format!( + "rustqueue_proxy_publish_backends {}\n\ +rustqueue_proxy_broker_backends {}\n\ +rustqueue_proxy_kodo_gateway {}\n", + state.pool.len(), + state.broker_pool.len(), + u8::from(state.kodo.is_some()), + )); + output +} + +async fn ping() -> &'static str { + "OK" } async fn health(State(state): State) -> Response { @@ -75,6 +221,13 @@ async fn health(State(state): State) -> Response { } async fn forward(State(state): State, request: Request) -> Response { + if state.kodo.is_some() { + return ( + StatusCode::NOT_FOUND, + "E_NOT_FOUND Kodo Gateway HTTP forwarding is disabled", + ) + .into_response(); + } let (parts, body) = request.into_parts(); let content_length = parts .headers @@ -114,7 +267,8 @@ async fn forward(State(state): State, request: Request) -> Res Err(_) => return body_timeout(), }; let path = metadata.path_and_query; - let backends = state.pool.shuffled(2); + let retry_safe = retry_safe_method(&parts.method); + let backends = state.pool.shuffled(if retry_safe { 2 } else { 1 }); if backends.is_empty() { return unavailable(); } @@ -132,24 +286,47 @@ async fn forward(State(state): State, request: Request) -> Res } } match outgoing.send().await { - Ok(response) => return backend_response(response).await, + Ok(response) => { + return backend_response(response, retry_safe, &state.metrics).await; + } Err(error) => last_error = Some(error), } } if let Some(error) = last_error { - tracing::debug!(%error, "all producer HTTP backends failed"); + tracing::debug!(%error, "producer HTTP backend failed"); + if !retry_safe { + state + .metrics + .producer_ambiguous_failures + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return ambiguous(); + } } unavailable() } -async fn backend_response(response: reqwest::Response) -> Response { +fn retry_safe_method(method: &Method) -> bool { + method == Method::GET || method == Method::HEAD || method == Method::OPTIONS +} + +async fn backend_response( + response: reqwest::Response, + retry_safe: bool, + metrics: &ProxyMetrics, +) -> Response { let status = response.status(); let headers = response.headers().clone(); let body = match read_body_bounded(response, MAX_BACKEND_RESPONSE_BYTES).await { Some(body) => body, None => { tracing::debug!("backend response body failed or exceeded its limit"); - return unavailable(); + if retry_safe { + return unavailable(); + } + metrics + .producer_ambiguous_failures + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return ambiguous(); } }; let mut output = Response::builder().status(status); @@ -193,6 +370,14 @@ fn unavailable() -> Response { response } +fn ambiguous() -> Response { + ( + StatusCode::BAD_GATEWAY, + "E_AMBIGUOUS backend may have committed the request; automatic retry is unsafe", + ) + .into_response() +} + fn body_too_large() -> Response { ( StatusCode::PAYLOAD_TOO_LARGE, @@ -225,17 +410,72 @@ fn throttled() -> Response { #[cfg(test)] mod tests { use super::*; + use crate::backend::Backend; use tower::ServiceExt; + #[test] + fn kodo_stats_ports_override_the_pod_ordinal() { + let state = ProxyState { + pool: BackendPool::default(), + broker_pool: BackendPool::default(), + client: reqwest::Client::new(), + max_body_bytes: 1024, + body_timeout: std::time::Duration::from_secs(1), + inflight_bytes: Arc::new(Semaphore::new(1024)), + metrics: ProxyMetrics::default(), + kodo: Some(KodoConfig { + ordinal: 2, + cleanup_enabled: false, + cleanup_token: None, + registry_token: None, + }), + }; + assert_eq!(state_for_stats_shard(&state, 0).kodo.unwrap().ordinal, 0); + assert_eq!(state_for_stats_shard(&state, 1).kodo.unwrap().ordinal, 1); + } + + #[test] + fn mutating_http_requests_never_fail_over_after_an_ambiguous_send() { + assert!(!retry_safe_method(&Method::POST)); + assert!(!retry_safe_method(&Method::PUT)); + assert!(retry_safe_method(&Method::GET)); + assert!(retry_safe_method(&Method::HEAD)); + } + + #[tokio::test] + async fn ping_does_not_depend_on_backend_health() { + assert_eq!(ping().await, "OK"); + } + + #[tokio::test] + async fn metrics_expose_backend_inventory() { + let state = ProxyState { + pool: BackendPool::default(), + broker_pool: BackendPool::default(), + client: reqwest::Client::new(), + max_body_bytes: 1024, + body_timeout: std::time::Duration::from_secs(1), + inflight_bytes: Arc::new(Semaphore::new(1024)), + metrics: ProxyMetrics::default(), + kodo: None, + }; + let body = prometheus(State(state)).await; + assert!(body.contains("rustqueue_proxy_publish_backends 0\n")); + assert!(body.contains("rustqueue_proxy_broker_backends 0\n")); + assert!(body.contains("rustqueue_proxy_kodo_gateway 0\n")); + } + #[tokio::test] async fn no_backend_is_retryable() { let state = ProxyState { pool: BackendPool::default(), + broker_pool: BackendPool::default(), client: reqwest::Client::new(), max_body_bytes: 1024, body_timeout: std::time::Duration::from_secs(1), inflight_bytes: Arc::new(Semaphore::new(1024)), metrics: ProxyMetrics::default(), + kodo: None, }; let response = Router::new() .fallback(any(forward)) @@ -252,15 +492,49 @@ mod tests { assert_eq!(response.headers().get(header::RETRY_AFTER).unwrap(), "1"); } + #[tokio::test] + async fn kodo_gateway_does_not_expose_the_http_publish_proxy() { + let state = ProxyState { + pool: BackendPool::default(), + broker_pool: BackendPool::default(), + client: reqwest::Client::new(), + max_body_bytes: 1024, + body_timeout: std::time::Duration::from_secs(1), + inflight_bytes: Arc::new(Semaphore::new(1024)), + metrics: ProxyMetrics::default(), + kodo: Some(KodoConfig { + ordinal: 0, + cleanup_enabled: false, + cleanup_token: None, + registry_token: None, + }), + }; + let response = Router::new() + .fallback(any(forward)) + .with_state(state) + .oneshot( + Request::builder() + .method("POST") + .uri("/pub?topic=x") + .body(Body::from("message")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + #[tokio::test] async fn rejects_a_request_that_exceeds_the_node_inflight_budget() { let state = ProxyState { pool: BackendPool::default(), + broker_pool: BackendPool::default(), client: reqwest::Client::new(), max_body_bytes: 1024, body_timeout: std::time::Duration::from_secs(1), inflight_bytes: Arc::new(Semaphore::new(1)), metrics: ProxyMetrics::default(), + kodo: None, }; let response = Router::new() .fallback(any(forward)) @@ -278,4 +552,55 @@ mod tests { assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); assert_eq!(response.headers().get(header::RETRY_AFTER).unwrap(), "1"); } + + #[tokio::test] + async fn non_kodo_channel_delete_is_forwarded_to_the_backend() { + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let address = listener.local_addr().unwrap(); + let backend = tokio::spawn(async move { + axum::serve( + listener, + Router::new().route( + "/channel/delete", + axum::routing::post(|| async { StatusCode::ACCEPTED }), + ), + ) + .await + .unwrap(); + }); + let pool = BackendPool::default(); + pool.replace(vec![Backend { + broadcast_address: address.ip().to_string(), + tcp_port: 4150, + http_port: address.port(), + node_id: 1, + }]); + let state = ProxyState { + pool, + broker_pool: BackendPool::default(), + client: reqwest::Client::new(), + max_body_bytes: 1024, + body_timeout: std::time::Duration::from_secs(1), + inflight_bytes: Arc::new(Semaphore::new(1024)), + metrics: ProxyMetrics::default(), + kodo: None, + }; + let response = Router::new() + .route("/channel/delete", axum::routing::post(kodo_delete_channel)) + .fallback(any(forward)) + .with_state(state) + .oneshot( + Request::builder() + .method("POST") + .uri("/channel/delete?topic=events&channel=workers") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::ACCEPTED); + backend.abort(); + } } diff --git a/crates/proxy/src/kodo.rs b/crates/proxy/src/kodo.rs new file mode 100644 index 0000000..81f7c10 --- /dev/null +++ b/crates/proxy/src/kodo.rs @@ -0,0 +1,571 @@ +use crate::backend::{Backend, BackendPool}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::Json; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +const EXPECTED_BROKERS: usize = 3; +const MAX_STATS_BYTES: usize = 16 * 1024 * 1024; +const MAX_REGISTRY_HEAD_BYTES: usize = 64 * 1024; + +#[derive(Clone)] +pub(crate) struct KodoConfig { + pub ordinal: usize, + pub cleanup_enabled: bool, + pub cleanup_token: Option>, + pub registry_token: Option>, +} + +#[derive(Default, Deserialize, Serialize)] +struct StatsResponse { + #[serde(default)] + version: String, + #[serde(default)] + health: String, + #[serde(default)] + start_time: i64, + #[serde(default)] + topics: Vec, +} + +#[derive(Default, Deserialize, Serialize)] +struct TopicStats { + #[serde(default)] + topic_name: String, + #[serde(default)] + name: String, + #[serde(default)] + depth: u64, + #[serde(default)] + memory_depth: u64, + #[serde(default)] + backend_depth: u64, + #[serde(default)] + message_count: u64, + #[serde(default)] + paused: bool, + #[serde(default)] + channels: Vec, +} + +#[derive(Default, Deserialize, Serialize)] +struct ChannelStats { + #[serde(default)] + channel_name: String, + #[serde(default)] + name: String, + #[serde(default)] + depth: u64, + #[serde(default)] + memory_depth: u64, + #[serde(default)] + backend_depth: u64, + #[serde(default)] + message_count: u64, + #[serde(default)] + in_flight_count: u64, + #[serde(default)] + deferred_count: u64, + #[serde(default)] + requeue_count: u64, + #[serde(default)] + timeout_count: u64, + #[serde(default)] + client_count: usize, + #[serde(default)] + clients: Vec, + #[serde(default)] + paused: bool, +} + +#[derive(Deserialize)] +struct Registry { + revision: u64, + node_id: u64, +} + +pub(crate) async fn stats( + config: &KodoConfig, + pool: &BackendPool, + client: &reqwest::Client, + path_and_query: &str, +) -> Response { + let backends = sharded_backends(pool, config.ordinal); + if backends.is_empty() { + return unavailable("gateway stats shard has no broker"); + } + + let mut aggregate = StatsResponse { + version: env!("CARGO_PKG_VERSION").into(), + health: "OK".into(), + start_time: i64::MAX, + topics: Vec::new(), + }; + let path = force_json_stats(path_and_query); + for backend in backends { + let response = match client + .get(format!("{}{path}", backend.http_origin())) + .send() + .await + .and_then(reqwest::Response::error_for_status) + { + Ok(response) => response, + Err(error) => { + tracing::warn!(%error, node_id = backend.node_id, "Kodo stats backend failed"); + return unavailable("broker stats are unavailable"); + } + }; + let stats: StatsResponse = match read_json_bounded(response, MAX_STATS_BYTES).await { + Ok(stats) => stats, + Err(error) => { + tracing::warn!(%error, node_id = backend.node_id, "Kodo stats backend was invalid"); + return unavailable("broker stats are invalid"); + } + }; + if stats.start_time > 0 { + aggregate.start_time = aggregate.start_time.min(stats.start_time); + } + merge_topics(&mut aggregate.topics, stats.topics); + } + if aggregate.start_time == i64::MAX { + aggregate.start_time = unix_seconds(); + } + Json(aggregate).into_response() +} + +pub(crate) async fn delete_channel( + config: &KodoConfig, + pool: &BackendPool, + client: &reqwest::Client, + topic: &str, + channel: &str, +) -> Response { + if !config.cleanup_enabled { + return ( + StatusCode::NOT_FOUND, + "E_NOT_FOUND Kodo cleanup compatibility is disabled", + ) + .into_response(); + } + if topic.is_empty() || channel.is_empty() { + return ( + StatusCode::BAD_REQUEST, + "E_BAD_REQUEST topic and channel are required", + ) + .into_response(); + } + let (Some(cleanup_token), Some(registry_token)) = ( + config.cleanup_token.as_deref(), + config.registry_token.as_deref(), + ) else { + return unavailable("Kodo cleanup tokens are unavailable"); + }; + let Some(backend) = complete_sharded_backend(pool, config.ordinal) else { + return unavailable("complete broker inventory is required for cleanup"); + }; + + if let Err(response) = delete_from_backend( + client, + cleanup_token, + registry_token, + &backend, + topic, + channel, + ) + .await + { + return response; + } + "OK".into_response() +} + +async fn delete_from_backend( + client: &reqwest::Client, + cleanup_token: &str, + registry_token: &str, + backend: &Backend, + topic: &str, + channel: &str, +) -> Result<(), Response> { + let registry = client + .get(format!("{}/v1/registry/head", backend.http_origin())) + .bearer_auth(registry_token) + .send() + .await + .and_then(reqwest::Response::error_for_status) + .map_err(|error| { + tracing::warn!(%error, node_id = backend.node_id, "Kodo cleanup registry read failed"); + unavailable("broker registry is unavailable") + })?; + let registry: Registry = read_json_bounded(registry, MAX_REGISTRY_HEAD_BYTES) + .await + .map_err(|error| { + tracing::warn!(%error, node_id = backend.node_id, "Kodo cleanup registry was invalid"); + unavailable("broker registry is invalid") + })?; + if registry.node_id != backend.node_id { + tracing::warn!( + expected_node_id = backend.node_id, + actual_node_id = registry.node_id, + "Kodo cleanup registry identity changed" + ); + return Err(unavailable("broker registry identity changed")); + } + let operation_id = format!( + "kodo-delete-{:016x}-{:016x}-{:016x}", + registry.node_id, + registry.revision, + stable_hash(topic, channel) + ); + let response = client + .post(format!( + "{}/v1/manage/channels/delete-if-idle", + backend.http_origin() + )) + .bearer_auth(cleanup_token) + .json(&json!({ + "operation_id": operation_id, + "topic": topic, + "channel": channel, + "expected_revision": registry.revision, + })) + .send() + .await + .map_err(|error| { + tracing::warn!(%error, node_id = backend.node_id, "Kodo cleanup request failed"); + unavailable("broker cleanup is unavailable") + })?; + if response.status().is_success() { + return Ok(()); + } + let status = response.status(); + let detail = response + .text() + .await + .unwrap_or_else(|_| "broker cleanup failed".into()); + let outward = if status == StatusCode::CONFLICT { + StatusCode::CONFLICT + } else if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + StatusCode::BAD_GATEWAY + } else { + StatusCode::SERVICE_UNAVAILABLE + }; + Err((outward, detail).into_response()) +} + +fn complete_broker_set(pool: &BackendPool) -> Option> { + let backends = pool.all(); + (backends.len() == EXPECTED_BROKERS).then_some(backends) +} + +fn sharded_backend(pool: &BackendPool, ordinal: usize) -> Option { + let mut matching = sharded_backends(pool, ordinal).into_iter(); + let backend = matching.next()?; + matching.next().is_none().then_some(backend) +} + +fn sharded_backends(pool: &BackendPool, ordinal: usize) -> Vec { + if ordinal >= EXPECTED_BROKERS { + return Vec::new(); + } + pool.all() + .into_iter() + .filter(|backend| { + backend.node_id.saturating_sub(1) % EXPECTED_BROKERS as u64 == ordinal as u64 + }) + .collect() +} + +fn complete_sharded_backend(pool: &BackendPool, ordinal: usize) -> Option { + complete_broker_set(pool)?; + sharded_backend(pool, ordinal) +} + +fn force_json_stats(path_and_query: &str) -> String { + let (path, query) = path_and_query + .split_once('?') + .unwrap_or((path_and_query, "")); + let mut parameters: Vec<_> = query + .split('&') + .filter(|parameter| { + !parameter.is_empty() + && parameter + .split_once('=') + .map_or(*parameter != "format", |(key, _)| key != "format") + }) + .collect(); + parameters.push("format=json"); + format!("{path}?{}", parameters.join("&")) +} + +fn merge_topics(output: &mut Vec, topics: Vec) { + let mut merged: BTreeMap = output + .drain(..) + .map(|topic| (topic.topic_name.clone(), topic)) + .collect(); + for mut topic in topics { + normalize_topic(&mut topic); + match merged.get_mut(&topic.topic_name) { + Some(existing) => merge_topic(existing, topic), + None => { + merged.insert(topic.topic_name.clone(), topic); + } + } + } + *output = merged.into_values().collect(); +} + +fn normalize_topic(topic: &mut TopicStats) { + if topic.topic_name.is_empty() { + topic.topic_name = topic.name.clone(); + } + topic.name = topic.topic_name.clone(); + for channel in &mut topic.channels { + if channel.channel_name.is_empty() { + channel.channel_name = channel.name.clone(); + } + channel.name = channel.channel_name.clone(); + } +} + +fn merge_topic(existing: &mut TopicStats, incoming: TopicStats) { + existing.depth = existing.depth.saturating_add(incoming.depth); + existing.memory_depth = existing.memory_depth.saturating_add(incoming.memory_depth); + existing.backend_depth = existing + .backend_depth + .saturating_add(incoming.backend_depth); + existing.message_count = existing + .message_count + .saturating_add(incoming.message_count); + existing.paused |= incoming.paused; + let mut channels: BTreeMap = existing + .channels + .drain(..) + .map(|channel| (channel.channel_name.clone(), channel)) + .collect(); + for channel in incoming.channels { + match channels.get_mut(&channel.channel_name) { + Some(existing) => merge_channel(existing, channel), + None => { + channels.insert(channel.channel_name.clone(), channel); + } + } + } + existing.channels = channels.into_values().collect(); +} + +fn merge_channel(existing: &mut ChannelStats, mut incoming: ChannelStats) { + existing.depth = existing.depth.saturating_add(incoming.depth); + existing.memory_depth = existing.memory_depth.saturating_add(incoming.memory_depth); + existing.backend_depth = existing + .backend_depth + .saturating_add(incoming.backend_depth); + existing.message_count = existing + .message_count + .saturating_add(incoming.message_count); + existing.in_flight_count = existing + .in_flight_count + .saturating_add(incoming.in_flight_count); + existing.deferred_count = existing + .deferred_count + .saturating_add(incoming.deferred_count); + existing.requeue_count = existing + .requeue_count + .saturating_add(incoming.requeue_count); + existing.timeout_count = existing + .timeout_count + .saturating_add(incoming.timeout_count); + existing.client_count = existing.client_count.saturating_add(incoming.client_count); + existing.clients.append(&mut incoming.clients); + existing.paused |= incoming.paused; +} + +async fn read_json_bounded( + mut response: reqwest::Response, + maximum: usize, +) -> anyhow::Result { + if response + .content_length() + .is_some_and(|length| length > maximum as u64) + { + anyhow::bail!("response body exceeds {maximum} bytes"); + } + let mut bytes = Vec::new(); + while let Some(chunk) = response.chunk().await? { + if bytes.len().saturating_add(chunk.len()) > maximum { + anyhow::bail!("response body exceeds {maximum} bytes"); + } + bytes.extend_from_slice(&chunk); + } + Ok(serde_json::from_slice(&bytes)?) +} + +fn stable_hash(topic: &str, channel: &str) -> u64 { + topic + .bytes() + .chain(std::iter::once(0)) + .chain(channel.bytes()) + .fold(0xcbf29ce484222325, |hash, byte| { + (hash ^ byte as u64).wrapping_mul(0x100000001b3) + }) +} + +fn unix_seconds() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + .min(i64::MAX as u64) as i64 +} + +fn unavailable(detail: &'static str) -> Response { + (StatusCode::SERVICE_UNAVAILABLE, detail).into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::http::{header, HeaderMap}; + use axum::routing::{get, post}; + use axum::Router; + + fn bearer_is(headers: &HeaderMap, token: &str) -> bool { + headers + .get(header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + == Some(token) + } + + async fn registry_with_registry_token(headers: HeaderMap) -> Response { + if !bearer_is(&headers, "registry-secret") { + return StatusCode::UNAUTHORIZED.into_response(); + } + Json(json!({"revision": 7, "node_id": 1})).into_response() + } + + async fn delete_with_cleanup_token(headers: HeaderMap) -> Response { + if !bearer_is(&headers, "cleanup-secret") { + return StatusCode::UNAUTHORIZED.into_response(); + } + StatusCode::OK.into_response() + } + + #[test] + fn stats_merge_keeps_standard_names_and_counters() { + let mut topics = Vec::new(); + merge_topics( + &mut topics, + vec![TopicStats { + name: "events".into(), + message_count: 3, + channels: vec![ChannelStats { + name: "workers".into(), + message_count: 3, + client_count: 1, + ..Default::default() + }], + ..Default::default() + }], + ); + assert_eq!(topics[0].topic_name, "events"); + assert_eq!(topics[0].channels[0].channel_name, "workers"); + assert_eq!(topics[0].channels[0].client_count, 1); + } + + #[test] + fn operation_hash_separates_topic_channel_pairs() { + assert_ne!(stable_hash("a:b", "c"), stable_hash("a", "b:c")); + } + + #[test] + fn stats_requests_always_use_json_without_dropping_filters() { + assert_eq!( + force_json_stats("/stats?format=text&topic=events"), + "/stats?topic=events&format=json" + ); + assert_eq!(force_json_stats("/stats"), "/stats?format=json"); + } + + #[test] + fn gateway_ordinal_shards_every_broker_during_scale_transitions() { + let pool = BackendPool::default(); + pool.replace( + (1..=3) + .map(|node_id| Backend { + broadcast_address: format!("broker-{node_id}"), + tcp_port: 4150, + http_port: 4151, + node_id, + }) + .collect(), + ); + assert_eq!(sharded_backend(&pool, 1).unwrap().node_id, 2); + pool.replace(pool.all().into_iter().take(2).collect()); + assert_eq!(sharded_backend(&pool, 1).unwrap().node_id, 2); + assert!(complete_sharded_backend(&pool, 1).is_none()); + + pool.replace( + [1, 4, 7] + .into_iter() + .map(|node_id| Backend { + broadcast_address: format!("broker-{node_id}"), + tcp_port: 4150, + http_port: 4151, + node_id, + }) + .collect(), + ); + assert_eq!( + sharded_backends(&pool, 0) + .into_iter() + .map(|backend| backend.node_id) + .collect::>(), + vec![1, 4, 7] + ); + assert!(sharded_backend(&pool, 0).is_none()); + } + + #[tokio::test] + async fn cleanup_uses_distinct_registry_and_cleanup_tokens() { + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + axum::serve( + listener, + Router::new() + .route("/v1/registry/head", get(registry_with_registry_token)) + .route( + "/v1/manage/channels/delete-if-idle", + post(delete_with_cleanup_token), + ), + ) + .await + .unwrap(); + }); + let backend = Backend { + broadcast_address: address.ip().to_string(), + tcp_port: 4150, + http_port: address.port(), + node_id: 1, + }; + + assert!(delete_from_backend( + &reqwest::Client::new(), + "cleanup-secret", + "registry-secret", + &backend, + "events", + "workers", + ) + .await + .is_ok()); + task.abort(); + } +} diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index 158c032..8c343ca 100644 --- a/crates/proxy/src/main.rs +++ b/crates/proxy/src/main.rs @@ -1,6 +1,7 @@ mod backend; mod discovery; mod http; +mod kodo; mod metrics; mod tcp; @@ -8,7 +9,11 @@ use backend::BackendPool; use clap::Parser; use metrics::ProxyMetrics; use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; use std::time::Duration; +use tokio::sync::watch; +use tracing::{info, warn}; use tracing_subscriber::EnvFilter; #[derive(Debug, Parser)] @@ -35,6 +40,8 @@ struct Cli { http_address: SocketAddr, #[arg(long, env = "RUSTQUEUE_PROXY_MAX_BODY_BYTES", default_value_t = 64 * 1024 * 1024)] max_body_bytes: usize, + #[arg(long, env = "RUSTQUEUE_PROXY_MAX_MESSAGE_BYTES", default_value_t = 20 * 1024 * 1024)] + max_message_bytes: usize, #[arg( long, env = "RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES", @@ -59,6 +66,32 @@ struct Cli { default_value_t = 300 )] tcp_max_connection_age_seconds: u64, + #[arg( + long, + env = "RUSTQUEUE_PROXY_TCP_COMMAND_TIMEOUT_MS", + default_value_t = 120_000 + )] + tcp_command_timeout_ms: u64, + #[arg( + long, + env = "RUSTQUEUE_PROXY_SHUTDOWN_GRACE_SECONDS", + default_value_t = 30 + )] + shutdown_grace_seconds: u64, + #[arg( + long, + env = "RUSTQUEUE_KODO_COMPATIBILITY_ENABLED", + default_value_t = false + )] + kodo_compatibility_enabled: bool, + #[arg(long, env = "RUSTQUEUE_KODO_GATEWAY_ORDINAL")] + kodo_gateway_ordinal: Option, + #[arg(long, env = "RUSTQUEUE_KODO_CLEANUP_ENABLED", default_value_t = false)] + kodo_cleanup_enabled: bool, + #[arg(long, env = "RUSTQUEUE_KODO_CLEANUP_TOKEN_FILE")] + kodo_cleanup_token_file: Option, + #[arg(long, env = "RUSTQUEUE_REGISTRY_TOKEN_FILE")] + registry_token_file: Option, } #[tokio::main] @@ -70,32 +103,250 @@ async fn main() -> anyhow::Result<()> { ) .init(); let cli = Cli::parse(); + validate_proxy_limits(&cli)?; + let publish_pool = BackendPool::default(); + let broker_pool = BackendPool::default(); + let metrics = ProxyMetrics::default(); + let kodo = build_kodo_config(&cli)?; + let terminate_producer_protocol = kodo.is_some(); + let inflight_bytes = Arc::new(tokio::sync::Semaphore::new(cli.max_inflight_bytes)); + let shutdown_grace = Duration::from_secs(cli.shutdown_grace_seconds); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let discovery = discovery::run( + publish_pool.clone(), + broker_pool.clone(), + cli.discovery_urls, + metrics.clone(), + ); + let tcp = tcp::serve( + cli.tcp_address, + publish_pool.clone(), + tcp::Limits { + max_connections: cli.max_connections, + max_connection_age: Duration::from_secs(cli.tcp_max_connection_age_seconds), + terminate_producer_protocol, + max_message_bytes: cli.max_message_bytes, + max_body_bytes: cli.max_body_bytes, + command_timeout: Duration::from_millis(cli.tcp_command_timeout_ms), + inflight_bytes: Arc::clone(&inflight_bytes), + }, + metrics.clone(), + shutdown_rx.clone(), + shutdown_grace, + ); + let http = http::serve( + cli.http_address, + publish_pool, + broker_pool, + http::Limits { + max_body_bytes: cli.max_body_bytes, + inflight_bytes, + body_timeout: Duration::from_millis(cli.http_body_timeout_ms), + }, + metrics, + kodo, + shutdown_rx, + ); + tokio::pin!(discovery, tcp, http); + let (completed_task, terminal_result) = tokio::select! { + result = &mut discovery => ( + Some(RuntimeTask::Discovery), + unexpected_task_exit("discovery", result), + ), + result = &mut tcp => ( + Some(RuntimeTask::Tcp), + unexpected_task_exit("TCP", result), + ), + result = &mut http => ( + Some(RuntimeTask::Http), + unexpected_task_exit("HTTP", result), + ), + _ = shutdown_signal() => { + info!("proxy shutdown signal received"); + (None, Ok(())) + } + }; + let _ = shutdown_tx.send(true); + let graceful = async { + match completed_task { + Some(RuntimeTask::Tcp) => (&mut http).await, + Some(RuntimeTask::Http) => (&mut tcp).await, + Some(RuntimeTask::Discovery) | None => { + tokio::try_join!(async { (&mut tcp).await }, async { (&mut http).await },)?; + Ok::<(), anyhow::Error>(()) + } + } + }; + let cleanup_result = + match tokio::time::timeout(shutdown_grace + Duration::from_secs(2), graceful).await { + Ok(result) => result, + Err(_) => { + warn!("proxy shutdown grace expired"); + Err(anyhow::anyhow!("proxy shutdown grace expired")) + } + }; + combine_results(terminal_result, cleanup_result) +} + +fn validate_proxy_limits(cli: &Cli) -> anyhow::Result<()> { if cli.max_body_bytes == 0 + || cli.max_message_bytes == 0 + || cli.max_message_bytes > cli.max_body_bytes + || cli.max_message_bytes > rustqueue_protocol::MAX_MESSAGE_BYTES + || cli.max_body_bytes > rustqueue_protocol::MAX_BATCH_BYTES || cli.max_inflight_bytes == 0 + || cli.max_inflight_bytes < cli.max_body_bytes || cli.max_inflight_bytes > u32::MAX as usize || cli.max_connections == 0 || cli.http_body_timeout_ms == 0 + || cli.tcp_command_timeout_ms == 0 + || cli.shutdown_grace_seconds == 0 { - anyhow::bail!("proxy limits must be non-zero and inflight bytes must fit u32"); + anyhow::bail!( + "proxy limits must be non-zero, fit the 100 MiB message and 128 MiB batch contract, and fit the inflight byte budget" + ); } - let pool = BackendPool::default(); - let metrics = ProxyMetrics::default(); + if cli.kodo_compatibility_enabled + && (cli.max_message_bytes != rustqueue_protocol::MAX_MESSAGE_BYTES + || cli.max_body_bytes != rustqueue_protocol::MAX_BATCH_BYTES) + { + anyhow::bail!("Kodo compatibility requires the 100 MiB message and 128 MiB batch limits"); + } + Ok(()) +} + +#[derive(Clone, Copy, Eq, PartialEq)] +enum RuntimeTask { + Discovery, + Tcp, + Http, +} + +fn unexpected_task_exit(name: &str, result: anyhow::Result<()>) -> anyhow::Result<()> { + match result { + Ok(()) => anyhow::bail!("{name} task stopped unexpectedly"), + Err(error) => Err(anyhow::anyhow!("{name} task failed: {error:#}")), + } +} + +fn combine_results( + terminal: anyhow::Result<()>, + cleanup: anyhow::Result<()>, +) -> anyhow::Result<()> { + match (terminal, cleanup) { + (Ok(()), Ok(())) => Ok(()), + (Err(error), Ok(())) | (Ok(()), Err(error)) => Err(error), + (Err(terminal), Err(cleanup)) => Err(anyhow::anyhow!( + "{terminal:#}; shutdown cleanup also failed: {cleanup:#}" + )), + } +} + +async fn shutdown_signal() { + let ctrl_c = async { + if tokio::signal::ctrl_c().await.is_err() { + std::future::pending::<()>().await; + } + }; + #[cfg(unix)] + let terminate = async { + match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) { + Ok(mut signal) => { + signal.recv().await; + } + Err(_) => std::future::pending::<()>().await, + } + }; + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); tokio::select! { - result = discovery::run(pool.clone(), cli.discovery_urls, metrics.clone()) => result, - result = tcp::serve( - cli.tcp_address, - pool.clone(), - cli.max_connections, - Duration::from_secs(cli.tcp_max_connection_age_seconds), - metrics.clone(), - ) => result, - result = http::serve( - cli.http_address, - pool, - cli.max_body_bytes, - cli.max_inflight_bytes, - Duration::from_millis(cli.http_body_timeout_ms), - metrics, - ) => result, + _ = ctrl_c => {} + _ = terminate => {} + } +} + +fn build_kodo_config(cli: &Cli) -> anyhow::Result> { + if cli.kodo_cleanup_enabled { + anyhow::bail!( + "Kodo automatic cleanup is disabled until cluster-wide atomic deletion is available" + ); + } + if !cli.kodo_compatibility_enabled { + if cli.kodo_cleanup_enabled || cli.kodo_gateway_ordinal.is_some() { + anyhow::bail!("Kodo cleanup and gateway ordinal require Kodo compatibility"); + } + return Ok(None); + } + let ordinal = cli + .kodo_gateway_ordinal + .or_else(|| { + std::env::var("POD_NAME") + .ok() + .and_then(|name| name.rsplit_once('-')?.1.parse().ok()) + }) + .ok_or_else(|| anyhow::anyhow!("Kodo gateway ordinal is required"))?; + if ordinal >= 3 { + anyhow::bail!("Kodo gateway ordinal must be in 0..3"); + } + let cleanup_token = read_token(cli.kodo_cleanup_token_file.as_deref())?; + let registry_token = read_token(cli.registry_token_file.as_deref())?; + if cli.kodo_cleanup_enabled && (cleanup_token.is_none() || registry_token.is_none()) { + anyhow::bail!("Kodo cleanup requires non-empty cleanup and registry token files"); + } + if cli.kodo_cleanup_enabled && cleanup_token == registry_token { + anyhow::bail!("Kodo cleanup and registry tokens must be distinct"); + } + Ok(Some(kodo::KodoConfig { + ordinal, + cleanup_enabled: cli.kodo_cleanup_enabled, + cleanup_token, + registry_token, + })) +} + +fn read_token(path: Option<&std::path::Path>) -> anyhow::Result>> { + path.map(std::fs::read_to_string) + .transpose() + .map(|token| { + token + .map(|token| token.trim().to_owned()) + .filter(|token| !token.is_empty()) + .map(Arc::::from) + }) + .map_err(Into::into) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn default_cli() -> Cli { + Cli::try_parse_from(["rustqueue-proxy"]).unwrap() + } + + #[test] + fn proxy_limits_enforce_the_wire_and_memory_contract() { + let mut cli = default_cli(); + assert!(validate_proxy_limits(&cli).is_ok()); + + cli.max_message_bytes = rustqueue_protocol::MAX_MESSAGE_BYTES + 1; + cli.max_body_bytes = rustqueue_protocol::MAX_BATCH_BYTES; + assert!(validate_proxy_limits(&cli).is_err()); + + let mut cli = default_cli(); + cli.max_inflight_bytes = cli.max_body_bytes - 1; + assert!(validate_proxy_limits(&cli).is_err()); + } + + #[test] + fn kodo_mode_requires_the_full_hundred_mebibyte_profile() { + let mut cli = default_cli(); + cli.kodo_compatibility_enabled = true; + cli.max_message_bytes = rustqueue_protocol::MAX_MESSAGE_BYTES; + cli.max_body_bytes = rustqueue_protocol::MAX_BATCH_BYTES; + assert!(validate_proxy_limits(&cli).is_ok()); + + cli.max_message_bytes -= 1; + assert!(validate_proxy_limits(&cli).is_err()); } } diff --git a/crates/proxy/src/metrics.rs b/crates/proxy/src/metrics.rs index 956d801..ec065c2 100644 --- a/crates/proxy/src/metrics.rs +++ b/crates/proxy/src/metrics.rs @@ -7,6 +7,8 @@ pub struct ProxyMetrics { pub backend: Arc, pub discovery_poll: Arc, pub tcp_connection_rotations: Arc, + pub producer_retries: Arc, + pub producer_ambiguous_failures: Arc, } impl ProxyMetrics { @@ -29,6 +31,22 @@ impl ProxyMetrics { "rustqueue_proxy_tcp_connection_rotations_total {}\n", self.tcp_connection_rotations.load(Ordering::Relaxed) )); + output.push_str( + "# HELP rustqueue_proxy_producer_retries_total Publish commands retried after an explicit pre-commit Broker rejection.\n\ +# TYPE rustqueue_proxy_producer_retries_total counter\n", + ); + output.push_str(&format!( + "rustqueue_proxy_producer_retries_total {}\n", + self.producer_retries.load(Ordering::Relaxed) + )); + output.push_str( + "# HELP rustqueue_proxy_producer_ambiguous_failures_total Publish commands failed without retry because the Broker commit outcome was ambiguous.\n\ +# TYPE rustqueue_proxy_producer_ambiguous_failures_total counter\n", + ); + output.push_str(&format!( + "rustqueue_proxy_producer_ambiguous_failures_total {}\n", + self.producer_ambiguous_failures.load(Ordering::Relaxed) + )); output } } diff --git a/crates/proxy/src/tcp.rs b/crates/proxy/src/tcp.rs index cff6bac..6818148 100644 --- a/crates/proxy/src/tcp.rs +++ b/crates/proxy/src/tcp.rs @@ -1,3 +1,6 @@ +#[path = "tcp/gateway.rs"] +mod gateway; + use crate::backend::BackendPool; use crate::metrics::ProxyMetrics; use rand::Rng; @@ -5,33 +8,77 @@ use std::net::SocketAddr; use std::sync::atomic::Ordering; use std::sync::Arc; use std::time::Duration; -use tokio::sync::Semaphore; +use tokio::sync::{watch, Semaphore}; +use tokio::task::JoinSet; + +pub struct Limits { + pub max_connections: usize, + pub max_connection_age: Duration, + pub terminate_producer_protocol: bool, + pub max_message_bytes: usize, + pub max_body_bytes: usize, + pub command_timeout: Duration, + pub inflight_bytes: Arc, +} pub async fn serve( address: SocketAddr, pool: BackendPool, - max_connections: usize, - max_connection_age: Duration, + limits: Limits, metrics: ProxyMetrics, + mut shutdown: watch::Receiver, + shutdown_grace: Duration, ) -> anyhow::Result<()> { let listener = tokio::net::TcpListener::bind(address).await?; - let connections = Arc::new(Semaphore::new(max_connections)); + let connections = Arc::new(Semaphore::new(limits.max_connections)); + let mut sessions = JoinSet::new(); tracing::info!(%address, "producer TCP proxy listening"); loop { - let (mut client, peer) = listener.accept().await?; + let accepted = tokio::select! { + result = listener.accept() => Some(result?), + result = sessions.join_next(), if !sessions.is_empty() => { + log_session_result(result); + None + } + _ = shutdown.changed() => break, + }; + let Some((mut client, peer)) = accepted else { + continue; + }; let Ok(permit) = Arc::clone(&connections).try_acquire_owned() else { tracing::debug!(%peer, "rejecting producer because proxy connection limit is reached"); continue; }; + let metrics = metrics.clone(); + if limits.terminate_producer_protocol { + let pool = pool.clone(); + let gateway_limits = gateway::Limits { + max_message_bytes: limits.max_message_bytes, + max_body_bytes: limits.max_body_bytes, + command_timeout: limits.command_timeout, + inflight_bytes: Arc::clone(&limits.inflight_bytes), + }; + let session_shutdown = shutdown.clone(); + sessions.spawn(async move { + let _permit = permit; + let _ = client.set_nodelay(true); + if let Err(error) = + gateway::run(client, pool, gateway_limits, metrics, session_shutdown).await + { + tracing::debug!(%peer, %error, "producer Gateway session closed"); + } + }); + continue; + } let Some(backend) = pool.lease() else { tracing::debug!(%peer, "rejecting producer because no broker is ready"); continue; }; - let metrics = metrics.clone(); - let connection_age = jittered_connection_age(max_connection_age); - tokio::spawn(async move { + let connection_age = jittered_connection_age(limits.max_connection_age); + sessions.spawn(async move { let _permit = permit; let _backend_lease = backend; + let mut invalidation = _backend_lease.invalidation(); let connect_timer = metrics.backend.timer(); match tokio::time::timeout( std::time::Duration::from_secs(2), @@ -43,33 +90,36 @@ pub async fn serve( drop(connect_timer); let _ = client.set_nodelay(true); let _ = broker.set_nodelay(true); - if let Some(connection_age) = connection_age { - match tokio::time::timeout( - connection_age, - tokio::io::copy_bidirectional(&mut client, &mut broker), - ) - .await - { - Ok(Ok(_)) => {} - Ok(Err(error)) => { - tracing::debug!(%peer, %error, "producer TCP proxy closed") + let tunnel = tokio::io::copy_bidirectional(&mut client, &mut broker); + tokio::select! { + result = tunnel => { + if let Err(error) = result { + tracing::debug!(%peer, %error, "producer TCP proxy closed"); } - Err(_) => { - metrics - .tcp_connection_rotations - .fetch_add(1, Ordering::Relaxed); - tracing::debug!( - %peer, - node_id = _backend_lease.node_id, - max_age_seconds = connection_age.as_secs(), - "rotating producer TCP connection" - ); + } + _ = invalidation.changed() => { + metrics.tcp_connection_rotations.fetch_add(1, Ordering::Relaxed); + tracing::debug!( + %peer, + node_id = _backend_lease.node_id, + "closing producer tunnel because backend left discovery" + ); + } + _ = async { + if let Some(connection_age) = connection_age { + tokio::time::sleep(connection_age).await; + } else { + std::future::pending::<()>().await; } + } => { + metrics.tcp_connection_rotations.fetch_add(1, Ordering::Relaxed); + tracing::debug!( + %peer, + node_id = _backend_lease.node_id, + max_age_seconds = connection_age.map_or(0, |age| age.as_secs()), + "rotating producer TCP connection" + ); } - } else if let Err(error) = - tokio::io::copy_bidirectional(&mut client, &mut broker).await - { - tracing::debug!(%peer, %error, "producer TCP proxy closed"); } } Ok(Err(error)) => { @@ -83,6 +133,31 @@ pub async fn serve( } }); } + tracing::info!("producer TCP proxy stopped accepting new connections"); + if tokio::time::timeout(shutdown_grace, drain_sessions(&mut sessions)) + .await + .is_err() + { + tracing::warn!( + active_sessions = sessions.len(), + "producer TCP proxy shutdown grace expired" + ); + sessions.abort_all(); + drain_sessions(&mut sessions).await; + } + Ok(()) +} + +async fn drain_sessions(sessions: &mut JoinSet<()>) { + while let Some(result) = sessions.join_next().await { + log_session_result(Some(result)); + } +} + +fn log_session_result(result: Option>) { + if let Some(Err(error)) = result { + tracing::warn!(%error, "producer proxy session task failed"); + } } fn jittered_connection_age(base: Duration) -> Option { diff --git a/crates/proxy/src/tcp/gateway.rs b/crates/proxy/src/tcp/gateway.rs new file mode 100644 index 0000000..cdbd91b --- /dev/null +++ b/crates/proxy/src/tcp/gateway.rs @@ -0,0 +1,1092 @@ +use crate::backend::{Backend, BackendPool}; +use crate::metrics::ProxyMetrics; +use axum::body::Bytes; +use rustqueue_protocol::{ + encode_frame, Command, FrameType, IdentifyRequest, IdentifyResponse, HEARTBEAT, MAGIC_V2, OK, +}; +use serde_json::{json, Value}; +use std::collections::BTreeSet; +use std::io; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore}; +use tokio::time::{Instant, Sleep}; + +const MAX_COMMAND_LINE_BYTES: usize = 1024; +const MAX_CONTROL_BODY_BYTES: usize = 1024 * 1024; +const MAX_BACKEND_FRAME_BYTES: usize = 1024 * 1024; +const BACKEND_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); +const BACKEND_IDENTIFY_TIMEOUT: Duration = Duration::from_secs(5); +// Kodo's go-nsq producer waits at most 60 seconds for the response. Two +// pre-commit attempts, including both handshakes and the final response, must +// remain inside that client envelope while still allowing a 100 MiB body to +// cross a temporarily congested cluster network. +const BACKEND_WRITE_TIMEOUT: Duration = Duration::from_secs(12); +const BACKEND_TIMEOUT: Duration = Duration::from_secs(15); +const MAX_PUBLISH_ATTEMPTS: usize = 3; +const PUBLISH_DEADLINE: Duration = Duration::from_secs(55); +const NEW_ATTEMPT_BUDGET: Duration = Duration::from_secs(34); +const DEFAULT_HEARTBEAT: Duration = Duration::from_secs(30); + +struct ClientCommand { + command: Command, + line: Vec, + body: Option, + _permit: Option, +} + +struct BrokerSession { + backend: Backend, + stream: TcpStream, +} + +pub(super) struct Limits { + pub max_message_bytes: usize, + pub max_body_bytes: usize, + pub command_timeout: Duration, + pub inflight_bytes: Arc, +} + +enum Attempt { + Success, + Retry(Vec), + Final(Vec), + Ambiguous(String), +} + +#[derive(Debug, Eq, PartialEq)] +enum PublishFailure { + Final(Vec), + RetryByReconnect, + Ambiguous(Vec), +} + +pub(super) async fn run( + mut client: TcpStream, + pool: BackendPool, + limits: Limits, + metrics: ProxyMetrics, + mut shutdown: watch::Receiver, +) -> anyhow::Result<()> { + let mut magic = [0; 4]; + tokio::select! { + biased; + _ = shutdown.changed() => return Ok(()), + result = tokio::time::timeout(Duration::from_secs(5), client.read_exact(&mut magic)) => { + result??; + } + } + if &magic != MAGIC_V2 { + return Ok(()); + } + let (mut reader, mut writer) = client.into_split(); + let (commands, mut incoming) = tokio::sync::mpsc::channel(1); + let reader_task = tokio::spawn(async move { + loop { + let command = read_command( + &mut reader, + limits.max_message_bytes, + limits.max_body_bytes, + limits.command_timeout, + &limits.inflight_bytes, + ) + .await; + let failed = command.is_err(); + if commands.send(command).await.is_err() || failed { + break; + } + } + }); + let _reader = AbortOnDrop(reader_task.abort_handle()); + let mut identified = false; + let mut heartbeat_interval = Some(DEFAULT_HEARTBEAT); + let mut heartbeat = Box::pin(tokio::time::sleep(DEFAULT_HEARTBEAT)); + let mut backend_identify = default_backend_identify(); + let mut backend_session = None; + + loop { + tokio::select! { + biased; + _ = shutdown.changed() => break, + command = incoming.recv() => { + let Some(command) = command else { + break; + }; + let command = match command { + Ok(command) => command, + Err(ReadError::RetryByReconnect(detail)) => { + tracing::debug!( + detail, + "closing Kodo producer connection for safe Gateway failover" + ); + break; + } + Err(error) => { + write_error(&mut writer, error.code(), &error.to_string()).await?; + break; + } + }; + reset_heartbeat(&mut heartbeat, heartbeat_interval); + match command.command { + Command::Identify if !identified => { + let body = command.body.as_deref().unwrap_or_default(); + let request: IdentifyRequest = match serde_json::from_slice(body) { + Ok(request) => request, + Err(_) => { + write_error(&mut writer, "E_BAD_BODY", "IDENTIFY body is invalid").await?; + break; + } + }; + if request.tls_v1 || request.snappy || request.deflate { + write_error( + &mut writer, + "E_IDENTIFY_FAILED", + "Kodo publish gateways do not negotiate TLS or compression", + ) + .await?; + break; + } + heartbeat_interval = match request.heartbeat_interval { + Some(-1) => None, + Some(value) if (1_000..=300_000).contains(&value) => { + Some(Duration::from_millis(value as u64)) + } + Some(_) => { + write_error( + &mut writer, + "E_IDENTIFY_FAILED", + "heartbeat_interval is outside 1000..=300000", + ) + .await?; + break; + } + None => Some(DEFAULT_HEARTBEAT), + }; + backend_identify = backend_identify_body(body)?; + identified = true; + reset_heartbeat(&mut heartbeat, heartbeat_interval); + let response = IdentifyResponse { + max_rdy_count: 2_500, + version: env!("CARGO_PKG_VERSION").into(), + max_msg_timeout: 900_000, + msg_timeout: request.msg_timeout.unwrap_or(60_000).max(1_000), + tls_v1: false, + deflate: false, + deflate_level: 0, + max_deflate_level: 0, + snappy: false, + sample_rate: 0, + auth_required: false, + output_buffer_size: 16_384, + output_buffer_timeout: 250, + }; + write_frame( + &mut writer, + FrameType::Response, + &serde_json::to_vec(&response)?, + ) + .await?; + } + Command::Identify => { + write_error(&mut writer, "E_INVALID", "IDENTIFY may only be sent once").await?; + break; + } + Command::Publish { .. } + | Command::MultiPublish { .. } + | Command::DeferredPublish { .. } => { + let response = publish( + &pool, + &backend_identify, + &command, + &mut backend_session, + &metrics, + ) + .await; + match response { + Ok(()) => write_frame(&mut writer, FrameType::Response, OK).await?, + Err(PublishFailure::Final(error)) => { + write_frame(&mut writer, FrameType::Error, &error).await?; + } + // Kodo retries another advertised Gateway only for a + // connection error. Every backend attempt represented + // here failed before commit, so closing is both safe + // and necessary for cross-Gateway failover. + Err(PublishFailure::RetryByReconnect) => break, + Err(PublishFailure::Ambiguous(error)) => { + write_frame(&mut writer, FrameType::Error, &error).await?; + break; + } + } + } + Command::Noop => {} + Command::Auth => { + write_error(&mut writer, "E_AUTH_DISABLED", "AUTH disabled").await?; + break; + } + _ => { + write_error( + &mut writer, + "E_INVALID", + "publish Gateway accepts only IDENTIFY, PUB, MPUB, DPUB, and NOP", + ) + .await?; + break; + } + } + } + _ = &mut heartbeat, if heartbeat_interval.is_some() => { + write_frame(&mut writer, FrameType::Response, HEARTBEAT).await?; + reset_heartbeat(&mut heartbeat, heartbeat_interval); + } + } + } + Ok(()) +} + +async fn publish( + pool: &BackendPool, + identify: &[u8], + command: &ClientCommand, + current: &mut Option, + metrics: &ProxyMetrics, +) -> Result<(), PublishFailure> { + let deadline = Instant::now() + PUBLISH_DEADLINE; + let available = pool.all(); + let mut attempted = BTreeSet::new(); + let mut last_retry = None; + + if current + .as_ref() + .is_some_and(|session| available.contains(&session.backend)) + { + let mut session = current.take().expect("checked current session"); + attempted.insert(session.backend.node_id); + match attempt(&mut session.stream, command).await { + Attempt::Success => { + *current = Some(session); + return Ok(()); + } + Attempt::Retry(error) => { + metrics.producer_retries.fetch_add(1, Ordering::Relaxed); + last_retry = Some(error); + } + Attempt::Final(error) => return Err(PublishFailure::Final(error)), + Attempt::Ambiguous(detail) => { + metrics + .producer_ambiguous_failures + .fetch_add(1, Ordering::Relaxed); + return Err(PublishFailure::Ambiguous( + format!("E_AMBIGUOUS {detail}").into_bytes(), + )); + } + } + } else { + *current = None; + } + + for backend in pool.shuffled(MAX_PUBLISH_ATTEMPTS) { + if attempted.len() >= MAX_PUBLISH_ATTEMPTS + || deadline.saturating_duration_since(Instant::now()) < NEW_ATTEMPT_BUDGET + { + break; + } + if !attempted.insert(backend.node_id) { + continue; + } + let mut session = match connect_backend(backend, identify, metrics).await { + Ok(session) => session, + Err(error) => { + tracing::debug!(%error, "publish Gateway backend handshake failed"); + continue; + } + }; + match attempt(&mut session.stream, command).await { + Attempt::Success => { + *current = Some(session); + return Ok(()); + } + Attempt::Retry(error) => { + metrics.producer_retries.fetch_add(1, Ordering::Relaxed); + last_retry = Some(error); + } + Attempt::Final(error) => return Err(PublishFailure::Final(error)), + Attempt::Ambiguous(detail) => { + metrics + .producer_ambiguous_failures + .fetch_add(1, Ordering::Relaxed); + return Err(PublishFailure::Ambiguous( + format!("E_AMBIGUOUS {detail}").into_bytes(), + )); + } + } + } + if let Some(error) = last_retry { + tracing::debug!( + error = %String::from_utf8_lossy(&error), + "publish Gateway exhausted pre-commit Broker retries" + ); + } else { + tracing::debug!("publish Gateway has no publish-ready Broker"); + } + Err(PublishFailure::RetryByReconnect) +} + +async fn connect_backend( + backend: Backend, + identify: &[u8], + metrics: &ProxyMetrics, +) -> anyhow::Result { + let _timer = metrics.backend.timer(); + let mut stream = tokio::time::timeout( + BACKEND_CONNECT_TIMEOUT, + TcpStream::connect(backend.tcp_address()), + ) + .await??; + stream.set_nodelay(true)?; + stream.write_all(MAGIC_V2).await?; + stream.write_all(b"IDENTIFY\n").await?; + stream.write_u32(identify.len() as u32).await?; + stream.write_all(identify).await?; + stream.flush().await?; + let (frame_type, body) = + tokio::time::timeout(BACKEND_IDENTIFY_TIMEOUT, read_frame(&mut stream)).await??; + if frame_type == FrameType::Error as i32 { + anyhow::bail!( + "backend IDENTIFY failed: {}", + String::from_utf8_lossy(&body) + ); + } + if frame_type != FrameType::Response as i32 { + anyhow::bail!("backend IDENTIFY returned an invalid frame"); + } + Ok(BrokerSession { backend, stream }) +} + +async fn attempt(stream: &mut TcpStream, command: &ClientCommand) -> Attempt { + if send_before_commit(stream, command, BACKEND_WRITE_TIMEOUT) + .await + .is_err() + { + return Attempt::Retry( + b"E_PUB_RETRY Broker connection failed before the publish body was complete".to_vec(), + ); + } + if !matches!( + tokio::time::timeout(BACKEND_WRITE_TIMEOUT, stream.flush()).await, + Ok(Ok(())) + ) { + return Attempt::Ambiguous("Broker connection failed after sending the body".into()); + } + match tokio::time::timeout(BACKEND_TIMEOUT, read_frame(stream)).await { + Ok(Ok((frame_type, body))) if frame_type == FrameType::Response as i32 && body == OK => { + Attempt::Success + } + Ok(Ok((frame_type, body))) if frame_type == FrameType::Error as i32 => { + if retriable_error(&body) { + Attempt::Retry(body) + } else { + Attempt::Final(body) + } + } + Ok(Ok(_)) => Attempt::Ambiguous("Broker returned an invalid publish response".into()), + Ok(Err(error)) => Attempt::Ambiguous(format!("Broker response failed: {error}")), + Err(_) => Attempt::Ambiguous("Broker publish response timed out".into()), + } +} + +async fn send_before_commit( + stream: &mut W, + command: &ClientCommand, + timeout: Duration, +) -> io::Result<()> +where + W: AsyncWrite + Unpin, +{ + tokio::time::timeout(timeout, async { + stream.write_all(&command.line).await?; + if let Some(body) = command.body.as_deref() { + stream.write_u32(body.len() as u32).await?; + stream.write_all(body).await?; + } + io::Result::Ok(()) + }) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "backend publish write timed out"))? +} + +fn retriable_error(body: &[u8]) -> bool { + let error = String::from_utf8_lossy(body); + error.starts_with("E_DRAINING") + || error.starts_with("E_CLOSING") + || error.starts_with("E_THROTTLED") + || error.starts_with("E_PUB_RETRY") +} + +fn backend_identify_body(body: &[u8]) -> anyhow::Result> { + let mut value: Value = serde_json::from_slice(body)?; + let object = value + .as_object_mut() + .ok_or_else(|| anyhow::anyhow!("IDENTIFY body must be an object"))?; + object.insert("feature_negotiation".into(), json!(true)); + object.insert("heartbeat_interval".into(), json!(-1)); + object.insert("tls_v1".into(), json!(false)); + object.insert("snappy".into(), json!(false)); + object.insert("deflate".into(), json!(false)); + Ok(serde_json::to_vec(&value)?) +} + +fn default_backend_identify() -> Vec { + serde_json::to_vec(&json!({ + "feature_negotiation": true, + "heartbeat_interval": -1, + "tls_v1": false, + "snappy": false, + "deflate": false, + "user_agent": "rustqueue-kodo-gateway" + })) + .expect("static IDENTIFY body") +} + +fn reset_heartbeat(heartbeat: &mut std::pin::Pin>, interval: Option) { + heartbeat.as_mut().reset( + Instant::now() + interval.unwrap_or_else(|| Duration::from_secs(365 * 24 * 60 * 60)), + ); +} + +async fn read_command( + reader: &mut R, + max_message_bytes: usize, + max_body_bytes: usize, + command_timeout: Duration, + inflight_bytes: &Arc, +) -> Result +where + R: AsyncRead + Unpin, +{ + tokio::time::timeout(command_timeout, async { + let first = reader.read_u8().await?; + read_started_command( + reader, + first, + max_message_bytes, + max_body_bytes, + inflight_bytes, + ) + .await + }) + .await + .map_err(|_| ReadError::Protocol("E_INVALID", "command read timed out".into()))? +} + +async fn read_started_command( + reader: &mut R, + first: u8, + max_message_bytes: usize, + max_body_bytes: usize, + inflight_bytes: &Arc, +) -> Result +where + R: AsyncRead + Unpin, +{ + let mut line = Vec::with_capacity(64); + line.push(first); + loop { + if line.last() == Some(&b'\n') { + break; + } + line.push(reader.read_u8().await?); + if line.len() > MAX_COMMAND_LINE_BYTES { + return Err(ReadError::Protocol( + "E_INVALID", + "command line exceeds limit".into(), + )); + } + } + let command = Command::parse(&line) + .map_err(|error| ReadError::Protocol("E_INVALID", error.to_string()))?; + let limit = match command { + Command::Identify | Command::Auth => Some(("E_BAD_BODY", MAX_CONTROL_BODY_BYTES)), + Command::Publish { .. } | Command::DeferredPublish { .. } => { + Some(("E_BAD_MESSAGE", max_message_bytes)) + } + Command::MultiPublish { .. } => Some(("E_BAD_BODY", max_body_bytes)), + _ => None, + }; + let Some((limit_code, maximum)) = limit else { + return Ok(ClientCommand { + command, + line, + body: None, + _permit: None, + }); + }; + let length = reader.read_u32().await? as usize; + if length == 0 || length > maximum { + return Err(ReadError::Protocol( + limit_code, + format!("command body size {length} is outside 1..={maximum}"), + )); + } + let permit = Arc::clone(inflight_bytes) + .try_acquire_many_owned(length as u32) + .map_err(|_| ReadError::RetryByReconnect("publish Gateway byte budget is exhausted"))?; + let mut body = vec![0; length]; + reader.read_exact(&mut body).await?; + let body = Bytes::from(body); + if matches!(command, Command::MultiPublish { .. }) { + rustqueue_protocol::parse_mpub_bytes(body.clone(), max_message_bytes) + .map_err(|error| ReadError::Protocol(error.code(), error.to_string()))?; + } + Ok(ClientCommand { + command, + line, + body: Some(body), + _permit: Some(permit), + }) +} + +async fn read_frame(reader: &mut R) -> io::Result<(i32, Vec)> +where + R: AsyncRead + Unpin, +{ + let size = reader.read_u32().await? as usize; + if !(4..=MAX_BACKEND_FRAME_BYTES).contains(&size) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "backend frame size is invalid", + )); + } + let frame_type = reader.read_i32().await?; + let mut body = vec![0; size - 4]; + reader.read_exact(&mut body).await?; + Ok((frame_type, body)) +} + +async fn write_frame(writer: &mut W, frame_type: FrameType, body: &[u8]) -> io::Result<()> +where + W: AsyncWrite + Unpin, +{ + writer.write_all(&encode_frame(frame_type, body)).await?; + writer.flush().await +} + +async fn write_error(writer: &mut W, code: &str, detail: &str) -> io::Result<()> +where + W: AsyncWrite + Unpin, +{ + let detail = detail.replace(['\r', '\n'], " "); + write_frame( + writer, + FrameType::Error, + format!("{code} {detail}").as_bytes(), + ) + .await +} + +struct AbortOnDrop(tokio::task::AbortHandle); + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } +} + +#[derive(Debug)] +enum ReadError { + Io(io::Error), + Protocol(&'static str, String), + RetryByReconnect(&'static str), +} + +impl ReadError { + fn code(&self) -> &'static str { + match self { + Self::Io(_) => "E_INVALID", + Self::Protocol(code, _) => code, + Self::RetryByReconnect(_) => "E_THROTTLED", + } + } +} + +impl std::fmt::Display for ReadError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Io(error) => error.fmt(formatter), + Self::Protocol(_, detail) => detail.fmt(formatter), + Self::RetryByReconnect(detail) => detail.fmt(formatter), + } + } +} + +impl From for ReadError { + fn from(error: io::Error) -> Self { + Self::Io(error) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn read_line(stream: &mut TcpStream) -> Vec { + let mut line = Vec::new(); + loop { + let byte = stream.read_u8().await.unwrap(); + line.push(byte); + if byte == b'\n' { + return line; + } + } + } + + async fn mock_backend(listener: tokio::net::TcpListener, response: (FrameType, &'static [u8])) { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut magic = [0; 4]; + stream.read_exact(&mut magic).await.unwrap(); + assert_eq!(&magic, MAGIC_V2); + assert_eq!(read_line(&mut stream).await, b"IDENTIFY\n"); + let length = stream.read_u32().await.unwrap() as usize; + let mut identify = vec![0; length]; + stream.read_exact(&mut identify).await.unwrap(); + write_frame(&mut stream, FrameType::Response, b"{}") + .await + .unwrap(); + assert_eq!(read_line(&mut stream).await, b"PUB events\n"); + let length = stream.read_u32().await.unwrap() as usize; + let mut body = vec![0; length]; + stream.read_exact(&mut body).await.unwrap(); + assert_eq!(body, b"payload"); + write_frame(&mut stream, response.0, response.1) + .await + .unwrap(); + } + + async fn mock_large_backend(listener: tokio::net::TcpListener, expected: usize) { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut magic = [0; 4]; + stream.read_exact(&mut magic).await.unwrap(); + assert_eq!(&magic, MAGIC_V2); + assert_eq!(read_line(&mut stream).await, b"IDENTIFY\n"); + let length = stream.read_u32().await.unwrap() as usize; + let mut identify = vec![0; length]; + stream.read_exact(&mut identify).await.unwrap(); + write_frame(&mut stream, FrameType::Response, b"{}") + .await + .unwrap(); + assert_eq!(read_line(&mut stream).await, b"PUB large\n"); + assert_eq!(stream.read_u32().await.unwrap() as usize, expected); + let mut remaining = expected; + let mut buffer = vec![0; 1024 * 1024]; + while remaining > 0 { + let count = remaining.min(buffer.len()); + stream.read_exact(&mut buffer[..count]).await.unwrap(); + assert!(buffer[..count].iter().all(|byte| *byte == 0x5a)); + remaining -= count; + } + write_frame(&mut stream, FrameType::Response, OK) + .await + .unwrap(); + } + + fn backend(listener: &tokio::net::TcpListener, node_id: u64) -> Backend { + Backend { + broadcast_address: listener.local_addr().unwrap().ip().to_string(), + tcp_port: listener.local_addr().unwrap().port(), + http_port: 4151, + node_id, + } + } + + #[test] + fn only_explicit_node_local_rejections_are_retried() { + assert!(retriable_error(b"E_DRAINING broker is draining")); + assert!(retriable_error(b"E_CLOSING broker is shutting down")); + assert!(retriable_error( + b"E_THROTTLED local disk is above its publish watermark" + )); + assert!(retriable_error( + b"E_PUB_RETRY active topic publish worker limit reached" + )); + assert!(!retriable_error( + b"E_PUB_FAILED local storage is isolated after an earlier failure" + )); + assert!(!retriable_error( + b"E_PUB_FAILED message exceeds configured limit" + )); + assert!(!retriable_error( + b"E_PUB_FAILED topic is protected by an active deletion tombstone" + )); + } + + #[tokio::test] + async fn no_backend_requests_a_safe_cross_gateway_reconnect() { + let pool = BackendPool::default(); + let command = ClientCommand { + command: Command::Publish { + topic: "events".into(), + }, + line: b"PUB events\n".to_vec(), + body: Some(Bytes::from_static(b"payload")), + _permit: None, + }; + let mut current = None; + assert_eq!( + publish( + &pool, + &default_backend_identify(), + &command, + &mut current, + &ProxyMetrics::default(), + ) + .await, + Err(PublishFailure::RetryByReconnect) + ); + } + + #[tokio::test] + async fn final_backend_rejections_remain_protocol_errors() { + let rejected = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let rejected_backend = backend(&rejected, 1); + let backend_task = tokio::spawn(mock_backend( + rejected, + ( + FrameType::Error, + b"E_PUB_FAILED message exceeds configured limit", + ), + )); + let pool = BackendPool::default(); + pool.replace(vec![rejected_backend]); + let command = ClientCommand { + command: Command::Publish { + topic: "events".into(), + }, + line: b"PUB events\n".to_vec(), + body: Some(Bytes::from_static(b"payload")), + _permit: None, + }; + let mut current = None; + assert_eq!( + publish( + &pool, + &default_backend_identify(), + &command, + &mut current, + &ProxyMetrics::default(), + ) + .await, + Err(PublishFailure::Final( + b"E_PUB_FAILED message exceeds configured limit".to_vec() + )) + ); + backend_task.await.unwrap(); + } + + #[tokio::test] + async fn idle_client_command_is_bounded_before_the_first_byte() { + let (_client, mut gateway) = tokio::io::duplex(1024); + let budget = Arc::new(Semaphore::new(4)); + let error = match read_command(&mut gateway, 1024, 1024, Duration::from_millis(10), &budget) + .await + { + Ok(_) => panic!("idle client command was accepted"), + Err(error) => error, + }; + assert!(matches!(error, ReadError::Protocol("E_INVALID", _))); + assert_eq!(budget.available_permits(), 4); + } + + #[tokio::test] + async fn partial_client_body_releases_its_budget_after_timeout() { + let (mut client, mut gateway) = tokio::io::duplex(1024); + client.write_all(b"PUB events\n\0\0\0\x04x").await.unwrap(); + let budget = Arc::new(Semaphore::new(4)); + let error = match read_command(&mut gateway, 1024, 1024, Duration::from_millis(10), &budget) + .await + { + Ok(_) => panic!("partial publish body was accepted"), + Err(error) => error, + }; + assert!(matches!(error, ReadError::Protocol("E_INVALID", _))); + assert_eq!(budget.available_permits(), 4); + } + + #[tokio::test] + async fn exhausted_ingress_budget_closes_without_a_protocol_error() { + let pool = BackendPool::default(); + let gateway_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let gateway_address = gateway_listener.local_addr().unwrap(); + let gateway_task = tokio::spawn(async move { + let (client, _) = gateway_listener.accept().await.unwrap(); + let (_shutdown_tx, shutdown) = watch::channel(false); + run( + client, + pool, + Limits { + max_message_bytes: 1024, + max_body_bytes: 1024, + command_timeout: Duration::from_secs(1), + inflight_bytes: Arc::new(Semaphore::new(1)), + }, + ProxyMetrics::default(), + shutdown, + ) + .await + .unwrap(); + }); + + let mut client = TcpStream::connect(gateway_address).await.unwrap(); + client.write_all(MAGIC_V2).await.unwrap(); + client.write_all(b"PUB events\n\0\0\0\x02").await.unwrap(); + let mut frame = [0]; + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), client.read(&mut frame)) + .await + .unwrap() + .unwrap(), + 0 + ); + gateway_task.await.unwrap(); + } + + #[tokio::test] + async fn single_message_limit_is_independent_from_the_batch_body_limit() { + let (mut client, mut gateway) = tokio::io::duplex(1024); + client + .write_all(b"PUB events\n\0\0\0\x05hello") + .await + .unwrap(); + let budget = Arc::new(Semaphore::new(64)); + let error = match read_command(&mut gateway, 4, 64, Duration::from_secs(1), &budget).await { + Ok(_) => panic!("oversized PUB was accepted"), + Err(error) => error, + }; + assert!(matches!(error, ReadError::Protocol("E_BAD_MESSAGE", _))); + assert_eq!(budget.available_permits(), 64); + } + + #[tokio::test] + async fn mpub_rejects_an_entry_above_the_single_message_limit() { + let mut body = 1u32.to_be_bytes().to_vec(); + body.extend_from_slice(&5u32.to_be_bytes()); + body.extend_from_slice(b"hello"); + let (mut client, mut gateway) = tokio::io::duplex(1024); + client.write_all(b"MPUB events\n").await.unwrap(); + client.write_u32(body.len() as u32).await.unwrap(); + client.write_all(&body).await.unwrap(); + let budget = Arc::new(Semaphore::new(64)); + let error = match read_command(&mut gateway, 4, 64, Duration::from_secs(1), &budget).await { + Ok(_) => panic!("MPUB entry above the message limit was accepted"), + Err(error) => error, + }; + assert!(matches!(error, ReadError::Protocol("E_BAD_MESSAGE", _))); + assert_eq!(budget.available_permits(), 64); + } + + #[tokio::test] + async fn stalled_backend_write_is_bounded_before_commit() { + let (mut gateway, _backend) = tokio::io::duplex(1); + let command = ClientCommand { + command: Command::Publish { + topic: "events".into(), + }, + line: b"PUB events\n".to_vec(), + body: Some(Bytes::from(vec![0x5a; 4096])), + _permit: None, + }; + let error = send_before_commit(&mut gateway, &command, Duration::from_millis(10)) + .await + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::TimedOut); + } + + #[tokio::test] + async fn explicit_pre_commit_rejection_retries_another_broker() { + let rejected = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let accepted = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let rejected_backend = backend(&rejected, 1); + let accepted_backend = backend(&accepted, 2); + let rejected_task = tokio::spawn(mock_backend( + rejected, + (FrameType::Error, b"E_DRAINING broker is draining"), + )); + let accepted_task = tokio::spawn(mock_backend(accepted, (FrameType::Response, OK))); + let metrics = ProxyMetrics::default(); + let identify = default_backend_identify(); + let first = connect_backend(rejected_backend.clone(), &identify, &metrics) + .await + .unwrap(); + let pool = BackendPool::default(); + pool.replace(vec![rejected_backend, accepted_backend]); + let command = ClientCommand { + command: Command::Publish { + topic: "events".into(), + }, + line: b"PUB events\n".to_vec(), + body: Some(Bytes::from_static(b"payload")), + _permit: None, + }; + let mut current = Some(first); + publish(&pool, &identify, &command, &mut current, &metrics) + .await + .unwrap(); + assert_eq!(current.as_ref().unwrap().backend.node_id, 2); + assert_eq!(metrics.producer_retries.load(Ordering::Relaxed), 1); + rejected_task.await.unwrap(); + accepted_task.await.unwrap(); + } + + #[tokio::test] + async fn shutdown_finishes_the_inflight_publish_but_rejects_the_queued_publish() { + let backend_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let backend_address = backend(&backend_listener, 1); + let (body_received_tx, body_received_rx) = tokio::sync::oneshot::channel(); + let (respond_tx, respond_rx) = tokio::sync::oneshot::channel(); + let backend_task = tokio::spawn(async move { + let (mut stream, _) = backend_listener.accept().await.unwrap(); + let mut magic = [0; 4]; + stream.read_exact(&mut magic).await.unwrap(); + assert_eq!(&magic, MAGIC_V2); + assert_eq!(read_line(&mut stream).await, b"IDENTIFY\n"); + let length = stream.read_u32().await.unwrap() as usize; + let mut identify = vec![0; length]; + stream.read_exact(&mut identify).await.unwrap(); + write_frame(&mut stream, FrameType::Response, b"{}") + .await + .unwrap(); + assert_eq!(read_line(&mut stream).await, b"PUB events\n"); + let length = stream.read_u32().await.unwrap() as usize; + let mut body = vec![0; length]; + stream.read_exact(&mut body).await.unwrap(); + assert_eq!(body, b"payload"); + body_received_tx.send(()).unwrap(); + respond_rx.await.unwrap(); + write_frame(&mut stream, FrameType::Response, OK) + .await + .unwrap(); + let mut next = [0]; + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), stream.read(&mut next)) + .await + .unwrap() + .unwrap(), + 0 + ); + }); + let pool = BackendPool::default(); + pool.replace(vec![backend_address]); + let gateway_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let gateway_address = gateway_listener.local_addr().unwrap(); + let (shutdown_tx, shutdown) = watch::channel(false); + let gateway_task = tokio::spawn(async move { + let (client, _) = gateway_listener.accept().await.unwrap(); + run( + client, + pool, + Limits { + max_message_bytes: 1024, + max_body_bytes: 1024, + command_timeout: Duration::from_secs(1), + inflight_bytes: Arc::new(Semaphore::new(1024)), + }, + ProxyMetrics::default(), + shutdown, + ) + .await + .unwrap(); + }); + + let mut client = TcpStream::connect(gateway_address).await.unwrap(); + client.write_all(MAGIC_V2).await.unwrap(); + client.write_all(b"IDENTIFY\n\0\0\0\x02{}").await.unwrap(); + assert_eq!(read_frame(&mut client).await.unwrap().0, 0); + client + .write_all(b"PUB events\n\0\0\0\x07payload") + .await + .unwrap(); + body_received_rx.await.unwrap(); + client + .write_all(b"PUB events\n\0\0\0\x06queued") + .await + .unwrap(); + shutdown_tx.send(true).unwrap(); + respond_tx.send(()).unwrap(); + let response = read_frame(&mut client).await.unwrap(); + assert_eq!((response.0, response.1.as_slice()), (0, OK)); + let mut closed = [0]; + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), client.read(&mut closed)) + .await + .unwrap() + .unwrap(), + 0 + ); + tokio::time::timeout(Duration::from_secs(1), gateway_task) + .await + .unwrap() + .unwrap(); + backend_task.await.unwrap(); + } + + #[tokio::test] + async fn gateway_accepts_a_hundred_mebibyte_publish() { + const MESSAGE_BYTES: usize = 100 * 1024 * 1024; + let backend_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let pool = BackendPool::default(); + pool.replace(vec![backend(&backend_listener, 1)]); + let backend_task = tokio::spawn(mock_large_backend(backend_listener, MESSAGE_BYTES)); + + let gateway_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let gateway_address = gateway_listener.local_addr().unwrap(); + let gateway_task = tokio::spawn(async move { + let (client, _) = gateway_listener.accept().await.unwrap(); + let (_shutdown_tx, shutdown) = watch::channel(false); + run( + client, + pool, + Limits { + max_message_bytes: MESSAGE_BYTES, + max_body_bytes: 128 * 1024 * 1024, + command_timeout: Duration::from_secs(120), + inflight_bytes: Arc::new(Semaphore::new(512 * 1024 * 1024)), + }, + ProxyMetrics::default(), + shutdown, + ) + .await + .unwrap(); + }); + + let mut client = TcpStream::connect(gateway_address).await.unwrap(); + client.write_all(MAGIC_V2).await.unwrap(); + let identify = serde_json::to_vec(&json!({ + "feature_negotiation": true, + "heartbeat_interval": -1 + })) + .unwrap(); + client.write_all(b"IDENTIFY\n").await.unwrap(); + client.write_u32(identify.len() as u32).await.unwrap(); + client.write_all(&identify).await.unwrap(); + assert_eq!(read_frame(&mut client).await.unwrap().0, 0); + client.write_all(b"PUB large\n").await.unwrap(); + client.write_u32(MESSAGE_BYTES as u32).await.unwrap(); + client.write_all(&vec![0x5a; MESSAGE_BYTES]).await.unwrap(); + let response = read_frame(&mut client).await.unwrap(); + assert_eq!((response.0, response.1.as_slice()), (0, OK)); + drop(client); + backend_task.await.unwrap(); + gateway_task.await.unwrap(); + } + + #[test] + fn two_slow_precommit_attempts_fit_the_kodo_read_timeout() { + let handshake = BACKEND_CONNECT_TIMEOUT + BACKEND_IDENTIFY_TIMEOUT; + let first_precommit_failure = handshake + BACKEND_WRITE_TIMEOUT; + let total = first_precommit_failure + NEW_ATTEMPT_BUDGET; + assert!(total < PUBLISH_DEADLINE); + assert!(PUBLISH_DEADLINE < Duration::from_secs(60)); + assert!(NEW_ATTEMPT_BUDGET >= handshake + BACKEND_WRITE_TIMEOUT + BACKEND_TIMEOUT); + } +} diff --git a/crates/queue/src/batch.rs b/crates/queue/src/batch.rs index ed3994a..2af1b0e 100644 --- a/crates/queue/src/batch.rs +++ b/crates/queue/src/batch.rs @@ -1,11 +1,14 @@ use crate::model::MessageMeta; use crate::BrokerError; use bytes::Bytes; -use rustqueue_storage::{PayloadRef, Record, RecordLocation, HEADER_LEN}; +#[cfg(test)] +use rustqueue_storage::Record; +use rustqueue_storage::{PayloadRef, RecordLocation, HEADER_LEN}; +use std::io::{self, Read}; use std::sync::Arc; const ITEM_HEADER: usize = 20; -pub(crate) const MAX_MESSAGES: usize = 10_000; +pub(crate) const MAX_MESSAGES: usize = rustqueue_protocol::MAX_MPUB_MESSAGES; pub(crate) struct EncodedBatch<'a> { count: [u8; 4], @@ -67,6 +70,7 @@ impl EncodedBatch<'_> { } } +#[cfg(test)] pub(crate) fn metas( record: &Record, location: &RecordLocation, @@ -143,6 +147,86 @@ pub(crate) fn metas( Ok(output) } +pub(crate) fn metas_from_reader( + header: rustqueue_storage::RecordHeader, + payload_len: usize, + reader: &mut dyn Read, + location: &RecordLocation, +) -> io::Result> { + if payload_len < 4 { + return Err(invalid_data("publish batch is truncated")); + } + let mut count = [0u8; 4]; + reader.read_exact(&mut count)?; + let count = u32::from_be_bytes(count) as usize; + let maximum_count = payload_len.saturating_sub(4) / ITEM_HEADER; + if count == 0 || count > MAX_MESSAGES || count > maximum_count { + return Err(invalid_data("publish batch count is invalid")); + } + let mut cursor = 4usize; + let mut first_position = None; + let mut output = Vec::with_capacity(count); + let mut body_buffer = [0u8; 64 * 1024]; + for ordinal in 0..count { + let mut item = [0u8; ITEM_HEADER]; + reader.read_exact(&mut item)?; + cursor = cursor.saturating_add(ITEM_HEADER); + let position = u64::from_be_bytes(item[0..8].try_into().unwrap()); + let id = u64::from_be_bytes(item[8..16].try_into().unwrap()); + let len = u32::from_be_bytes(item[16..20].try_into().unwrap()); + let end = cursor + .checked_add(len as usize) + .ok_or_else(|| invalid_data("publish batch length overflow"))?; + if end > payload_len { + return Err(invalid_data("publish batch body is truncated")); + } + let base_position = *first_position.get_or_insert(position); + let expected_position = base_position + .checked_add(ordinal as u64) + .ok_or_else(|| invalid_data("publish batch position overflow"))?; + let expected_id = header + .message_id + .checked_add(ordinal as u64) + .ok_or_else(|| invalid_data("publish batch message ID overflow"))?; + if position != expected_position || id != expected_id { + return Err(invalid_data( + "publish batch positions or message IDs are not contiguous", + )); + } + let body_offset = cursor; + let mut remaining = len as usize; + let mut body_crc = 0u32; + while remaining > 0 { + let wanted = remaining.min(body_buffer.len()); + reader.read_exact(&mut body_buffer[..wanted])?; + body_crc = crc32c::crc32c_append(body_crc, &body_buffer[..wanted]); + remaining -= wanted; + } + output.push(MessageMeta { + position, + id, + timestamp_ns: header.timestamp_ns, + available_at_ms: header.available_at_ms, + log_index: location.index, + payload: PayloadRef { + path: Arc::clone(&location.segment), + offset: location.offset + HEADER_LEN as u64 + body_offset as u64, + len, + crc32c: body_crc, + }, + }); + cursor = end; + } + if cursor != payload_len { + return Err(invalid_data("publish batch has trailing bytes")); + } + Ok(output) +} + +fn invalid_data(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidData, message.into()) +} + pub(crate) fn metas_after_append( timestamp_ns: i64, available_at_ms: i64, diff --git a/crates/queue/src/broker.rs b/crates/queue/src/broker.rs index e515907..ccee8bd 100644 --- a/crates/queue/src/broker.rs +++ b/crates/queue/src/broker.rs @@ -12,6 +12,8 @@ mod maintenance; mod management; #[path = "broker/metadata_budget.rs"] mod metadata_budget; +#[path = "broker/stats.rs"] +mod stats; #[path = "broker/topics.rs"] mod topics; @@ -19,7 +21,6 @@ use crate::delivery_budget::DeliveryBudget; use crate::management::FenceCatalog; use crate::management_ops::OperationCatalog; use crate::metadata::{load_optional, load_topic_manifest, store_atomic, BrokerMeta}; -use crate::model::{BrokerStats, QueueAggregateStats}; use crate::payload_reader::PayloadReader; use crate::telemetry::QueueMetrics; use crate::topic::index::MessageIndexCache; @@ -37,6 +38,8 @@ use std::sync::Arc; use std::time::Duration; use thiserror::Error; +const FEATURE_LEVEL_2_MAX_MESSAGE_BYTES: usize = 100 * 1024 * 1024; + #[derive(Clone, Debug)] pub struct BrokerConfig { pub data_path: PathBuf, @@ -100,6 +103,12 @@ pub enum BrokerError { ChannelNotFound, #[error("channel is protected by an active deletion tombstone")] ChannelTombstoned, + #[error("channel is not idle: depth={depth}, in_flight={in_flight}, deferred={deferred}")] + ChannelNotIdle { + depth: u64, + in_flight: u64, + deferred: u64, + }, #[error("management fence catalog has not been synchronized")] ManagementUnavailable, #[error("registry revision conflict: expected {expected}, actual {actual}")] @@ -186,14 +195,18 @@ impl Broker { "topic and publish worker limits must be greater than zero".into(), )); } - if config - .max_message_bytes + let readable_message_bytes = if config.storage_feature_level >= 2 { + FEATURE_LEVEL_2_MAX_MESSAGE_BYTES + } else { + config.max_message_bytes + }; + if readable_message_bytes .checked_mul(2) .is_none_or(|minimum| config.delivery_inflight_bytes < minimum) || config.delivery_inflight_bytes > u32::MAX as usize { return Err(BrokerError::InvalidRecord( - "delivery byte budget must fit u32 and the payload read working set".into(), + "delivery byte budget must fit u32 and every message readable at the active storage feature level".into(), )); } ensure_data_format(&config.data_path) @@ -453,56 +466,6 @@ impl Broker { Ok(self.topic(topic)?.state.lock().channel_names()) } - pub fn stats(&self) -> BrokerStats { - let mut topics: Vec<_> = self - .inner - .topics - .read() - .values() - .map(|topic| topic.state.lock().stats()) - .collect(); - topics.sort_by(|left, right| left.name.cmp(&right.name)); - let mut aggregate = QueueAggregateStats::default(); - for topic in &topics { - aggregate.add_topic(topic); - } - BrokerStats { - publish_group_commit: self.inner.publish_groups.stats(), - channel_group_commit: self.inner.channel_groups.stats(), - latency: self.inner.metrics.snapshot(), - delivery_budget: self.inner.delivery_budget.snapshot(), - aggregate, - topics, - } - } - - pub fn metrics_stats(&self, detailed: bool, max_series: usize) -> BrokerStats { - let mut aggregate = QueueAggregateStats::default(); - let mut topics = Vec::new(); - let mut remaining = max_series; - for handle in self.inner.topics.read().values() { - let topic = handle.state.lock(); - topic.add_aggregate_stats(&mut aggregate); - if detailed && remaining > 0 { - let mut stats = topic.stats(); - remaining = remaining.saturating_sub(1); - let channel_limit = remaining / 4; - stats.channels.truncate(channel_limit); - remaining = remaining.saturating_sub(stats.channels.len().saturating_mul(4)); - topics.push(stats); - } - } - topics.sort_by(|left, right| left.name.cmp(&right.name)); - BrokerStats { - publish_group_commit: self.inner.publish_groups.stats(), - channel_group_commit: self.inner.channel_groups.stats(), - latency: self.inner.metrics.snapshot(), - delivery_budget: self.inner.delivery_budget.snapshot(), - aggregate, - topics, - } - } - pub fn registry_revision(&self) -> u64 { self.inner.registry_revision.load(Ordering::Acquire) } diff --git a/crates/queue/src/broker/delivery.rs b/crates/queue/src/broker/delivery.rs index 8d915aa..fdb91d8 100644 --- a/crates/queue/src/broker/delivery.rs +++ b/crates/queue/src/broker/delivery.rs @@ -31,6 +31,7 @@ impl Broker { ) -> Result { self.ensure_storage_healthy()?; self.ensure_management_access(topic, Some(channel))?; + self.expire_channel_in_flight(topic, channel).await?; let handle = self.topic(topic)?; let mut wake = handle.wake.subscribe(); let timeout = timeout.unwrap_or(self.inner.config.message_timeout); diff --git a/crates/queue/src/broker/group_commit.rs b/crates/queue/src/broker/group_commit.rs index aa273eb..1a7bb1d 100644 --- a/crates/queue/src/broker/group_commit.rs +++ b/crates/queue/src/broker/group_commit.rs @@ -415,6 +415,15 @@ pub(super) fn copy_error(error: &BrokerError) -> BrokerError { BrokerError::TopicTombstoned => BrokerError::TopicTombstoned, BrokerError::ChannelNotFound => BrokerError::ChannelNotFound, BrokerError::ChannelTombstoned => BrokerError::ChannelTombstoned, + BrokerError::ChannelNotIdle { + depth, + in_flight, + deferred, + } => BrokerError::ChannelNotIdle { + depth: *depth, + in_flight: *in_flight, + deferred: *deferred, + }, BrokerError::ManagementUnavailable => BrokerError::ManagementUnavailable, BrokerError::RevisionConflict { expected, actual } => BrokerError::RevisionConflict { expected: *expected, diff --git a/crates/queue/src/broker/io.rs b/crates/queue/src/broker/io.rs index b6b368e..6ab8a32 100644 --- a/crates/queue/src/broker/io.rs +++ b/crates/queue/src/broker/io.rs @@ -60,7 +60,8 @@ impl Broker { self.ensure_storage_healthy()?; self.ensure_management_access(topic, Some(channel))?; let available = now_ms().saturating_add(delay.as_millis().min(i64::MAX as u128) as i64); - self.inner + let result = self + .inner .channel_groups .submit( self, @@ -71,7 +72,8 @@ impl Broker { available_at_ms: available, }, ) - .await + .await; + result } pub fn touch( @@ -118,11 +120,12 @@ impl Broker { let path = self .storage_task(move || crate::outbox::store(&directory, &entry_for_write)) .await?; - self.publish( - &entry.target_topic, - vec![entry.body.clone()], - Duration::ZERO, - ) + let broker = self.clone(); + let target_topic = entry.target_topic.clone(); + let body = entry.body.clone(); + self.storage_task(move || { + broker.publish_durable_body_sync(&target_topic, &[body], Duration::ZERO) + }) .await?; self.finish_inner( &entry.source_topic, @@ -135,13 +138,13 @@ impl Broker { .await } - pub(super) fn publish_sync( + fn publish_durable_body_sync( &self, topic: &str, bodies: &[Bytes], delay: Duration, ) -> Result, BrokerError> { - self.validate_publish_request(topic, bodies)?; + self.validate_publish_request_with_limit(topic, bodies, self.durable_message_read_limit())?; let mut metadata = self.reserve_message_metadata(bodies.len())?; let handle = self.get_or_create_topic(topic)?; let mut state = handle.state.lock(); @@ -158,6 +161,15 @@ impl Broker { &self, topic: &str, bodies: &[Bytes], + ) -> Result { + self.validate_publish_request_with_limit(topic, bodies, self.inner.config.max_message_bytes) + } + + fn validate_publish_request_with_limit( + &self, + topic: &str, + bodies: &[Bytes], + max_message_bytes: usize, ) -> Result { validate_name(topic).map_err(|_| BrokerError::InvalidTopic)?; self.ensure_management_access(topic, None)?; @@ -166,7 +178,7 @@ impl Broker { } if bodies .iter() - .any(|body| body.is_empty() || body.len() > self.inner.config.max_message_bytes) + .any(|body| body.is_empty() || body.len() > max_message_bytes) { return Err(BrokerError::MessageTooLarge); } @@ -179,6 +191,14 @@ impl Broker { Ok(encoded_bytes.expect("validated encoded length")) } + fn durable_message_read_limit(&self) -> usize { + if self.inner.compatibility.minimum_reader_feature_level >= 2 { + 100 * 1024 * 1024 + } else { + rustqueue_storage::LEGACY_MAX_RECORD_BYTES.saturating_sub(24) + } + } + pub(super) fn append_publish_to_topic( &self, state: &mut Topic, @@ -220,7 +240,7 @@ impl Broker { pub(super) fn recover_outbox(&self) -> Result<(), BrokerError> { for path in crate::outbox::paths(&self.inner.config.data_path.join("dlq-outbox"))? { let entry = crate::outbox::load(&path)?; - self.publish_sync(&entry.target_topic, &[entry.body], Duration::ZERO)?; + self.publish_durable_body_sync(&entry.target_topic, &[entry.body], Duration::ZERO)?; let finish = self.topic(&entry.source_topic).and_then(|topic| { topic .state diff --git a/crates/queue/src/broker/maintenance.rs b/crates/queue/src/broker/maintenance.rs index e84b322..2719660 100644 --- a/crates/queue/src/broker/maintenance.rs +++ b/crates/queue/src/broker/maintenance.rs @@ -120,13 +120,56 @@ impl Broker { .sum() } - pub fn expire_in_flight(&self) -> usize { - self.inner + pub async fn expire_in_flight(&self) -> Result { + self.ensure_storage_healthy()?; + if !self + .inner .topics .read() .values() - .map(|topic| topic.state.lock().expire_in_flight()) - .sum() + .any(|topic| topic.state.lock().has_expired_in_flight()) + { + return Ok(0); + } + let broker = self.clone(); + self.storage_task(move || { + broker + .inner + .topics + .read() + .values() + .try_fold(0usize, |total, topic| { + Ok(total.saturating_add(topic.state.lock().expire_in_flight()?)) + }) + }) + .await + } + + pub async fn expire_channel_in_flight( + &self, + topic: &str, + channel: &str, + ) -> Result { + self.ensure_storage_healthy()?; + if !self + .topic(topic)? + .state + .lock() + .channel_has_expired_in_flight(channel)? + { + return Ok(0); + } + let broker = self.clone(); + let topic = topic.to_owned(); + let channel = channel.to_owned(); + self.storage_task(move || { + broker + .topic(&topic)? + .state + .lock() + .expire_channel_in_flight(&channel) + }) + .await } pub async fn flush(&self) -> Result<(), BrokerError> { diff --git a/crates/queue/src/broker/management.rs b/crates/queue/src/broker/management.rs index b6247cc..420cd22 100644 --- a/crates/queue/src/broker/management.rs +++ b/crates/queue/src/broker/management.rs @@ -1,6 +1,7 @@ use super::*; use crate::management::{ - ChannelManagementAction, ManagementFenceSnapshot, ManagementResult, TopicManagementAction, + ChannelManagementAction, ChannelManagementCommand, ManagementFenceSnapshot, ManagementResult, + TopicManagementAction, }; use crate::management_ops::OperationLookup; @@ -43,15 +44,15 @@ impl Broker { let _lifecycle = broker.inner.topic_lifecycle.lock(); let fingerprint = serde_json::to_string(&("topic", &topic, action)) .map_err(|error| BrokerError::InvalidRecord(error.to_string()))?; - let replayed = match broker.prepare_management_operation( + let (replayed, operation_id) = match broker.prepare_management_operation( &operation_id, &fingerprint, &topic, expected_revision, )? { PreparedOperation::Completed(result) => return Ok(result), - PreparedOperation::New => false, - PreparedOperation::Pending => true, + PreparedOperation::New(id) => (false, id), + PreparedOperation::Pending(id) => (true, id), }; let mut changed = false; match action { @@ -106,13 +107,17 @@ impl Broker { pub async fn manage_channel( &self, - operation_id: &str, - topic: &str, - channel: &str, - action: ChannelManagementAction, - expected_revision: u64, - tombstone_until_ms: Option, + command: ChannelManagementCommand<'_>, ) -> Result { + let ChannelManagementCommand { + operation_id, + topic, + channel, + action, + expected_revision, + tombstone_until_ms, + require_idle, + } = command; validate_name(topic).map_err(|_| BrokerError::InvalidTopic)?; validate_channel(channel)?; let broker = self.clone(); @@ -121,17 +126,61 @@ impl Broker { let channel = channel.to_owned(); self.storage_task(move || { let _lifecycle = broker.inner.topic_lifecycle.lock(); - let fingerprint = serde_json::to_string(&("channel", &topic, &channel, action)) - .map_err(|error| BrokerError::InvalidRecord(error.to_string()))?; - let replayed = match broker.prepare_management_operation( - &operation_id, - &fingerprint, - &topic, - expected_revision, - )? { - PreparedOperation::Completed(result) => return Ok(result), - PreparedOperation::New => false, - PreparedOperation::Pending => true, + let fingerprint = + serde_json::to_string(&("channel", &topic, &channel, action, require_idle)) + .map_err(|error| BrokerError::InvalidRecord(error.to_string()))?; + let pending_operation = { + let operations = broker.inner.management_ops.lock(); + match operations + .lookup(&operation_id, &fingerprint) + .map_err(operation_catalog_error)? + { + OperationLookup::Completed(result) => return Ok(result), + OperationLookup::Pending => Some(operation_id.clone()), + OperationLookup::New => operations.pending_id(&topic, &fingerprint), + } + }; + let idle_handle = if action == ChannelManagementAction::Delete && require_idle { + match broker.topic(&topic) { + Ok(handle) => Some(handle), + Err(BrokerError::TopicNotFound) if pending_operation.is_some() => None, + Err(error) => return Err(error), + } + } else { + None + }; + let mut idle_state = idle_handle.as_ref().map(|handle| handle.state.lock()); + if let Some(topic_state) = idle_state.as_mut() { + let channel_exists = topic_state + .channel_names() + .iter() + .any(|name| name == &channel); + if !channel_exists && pending_operation.is_none() { + return Err(BrokerError::ChannelNotFound); + } + if channel_exists { + let (depth, in_flight, deferred) = topic_state.channel_counts(&channel)?; + if depth != 0 || in_flight != 0 || deferred != 0 { + return Err(BrokerError::ChannelNotIdle { + depth, + in_flight, + deferred, + }); + } + } + } + let (replayed, operation_id) = match pending_operation { + Some(id) => (true, id), + None => match broker.prepare_management_operation( + &operation_id, + &fingerprint, + &topic, + expected_revision, + )? { + PreparedOperation::Completed(result) => return Ok(result), + PreparedOperation::New(id) => (false, id), + PreparedOperation::Pending(id) => (true, id), + }, }; let mut changed = false; match action { @@ -162,24 +211,44 @@ impl Broker { } ChannelManagementAction::Delete | ChannelManagementAction::Tombstone => { let until = valid_tombstone(tombstone_until_ms, replayed)?; - let mut fences = broker.inner.fences.lock(); - if fences.set_channel(&topic, &channel, until) { - fences.store(&broker.inner.fences_path)?; - changed = true; - } - drop(fences); - if action == ChannelManagementAction::Delete { - if let Ok(handle) = broker.topic(&topic) { - if handle - .state - .lock() + if let Some(mut topic_state) = idle_state.take() { + let mut fences = broker.inner.fences.lock(); + if set_channel_fence(&mut fences, &topic, &channel, until, require_idle) { + fences.store(&broker.inner.fences_path)?; + changed = true; + } + drop(fences); + if action == ChannelManagementAction::Delete + && topic_state + .channel_names() + .iter() + .any(|name| name == &channel) + { + topic_state.delete_channel(&channel)?; + changed = true; + } + } else if let Ok(handle) = broker.topic(&topic) { + let mut topic_state = handle.state.lock(); + let mut fences = broker.inner.fences.lock(); + if set_channel_fence(&mut fences, &topic, &channel, until, require_idle) { + fences.store(&broker.inner.fences_path)?; + changed = true; + } + drop(fences); + if action == ChannelManagementAction::Delete + && topic_state .channel_names() .iter() .any(|name| name == &channel) - { - handle.state.lock().delete_channel(&channel)?; - changed = true; - } + { + topic_state.delete_channel(&channel)?; + changed = true; + } + } else { + let mut fences = broker.inner.fences.lock(); + if set_channel_fence(&mut fences, &topic, &channel, until, require_idle) { + fences.store(&broker.inner.fences_path)?; + changed = true; } } } @@ -241,8 +310,11 @@ impl Broker { .map_err(operation_catalog_error)? { OperationLookup::Completed(result) => Ok(PreparedOperation::Completed(result)), - OperationLookup::Pending => Ok(PreparedOperation::Pending), + OperationLookup::Pending => Ok(PreparedOperation::Pending(operation_id.to_owned())), OperationLookup::New => { + if let Some(pending_id) = operations.pending_id(topic, fingerprint) { + return Ok(PreparedOperation::Pending(pending_id)); + } self.check_revision(expected_revision)?; operations .prepare( @@ -252,7 +324,7 @@ impl Broker { topic.to_owned(), ) .map_err(operation_catalog_error)?; - Ok(PreparedOperation::New) + Ok(PreparedOperation::New(operation_id.to_owned())) } } } @@ -272,8 +344,8 @@ impl Broker { } enum PreparedOperation { - New, - Pending, + New(String), + Pending(String), Completed(ManagementResult), } @@ -283,6 +355,20 @@ fn valid_tombstone(value: Option, replayed: bool) -> Result bool { + if local { + fences.set_local_channel(topic, channel, until_ms) + } else { + fences.set_channel(topic, channel, until_ms) + } +} + fn operation_catalog_error(error: std::io::Error) -> BrokerError { if error.kind() == std::io::ErrorKind::InvalidInput { BrokerError::OperationConflict diff --git a/crates/queue/src/broker/stats.rs b/crates/queue/src/broker/stats.rs new file mode 100644 index 0000000..5b91e45 --- /dev/null +++ b/crates/queue/src/broker/stats.rs @@ -0,0 +1,108 @@ +use super::*; +use crate::model::{BrokerStats, QueueAggregateStats, TopicStats}; + +impl Broker { + pub fn stats(&self) -> BrokerStats { + self.filtered_stats(None, None) + } + + pub fn filtered_stats(&self, topic: Option<&str>, channel: Option<&str>) -> BrokerStats { + let mut topics = self.topic_stats(topic); + if let Some(channel) = channel { + for topic in &mut topics { + topic.channels.retain(|candidate| candidate.name == channel); + } + } + topics.sort_by(|left, right| left.name.cmp(&right.name)); + self.snapshot(topics) + } + + pub fn metrics_stats(&self, detailed: bool, max_series: usize) -> BrokerStats { + let handles: Vec<_> = self.inner.topics.read().values().cloned().collect(); + let mut aggregate = QueueAggregateStats::default(); + let mut topics = Vec::new(); + let mut remaining = max_series; + for handle in handles { + let mut topic = handle.state.lock(); + topic.add_aggregate_stats(&mut aggregate); + if detailed && remaining > 0 { + let mut stats = topic.stats(); + remaining = remaining.saturating_sub(1); + let channel_limit = remaining / 4; + stats.channels.truncate(channel_limit); + remaining = remaining.saturating_sub(stats.channels.len().saturating_mul(4)); + topics.push(stats); + } + } + topics.sort_by(|left, right| left.name.cmp(&right.name)); + BrokerStats { + aggregate, + ..self.snapshot(topics) + } + } + + fn topic_stats(&self, topic: Option<&str>) -> Vec { + let handles: Vec<_> = match topic { + Some(name) => self + .inner + .topics + .read() + .get(name) + .cloned() + .into_iter() + .collect(), + None => self.inner.topics.read().values().cloned().collect(), + }; + handles + .into_iter() + .map(|topic| topic.state.lock().stats()) + .collect() + } + + fn snapshot(&self, topics: Vec) -> BrokerStats { + let mut aggregate = QueueAggregateStats::default(); + for topic in &topics { + aggregate.add_topic(topic); + } + BrokerStats { + publish_group_commit: self.inner.publish_groups.stats(), + channel_group_commit: self.inner.channel_groups.stats(), + latency: self.inner.metrics.snapshot(), + delivery_budget: self.inner.delivery_budget.snapshot(), + aggregate, + topics, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[tokio::test] + async fn filtered_stats_only_return_the_requested_topic_and_channel() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + broker.create_channel("events", "audit").await.unwrap(); + broker.create_channel("other", "workers").await.unwrap(); + + let stats = broker.filtered_stats(Some("events"), Some("workers")); + assert_eq!(stats.aggregate.topic_count, 1); + assert_eq!(stats.aggregate.channel_count, 1); + assert_eq!(stats.topics.len(), 1); + assert_eq!(stats.topics[0].name, "events"); + assert_eq!(stats.topics[0].channels.len(), 1); + assert_eq!(stats.topics[0].channels[0].name, "workers"); + + assert!(broker + .filtered_stats(Some("missing"), None) + .topics + .is_empty()); + } +} diff --git a/crates/queue/src/broker_tests.rs b/crates/queue/src/broker_tests.rs index 0c0eab8..5a5629e 100644 --- a/crates/queue/src/broker_tests.rs +++ b/crates/queue/src/broker_tests.rs @@ -1,7 +1,10 @@ use super::io::SEQUENCE_RESERVATION; use super::*; use crate::outbox::OutboxEntry; -use crate::{ManagementFenceSnapshot, TopicManagementAction}; +use crate::{ + ChannelManagementAction, ChannelManagementCommand, ManagementFenceSnapshot, + TopicManagementAction, +}; use futures::future::join_all; use std::collections::{BTreeMap, HashSet}; use std::sync::atomic::{AtomicBool, Ordering}; @@ -127,6 +130,174 @@ async fn startup_replays_dlq_outbox_before_finishing_the_source() { assert_eq!(&*dlq.body, b"poison"); } +#[tokio::test] +async fn stats_settle_expired_in_flight_messages_without_another_fetch() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + broker + .publish("events", vec![b"body".to_vec()], Duration::ZERO) + .await + .unwrap(); + assert!(broker + .next_message("events", "workers", Some(Duration::ZERO)) + .await + .unwrap() + .is_some()); + + broker.expire_in_flight().await.unwrap(); + let stats = broker.stats(); + let channel = &stats.topics[0].channels[0]; + assert_eq!(channel.in_flight_count, 0); + assert_eq!(channel.timeout_count, 1); + assert_eq!(channel.depth, 1); +} + +#[tokio::test] +async fn kodo_channel_counters_are_monotonic_across_empty_and_restart() { + if rustqueue_storage::MAX_WRITER_FEATURE_LEVEL < 2 { + return; + } + let root = tempdir().unwrap(); + let config = BrokerConfig { + data_path: root.path().into(), + storage_feature_level: 2, + ..BrokerConfig::default() + }; + let broker = Broker::open(config.clone()).unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + broker + .publish( + "events", + vec![b"one".to_vec(), b"two".to_vec()], + Duration::ZERO, + ) + .await + .unwrap(); + let delivery = broker + .next_message("events", "workers", None) + .await + .unwrap() + .unwrap(); + broker + .requeue("events", "workers", delivery.id, Duration::ZERO) + .await + .unwrap(); + assert!(broker + .next_message("events", "workers", Some(Duration::ZERO)) + .await + .unwrap() + .is_some()); + broker.expire_in_flight().await.unwrap(); + broker.empty_channel("events", "workers").await.unwrap(); + broker + .publish("events", vec![b"three".to_vec()], Duration::ZERO) + .await + .unwrap(); + + let channel = &broker.stats().topics[0].channels[0]; + assert_eq!(channel.message_count, 3); + assert_eq!(channel.requeue_count, 1); + assert_eq!(channel.timeout_count, 1); + drop(broker); + + let reopened = Broker::open(config).unwrap(); + let channel = &reopened.stats().topics[0].channels[0]; + assert_eq!(channel.message_count, 3); + assert_eq!(channel.requeue_count, 1); + assert_eq!(channel.timeout_count, 1); +} + +#[tokio::test] +async fn durable_batch_accepts_the_full_protocol_message_count() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + let messages = vec![vec![b'x']; rustqueue_protocol::MAX_MPUB_MESSAGES]; + let ids = broker + .publish("events", messages, Duration::ZERO) + .await + .unwrap(); + assert_eq!(ids.len(), rustqueue_protocol::MAX_MPUB_MESSAGES); + assert_eq!( + broker.stats().topics[0].published_count, + rustqueue_protocol::MAX_MPUB_MESSAGES as u64 + ); +} + +#[tokio::test] +async fn startup_replays_a_large_dlq_outbox_after_lowering_the_publish_limit() { + if rustqueue_storage::MAX_WRITER_FEATURE_LEVEL < 2 { + return; + } + let root = tempdir().unwrap(); + let original = BrokerConfig { + data_path: root.path().into(), + max_segment_bytes: 256 * 1024 * 1024, + max_message_bytes: 100 * 1024 * 1024, + storage_feature_level: 2, + ..BrokerConfig::default() + }; + let broker = Broker::open(original.clone()).unwrap(); + let body = bytes::Bytes::from(vec![0x6b; 100 * 1024 * 1024]); + let entry = OutboxEntry { + source_topic: "missing-source".into(), + source_channel: "missing-channel".into(), + message_id: 1, + target_topic: "events.DLQ".into(), + body: body.clone(), + }; + crate::outbox::store(&root.path().join("dlq-outbox"), &entry).unwrap(); + drop(broker); + + let broker = Broker::open(BrokerConfig { + max_message_bytes: 20 * 1024 * 1024, + ..original + }) + .unwrap(); + broker + .create_channel("events.DLQ", "inspect") + .await + .unwrap(); + let delivery = broker + .next_message("events.DLQ", "inspect", None) + .await + .unwrap() + .unwrap(); + assert_eq!(delivery.body.as_ref(), body.as_ref()); + assert_eq!( + std::fs::read_dir(root.path().join("dlq-outbox")) + .unwrap() + .count(), + 0 + ); +} + +#[test] +fn feature_level_two_rejects_a_delivery_budget_that_cannot_read_retained_messages() { + let root = tempdir().unwrap(); + let error = match Broker::open(BrokerConfig { + data_path: root.path().into(), + max_message_bytes: 20 * 1024 * 1024, + delivery_inflight_bytes: 40 * 1024 * 1024, + storage_feature_level: 2, + ..BrokerConfig::default() + }) { + Ok(_) => panic!("undersized delivery budget was accepted"), + Err(error) => error, + }; + assert!(error + .to_string() + .contains("every message readable at the active storage feature level")); +} + #[tokio::test] async fn restart_never_reuses_ids_from_a_reserved_sequence_block() { let root = tempdir().unwrap(); @@ -463,6 +634,64 @@ async fn idle_gc_does_not_seal_a_channel_blocked_active_segment() { assert_eq!(stats.topics[0].message_count, 1); } +#[tokio::test] +async fn deferred_stats_are_exact_without_a_consumer_and_after_restart() { + let root = tempdir().unwrap(); + let config = BrokerConfig { + data_path: root.path().into(), + max_ack_gap: 2, + ..BrokerConfig::default() + }; + let broker = Broker::open(config.clone()).unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + broker + .publish( + "events", + vec![b"one".to_vec(), b"two".to_vec(), b"three".to_vec()], + Duration::ZERO, + ) + .await + .unwrap(); + broker + .publish( + "events", + vec![b"deferred".to_vec()], + Duration::from_secs(60), + ) + .await + .unwrap(); + let channel = &broker.stats().topics[0].channels[0]; + assert_eq!((channel.depth, channel.deferred_count), (4, 1)); + drop(broker); + + let reopened = Broker::open(config).unwrap(); + let channel = &reopened.stats().topics[0].channels[0]; + assert_eq!((channel.depth, channel.deferred_count), (4, 1)); +} + +#[tokio::test] +async fn deferred_stats_promote_messages_when_the_deadline_passes() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + broker + .publish( + "events", + vec![b"deferred".to_vec()], + Duration::from_millis(20), + ) + .await + .unwrap(); + assert_eq!(broker.stats().topics[0].channels[0].deferred_count, 1); + tokio::time::sleep(Duration::from_millis(40)).await; + let channel = &broker.stats().topics[0].channels[0]; + assert_eq!((channel.depth, channel.deferred_count), (1, 0)); +} + #[tokio::test] async fn bounded_gc_rotates_across_topics() { let root = tempdir().unwrap(); @@ -789,6 +1018,121 @@ async fn management_rejects_stale_registry_revisions() { assert!(!broker.stats().topics[0].paused); } +#[tokio::test] +async fn delete_if_idle_rejects_backlog_without_leaving_a_management_block() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + broker + .publish("events", vec![b"one".to_vec()], Duration::ZERO) + .await + .unwrap(); + let revision = broker.registry_revision(); + let deadline = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() + .min(i64::MAX as u128) as i64 + + 60_000; + let error = broker + .manage_channel(ChannelManagementCommand { + operation_id: "delete-idle-test-0001", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: revision, + tombstone_until_ms: Some(deadline), + require_idle: true, + }) + .await + .unwrap_err(); + assert!(matches!( + error, + BrokerError::ChannelNotIdle { + depth: 1, + in_flight: 0, + deferred: 0 + } + )); + + broker + .publish("events", vec![b"two".to_vec()], Duration::ZERO) + .await + .unwrap(); + broker.empty_channel("events", "workers").await.unwrap(); + broker + .manage_channel(ChannelManagementCommand { + operation_id: "delete-idle-test-0002", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: broker.registry_revision(), + tombstone_until_ms: Some(deadline), + require_idle: true, + }) + .await + .unwrap(); + assert!(broker.channel_names("events").unwrap().is_empty()); +} + +#[tokio::test] +async fn delete_if_idle_reports_missing_topics_and_channels() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + let deadline = now_ms() + 60_000; + + let missing_topic = broker + .manage_channel(ChannelManagementCommand { + operation_id: "delete-missing-topic-0001", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: broker.registry_revision(), + tombstone_until_ms: Some(deadline), + require_idle: true, + }) + .await + .unwrap_err(); + assert!(matches!(missing_topic, BrokerError::TopicNotFound)); + + broker.create_channel("events", "workers").await.unwrap(); + broker.empty_channel("events", "workers").await.unwrap(); + broker + .manage_channel(ChannelManagementCommand { + operation_id: "delete-existing-channel-0001", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: broker.registry_revision(), + tombstone_until_ms: Some(deadline), + require_idle: true, + }) + .await + .unwrap(); + + let missing_channel = broker + .manage_channel(ChannelManagementCommand { + operation_id: "delete-missing-channel-0001", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: broker.registry_revision(), + tombstone_until_ms: Some(deadline), + require_idle: true, + }) + .await + .unwrap_err(); + assert!(matches!(missing_channel, BrokerError::ChannelNotFound)); +} + #[tokio::test] async fn management_operation_results_are_idempotent_across_restart() { let root = tempdir().unwrap(); @@ -889,6 +1233,110 @@ async fn expired_pending_tombstone_can_finish_and_unblock_the_topic() { .unwrap(); } +#[tokio::test] +async fn changed_cleanup_operation_id_adopts_a_pending_delete_after_restart() { + let root = tempdir().unwrap(); + let config = BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }; + let broker = Broker::open(config.clone()).unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + let fingerprint = serde_json::to_string(&( + "channel", + "events", + "workers", + ChannelManagementAction::Delete, + true, + )) + .unwrap(); + broker + .inner + .management_ops + .lock() + .prepare( + &broker.inner.management_ops_path, + "kodo-delete-old-0001", + fingerprint, + "events".into(), + ) + .unwrap(); + drop(broker); + + let broker = Broker::open(config).unwrap(); + broker + .manage_channel(ChannelManagementCommand { + operation_id: "kodo-delete-new-0001", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: 0, + tombstone_until_ms: Some(now_ms() + 60_000), + require_idle: true, + }) + .await + .unwrap(); + assert!(broker.channel_names("events").unwrap().is_empty()); + assert!(!broker.inner.management_ops.lock().blocks_topic("events")); +} + +#[tokio::test] +async fn pending_idle_delete_completes_when_the_channel_was_already_removed() { + let root = tempdir().unwrap(); + let config = BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }; + let broker = Broker::open(config.clone()).unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + let fingerprint = serde_json::to_string(&( + "channel", + "events", + "workers", + ChannelManagementAction::Delete, + true, + )) + .unwrap(); + broker + .inner + .management_ops + .lock() + .prepare( + &broker.inner.management_ops_path, + "kodo-delete-old-0002", + fingerprint, + "events".into(), + ) + .unwrap(); + broker + .topic("events") + .unwrap() + .state + .lock() + .delete_channel("workers") + .unwrap(); + drop(broker); + + let broker = Broker::open(config).unwrap(); + broker + .manage_channel(ChannelManagementCommand { + operation_id: "kodo-delete-new-0002", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: 0, + tombstone_until_ms: Some(now_ms() + 60_000), + require_idle: true, + }) + .await + .unwrap(); + assert!(!broker.inner.management_ops.lock().blocks_topic("events")); + broker + .publish("events", vec![b"accepted".to_vec()], Duration::ZERO) + .await + .unwrap(); +} + #[tokio::test] async fn concurrent_registry_updates_persist_the_latest_revision() { let root = tempdir().unwrap(); diff --git a/crates/queue/src/channel.rs b/crates/queue/src/channel.rs index e5e2465..bbc0b2b 100644 --- a/crates/queue/src/channel.rs +++ b/crates/queue/src/channel.rs @@ -16,6 +16,8 @@ pub(crate) enum ChannelCommand { message_id: u64, available_at_ms: i64, attempts: u16, + #[serde(default)] + cumulative_count: Option, }, Pause { paused: bool, @@ -26,6 +28,9 @@ pub(crate) enum ChannelCommand { Evict { through_position: u64, }, + Timeout { + cumulative_count: u64, + }, } #[derive(Clone, Debug, Deserialize, Serialize)] @@ -38,6 +43,12 @@ pub(crate) struct ChannelCheckpoint { pub requeued_until: BTreeMap, #[serde(default)] pub attempts: BTreeMap, + #[serde(default)] + pub message_count_origin_position: Option, + #[serde(default)] + pub requeue_count: u64, + #[serde(default)] + pub timeout_count: u64, pub paused: bool, pub ephemeral: bool, } @@ -50,6 +61,7 @@ struct InFlight { pub(crate) struct ChannelRuntime { pub state: ChannelState, pub store: Option, + pub durable_counters: bool, } pub(crate) struct ChannelState { @@ -60,6 +72,7 @@ pub(crate) struct ChannelState { pub requeued_until: BTreeMap, pub paused: bool, pub ephemeral: bool, + message_count_origin_position: u64, next_position: u64, in_flight: HashMap, in_flight_ids: HashMap, @@ -67,6 +80,8 @@ pub(crate) struct ChannelState { attempts: HashMap, next_token: u64, max_ack_gap: usize, + requeue_count: u64, + timeout_count: u64, } pub(crate) enum MessageAvailability { @@ -91,6 +106,7 @@ impl ChannelState { requeued_until: BTreeMap::new(), paused: false, ephemeral, + message_count_origin_position: barrier_position, next_position: barrier_position.saturating_add(1), in_flight: HashMap::new(), in_flight_ids: HashMap::new(), @@ -98,6 +114,8 @@ impl ChannelState { attempts: HashMap::new(), next_token: 1, max_ack_gap: max_ack_gap.max(1), + requeue_count: 0, + timeout_count: 0, } } @@ -105,12 +123,19 @@ impl ChannelState { checkpoint: ChannelCheckpoint, max_ack_gap: usize, ) -> Result { - if checkpoint.format != 7 || checkpoint.ack_floor_position < checkpoint.barrier_position { + let message_count_origin_position = checkpoint + .message_count_origin_position + .unwrap_or(checkpoint.barrier_position); + if checkpoint.format != 7 + || checkpoint.ack_floor_position < checkpoint.barrier_position + || message_count_origin_position > checkpoint.barrier_position + { return Err(BrokerError::InvalidRecord( "invalid channel checkpoint".into(), )); } let next_position = checkpoint.ack_floor_position.saturating_add(1); + let redelivery = checkpoint.requeued_until.keys().copied().collect(); Ok(Self { name: checkpoint.name, barrier_position: checkpoint.barrier_position, @@ -119,13 +144,16 @@ impl ChannelState { requeued_until: checkpoint.requeued_until, paused: checkpoint.paused, ephemeral: checkpoint.ephemeral, + message_count_origin_position, next_position, in_flight: HashMap::new(), in_flight_ids: HashMap::new(), - redelivery: BTreeSet::new(), + redelivery, attempts: checkpoint.attempts.into_iter().collect(), next_token: 1, max_ack_gap: max_ack_gap.max(1), + requeue_count: checkpoint.requeue_count, + timeout_count: checkpoint.timeout_count, }) } @@ -142,6 +170,9 @@ impl ChannelState { .iter() .map(|(position, attempts)| (*position, *attempts)) .collect(), + message_count_origin_position: Some(self.message_count_origin_position), + requeue_count: self.requeue_count, + timeout_count: self.timeout_count, paused: self.paused, ephemeral: self.ephemeral, } @@ -154,12 +185,20 @@ impl ChannelState { position, available_at_ms, attempts, + cumulative_count, .. } => { + let already_applied = self.requeued_until.get(&position) == Some(&available_at_ms) + && self.attempts.get(&position) == Some(&attempts); self.remove_in_flight(position); self.redelivery.insert(position); self.requeued_until.insert(position, available_at_ms); self.attempts.insert(position, attempts); + if let Some(count) = cumulative_count { + self.requeue_count = self.requeue_count.max(count); + } else if !already_applied { + self.requeue_count = self.requeue_count.saturating_add(1); + } } ChannelCommand::Pause { paused } => self.paused = paused, ChannelCommand::Empty { through_position } @@ -184,6 +223,9 @@ impl ChannelState { .retain(|position| *position > through_position); self.next_position = self.next_position.max(through_position.saturating_add(1)); } + ChannelCommand::Timeout { cumulative_count } => { + self.timeout_count = self.timeout_count.max(cumulative_count); + } } } @@ -196,7 +238,6 @@ impl ChannelState { if self.paused { return NextCandidate::None; } - self.expire_in_flight(); let mut absent = Vec::new(); let redelivery: Vec<_> = self.redelivery.iter().copied().collect(); for position in redelivery { @@ -236,6 +277,7 @@ impl ChannelState { if position <= self.ack_floor_position || self.acknowledged.contains(&position) || self.in_flight.contains_key(&position) + || self.redelivery.contains(&position) { self.next_position = self.next_position.saturating_add(1); continue; @@ -304,6 +346,10 @@ impl ChannelState { self.attempts.get(&position).copied().unwrap_or_default() } + pub fn next_requeue_count(&self) -> u64 { + self.requeue_count.saturating_add(1) + } + pub fn touch(&mut self, position: u64, timeout: Duration) -> bool { let Some(flight) = self.in_flight.get_mut(&position) else { return false; @@ -328,32 +374,40 @@ impl ChannelState { count } - pub fn expire_in_flight(&mut self) -> usize { - let now = Instant::now(); - let expired: Vec<_> = self - .in_flight + fn expired_positions(&self, now: Instant) -> Vec { + self.in_flight .iter() .filter_map(|(position, flight)| (flight.deadline <= now).then_some(*position)) - .collect(); - let count = expired.len(); - for position in expired { - self.remove_in_flight(position); - self.redelivery.insert(position); + .collect() + } + + fn expire_positions(&mut self, positions: &[u64]) { + for position in positions { + self.remove_in_flight(*position); + self.redelivery.insert(*position); } - count } pub fn first_in_flight_position(&self) -> Option { self.in_flight.keys().copied().min() } - pub fn stats(&self, last_position: u64) -> ChannelStats { - let (depth, in_flight_count, deferred_count, ack_gap) = self.metric_counts(last_position); + pub fn stats( + &self, + last_position: u64, + scheduled: &BTreeSet, + now_ms: i64, + ) -> ChannelStats { + let (depth, in_flight_count, deferred_count, ack_gap) = + self.metric_counts(last_position, scheduled, now_ms); ChannelStats { name: self.name.clone(), depth, + message_count: last_position.saturating_sub(self.message_count_origin_position), in_flight_count, deferred_count, + requeue_count: self.requeue_count, + timeout_count: self.timeout_count, paused: self.paused, ephemeral: self.ephemeral, ack_cursor: self.ack_floor_position, @@ -361,16 +415,41 @@ impl ChannelState { } } - pub fn metric_counts(&self, last_position: u64) -> (u64, u64, u64, u64) { + pub fn metric_counts( + &self, + last_position: u64, + scheduled: &BTreeSet, + now_ms: i64, + ) -> (u64, u64, u64, u64) { let total = last_position.saturating_sub(self.ack_floor_position); + let scheduled_count = scheduled + .iter() + .filter(|position| self.is_outstanding(**position, last_position)) + .count() as u64; + let requeued_count = self + .requeued_until + .iter() + .filter(|(position, until)| { + **until > now_ms + && !scheduled.contains(position) + && self.is_outstanding(**position, last_position) + }) + .count() as u64; ( total.saturating_sub(self.acknowledged.len() as u64), self.in_flight.len() as u64, - self.requeued_until.len() as u64, + scheduled_count.saturating_add(requeued_count), self.acknowledged.len() as u64, ) } + fn is_outstanding(&self, position: u64, last_position: u64) -> bool { + position > self.ack_floor_position + && position <= last_position + && !self.acknowledged.contains(&position) + && !self.in_flight.contains_key(&position) + } + fn acknowledge(&mut self, position: u64) { self.remove_in_flight(position); self.redelivery.remove(&position); @@ -398,6 +477,38 @@ impl ChannelState { } } +impl ChannelRuntime { + pub fn has_expired_in_flight(&self) -> bool { + !self.state.expired_positions(Instant::now()).is_empty() + } + + pub fn expire_in_flight(&mut self) -> Result { + let positions = self.state.expired_positions(Instant::now()); + if positions.is_empty() { + return Ok(0); + } + let cumulative_count = self + .state + .timeout_count + .saturating_add(positions.len() as u64); + if self.durable_counters { + let command = ChannelCommand::Timeout { cumulative_count }; + if let Some(store) = self.store.as_mut() { + store.append(&command)?; + } + self.state.expire_positions(&positions); + self.state.apply(&command); + } else { + self.state.expire_positions(&positions); + self.state.timeout_count = cumulative_count; + } + if let Some(store) = self.store.as_mut() { + store.checkpoint_if_needed(&self.state)?; + } + Ok(positions.len()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -418,7 +529,71 @@ mod tests { fn in_flight_expiry_uses_a_monotonic_deadline() { let mut channel = ChannelState::new("workers".into(), 0, false, 16); channel.reserve(1, 10, Duration::ZERO); - assert_eq!(channel.expire_in_flight(), 1); + let mut runtime = ChannelRuntime { + state: channel, + store: None, + durable_counters: true, + }; + assert_eq!(runtime.expire_in_flight().unwrap(), 1); + let channel = runtime.state; assert_eq!(channel.in_flight_position(10), None); + assert_eq!( + channel.stats(1, &BTreeSet::new(), i64::MAX).timeout_count, + 1 + ); + } + + #[test] + fn cumulative_counters_and_message_count_survive_checkpoint_and_empty() { + let mut channel = ChannelState::new("workers".into(), 0, false, 16); + channel.reserve(1, 10, Duration::ZERO); + let mut runtime = ChannelRuntime { + state: channel, + store: None, + durable_counters: true, + }; + runtime.state.apply(&ChannelCommand::Requeue { + position: 1, + message_id: 10, + available_at_ms: 0, + attempts: 1, + cumulative_count: Some(1), + }); + runtime.state.reserve(2, 11, Duration::ZERO); + runtime.expire_in_flight().unwrap(); + runtime.state.apply(&ChannelCommand::Empty { + through_position: 2, + }); + let current = runtime.state.stats(2, &BTreeSet::new(), i64::MAX); + assert_eq!(current.message_count, 2); + assert_eq!(current.requeue_count, 1); + assert_eq!(current.timeout_count, 1); + + let recovered = ChannelState::from_checkpoint(runtime.state.checkpoint(), 16).unwrap(); + let restarted = recovered.stats(4, &BTreeSet::new(), i64::MAX); + assert_eq!(restarted.message_count, 4); + assert_eq!(restarted.requeue_count, 1); + assert_eq!(restarted.timeout_count, 1); + } + + #[test] + fn legacy_checkpoint_starts_monotonic_count_at_its_persisted_barrier() { + let checkpoint: ChannelCheckpoint = serde_json::from_value(serde_json::json!({ + "format": 7, + "name": "workers", + "barrier_position": 5, + "ack_floor_position": 5, + "acknowledged": [], + "requeued_until": {}, + "paused": false, + "ephemeral": false + })) + .unwrap(); + let mut state = ChannelState::from_checkpoint(checkpoint, 16).unwrap(); + assert_eq!(state.stats(7, &BTreeSet::new(), i64::MAX).message_count, 2); + state.apply(&ChannelCommand::Empty { + through_position: 7, + }); + assert_eq!(state.stats(7, &BTreeSet::new(), i64::MAX).message_count, 2); } } diff --git a/crates/queue/src/channel_store.rs b/crates/queue/src/channel_store.rs index b89a8eb..226ef0e 100644 --- a/crates/queue/src/channel_store.rs +++ b/crates/queue/src/channel_store.rs @@ -373,12 +373,16 @@ fn encode_command(command: &ChannelCommand) -> Vec { message_id, available_at_ms, attempts, + cumulative_count, } => { - body.push(2); + body.push(if cumulative_count.is_some() { 7 } else { 2 }); body.extend_from_slice(&position.to_be_bytes()); body.extend_from_slice(&message_id.to_be_bytes()); body.extend_from_slice(&available_at_ms.to_be_bytes()); body.extend_from_slice(&attempts.to_be_bytes()); + if let Some(count) = cumulative_count { + body.extend_from_slice(&count.to_be_bytes()); + } } ChannelCommand::Pause { paused } => { body.extend_from_slice(&[3, paused as u8]); @@ -391,6 +395,10 @@ fn encode_command(command: &ChannelCommand) -> Vec { body.push(5); body.extend_from_slice(&through_position.to_be_bytes()); } + ChannelCommand::Timeout { cumulative_count } => { + body.push(6); + body.extend_from_slice(&cumulative_count.to_be_bytes()); + } } body } @@ -407,6 +415,7 @@ fn decode_command(body: &[u8]) -> Result { message_id: u64::from_be_bytes(body[9..17].try_into().unwrap()), available_at_ms: i64::from_be_bytes(body[17..25].try_into().unwrap()), attempts: u16::from_be_bytes(body[25..27].try_into().unwrap()), + cumulative_count: None, }), Some(3) if body.len() == 2 && body[1] <= 1 => Ok(ChannelCommand::Pause { paused: body[1] == 1, @@ -417,6 +426,16 @@ fn decode_command(body: &[u8]) -> Result { Some(5) if body.len() == 9 => Ok(ChannelCommand::Evict { through_position: u64::from_be_bytes(body[1..9].try_into().unwrap()), }), + Some(6) if body.len() == 9 => Ok(ChannelCommand::Timeout { + cumulative_count: u64::from_be_bytes(body[1..9].try_into().unwrap()), + }), + Some(7) if body.len() == 35 => Ok(ChannelCommand::Requeue { + position: u64::from_be_bytes(body[1..9].try_into().unwrap()), + message_id: u64::from_be_bytes(body[9..17].try_into().unwrap()), + available_at_ms: i64::from_be_bytes(body[17..25].try_into().unwrap()), + attempts: u16::from_be_bytes(body[25..27].try_into().unwrap()), + cumulative_count: Some(u64::from_be_bytes(body[27..35].try_into().unwrap())), + }), _ => Err(invalid()), } } @@ -453,6 +472,80 @@ mod tests { assert_eq!(fs::metadata(wal).unwrap().len(), (HEADER_LEN + 17) as u64); } + #[test] + fn requeue_and_timeout_counters_survive_wal_recovery() { + let root = tempdir().unwrap(); + let state = ChannelState::new("workers".into(), 0, false, 65_536); + let mut store = ChannelStore::create(root.path(), &state).unwrap(); + store + .append(&ChannelCommand::Requeue { + position: 1, + message_id: 7, + available_at_ms: 0, + attempts: 1, + cumulative_count: Some(1), + }) + .unwrap(); + store + .append(&ChannelCommand::Timeout { + cumulative_count: 2, + }) + .unwrap(); + drop(store); + + let checkpoint = root + .path() + .join(format!("{}.checkpoint", hex::encode("workers"))); + let (state, _) = ChannelStore::open(&checkpoint, 65_536).unwrap(); + let stats = state.stats(1, &Default::default(), i64::MAX); + assert_eq!(stats.requeue_count, 1); + assert_eq!(stats.timeout_count, 2); + } + + #[test] + fn checkpoint_before_wal_reset_does_not_double_absolute_counters() { + let root = tempdir().unwrap(); + let mut state = ChannelState::new("workers".into(), 0, false, 65_536); + let mut store = ChannelStore::create(root.path(), &state).unwrap(); + let requeue = ChannelCommand::Requeue { + position: 1, + message_id: 7, + available_at_ms: 1_000, + attempts: 1, + cumulative_count: Some(1), + }; + store.append(&requeue).unwrap(); + state.apply(&requeue); + let timeout = ChannelCommand::Timeout { + cumulative_count: 2, + }; + store.append(&timeout).unwrap(); + state.apply(&timeout); + + // Simulate a crash after the checkpoint rename but before the old WAL + // is reset. Recovery must tolerate replaying both absolute commands. + write_checkpoint(&store.checkpoint_path, &state.checkpoint()).unwrap(); + drop(store); + + let checkpoint = root + .path() + .join(format!("{}.checkpoint", hex::encode("workers"))); + let (mut recovered, _) = ChannelStore::open(&checkpoint, 65_536).unwrap(); + let stats = recovered.stats(1, &Default::default(), 0); + assert_eq!(stats.requeue_count, 1); + assert_eq!(stats.timeout_count, 2); + assert!(matches!( + recovered.next_candidate(0, 1, |_| crate::channel::MessageAvailability::Ready(0)), + crate::channel::NextCandidate::None + )); + assert!(matches!( + recovered.next_candidate(1_000, 1, |_| { + crate::channel::MessageAvailability::Ready(0) + }), + crate::channel::NextCandidate::Ready(1) + )); + } + #[test] fn refuses_a_checksum_corrupt_complete_wal_tail() { let root = tempdir().unwrap(); @@ -495,6 +588,38 @@ mod tests { assert!(decode_command(&encoded).is_err()); } + #[test] + fn feature_one_requeue_keeps_the_legacy_wal_encoding() { + let legacy = ChannelCommand::Requeue { + position: 1, + message_id: 7, + available_at_ms: 10, + attempts: 1, + cumulative_count: None, + }; + let encoded = encode_command(&legacy); + assert_eq!(encoded[0], 2); + assert_eq!(encoded.len(), 27); + assert!(matches!( + decode_command(&encoded).unwrap(), + ChannelCommand::Requeue { + cumulative_count: None, + .. + } + )); + + let durable = ChannelCommand::Requeue { + position: 1, + message_id: 7, + available_at_ms: 10, + attempts: 1, + cumulative_count: Some(1), + }; + let encoded = encode_command(&durable); + assert_eq!(encoded[0], 7); + assert_eq!(encoded.len(), 35); + } + #[test] fn recreating_a_channel_discards_an_orphaned_old_wal() { let root = tempdir().unwrap(); diff --git a/crates/queue/src/delivery_budget.rs b/crates/queue/src/delivery_budget.rs index 4b605a6..a30e635 100644 --- a/crates/queue/src/delivery_budget.rs +++ b/crates/queue/src/delivery_budget.rs @@ -49,6 +49,12 @@ impl DeliveryBudget { .ok_or_else(|| { BrokerError::InvalidRecord("delivery working set overflows usize".into()) })?; + if bytes > self.capacity { + return Err(BrokerError::InvalidRecord(format!( + "delivery working set {bytes} exceeds configured byte budget {}", + self.capacity + ))); + } let permits = u32::try_from(bytes).map_err(|_| { BrokerError::InvalidRecord("delivery batch exceeds the byte budget contract".into()) })?; @@ -102,3 +108,19 @@ impl Drop for Waiter { self.0.fetch_sub(1, Ordering::AcqRel); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn oversized_request_fails_instead_of_waiting_forever() { + let budget = DeliveryBudget::new(8); + let error = match budget.acquire(5).await { + Ok(_) => panic!("oversized delivery request was accepted"), + Err(error) => error, + }; + assert!(error.to_string().contains("exceeds configured byte budget")); + assert_eq!(budget.snapshot().waiters, 0); + } +} diff --git a/crates/queue/src/lib.rs b/crates/queue/src/lib.rs index 30a0b3b..1bc31e4 100644 --- a/crates/queue/src/lib.rs +++ b/crates/queue/src/lib.rs @@ -19,8 +19,8 @@ pub use delivery_budget::DeliveryHold; pub use delivery_guard::DeliveryGuard; pub use eviction::ProtectiveEviction; pub use management::{ - ChannelFence, ChannelManagementAction, ManagementFenceSnapshot, ManagementResult, - TopicManagementAction, + ChannelFence, ChannelManagementAction, ChannelManagementCommand, ManagementFenceSnapshot, + ManagementResult, TopicManagementAction, }; pub use model::{ BrokerLatencyStats, BrokerStats, ChannelGroupCommitStats, ChannelStats, Delivery, diff --git a/crates/queue/src/management.rs b/crates/queue/src/management.rs index 0acbcbd..639fcdc 100644 --- a/crates/queue/src/management.rs +++ b/crates/queue/src/management.rs @@ -28,6 +28,17 @@ pub enum ChannelManagementAction { Tombstone, } +#[derive(Clone, Copy, Debug)] +pub struct ChannelManagementCommand<'a> { + pub operation_id: &'a str, + pub topic: &'a str, + pub channel: &'a str, + pub action: ChannelManagementAction, + pub expected_revision: u64, + pub tombstone_until_ms: Option, + pub require_idle: bool, +} + #[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] pub struct ManagementFenceSnapshot { #[serde(default)] @@ -60,6 +71,8 @@ pub(crate) struct FenceCatalog { topics: BTreeMap, #[serde(default)] channels: BTreeMap, + #[serde(default)] + local_channels: BTreeMap, } impl Default for FenceCatalog { @@ -69,6 +82,7 @@ impl Default for FenceCatalog { revision: String::new(), topics: BTreeMap::new(), channels: BTreeMap::new(), + local_channels: BTreeMap::new(), } } } @@ -95,9 +109,12 @@ impl FenceCatalog { } pub(crate) fn channel_blocked(&self, topic: &str, channel: &str, now_ms: i64) -> bool { - self.channels - .get(&channel_key(topic, channel)) - .is_some_and(|until| *until > now_ms) + let key = channel_key(topic, channel); + self.channels.get(&key).is_some_and(|until| *until > now_ms) + || self + .local_channels + .get(&key) + .is_some_and(|until| *until > now_ms) } pub(crate) fn set_topic(&mut self, topic: &str, until_ms: i64) -> bool { @@ -112,11 +129,22 @@ impl FenceCatalog { self.channels.insert(channel_key(topic, channel), until_ms) != Some(until_ms) } + pub(crate) fn set_local_channel(&mut self, topic: &str, channel: &str, until_ms: i64) -> bool { + self.local_channels + .insert(channel_key(topic, channel), until_ms) + != Some(until_ms) + } + pub(crate) fn clear_channel(&mut self, topic: &str, channel: &str) -> bool { - self.channels.remove(&channel_key(topic, channel)).is_some() + let key = channel_key(topic, channel); + let removed_global = self.channels.remove(&key).is_some(); + let removed_local = self.local_channels.remove(&key).is_some(); + removed_global || removed_local } pub(crate) fn replace(&mut self, snapshot: ManagementFenceSnapshot) { + let now_ms = unix_millis(); + self.local_channels.retain(|_, until_ms| *until_ms > now_ms); self.revision = snapshot.revision; self.topics = snapshot.topics; self.channels = snapshot @@ -127,6 +155,14 @@ impl FenceCatalog { } } +fn unix_millis() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(i64::MAX as u128) as i64 +} + fn channel_key(topic: &str, channel: &str) -> String { format!("{}:{}:{}", topic.len(), topic, channel) } @@ -146,4 +182,32 @@ mod tests { catalog.set_topic("orders", 9); assert!(!catalog.topic_blocked("orders", 9)); } + + #[test] + fn snapshot_sync_preserves_only_active_kodo_fences() { + let mut catalog = FenceCatalog::default(); + let active = unix_millis() + 60_000; + catalog.set_local_channel("events", "kodo", active); + catalog.set_local_channel("events", "expired", unix_millis() - 1); + catalog.set_channel("events", "console", active); + catalog.replace(ManagementFenceSnapshot { + revision: "console-2".into(), + topics: BTreeMap::new(), + channels: Vec::new(), + }); + assert!(catalog.channel_blocked("events", "kodo", unix_millis())); + assert!(!catalog.channel_blocked("events", "expired", unix_millis())); + assert!(!catalog.channel_blocked("events", "console", unix_millis())); + } + + #[test] + fn kodo_fences_survive_catalog_reopen() { + let root = tempfile::tempdir().unwrap(); + let (path, mut catalog) = FenceCatalog::load(root.path()).unwrap(); + catalog.set_local_channel("events", "workers", unix_millis() + 60_000); + catalog.store(&path).unwrap(); + + let (_, reopened) = FenceCatalog::load(root.path()).unwrap(); + assert!(reopened.channel_blocked("events", "workers", unix_millis())); + } } diff --git a/crates/queue/src/management_ops.rs b/crates/queue/src/management_ops.rs index 7bce501..86dbc44 100644 --- a/crates/queue/src/management_ops.rs +++ b/crates/queue/src/management_ops.rs @@ -112,6 +112,17 @@ impl OperationCatalog { Ok(()) } + pub(crate) fn pending_id(&self, topic: &str, fingerprint: &str) -> Option { + self.records + .iter() + .find(|(_, record)| { + record.topic == topic + && record.fingerprint == fingerprint + && record.state == OperationState::Pending + }) + .map(|(id, _)| id.clone()) + } + pub(crate) fn complete( &mut self, path: &Path, @@ -119,12 +130,30 @@ impl OperationCatalog { result: ManagementResult, ) -> io::Result<()> { let mut next = self.clone(); - let record = next.records.get_mut(id).ok_or_else(|| { + let record = next.records.get(id).ok_or_else(|| { io::Error::new(io::ErrorKind::NotFound, "management operation disappeared") })?; - record.state = OperationState::Completed; - record.result = Some(result); - next.completed_order.push_back(id.to_owned()); + let fingerprint = record.fingerprint.clone(); + let topic = record.topic.clone(); + let matching: Vec<_> = next + .records + .iter() + .filter(|(_, record)| { + record.topic == topic + && record.fingerprint == fingerprint + && record.state == OperationState::Pending + }) + .map(|(id, _)| id.clone()) + .collect(); + for id in matching { + let record = next + .records + .get_mut(&id) + .expect("matching operation still exists"); + record.state = OperationState::Completed; + record.result = Some(result.clone()); + next.completed_order.push_back(id); + } while next.completed_order.len() > MAX_COMPLETED { if let Some(expired) = next.completed_order.pop_front() { if next @@ -197,4 +226,22 @@ mod tests { .unwrap(); assert!(!catalog.blocks_topic("orders")); } + + #[test] + fn matching_pending_operation_can_be_adopted_after_the_caller_id_changes() { + let root = tempdir().unwrap(); + let (path, mut catalog) = OperationCatalog::load(root.path()).unwrap(); + catalog + .prepare( + &path, + "operation-00000001", + "channel:delete:orders:worker".into(), + "orders".into(), + ) + .unwrap(); + assert_eq!( + catalog.pending_id("orders", "channel:delete:orders:worker"), + Some("operation-00000001".into()) + ); + } } diff --git a/crates/queue/src/model.rs b/crates/queue/src/model.rs index 8b6c243..7fb5730 100644 --- a/crates/queue/src/model.rs +++ b/crates/queue/src/model.rs @@ -139,6 +139,8 @@ pub struct ChannelGroupCommitStats { pub struct TopicStats { pub name: String, pub paused: bool, + #[serde(default)] + pub published_count: u64, pub message_count: u64, #[serde(default)] pub segment_count: u64, @@ -151,8 +153,14 @@ pub struct TopicStats { pub struct ChannelStats { pub name: String, pub depth: u64, + #[serde(default)] + pub message_count: u64, pub in_flight_count: u64, pub deferred_count: u64, + #[serde(default)] + pub requeue_count: u64, + #[serde(default)] + pub timeout_count: u64, pub paused: bool, pub ephemeral: bool, pub ack_cursor: u64, diff --git a/crates/queue/src/topic.rs b/crates/queue/src/topic.rs index 35281f4..14af659 100644 --- a/crates/queue/src/topic.rs +++ b/crates/queue/src/topic.rs @@ -39,6 +39,8 @@ pub(crate) struct Topic { messages: MessageIndex, channels: HashMap, max_ack_gap: usize, + durable_channel_counters: bool, + published_count: u64, } impl TopicHandle { @@ -106,6 +108,8 @@ impl TopicHandle { messages: MessageIndex::new(index_cache), channels: HashMap::new(), max_ack_gap, + durable_channel_counters: storage_feature_level >= 2, + published_count: 0, }), wake, })) @@ -160,10 +164,16 @@ impl Topic { } let mut path_messages = Vec::new(); for location in log.locations_for_segment(&path)? { - let record = log.read_location(&location)?; - if record.kind == RecordKind::PublishBatch { - path_messages.extend(batch::metas(&record, &location)?); - } + let messages = + log.read_location_with(&location, |header, payload_len, reader| match header + .kind + { + RecordKind::PublishBatch => { + batch::metas_from_reader(header, payload_len, reader, &location) + } + _ => Ok(Vec::new()), + })?; + path_messages.extend(messages); } if immutable { log.persist_recovery_index(&path, recovery::encode(path_messages.iter()))?; @@ -182,6 +192,7 @@ impl Topic { manifest.next_position = recovered_next; store_atomic(&manifest_path, &manifest)?; } + let published_count = manifest.next_position.saturating_sub(1); Ok(Self { name: manifest.name.clone(), directory: directory.into(), @@ -191,6 +202,8 @@ impl Topic { messages, channels: HashMap::new(), max_ack_gap, + durable_channel_counters: storage_feature_level >= 2, + published_count, }) } @@ -212,6 +225,7 @@ impl Topic { ChannelRuntime { state, store: Some(store), + durable_counters: self.durable_channel_counters, }, ); } @@ -250,6 +264,9 @@ impl Topic { let metas = batch::metas_after_append(timestamp_ns, available_at_ms, &location, &batch); let ids = metas.iter().map(|message| message.id).collect(); self.messages.append(metas, reservation)?; + self.published_count = self + .published_count + .saturating_add(batch.entries.len() as u64); if previous_segment != self.log.current_segment_path() { self.persist_segment_index(&previous_segment)?; } @@ -348,8 +365,14 @@ impl Topic { &state, )?) }; - self.channels - .insert(name.into(), ChannelRuntime { state, store }); + self.channels.insert( + name.into(), + ChannelRuntime { + state, + store, + durable_counters: self.durable_channel_counters, + }, + ); Ok(true) } @@ -370,6 +393,21 @@ impl Topic { names } + pub fn channel_counts(&mut self, channel: &str) -> Result<(u64, u64, u64), BrokerError> { + let now_ms = now_ms(); + let scheduled = self.messages.deferred_positions(now_ms); + let last_position = self.last_position(); + let channel = self + .channels + .get(channel) + .ok_or(BrokerError::ChannelNotFound)?; + let (depth, in_flight, deferred, _) = + channel + .state + .metric_counts(last_position, &scheduled, now_ms); + Ok((depth, in_flight, deferred)) + } + #[cfg(test)] pub(crate) fn index_residency(&self) -> (usize, usize) { (self.messages.active_count(), self.messages.sealed_count()) diff --git a/crates/queue/src/topic/delivery.rs b/crates/queue/src/topic/delivery.rs index 1dde40f..e719116 100644 --- a/crates/queue/src/topic/delivery.rs +++ b/crates/queue/src/topic/delivery.rs @@ -154,6 +154,9 @@ impl Topic { message_id: id, available_at_ms, attempts, + cumulative_count: runtime + .durable_counters + .then(|| runtime.state.next_requeue_count()), }) } diff --git a/crates/queue/src/topic/index.rs b/crates/queue/src/topic/index.rs index 29c6312..355fcf4 100644 --- a/crates/queue/src/topic/index.rs +++ b/crates/queue/src/topic/index.rs @@ -4,7 +4,7 @@ use super::recovery; use crate::model::MessageMeta; use crate::BrokerError; use rustqueue_storage::RecoveryMetadataRef; -use std::collections::{BTreeSet, VecDeque}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -32,6 +32,7 @@ struct SealedMessages { pub(crate) struct MessageIndex { sealed: VecDeque, active: VecDeque, + scheduled: BTreeMap>, total_count: u64, cache: Arc, } @@ -41,6 +42,7 @@ impl MessageIndex { Self { sealed: VecDeque::new(), active: VecDeque::new(), + scheduled: BTreeMap::new(), total_count: 0, cache, } @@ -50,8 +52,11 @@ impl MessageIndex { &mut self, metadata: RecoveryMetadataRef, ) -> Result<(), BrokerError> { - let summary = recovery::inspect(&metadata)?; - self.push_sealed(SealedMessages::from_summary(metadata, summary)) + let summary = recovery::inspect(&metadata, unix_ms())?; + let scheduled = summary.scheduled.clone(); + self.push_sealed(SealedMessages::from_summary(metadata, summary))?; + self.extend_scheduled(scheduled); + Ok(()) } pub(crate) fn recover_active(&mut self, messages: Vec) -> Result<(), BrokerError> { @@ -194,6 +199,14 @@ impl MessageIndex { self.total_count } + pub(crate) fn deferred_positions(&mut self, now_ms: i64) -> BTreeSet { + self.scheduled = self.scheduled.split_off(&now_ms.saturating_add(1)); + self.scheduled + .values() + .flat_map(|positions| positions.iter().copied()) + .collect() + } + pub(crate) fn active_count(&self) -> usize { self.active.len() } @@ -302,12 +315,27 @@ impl MessageIndex { .sealed .iter() .filter(|segment| !existing.contains(segment.metadata.segment_path())) - .map(|segment| (segment.metadata.segment_path().to_path_buf(), segment.count)) + .map(|segment| { + ( + segment.metadata.segment_path().to_path_buf(), + segment.count, + segment.first_position, + segment.last_position, + ) + }) .collect(); - for (path, count) in &removed { + for (path, count, _, _) in &removed { self.total_count = self.total_count.saturating_sub(*count); self.cache.invalidate(path); } + self.scheduled.retain(|_, positions| { + positions.retain(|position| { + !removed + .iter() + .any(|(_, _, first, last)| (*first..=*last).contains(position)) + }); + !positions.is_empty() + }); self.sealed .retain(|segment| existing.contains(segment.metadata.segment_path())); } @@ -340,10 +368,26 @@ impl MessageIndex { expected = Some(message.position.saturating_add(1)); } self.total_count = self.total_count.saturating_add(messages.len() as u64); + let now_ms = unix_ms(); + self.extend_scheduled( + messages + .iter() + .filter(|message| message.available_at_ms > now_ms) + .map(|message| (message.available_at_ms, message.position)), + ); self.active.extend(messages); Ok(()) } + fn extend_scheduled(&mut self, scheduled: impl IntoIterator) { + for (available_at_ms, position) in scheduled { + self.scheduled + .entry(available_at_ms) + .or_default() + .insert(position); + } + } + fn load_ordinal( &self, segment: &SealedMessages, @@ -389,6 +433,14 @@ impl MessageIndex { } } +fn unix_ms() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(i64::MAX as u128) as i64 +} + impl Drop for MessageIndex { fn drop(&mut self) { for segment in &self.sealed { diff --git a/crates/queue/src/topic/maintenance.rs b/crates/queue/src/topic/maintenance.rs index 6807e15..ddf405f 100644 --- a/crates/queue/src/topic/maintenance.rs +++ b/crates/queue/src/topic/maintenance.rs @@ -32,18 +32,21 @@ impl Topic { self.persist_channel(channel, ChannelCommand::Pause { paused }) } - pub fn stats(&self) -> TopicStats { + pub fn stats(&mut self) -> TopicStats { let last = self.last_position(); + let now_ms = now_ms(); + let scheduled = self.messages.deferred_positions(now_ms); let (segment_count, segment_bytes) = self.log.storage_usage(); let mut channels: Vec<_> = self .channels - .values() - .map(|channel| channel.state.stats(last)) + .values_mut() + .map(|channel| channel.state.stats(last, &scheduled, now_ms)) .collect(); channels.sort_by(|left, right| left.name.cmp(&right.name)); TopicStats { name: self.name.clone(), paused: self.manifest.paused, + published_count: self.published_count, message_count: self.messages.total_count(), segment_count, segment_bytes, @@ -51,8 +54,10 @@ impl Topic { } } - pub fn add_aggregate_stats(&self, aggregate: &mut QueueAggregateStats) { + pub fn add_aggregate_stats(&mut self, aggregate: &mut QueueAggregateStats) { let last = self.last_position(); + let now_ms = now_ms(); + let scheduled = self.messages.deferred_positions(now_ms); let (segment_count, segment_bytes) = self.log.storage_usage(); aggregate.topic_count = aggregate.topic_count.saturating_add(1); aggregate.message_count = aggregate @@ -60,8 +65,9 @@ impl Topic { .saturating_add(self.messages.total_count()); aggregate.segment_count = aggregate.segment_count.saturating_add(segment_count); aggregate.segment_bytes = aggregate.segment_bytes.saturating_add(segment_bytes); - for channel in self.channels.values() { - let (depth, in_flight, deferred, ack_gap) = channel.state.metric_counts(last); + for channel in self.channels.values_mut() { + let (depth, in_flight, deferred, ack_gap) = + channel.state.metric_counts(last, &scheduled, now_ms); aggregate.channel_count = aggregate.channel_count.saturating_add(1); aggregate.channel_depth = aggregate.channel_depth.saturating_add(depth); aggregate.channel_in_flight = aggregate.channel_in_flight.saturating_add(in_flight); @@ -186,17 +192,40 @@ impl Topic { Ok(self.log.scrub_targets(false)?) } - pub fn sync(&self) -> Result<(), BrokerError> { + pub fn sync(&mut self) -> Result<(), BrokerError> { self.log.sync()?; + self.checkpoint_channels()?; store_atomic(&self.manifest_path, &self.manifest)?; Ok(()) } - pub fn expire_in_flight(&mut self) -> usize { + pub fn expire_in_flight(&mut self) -> Result { self.channels .values_mut() - .map(|channel| channel.state.expire_in_flight()) - .sum() + .try_fold(0usize, |total, channel| { + Ok(total.saturating_add(channel.expire_in_flight()?)) + }) + } + + pub fn expire_channel_in_flight(&mut self, channel: &str) -> Result { + self.channels + .get_mut(channel) + .ok_or(BrokerError::ChannelNotFound)? + .expire_in_flight() + } + + pub fn has_expired_in_flight(&self) -> bool { + self.channels + .values() + .any(|channel| channel.has_expired_in_flight()) + } + + pub fn channel_has_expired_in_flight(&self, channel: &str) -> Result { + Ok(self + .channels + .get(channel) + .ok_or(BrokerError::ChannelNotFound)? + .has_expired_in_flight()) } pub fn checkpoint_channels(&mut self) -> Result<(), BrokerError> { diff --git a/crates/queue/src/topic/recovery.rs b/crates/queue/src/topic/recovery.rs index 7f7f990..867996e 100644 --- a/crates/queue/src/topic/recovery.rs +++ b/crates/queue/src/topic/recovery.rs @@ -14,6 +14,7 @@ pub(super) struct Summary { pub count: u64, pub first: MessageMeta, pub last: MessageMeta, + pub scheduled: Vec<(i64, u64)>, } pub(super) fn encode<'a>(messages: impl Iterator + Clone) -> Vec { @@ -37,7 +38,10 @@ pub(super) fn encode<'a>(messages: impl Iterator + Clone bytes } -pub(super) fn inspect(reference: &RecoveryMetadataRef) -> Result { +pub(super) fn inspect( + reference: &RecoveryMetadataRef, + now_ms: i64, +) -> Result { let header = reference.read_range(0, HEADER_LEN)?; let count = decode_header(&header)?; if count == 0 { @@ -54,14 +58,34 @@ pub(super) fn inspect(reference: &RecoveryMetadataRef) -> Result = None; + let mut ordinal = 0; + while ordinal < count { + let page = read_page(reference, ordinal, (count - ordinal).min(1024) as usize)?; + for message in page { + if last.as_ref().is_some_and(|previous| { + message.position != previous.position.saturating_add(1) || message.id < previous.id + }) { + return Err(BrokerError::InvalidRecord( + "topic recovery index range is invalid".into(), + )); + } + if message.available_at_ms > now_ms { + scheduled.push((message.available_at_ms, message.position)); + } + first.get_or_insert_with(|| message.clone()); + last = Some(message); + } + ordinal += (count - ordinal).min(1024); } - Ok(Summary { count, first, last }) + Ok(Summary { + count, + first: first.expect("non-empty recovery index"), + last: last.expect("non-empty recovery index"), + scheduled, + }) } pub(super) fn read_page( @@ -92,19 +116,6 @@ pub(super) fn read_page( Ok(messages) } -fn read_entry(reference: &RecoveryMetadataRef, ordinal: u64) -> Result { - let offset = HEADER_LEN as u64 - + ordinal - .checked_mul(MESSAGE_LEN as u64) - .ok_or_else(|| BrokerError::InvalidRecord("topic index offset overflow".into()))?; - let bytes = reference.read_range(offset, MESSAGE_LEN)?; - decode_entry( - Arc::new(reference.segment_path().to_path_buf()), - reference.segment_len(), - &bytes, - ) -} - fn decode_header(bytes: &[u8]) -> Result { if bytes.len() < HEADER_LEN || &bytes[0..4] != MAGIC diff --git a/crates/queue/tests/large_message.rs b/crates/queue/tests/large_message.rs index 5f5b524..fee01a2 100644 --- a/crates/queue/tests/large_message.rs +++ b/crates/queue/tests/large_message.rs @@ -3,22 +3,32 @@ use rustqueue_queue::{Broker, BrokerConfig}; use std::time::Duration; #[tokio::test] -async fn persists_and_delivers_a_twenty_mebibyte_message() { +async fn persists_and_delivers_a_hundred_mebibyte_message() { let directory = tempfile::tempdir().unwrap(); - let broker = Broker::open(BrokerConfig { + let config = BrokerConfig { data_path: directory.path().into(), - max_message_bytes: 32 * 1024 * 1024, + max_segment_bytes: 256 * 1024 * 1024, + max_message_bytes: 100 * 1024 * 1024, + storage_feature_level: 2, ..BrokerConfig::default() - }) - .unwrap(); + }; + let broker = Broker::open(config.clone()).unwrap(); broker.create_topic("large").await.unwrap(); broker.create_channel("large", "workers").await.unwrap(); - let body = Bytes::from(vec![0x5a; 20 * 1024 * 1024]); + let body = Bytes::from(vec![0x5a; 100 * 1024 * 1024]); broker .publish("large", vec![body.clone()], Duration::ZERO) .await .unwrap(); + broker.flush().await.unwrap(); + drop(broker); + + let broker = Broker::open(BrokerConfig { + max_message_bytes: 20 * 1024 * 1024, + ..config + }) + .unwrap(); let delivery = broker .next_message("large", "workers", None) .await diff --git a/crates/queue/tests/recovery.rs b/crates/queue/tests/recovery.rs index 5c66df8..fa64e7a 100644 --- a/crates/queue/tests/recovery.rs +++ b/crates/queue/tests/recovery.rs @@ -261,7 +261,9 @@ async fn protective_eviction_persists_channel_gap_and_audit_before_deleting() { assert_eq!(report.topic, "events"); assert_eq!(report.messages, 1); assert_eq!(report.through_position, 1); - assert_eq!(broker.stats().topics[0].channels[0].ack_cursor, 1); + let channel = &broker.stats().topics[0].channels[0]; + assert_eq!(channel.ack_cursor, 1); + assert_eq!(channel.message_count, 3); assert_eq!( std::fs::read_dir(root.path().join("audit")) .unwrap() @@ -271,6 +273,7 @@ async fn protective_eviction_persists_channel_gap_and_audit_before_deleting() { drop(broker); let broker = Broker::open(config(root.path())).unwrap(); + assert_eq!(broker.stats().topics[0].channels[0].message_count, 3); let next = broker .next_message("events", "workers", None) .await diff --git a/crates/server/src/config.rs b/crates/server/src/config.rs index 77f7db4..678c0fd 100644 --- a/crates/server/src/config.rs +++ b/crates/server/src/config.rs @@ -3,14 +3,30 @@ mod environment; mod validation; use environment::read_optional_secret; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use std::fs; use std::net::SocketAddr; use std::path::{Path, PathBuf}; use std::time::Duration; -pub const MAX_SUPPORTED_MESSAGE_BYTES: usize = 32 * 1024 * 1024; -pub const MAX_SUPPORTED_BATCH_BYTES: usize = 64 * 1024 * 1024; +pub const MAX_SUPPORTED_MESSAGE_BYTES: usize = rustqueue_protocol::MAX_MESSAGE_BYTES; +pub const MAX_SUPPORTED_BATCH_BYTES: usize = rustqueue_protocol::MAX_BATCH_BYTES; + +#[derive(Clone, Debug, Serialize)] +pub struct RuntimeCapabilities { + #[serde(flatten)] + storage: rustqueue_storage::BinaryCapabilities, + pub maximum_message_bytes: usize, + pub maximum_batch_bytes: usize, +} + +pub fn runtime_capabilities() -> RuntimeCapabilities { + RuntimeCapabilities { + storage: rustqueue_storage::binary_capabilities(), + maximum_message_bytes: MAX_SUPPORTED_MESSAGE_BYTES, + maximum_batch_bytes: MAX_SUPPORTED_BATCH_BYTES, + } +} #[derive(Clone, Debug, Deserialize)] #[serde(default, deny_unknown_fields)] @@ -38,6 +54,7 @@ pub struct NodeConfig { pub struct NetworkConfig { pub tcp_address: SocketAddr, pub http_address: SocketAddr, + pub kodo_http_address: Option, pub advertised_tcp_port: u16, pub advertised_http_port: u16, pub snappy_enabled: bool, @@ -91,7 +108,9 @@ pub struct SecurityConfig { pub publish_token_file: Option, pub registry_token_file: Option, pub console_token_file: Option, + pub kodo_cleanup_token_file: Option, pub console_management_enabled: bool, + pub kodo_cleanup_enabled: bool, } #[derive(Clone, Debug, Deserialize)] @@ -128,6 +147,7 @@ pub struct LimitsConfig { pub auth_max_ttl_seconds: u64, pub auth_cache_max_entries: usize, pub http_body_timeout_ms: u64, + pub disconnect_on_retriable_publish_error: bool, } #[derive(Clone, Debug, Deserialize)] @@ -173,6 +193,7 @@ impl Default for NetworkConfig { Self { tcp_address: "0.0.0.0:4150".parse().unwrap(), http_address: "0.0.0.0:4151".parse().unwrap(), + kodo_http_address: None, advertised_tcp_port: 4150, advertised_http_port: 4151, snappy_enabled: true, @@ -198,7 +219,7 @@ impl Default for StorageConfig { disk_high_watermark_percent: 85, disk_low_watermark_percent: 75, min_free_bytes: 10 * 1024 * 1024 * 1024, - protective_eviction_enabled: true, + protective_eviction_enabled: false, disk_pressure_grace_seconds: 60, } } @@ -238,7 +259,7 @@ impl Default for TlsConfig { impl Default for LimitsConfig { fn default() -> Self { Self { - max_body_bytes: MAX_SUPPORTED_BATCH_BYTES, + max_body_bytes: 64 * 1024 * 1024, node_publish_inflight_bytes: 512 * 1024 * 1024, connection_publish_inflight_bytes: 80 * 1024 * 1024, node_delivery_inflight_bytes: 512 * 1024 * 1024, @@ -259,6 +280,7 @@ impl Default for LimitsConfig { auth_max_ttl_seconds: 3600, auth_cache_max_entries: 10_000, http_body_timeout_ms: 30_000, + disconnect_on_retriable_publish_error: false, } } } @@ -366,6 +388,20 @@ impl Config { if !(1..=9).contains(&self.network.max_deflate_level) { bail!("network.max_deflate_level must be in 1..=9"); } + if self.security.kodo_cleanup_enabled { + bail!( + "Kodo automatic cleanup is disabled until cluster-wide atomic deletion is available" + ); + } + if self.security.kodo_cleanup_enabled != self.network.kodo_http_address.is_some() { + bail!("network.kodo_http_address is required exactly when Kodo cleanup is enabled"); + } + if self.network.kodo_http_address.is_some_and(|address| { + address.port() == self.network.tcp_address.port() + || address.port() == self.network.http_address.port() + }) { + bail!("network.kodo_http_address must use a dedicated port"); + } if self.shutdown.grace_seconds == 0 { bail!("shutdown.grace_seconds must be greater than zero"); } @@ -391,6 +427,7 @@ impl Config { self.security.publish_token_file.as_ref(), self.security.registry_token_file.as_ref(), self.security.console_token_file.as_ref(), + self.security.kodo_cleanup_token_file.as_ref(), ] .into_iter() .flatten() @@ -402,6 +439,24 @@ impl Config { if self.security.console_management_enabled && self.security.console_token_file.is_none() { bail!("security.console_token_file is required when console management is enabled"); } + if self.security.kodo_cleanup_enabled && self.security.kodo_cleanup_token_file.is_none() { + bail!("security.kodo_cleanup_token_file is required when Kodo cleanup is enabled"); + } + if self.security.kodo_cleanup_enabled { + let cleanup = self + .read_kodo_cleanup_token()? + .context("Kodo cleanup token is unavailable")?; + for (scope, token) in [ + ("admin", self.read_admin_token()?), + ("publish", self.read_publish_token()?), + ("registry", self.read_registry_token()?), + ("console", self.read_console_token()?), + ] { + if token.as_deref() == Some(cleanup.as_str()) { + bail!("Kodo cleanup token must be distinct from the {scope} token"); + } + } + } Ok(()) } @@ -420,6 +475,9 @@ impl Config { pub fn read_console_token(&self) -> anyhow::Result> { read_optional_secret(self.security.console_token_file.as_deref()) } + pub fn read_kodo_cleanup_token(&self) -> anyhow::Result> { + read_optional_secret(self.security.kodo_cleanup_token_file.as_deref()) + } } #[cfg(test)] diff --git a/crates/server/src/config/environment.rs b/crates/server/src/config/environment.rs index fc5beed..a274180 100644 --- a/crates/server/src/config/environment.rs +++ b/crates/server/src/config/environment.rs @@ -27,6 +27,13 @@ impl Config { )?; set_from_env("RUSTQUEUE_TCP_ADDRESS", &mut self.network.tcp_address)?; set_from_env("RUSTQUEUE_HTTP_ADDRESS", &mut self.network.http_address)?; + if let Ok(value) = env::var("RUSTQUEUE_KODO_HTTP_ADDRESS") { + self.network.kodo_http_address = Some( + value + .parse() + .context("parse environment variable RUSTQUEUE_KODO_HTTP_ADDRESS")?, + ); + } if let Ok(value) = env::var("RUSTQUEUE_DATA_PATH") { self.storage.data_path = value.into(); } @@ -68,6 +75,10 @@ impl Config { "RUSTQUEUE_HTTP_BODY_TIMEOUT_MS", &mut self.limits.http_body_timeout_ms, )?; + set_from_env( + "RUSTQUEUE_DISCONNECT_ON_RETRIABLE_PUBLISH_ERROR", + &mut self.limits.disconnect_on_retriable_publish_error, + )?; set_from_env( "RUSTQUEUE_TCP_COMMAND_TIMEOUT_MS", &mut self.limits.tcp_command_timeout_ms, diff --git a/crates/server/src/config/tests.rs b/crates/server/src/config/tests.rs index dd2d94a..d84358e 100644 --- a/crates/server/src/config/tests.rs +++ b/crates/server/src/config/tests.rs @@ -24,8 +24,88 @@ fn maintenance_has_a_startup_quiet_period() { #[test] fn accepts_the_documented_large_message_limits() { let mut config = Config::default(); + config.storage.feature_level = 2; + config.storage.max_segment_bytes = 256 * 1024 * 1024; config.queue.max_message_bytes = MAX_SUPPORTED_MESSAGE_BYTES; config.limits.max_body_bytes = MAX_SUPPORTED_BATCH_BYTES; + config.limits.connection_publish_inflight_bytes = 160 * 1024 * 1024; + config.limits.node_publish_inflight_bytes = 1024 * 1024 * 1024; + config.limits.connection_delivery_inflight_bytes = 128 * 1024 * 1024; + config.limits.node_delivery_inflight_bytes = 1024 * 1024 * 1024; + config.validate().unwrap(); +} + +#[test] +fn binary_capabilities_advertise_protocol_limits() { + let capabilities = serde_json::to_value(runtime_capabilities()).unwrap(); + assert_eq!( + capabilities["maximum_message_bytes"], + MAX_SUPPORTED_MESSAGE_BYTES + ); + assert_eq!( + capabilities["maximum_batch_bytes"], + MAX_SUPPORTED_BATCH_BYTES + ); +} + +#[test] +fn kodo_cleanup_is_disabled_until_cluster_wide_deletion_is_atomic() { + let root = tempfile::tempdir().unwrap(); + let cleanup = root.path().join("cleanup-token"); + let admin = root.path().join("admin-token"); + std::fs::write(&cleanup, "cleanup-secret").unwrap(); + std::fs::write(&admin, "cleanup-secret").unwrap(); + let mut config = Config::default(); + config.security.kodo_cleanup_enabled = true; + assert!(config.validate().is_err()); + config.network.kodo_http_address = Some("0.0.0.0:4152".parse().unwrap()); + config.security.kodo_cleanup_token_file = Some(cleanup); + config.security.admin_token_file = Some(admin.clone()); + assert!(config.validate().is_err()); + std::fs::write(admin, "admin-secret").unwrap(); + let error = config.validate().unwrap_err(); + assert!(error.to_string().contains("cluster-wide atomic deletion")); +} + +#[test] +fn large_record_contract_requires_storage_feature_level_two() { + let mut config = Config::default(); + config.storage.max_segment_bytes = 256 * 1024 * 1024; + config.queue.max_message_bytes = MAX_SUPPORTED_MESSAGE_BYTES; + config.limits.max_body_bytes = MAX_SUPPORTED_BATCH_BYTES; + config.limits.connection_publish_inflight_bytes = 160 * 1024 * 1024; + config.limits.node_publish_inflight_bytes = 1024 * 1024 * 1024; + config.limits.connection_delivery_inflight_bytes = 128 * 1024 * 1024; + config.limits.node_delivery_inflight_bytes = 1024 * 1024 * 1024; + assert!(config.validate().is_err()); + + config.storage.feature_level = 2; + config.validate().unwrap(); +} + +#[test] +fn feature_level_two_keeps_enough_budget_to_read_retained_large_messages() { + let mut config = Config::default(); + config.storage.feature_level = 2; + assert!(config.validate().is_err()); + config.limits.connection_delivery_inflight_bytes = MAX_SUPPORTED_MESSAGE_BYTES; + config.limits.node_delivery_inflight_bytes = 2 * MAX_SUPPORTED_MESSAGE_BYTES; + config.validate().unwrap(); +} + +#[test] +fn large_publish_body_requires_matching_inflight_budget() { + let mut config = Config::default(); + config.storage.feature_level = 2; + config.storage.max_segment_bytes = 256 * 1024 * 1024; + config.queue.max_message_bytes = MAX_SUPPORTED_MESSAGE_BYTES; + config.limits.max_body_bytes = MAX_SUPPORTED_MESSAGE_BYTES; + config.limits.connection_delivery_inflight_bytes = MAX_SUPPORTED_MESSAGE_BYTES; + config.limits.node_delivery_inflight_bytes = 2 * MAX_SUPPORTED_MESSAGE_BYTES; + assert!(config.validate().is_err()); + + config.limits.connection_publish_inflight_bytes = 160 * 1024 * 1024; + config.limits.node_publish_inflight_bytes = 512 * 1024 * 1024; config.validate().unwrap(); } diff --git a/crates/server/src/config/validation.rs b/crates/server/src/config/validation.rs index 611ae4f..0b3f294 100644 --- a/crates/server/src/config/validation.rs +++ b/crates/server/src/config/validation.rs @@ -1,4 +1,4 @@ -use super::{Config, MAX_SUPPORTED_BATCH_BYTES}; +use super::{Config, MAX_SUPPORTED_BATCH_BYTES, MAX_SUPPORTED_MESSAGE_BYTES}; use crate::admission::{working_set_bytes, PublishShape}; use anyhow::bail; @@ -8,14 +8,37 @@ impl Config { { bail!("limits.max_body_bytes must be in 1..={MAX_SUPPORTED_BATCH_BYTES}"); } - if self.limits.connection_publish_inflight_bytes - < working_set_bytes(self.limits.max_body_bytes, PublishShape::Multi) + let publish_working_set = + working_set_bytes(self.limits.max_body_bytes, PublishShape::Multi).max( + working_set_bytes(self.queue.max_message_bytes, PublishShape::Single), + ); + if self.limits.connection_publish_inflight_bytes < publish_working_set || self.limits.node_publish_inflight_bytes < self.limits.connection_publish_inflight_bytes { bail!("publish inflight limits must fit the encoded working set and satisfy connection <= node"); } - if self.limits.connection_delivery_inflight_bytes < self.queue.max_message_bytes + let maximum_record_payload = self.queue.max_message_bytes.saturating_add(24).max( + self.limits + .max_body_bytes + .saturating_add(16 * rustqueue_protocol::MAX_MPUB_MESSAGES), + ); + if maximum_record_payload > rustqueue_storage::MAX_RECORD_BYTES { + bail!("configured publish limits can exceed the durable record contract"); + } + if maximum_record_payload > rustqueue_storage::LEGACY_MAX_RECORD_BYTES + && self.storage.feature_level < 2 + { + bail!( + "publish limits above the v7 legacy record bound require storage.feature_level = 2" + ); + } + let retained_message_bound = if self.storage.feature_level >= 2 { + MAX_SUPPORTED_MESSAGE_BYTES + } else { + self.queue.max_message_bytes + }; + if self.limits.connection_delivery_inflight_bytes < retained_message_bound || self .limits .connection_delivery_inflight_bytes @@ -23,7 +46,7 @@ impl Config { .is_none_or(|minimum| self.limits.node_delivery_inflight_bytes < minimum) || self.limits.node_delivery_inflight_bytes > u32::MAX as usize { - bail!("delivery inflight limits must fit one message plus the payload-read working set and fit u32"); + bail!("delivery inflight limits must fit every readable durable message plus the payload-read working set and fit u32"); } if self.limits.client_handshake_timeout_ms == 0 || self.limits.tcp_command_timeout_ms == 0 diff --git a/crates/server/src/http.rs b/crates/server/src/http.rs index 256589a..bcf36a7 100644 --- a/crates/server/src/http.rs +++ b/crates/server/src/http.rs @@ -1,16 +1,22 @@ mod compat; mod helpers; +mod kodo_compat; mod manage; mod native; +mod nsq_stats; +mod tokens; use compat::*; use helpers::*; use manage::*; use native::*; +use nsq_stats::*; +use tokens::TokenSet; use crate::admission::PublishAdmission; use crate::config::Config; use crate::metrics::Metrics; +use crate::subscriptions::SubscriptionRegistry; use axum::body::{Body, Bytes}; use axum::extract::{Query, State}; use axum::http::{header, HeaderMap, HeaderValue, Request, StatusCode}; @@ -27,18 +33,19 @@ use tokio::net::TcpListener; use tower_http::timeout::RequestBodyTimeoutLayer; use tracing::info; +pub(crate) use kodo_compat::serve_kodo_compat; + #[derive(Clone)] struct AppState { config: Arc, broker: Arc, metrics: Arc, - admin_token: Option>, - publish_token: Option>, - registry_token: Option>, - console_token: Option>, + tokens: TokenSet, accepting: Arc, delivering: Arc, publish_admission: Arc, + subscriptions: SubscriptionRegistry, + started_at: i64, } #[derive(Debug)] @@ -80,6 +87,7 @@ struct StatsQuery { format: Option, topic: Option, channel: Option, + include_clients: Option, } #[derive(Serialize)] @@ -98,6 +106,7 @@ struct Producer { version: &'static str, } +#[allow(clippy::too_many_arguments)] pub async fn serve( config: Arc, broker: Arc, @@ -105,19 +114,18 @@ pub async fn serve( accepting: Arc, delivering: Arc, publish_admission: Arc, + subscriptions: SubscriptionRegistry, + shutdown: tokio::sync::watch::Receiver, ) -> anyhow::Result<()> { - let state = AppState { - admin_token: config.read_admin_token()?.map(Arc::from), - publish_token: config.read_publish_token()?.map(Arc::from), - registry_token: config.read_registry_token()?.map(Arc::from), - console_token: config.read_console_token()?.map(Arc::from), - config: Arc::clone(&config), + let state = app_state( + config, broker, metrics, accepting, delivering, publish_admission, - }; + subscriptions, + )?; let mut router = Router::new() .route("/ping", get(ping)) .route("/info", get(info_handler)) @@ -149,21 +157,76 @@ pub async fn serve( .route("/v1/observe/head", get(observe_head)) .route("/v1/storage/scrub", post(scrub)) .layer(middleware::from_fn(nsq_content_negotiation)); - if config.security.console_management_enabled { + if state.config.security.console_management_enabled { router = router .route("/v1/manage/topics/{action}", post(manage_topic)) .route("/v1/manage/channels/{action}", post(manage_channel)) .route("/v1/manage/fences/sync", post(sync_fences)); + } else if state.config.security.kodo_cleanup_enabled { + router = router.route( + "/v1/manage/channels/delete-if-idle", + post(delete_idle_channel), + ); } + let timeout = state.config.limits.http_body_timeout_ms; + let address = state.config.network.http_address; + let tokens = state.tokens.clone(); + let token_shutdown = shutdown.clone(); let router = router.with_state(state).layer(RequestBodyTimeoutLayer::new( - std::time::Duration::from_millis(config.limits.http_body_timeout_ms), + std::time::Duration::from_millis(timeout), )); - let listener = TcpListener::bind(config.network.http_address).await?; - info!(address = %config.network.http_address, "HTTP API listening"); - axum::serve(listener, router).await?; + let listener = TcpListener::bind(address).await?; + info!(%address, "HTTP API listening"); + let reloader = tokio::spawn(tokens.reload(token_shutdown)); + let result = axum::serve(listener, router) + .with_graceful_shutdown(wait_for_shutdown(shutdown)) + .await; + reloader.abort(); + result?; Ok(()) } +async fn wait_for_shutdown(mut shutdown: tokio::sync::watch::Receiver) { + if *shutdown.borrow() { + return; + } + while shutdown.changed().await.is_ok() { + if *shutdown.borrow() { + return; + } + } +} + +fn app_state( + config: Arc, + broker: Arc, + metrics: Arc, + accepting: Arc, + delivering: Arc, + publish_admission: Arc, + subscriptions: SubscriptionRegistry, +) -> anyhow::Result { + Ok(AppState { + tokens: TokenSet::from_config(&config)?, + config: Arc::clone(&config), + broker, + metrics, + accepting, + delivering, + publish_admission, + subscriptions, + started_at: unix_seconds(), + }) +} + +fn unix_seconds() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + .min(i64::MAX as u64) as i64 +} + async fn nsq_content_negotiation(request: Request, next: Next) -> Response { let v1 = request .headers() @@ -202,6 +265,13 @@ impl ApiError { detail: detail.into(), } } + fn conflict(code: &'static str, detail: impl Into) -> Self { + Self { + status: StatusCode::CONFLICT, + code, + detail: detail.into(), + } + } fn timeout(code: &'static str, detail: impl Into) -> Self { Self { status: StatusCode::REQUEST_TIMEOUT, @@ -220,6 +290,7 @@ impl From for ApiError { BrokerError::InvalidTopic => (StatusCode::BAD_REQUEST, "E_BAD_TOPIC"), BrokerError::ChannelNotFound => (StatusCode::NOT_FOUND, "E_BAD_CHANNEL"), BrokerError::ChannelTombstoned => (StatusCode::CONFLICT, "E_CHANNEL_TOMBSTONED"), + BrokerError::ChannelNotIdle { .. } => (StatusCode::CONFLICT, "E_CHANNEL_NOT_IDLE"), BrokerError::ManagementUnavailable => { (StatusCode::SERVICE_UNAVAILABLE, "E_MANAGEMENT_UNAVAILABLE") } @@ -233,9 +304,10 @@ impl From for ApiError { | BrokerError::PublishWorkerLimit | BrokerError::ChannelWorkerLimit | BrokerError::ChannelLimit => (StatusCode::TOO_MANY_REQUESTS, "E_THROTTLED"), - BrokerError::StorageUnavailable | BrokerError::Storage(_) | BrokerError::Io(_) => { - (StatusCode::SERVICE_UNAVAILABLE, "E_STORAGE") - } + BrokerError::StorageUnavailable + | BrokerError::Storage(_) + | BrokerError::Io(_) + | BrokerError::InvalidRecord(_) => (StatusCode::SERVICE_UNAVAILABLE, "E_STORAGE"), _ => (StatusCode::INTERNAL_SERVER_ERROR, "E_INTERNAL"), }; Self { @@ -285,6 +357,13 @@ mod tests { assert!(parse_binary_mpub(Bytes::from(body), 10).is_err()); } + #[test] + fn durable_record_failures_are_reported_as_storage_unavailable() { + let error = ApiError::from(BrokerError::InvalidRecord("corrupt local state".into())); + assert_eq!(error.status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(error.code, "E_STORAGE"); + } + #[test] fn registry_exposes_channel_names_instead_of_internal_stats() { let topics = registry_topics(&BrokerStats { @@ -296,14 +375,18 @@ mod tests { topics: vec![TopicStats { name: "events".into(), paused: false, + published_count: 3, message_count: 3, segment_count: 1, segment_bytes: 256, channels: vec![ChannelStats { name: "workers".into(), depth: 2, + message_count: 3, in_flight_count: 1, deferred_count: 0, + requeue_count: 0, + timeout_count: 0, paused: false, ephemeral: false, ack_cursor: 1, @@ -321,13 +404,24 @@ mod tests { header::AUTHORIZATION, HeaderValue::from_static("Bearer console-secret"), ); - assert!(authorize(&headers, Some("console-secret"), "console").is_ok()); - assert!(authorize(&headers, Some("admin-secret"), "admin").is_err()); - assert!(authorize_any( - &headers, - &[Some("registry-secret"), Some("console-secret")], - "read-only" - ) - .is_ok()); + let console = tokens::TokenSource::fixed("console", "console-secret"); + let admin = tokens::TokenSource::fixed("admin", "admin-secret"); + let registry = tokens::TokenSource::fixed("registry", "registry-secret"); + assert!(authorize(&headers, &console, "console").is_ok()); + assert!(authorize(&headers, &admin, "admin").is_err()); + assert!(authorize_any(&headers, &[®istry, &console], "read-only").is_ok()); + } + + #[test] + fn kodo_cleanup_token_is_distinct_from_admin() { + let mut headers = HeaderMap::new(); + headers.insert( + header::AUTHORIZATION, + HeaderValue::from_static("Bearer cleanup-secret"), + ); + let cleanup = tokens::TokenSource::fixed("Kodo cleanup", "cleanup-secret"); + let admin = tokens::TokenSource::fixed("admin", "admin-secret"); + assert!(authorize(&headers, &cleanup, "Kodo cleanup").is_ok()); + assert!(authorize(&headers, &admin, "admin").is_err()); } } diff --git a/crates/server/src/http/compat.rs b/crates/server/src/http/compat.rs index e8abcd2..aac00b6 100644 --- a/crates/server/src/http/compat.rs +++ b/crates/server/src/http/compat.rs @@ -23,7 +23,7 @@ pub(super) async fn publish( headers: HeaderMap, request: Request, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.publish_token.as_deref(), "publish")?; + authorize(&headers, &state.tokens.publish, "publish")?; validate_defer(query.defer, &state.config)?; let (body, reservation) = read_publish_body( &state, @@ -54,7 +54,7 @@ pub(super) async fn multi_publish( headers: HeaderMap, request: Request, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.publish_token.as_deref(), "publish")?; + authorize(&headers, &state.tokens.publish, "publish")?; validate_defer(query.defer, &state.config)?; let (body, reservation) = read_publish_body( &state, @@ -87,31 +87,6 @@ pub(super) async fn multi_publish( Ok("OK") } -pub(super) async fn stats( - State(state): State, - Query(query): Query, -) -> Response { - let filtered = filter_stats( - state.broker.stats(), - query.topic.as_deref(), - query.channel.as_deref(), - ); - if query.format.as_deref() == Some("json") { - return Json(json!({"version": env!("CARGO_PKG_VERSION"), "health": "OK", "topics": filtered.topics})).into_response(); - } - let mut output = String::new(); - for topic in filtered.topics { - output.push_str(&format!("[{}] depth={}\n", topic.name, topic.message_count)); - for channel in topic.channels { - output.push_str(&format!( - " [{}] depth={} in-flight={} deferred={}\n", - channel.name, channel.depth, channel.in_flight_count, channel.deferred_count - )); - } - } - output.into_response() -} - pub(super) async fn lookup( State(state): State, Query(query): Query, @@ -153,6 +128,9 @@ pub(super) async fn nodes(State(state): State) -> Json { } pub(super) async fn metrics_handler(State(state): State) -> Response { + if let Err(error) = state.broker.expire_in_flight().await { + return ApiError::from(error).into_response(); + } let mut output = state.metrics.render(); let queue_stats = state.broker.metrics_stats( state.config.metrics.detailed_queue_metrics, @@ -174,7 +152,7 @@ pub(super) async fn create_topic( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state.broker.create_topic(&query.topic).await?; Ok("OK") } @@ -184,7 +162,7 @@ pub(super) async fn delete_topic( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state.broker.delete_topic(&query.topic).await?; Ok("OK") } @@ -194,7 +172,7 @@ pub(super) async fn empty_topic( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state.broker.empty_topic(&query.topic).await?; Ok("OK") } @@ -204,7 +182,7 @@ pub(super) async fn pause_topic( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state.broker.set_topic_paused(&query.topic, true).await?; Ok("OK") } @@ -214,7 +192,7 @@ pub(super) async fn unpause_topic( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state.broker.set_topic_paused(&query.topic, false).await?; Ok("OK") } @@ -224,7 +202,7 @@ pub(super) async fn create_channel( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state .broker .create_channel(&query.topic, &query.channel) @@ -237,7 +215,7 @@ pub(super) async fn delete_channel( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state .broker .delete_channel(&query.topic, &query.channel) @@ -250,7 +228,7 @@ pub(super) async fn empty_channel( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state .broker .empty_channel(&query.topic, &query.channel) @@ -263,7 +241,7 @@ pub(super) async fn pause_channel( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state .broker .set_channel_paused(&query.topic, &query.channel, true) @@ -276,7 +254,7 @@ pub(super) async fn unpause_channel( Query(query): Query, headers: HeaderMap, ) -> Result<&'static str, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state .broker .set_channel_paused(&query.topic, &query.channel, false) diff --git a/crates/server/src/http/helpers.rs b/crates/server/src/http/helpers.rs index 9d28b59..b43d6a7 100644 --- a/crates/server/src/http/helpers.rs +++ b/crates/server/src/http/helpers.rs @@ -1,3 +1,4 @@ +use super::tokens::TokenSource; use super::*; use std::time::Duration; @@ -26,6 +27,7 @@ pub(super) async fn publish_write( return Err(ApiError::unavailable("E_DRAINING", "broker is draining")); } if !state.broker.storage_healthy() { + state.metrics.storage_errors.fetch_add(1, Ordering::Relaxed); return Err(ApiError::unavailable( "E_STORAGE", "local storage is isolated; broker restart is required", @@ -36,7 +38,7 @@ pub(super) async fn publish_write( "local disk is above its publish watermark", )); } - Ok(state + let result = state .broker .publish_guarded( topic, @@ -44,23 +46,11 @@ pub(super) async fn publish_write( Duration::from_millis(defer_ms), reservation, ) - .await?) -} - -pub(super) fn filter_stats( - mut stats: BrokerStats, - topic: Option<&str>, - channel: Option<&str>, -) -> BrokerStats { - if let Some(topic) = topic { - stats.topics.retain(|candidate| candidate.name == topic); - } - if let Some(channel) = channel { - for topic in &mut stats.topics { - topic.channels.retain(|candidate| candidate.name == channel); - } + .await; + if result.as_ref().is_err_and(crate::tcp::broker_storage_error) { + state.metrics.storage_errors.fetch_add(1, Ordering::Relaxed); } - stats + Ok(result?) } pub(super) async fn read_publish_body( @@ -136,9 +126,15 @@ pub(super) fn parse_text_mpub(body: Bytes, max: usize) -> Result, Api pub(super) fn authorize( headers: &HeaderMap, - expected: Option<&str>, + expected: &TokenSource, scope: &'static str, ) -> Result<(), ApiError> { + let expected = expected.expected().map_err(|unavailable| { + ApiError::unavailable( + "E_AUTH_UNAVAILABLE", + format!("{unavailable} authorization token is unavailable"), + ) + })?; let Some(expected) = expected else { return Ok(()); }; @@ -164,10 +160,23 @@ pub(super) fn authorize( pub(super) fn authorize_any( headers: &HeaderMap, - expected: &[Option<&str>], + expected: &[&TokenSource], scope: &'static str, ) -> Result<(), ApiError> { - let expected: Vec<_> = expected.iter().flatten().copied().collect(); + let expected: Vec<_> = expected + .iter() + .map(|source| { + source.expected().map_err(|unavailable| { + ApiError::unavailable( + "E_AUTH_UNAVAILABLE", + format!("{unavailable} authorization token is unavailable"), + ) + }) + }) + .collect::, _>>()? + .into_iter() + .flatten() + .collect(); if expected.is_empty() { return Ok(()); } diff --git a/crates/server/src/http/kodo_compat.rs b/crates/server/src/http/kodo_compat.rs new file mode 100644 index 0000000..2e31846 --- /dev/null +++ b/crates/server/src/http/kodo_compat.rs @@ -0,0 +1,101 @@ +use super::*; + +#[allow(clippy::too_many_arguments)] +pub(crate) async fn serve_kodo_compat( + config: Arc, + broker: Arc, + metrics: Arc, + accepting: Arc, + delivering: Arc, + publish_admission: Arc, + subscriptions: SubscriptionRegistry, + shutdown: tokio::sync::watch::Receiver, +) -> anyhow::Result<()> { + let Some(address) = config.network.kodo_http_address else { + wait_for_shutdown(shutdown).await; + return Ok(()); + }; + let state = app_state( + config, + broker, + metrics, + accepting, + delivering, + publish_admission, + subscriptions, + )?; + let tokens = state.tokens.clone(); + let token_shutdown = shutdown.clone(); + let router = Router::new() + .route("/ping", get(ping)) + .route("/stats", get(stats)) + .route("/channel/delete", post(delete_idle_channel_compat)) + .layer(middleware::from_fn(nsq_content_negotiation)) + .with_state(state); + let listener = TcpListener::bind(address).await?; + info!(%address, "Kodo compatibility HTTP API listening"); + let reloader = tokio::spawn(tokens.reload(token_shutdown)); + let result = axum::serve(listener, router) + .with_graceful_shutdown(wait_for_shutdown(shutdown)) + .await; + reloader.abort(); + result?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use rustqueue_queue::BrokerConfig; + + #[tokio::test] + async fn compatibility_delete_is_idle_only_without_an_http_token() { + let root = tempfile::tempdir().unwrap(); + let broker = Arc::new( + Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(), + ); + broker.create_channel("events", "workers").await.unwrap(); + let metrics = Arc::new(Metrics::default()); + let state = AppState { + config: Arc::new(Config::default()), + broker: Arc::clone(&broker), + metrics: Arc::clone(&metrics), + tokens: TokenSet::from_config(&Config::default()).unwrap(), + accepting: Arc::new(AtomicBool::new(true)), + delivering: Arc::new(AtomicBool::new(true)), + publish_admission: Arc::new(PublishAdmission::new(1024, metrics)), + subscriptions: SubscriptionRegistry::default(), + started_at: 1, + }; + + let response = delete_idle_channel_compat( + State(state.clone()), + Query(ChannelQuery { + topic: "events".into(), + channel: "workers".into(), + }), + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + assert!(broker.stats().topics[0].channels.is_empty()); + + let missing = delete_idle_channel_compat( + State(state), + Query(ChannelQuery { + topic: "events".into(), + channel: "workers".into(), + }), + ) + .await; + assert_eq!(missing.status(), StatusCode::NOT_FOUND); + let body = axum::body::to_bytes(missing.into_body(), usize::MAX) + .await + .unwrap(); + assert!(String::from_utf8_lossy(&body).contains("CHANNEL_NOT_FOUND")); + } +} diff --git a/crates/server/src/http/manage.rs b/crates/server/src/http/manage.rs index aa5e1de..3a75a01 100644 --- a/crates/server/src/http/manage.rs +++ b/crates/server/src/http/manage.rs @@ -1,6 +1,9 @@ use super::*; use axum::extract::Path; -use rustqueue_queue::{ChannelManagementAction, ManagementFenceSnapshot, TopicManagementAction}; +use rustqueue_queue::{ + ChannelManagementAction, ChannelManagementCommand, ManagementFenceSnapshot, + TopicManagementAction, +}; #[derive(Debug, Deserialize)] pub(super) struct TopicManageRequest { @@ -25,7 +28,7 @@ pub(super) async fn manage_topic( headers: HeaderMap, Json(request): Json, ) -> Result, ApiError> { - authorize(&headers, state.console_token.as_deref(), "console")?; + authorize(&headers, &state.tokens.console, "console")?; let action = parse_topic_action(&action)?; let result = state .broker @@ -53,35 +56,151 @@ pub(super) async fn manage_channel( headers: HeaderMap, Json(request): Json, ) -> Result, ApiError> { - authorize(&headers, state.console_token.as_deref(), "console")?; - let action = parse_channel_action(&action)?; + let cleanup_enabled = state.config.security.kodo_cleanup_enabled; + apply_channel_management(state, &action, headers, request, cleanup_enabled).await +} + +pub(super) async fn delete_idle_channel( + State(state): State, + headers: HeaderMap, + Json(request): Json, +) -> Result, ApiError> { + apply_channel_management(state, "delete-if-idle", headers, request, true).await +} + +pub(super) async fn delete_idle_channel_compat( + State(state): State, + Query(query): Query, +) -> Response { + if query.topic.is_empty() || query.channel.is_empty() { + return ( + StatusCode::BAD_REQUEST, + "E_BAD_REQUEST topic and channel are required", + ) + .into_response(); + } + let _permit = match state + .subscriptions + .begin_delete(&query.topic, &query.channel) + { + Ok(permit) => permit, + Err(blocked) => { + return ApiError::conflict( + "E_CHANNEL_NOT_IDLE", + format!("channel deletion is blocked: {blocked:?}"), + ) + .into_response() + } + }; + let revision = state.broker.registry_revision(); + let operation_id = kodo_compat_operation_id(revision, &query.topic, &query.channel); let result = state .broker - .manage_channel( - &request.operation_id, - &request.topic, - &request.channel, - action, - request.expected_revision, - request.tombstone_until_ms, + .manage_channel(ChannelManagementCommand { + operation_id: &operation_id, + topic: &query.topic, + channel: &query.channel, + action: ChannelManagementAction::Delete, + expected_revision: revision, + tombstone_until_ms: Some(kodo_cleanup_deadline()), + require_idle: true, + }) + .await; + match result { + Ok(result) => { + tracing::info!( + target = %format!("{}/{}", query.topic, query.channel), + changed = result.changed, + revision = result.revision, + "Kodo compatibility channel cleanup applied" + ); + "OK".into_response() + } + Err(BrokerError::TopicNotFound | BrokerError::ChannelNotFound) => { + (StatusCode::NOT_FOUND, "E_NOT_FOUND CHANNEL_NOT_FOUND").into_response() + } + Err(error) => ApiError::from(error).into_response(), + } +} + +async fn apply_channel_management( + state: AppState, + action: &str, + headers: HeaderMap, + request: ChannelManageRequest, + cleanup_enabled: bool, +) -> Result, ApiError> { + let (action, require_idle) = parse_channel_action(action, cleanup_enabled)?; + if require_idle { + authorize(&headers, &state.tokens.kodo_cleanup, "Kodo cleanup")?; + } else { + authorize(&headers, &state.tokens.console, "console")?; + } + let _delete_permit = if require_idle { + Some( + state + .subscriptions + .begin_delete(&request.topic, &request.channel) + .map_err(|blocked| { + ApiError::conflict( + "E_CHANNEL_NOT_IDLE", + format!("channel deletion is blocked: {blocked:?}"), + ) + })?, ) + } else { + None + }; + let tombstone_until_ms = if require_idle { + Some(kodo_cleanup_deadline()) + } else { + request.tombstone_until_ms + }; + let result = state + .broker + .manage_channel(ChannelManagementCommand { + operation_id: &request.operation_id, + topic: &request.topic, + channel: &request.channel, + action, + expected_revision: request.expected_revision, + tombstone_until_ms, + require_idle, + }) .await?; tracing::info!( target = %format!("{}/{}", request.topic, request.channel), action = ?action, changed = result.changed, revision = result.revision, - "console channel management applied" + "channel management applied" ); Ok(Json(json!(result))) } +fn kodo_cleanup_deadline() -> i64 { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(i64::MAX as u128) as i64; + now.saturating_add(10 * 60 * 1_000) +} + +fn kodo_compat_operation_id(revision: u64, topic: &str, channel: &str) -> String { + let target = format!("{topic}\0{channel}"); + format!( + "kodo-compat-{revision:016x}-{:08x}", + crc32c::crc32c(target.as_bytes()) + ) +} + pub(super) async fn sync_fences( State(state): State, headers: HeaderMap, Json(snapshot): Json, ) -> Result, ApiError> { - authorize(&headers, state.console_token.as_deref(), "console")?; + authorize(&headers, &state.tokens.console, "console")?; state.broker.sync_management_fences(snapshot).await?; Ok(Json(json!({"status": "ok"}))) } @@ -101,17 +220,54 @@ fn parse_topic_action(value: &str) -> Result { } } -fn parse_channel_action(value: &str) -> Result { +fn parse_channel_action( + value: &str, + cleanup_enabled: bool, +) -> Result<(ChannelManagementAction, bool), ApiError> { match value { - "create" => Ok(ChannelManagementAction::Create), - "pause" => Ok(ChannelManagementAction::Pause), - "unpause" => Ok(ChannelManagementAction::Unpause), - "empty" => Ok(ChannelManagementAction::Empty), - "delete" => Ok(ChannelManagementAction::Delete), - "tombstone" => Ok(ChannelManagementAction::Tombstone), + "create" => Ok((ChannelManagementAction::Create, false)), + "pause" => Ok((ChannelManagementAction::Pause, false)), + "unpause" => Ok((ChannelManagementAction::Unpause, false)), + "empty" => Ok((ChannelManagementAction::Empty, false)), + "delete" => Ok((ChannelManagementAction::Delete, false)), + "delete-if-idle" if cleanup_enabled => Ok((ChannelManagementAction::Delete, true)), + "tombstone" => Ok((ChannelManagementAction::Tombstone, false)), _ => Err(ApiError::bad_request( "E_BAD_ACTION", "unknown channel action", )), } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn idle_delete_action_is_unavailable_until_kodo_cleanup_is_enabled() { + assert!(parse_channel_action("delete-if-idle", false).is_err()); + assert!(matches!( + parse_channel_action("delete-if-idle", true), + Ok((ChannelManagementAction::Delete, true)) + )); + } + + #[test] + fn cleanup_deadline_is_server_bounded() { + let deadline = kodo_cleanup_deadline(); + let now = super::unix_seconds() * 1_000; + assert!((now + 9 * 60 * 1_000..=now + 11 * 60 * 1_000).contains(&deadline)); + } + + #[test] + fn compatibility_cleanup_operation_ids_include_the_full_target() { + assert_ne!( + kodo_compat_operation_id(7, "a:b", "c"), + kodo_compat_operation_id(7, "a", "b:c") + ); + assert_ne!( + kodo_compat_operation_id(7, "events", "workers"), + kodo_compat_operation_id(8, "events", "workers") + ); + } +} diff --git a/crates/server/src/http/native.rs b/crates/server/src/http/native.rs index ae4fc73..2774573 100644 --- a/crates/server/src/http/native.rs +++ b/crates/server/src/http/native.rs @@ -62,12 +62,10 @@ pub(super) async fn registry( ) -> Result, ApiError> { authorize_any( &headers, - &[ - state.registry_token.as_deref(), - state.console_token.as_deref(), - ], + &[&state.tokens.registry, &state.tokens.console], "registry or console", )?; + state.broker.expire_in_flight().await?; let stats = state.broker.stats(); let topics = registry_topics(&stats); let (stored_messages, depth, in_flight) = backlog(&stats); @@ -83,7 +81,8 @@ pub(super) async fn registry( && management_ready && storage_ready && (process_ready || stored_messages > 0 || depth > 0 || in_flight > 0); - let (binary, storage) = state.broker.capabilities(); + let (_, storage) = state.broker.capabilities(); + let binary = crate::config::runtime_capabilities(); Ok(Json(json!({ "format": 7, "revision": state.broker.registry_revision(), @@ -108,10 +107,7 @@ pub(super) async fn registry_head( ) -> Result, ApiError> { authorize_any( &headers, - &[ - state.registry_token.as_deref(), - state.console_token.as_deref(), - ], + &[&state.tokens.registry, &state.tokens.console], "registry or console", )?; let process_ready = state.accepting.load(Ordering::Acquire); @@ -138,13 +134,11 @@ pub(super) async fn capabilities( ) -> Result, ApiError> { authorize_any( &headers, - &[ - state.registry_token.as_deref(), - state.console_token.as_deref(), - ], + &[&state.tokens.registry, &state.tokens.console], "registry or console", )?; - let (binary, storage) = state.broker.capabilities(); + let (_, storage) = state.broker.capabilities(); + let binary = crate::config::runtime_capabilities(); Ok(Json(json!({"binary": binary, "storage": storage}))) } @@ -174,13 +168,10 @@ pub(super) async fn drain_status( ) -> Result, ApiError> { authorize_any( &headers, - &[ - state.registry_token.as_deref(), - state.console_token.as_deref(), - ], + &[&state.tokens.registry, &state.tokens.console], "registry or console", )?; - state.broker.expire_in_flight(); + state.broker.expire_in_flight().await?; let stats = state.broker.metrics_stats(false, 0); let (stored_messages, depth, in_flight) = backlog(&stats); let draining = !state.accepting.load(Ordering::Acquire); @@ -212,7 +203,7 @@ pub(super) async fn set_drain( headers: HeaderMap, Json(request): Json, ) -> Result, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; state.accepting.store(!request.enabled, Ordering::Release); state.delivering.store( !request.enabled || !request.freeze_deliveries, @@ -223,7 +214,7 @@ pub(super) async fn set_drain( freeze_deliveries = request.freeze_deliveries, "broker drain state changed" ); - state.broker.expire_in_flight(); + state.broker.expire_in_flight().await?; let stats = state.broker.metrics_stats(false, 0); let (stored_messages, depth, in_flight) = backlog(&stats); let publish_inflight_bytes = state.metrics.publish_inflight_bytes.load(Ordering::Acquire); @@ -252,25 +243,25 @@ pub(super) async fn set_drain( pub(super) async fn native_stats( State(state): State, Query(query): Query, -) -> Json { - let filtered = filter_stats( - state.broker.stats(), - query.topic.as_deref(), - query.channel.as_deref(), - ); - Json(json!({ +) -> Result, ApiError> { + state.broker.expire_in_flight().await?; + let filtered = state + .broker + .filtered_stats(query.topic.as_deref(), query.channel.as_deref()); + Ok(Json(json!({ "complete": true, "node_id": state.config.node.id, "collected_at_ms": now_ms(), "topics": filtered.topics, - })) + }))) } pub(super) async fn observe( State(state): State, headers: HeaderMap, ) -> Result, ApiError> { - authorize(&headers, state.console_token.as_deref(), "console")?; + authorize(&headers, &state.tokens.console, "console")?; + state.broker.expire_in_flight().await?; let stats = state.broker.stats(); let segment_count = stats.aggregate.segment_count; let segment_bytes = stats.aggregate.segment_bytes; @@ -294,7 +285,7 @@ pub(super) async fn observe_head( State(state): State, headers: HeaderMap, ) -> Result, ApiError> { - authorize(&headers, state.console_token.as_deref(), "console")?; + authorize(&headers, &state.tokens.console, "console")?; Ok(Json(observation_head(&state))) } @@ -304,7 +295,8 @@ fn observation_head(state: &AppState) -> Value { let storage_healthy = state.broker.storage_healthy(); let disk_ready = state.publish_admission.storage_ready(); let management_ready = state.broker.management_fences_ready(); - let (binary, storage) = state.broker.capabilities(); + let (_, storage) = state.broker.capabilities(); + let binary = crate::config::runtime_capabilities(); let runtime = state.metrics.snapshot(); json!({ "schema_version": 1, @@ -350,7 +342,7 @@ pub(super) async fn scrub( State(state): State, headers: HeaderMap, ) -> Result, ApiError> { - authorize(&headers, state.admin_token.as_deref(), "admin")?; + authorize(&headers, &state.tokens.admin, "admin")?; let records = state.broker.scrub().await?; Ok(Json(json!({"status": "ok", "records_checked": records}))) } diff --git a/crates/server/src/http/nsq_stats.rs b/crates/server/src/http/nsq_stats.rs new file mode 100644 index 0000000..f098724 --- /dev/null +++ b/crates/server/src/http/nsq_stats.rs @@ -0,0 +1,242 @@ +use super::*; +use crate::subscriptions::ClientSnapshot; +use rustqueue_queue::{ChannelStats, TopicStats}; + +#[derive(Serialize)] +struct NsqStatsResponse { + version: &'static str, + health: &'static str, + start_time: i64, + topics: Vec, +} + +#[derive(Serialize)] +struct NsqTopicStats { + topic_name: String, + name: String, + depth: u64, + memory_depth: u64, + backend_depth: u64, + message_count: u64, + paused: bool, + channels: Vec, +} + +#[derive(Serialize)] +struct NsqChannelStats { + channel_name: String, + name: String, + depth: u64, + memory_depth: u64, + backend_depth: u64, + message_count: u64, + in_flight_count: u64, + deferred_count: u64, + requeue_count: u64, + timeout_count: u64, + client_count: usize, + clients: Vec, + paused: bool, +} + +pub(super) async fn stats( + State(state): State, + Query(query): Query, +) -> Response { + if let Err(error) = state.broker.expire_in_flight().await { + return ApiError::from(error).into_response(); + } + let filtered = state + .broker + .filtered_stats(query.topic.as_deref(), query.channel.as_deref()); + let include_clients = query.include_clients.unwrap_or(true); + if query.format.as_deref() == Some("json") { + let topics = filtered + .topics + .into_iter() + .map(|topic| convert_topic(&state, topic, include_clients)) + .collect(); + return Json(NsqStatsResponse { + version: env!("CARGO_PKG_VERSION"), + health: "OK", + start_time: state.started_at, + topics, + }) + .into_response(); + } + text_stats(filtered).into_response() +} + +fn convert_topic(state: &AppState, topic: TopicStats, include_clients: bool) -> NsqTopicStats { + let topic_name = topic.name; + let channels = topic + .channels + .into_iter() + .map(|channel| convert_channel(state, &topic_name, channel, include_clients)) + .collect(); + NsqTopicStats { + topic_name: topic_name.clone(), + name: topic_name, + depth: nsq_topic_depth(), + memory_depth: 0, + backend_depth: 0, + message_count: topic.published_count, + paused: topic.paused, + channels, + } +} + +fn convert_channel( + state: &AppState, + topic: &str, + channel: ChannelStats, + include_clients: bool, +) -> NsqChannelStats { + let queued_depth = nsq_channel_depth(&channel); + let channel_name = channel.name; + let (client_count, clients) = if include_clients { + let clients = state.subscriptions.clients(topic, &channel_name); + (clients.len(), clients) + } else { + ( + state.subscriptions.client_count(topic, &channel_name), + Vec::new(), + ) + }; + NsqChannelStats { + channel_name: channel_name.clone(), + name: channel_name, + depth: queued_depth, + memory_depth: 0, + backend_depth: queued_depth, + message_count: channel.message_count, + in_flight_count: channel.in_flight_count, + deferred_count: channel.deferred_count, + requeue_count: channel.requeue_count, + timeout_count: channel.timeout_count, + client_count, + clients, + paused: channel.paused, + } +} + +fn nsq_channel_depth(channel: &ChannelStats) -> u64 { + channel + .depth + .saturating_sub(channel.in_flight_count) + .saturating_sub(channel.deferred_count) +} + +fn nsq_topic_depth() -> u64 { + // RustQueue exposes every durable log position directly to each Channel; + // there is no intermediate NSQ Topic queue waiting to fan messages out. + 0 +} + +fn text_stats(stats: BrokerStats) -> String { + let mut output = String::new(); + for topic in stats.topics { + output.push_str(&format!("[{}] depth={}\n", topic.name, nsq_topic_depth())); + for channel in topic.channels { + output.push_str(&format!( + " [{}] depth={} in-flight={} deferred={}\n", + channel.name, + nsq_channel_depth(&channel), + channel.in_flight_count, + channel.deferred_count + )); + } + } + output +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn standard_channel_fields_use_nsq_names_and_keep_name_alias() { + let value = serde_json::to_value(NsqChannelStats { + channel_name: "workers".into(), + name: "workers".into(), + depth: 4, + memory_depth: 0, + backend_depth: 4, + message_count: 7, + in_flight_count: 1, + deferred_count: 2, + requeue_count: 3, + timeout_count: 5, + client_count: 0, + clients: Vec::new(), + paused: false, + }) + .unwrap(); + assert_eq!(value["channel_name"], "workers"); + assert_eq!(value["name"], "workers"); + assert_eq!(value["message_count"], 7); + assert_eq!(value["client_count"], 0); + } + + #[test] + fn nsq_depth_excludes_in_flight_and_deferred_messages() { + let channel = ChannelStats { + name: "workers".into(), + depth: 7, + message_count: 9, + in_flight_count: 2, + deferred_count: 3, + requeue_count: 0, + timeout_count: 0, + paused: false, + ephemeral: false, + ack_cursor: 0, + ack_gap: 0, + }; + assert_eq!(nsq_channel_depth(&channel), 2); + } + + #[test] + fn retained_topic_messages_are_not_reported_as_nsq_fanout_depth() { + assert_eq!(nsq_topic_depth(), 0); + } + + #[test] + fn kodo_replay_fixture_matches_the_rustqueue_stats_contract() { + let actual = serde_json::to_value(NsqStatsResponse { + version: "0.8.0", + health: "OK", + start_time: 1_700_000_001, + topics: vec![NsqTopicStats { + topic_name: "events".into(), + name: "events".into(), + depth: 0, + memory_depth: 0, + backend_depth: 0, + message_count: 10, + paused: false, + channels: vec![NsqChannelStats { + channel_name: "workers".into(), + name: "workers".into(), + depth: 4, + memory_depth: 0, + backend_depth: 4, + message_count: 10, + in_flight_count: 1, + deferred_count: 2, + requeue_count: 3, + timeout_count: 4, + client_count: 0, + clients: Vec::new(), + paused: false, + }], + }], + }) + .unwrap(); + let expected: serde_json::Value = serde_json::from_str(include_str!( + "../../../../tests/kodo-replay/fixtures/stats-1.json" + )) + .unwrap(); + assert_eq!(actual, expected); + } +} diff --git a/crates/server/src/http/tokens.rs b/crates/server/src/http/tokens.rs new file mode 100644 index 0000000..e3f57cd --- /dev/null +++ b/crates/server/src/http/tokens.rs @@ -0,0 +1,214 @@ +use crate::config::Config; +use parking_lot::RwLock; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::watch; + +#[derive(Clone)] +pub(super) struct TokenSource { + scope: &'static str, + path: Option>, + state: Arc>, +} + +#[derive(Clone)] +enum TokenState { + Unconfigured, + Ready(Arc), + Unavailable, +} + +#[derive(Clone)] +pub(super) struct TokenSet { + pub admin: TokenSource, + pub publish: TokenSource, + pub registry: TokenSource, + pub console: TokenSource, + pub kodo_cleanup: TokenSource, +} + +impl TokenSource { + fn from_path(scope: &'static str, path: Option<&Path>) -> anyhow::Result { + let path = path.map(|path| Arc::new(path.to_path_buf())); + let state = match path.as_deref() { + Some(path) => TokenState::Ready(read_token(path)?), + None => TokenState::Unconfigured, + }; + Ok(Self { + scope, + path, + state: Arc::new(RwLock::new(state)), + }) + } + + pub(super) fn expected(&self) -> Result>, &'static str> { + match &*self.state.read() { + TokenState::Unconfigured => Ok(None), + TokenState::Ready(token) => Ok(Some(Arc::clone(token))), + TokenState::Unavailable => Err(self.scope), + } + } + + fn refresh(&self) { + let Some(path) = self.path.as_deref() else { + return; + }; + match read_token(path) { + Ok(token) => { + let mut state = self.state.write(); + let changed = !matches!(&*state, TokenState::Ready(current) if current.as_ref() == token.as_ref()); + *state = TokenState::Ready(token); + if changed { + tracing::info!(scope = self.scope, "reloaded HTTP authorization token"); + } + } + Err(error) => { + let mut state = self.state.write(); + let changed = !matches!(&*state, TokenState::Unavailable); + *state = TokenState::Unavailable; + if changed { + tracing::warn!( + scope = self.scope, + %error, + "HTTP authorization token became unavailable" + ); + } + } + } + } + + #[cfg(test)] + pub(super) fn fixed(scope: &'static str, token: &str) -> Self { + Self { + scope, + path: None, + state: Arc::new(RwLock::new(TokenState::Ready(Arc::from(token)))), + } + } +} + +impl TokenSet { + pub(super) fn from_config(config: &Config) -> anyhow::Result { + Ok(Self { + admin: TokenSource::from_path("admin", config.security.admin_token_file.as_deref())?, + publish: TokenSource::from_path( + "publish", + config.security.publish_token_file.as_deref(), + )?, + registry: TokenSource::from_path( + "registry", + config.security.registry_token_file.as_deref(), + )?, + console: TokenSource::from_path( + "console", + config.security.console_token_file.as_deref(), + )?, + kodo_cleanup: TokenSource::from_path( + "Kodo cleanup", + config.security.kodo_cleanup_token_file.as_deref(), + )?, + }) + } + + pub(super) async fn reload(self, mut shutdown: watch::Receiver) { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + biased; + changed = shutdown.changed() => { + if changed.is_err() || *shutdown.borrow() { + return; + } + } + _ = interval.tick() => { + self.admin.refresh(); + self.publish.refresh(); + self.registry.refresh(); + self.console.refresh(); + self.kodo_cleanup.refresh(); + } + } + } + } +} + +fn read_token(path: &Path) -> anyhow::Result> { + let token = std::fs::read_to_string(path)?; + let token = token.trim(); + anyhow::ensure!(!token.is_empty(), "token file {} is empty", path.display()); + Ok(Arc::from(token)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn configured_tokens_reload_and_fail_closed() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("token"); + std::fs::write(&path, "first\n").unwrap(); + let source = TokenSource::from_path("registry", Some(&path)).unwrap(); + assert_eq!(source.expected().unwrap().unwrap().as_ref(), "first"); + + std::fs::write(&path, "second\n").unwrap(); + source.refresh(); + assert_eq!(source.expected().unwrap().unwrap().as_ref(), "second"); + + std::fs::write(&path, "\n").unwrap(); + source.refresh(); + assert_eq!(source.expected(), Err("registry")); + } + + #[tokio::test] + async fn reload_loop_observes_rotation_and_recovers() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("token"); + std::fs::write(&path, "first\n").unwrap(); + let registry = TokenSource::from_path("registry", Some(&path)).unwrap(); + let unused = TokenSource::from_path("unused", None).unwrap(); + let tokens = TokenSet { + admin: unused.clone(), + publish: unused.clone(), + registry: registry.clone(), + console: unused.clone(), + kodo_cleanup: unused, + }; + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let reload = tokio::spawn(tokens.reload(shutdown_rx)); + + std::fs::write(&path, "second\n").unwrap(); + wait_for(®istry, Ok(Some("second"))).await; + + std::fs::write(&path, "\n").unwrap(); + wait_for(®istry, Err("registry")).await; + + std::fs::write(&path, "third\n").unwrap(); + wait_for(®istry, Ok(Some("third"))).await; + + shutdown_tx.send(true).unwrap(); + tokio::time::timeout(Duration::from_secs(1), reload) + .await + .unwrap() + .unwrap(); + } + + async fn wait_for(source: &TokenSource, expected: Result, &'static str>) { + tokio::time::timeout(Duration::from_secs(3), async { + loop { + let current = source + .expected() + .map(|token| token.as_deref().map(str::to_owned)); + let expected = expected.map(|token| token.map(str::to_owned)); + if current == expected { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + } +} diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 30ea402..7693a1f 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -5,6 +5,7 @@ mod config; mod disk_guard; mod http; mod metrics; +mod subscriptions; mod tcp; mod tls; @@ -18,7 +19,8 @@ use rustqueue_storage::DataDirectoryLock; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use tracing::{error, info}; +use subscriptions::SubscriptionRegistry; +use tracing::{error, info, warn}; use tracing_subscriber::EnvFilter; #[derive(Parser, Debug)] @@ -86,7 +88,9 @@ async fn main() -> anyhow::Result<()> { config.limits.node_publish_inflight_bytes, Arc::clone(&metrics), )); + let subscriptions = SubscriptionRegistry::default(); let initially_pressured = disk_guard::initialize(&config, &publish_admission, &metrics)?; + let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); info!( node_id = config.node.id, @@ -96,63 +100,229 @@ async fn main() -> anyhow::Result<()> { "starting share-nothing RustQueue broker" ); - let tcp_task = tokio::spawn(tcp::serve( + let mut tcp_task = tokio::spawn(tcp::serve( Arc::clone(&config), Arc::clone(&broker), Arc::clone(&metrics), Arc::clone(&accepting), Arc::clone(&delivering), Arc::clone(&publish_admission), + subscriptions.clone(), + shutdown_rx.clone(), + std::time::Duration::from_secs(config.shutdown.grace_seconds) + .saturating_sub(std::time::Duration::from_millis(250)), )); - let http_task = tokio::spawn(http::serve( + let mut http_task = tokio::spawn(http::serve( Arc::clone(&config), Arc::clone(&broker), Arc::clone(&metrics), Arc::clone(&accepting), Arc::clone(&delivering), Arc::clone(&publish_admission), + subscriptions.clone(), + shutdown_rx.clone(), )); - let scrub_task = tokio::spawn(run_scrubber( + let mut kodo_http_task = tokio::spawn(http::serve_kodo_compat( + Arc::clone(&config), + Arc::clone(&broker), + Arc::clone(&metrics), + Arc::clone(&accepting), + Arc::clone(&delivering), + Arc::clone(&publish_admission), + subscriptions, + shutdown_rx, + )); + let mut scrub_task = tokio::spawn(run_scrubber( Arc::clone(&broker), std::time::Duration::from_secs(config.storage.maintenance_startup_delay_seconds), std::time::Duration::from_secs(config.storage.scrub_interval_seconds), )); - let disk_task = tokio::spawn(disk_guard::run( + let mut disk_task = tokio::spawn(disk_guard::run( Arc::clone(&config), Arc::clone(&broker), Arc::clone(&publish_admission), Arc::clone(&metrics), initially_pressured, )); - let storage_health_task = tokio::spawn(monitor_storage_health(Arc::clone(&broker))); - - tokio::select! { - result = tcp_task => result.context("TCP task join")??, - result = http_task => result.context("HTTP task join")??, - result = scrub_task => result.context("storage scrub task join")??, - result = disk_task => result.context("disk guard task join")??, - result = storage_health_task => result.context("storage health task join")??, + let mut storage_health_task = tokio::spawn(monitor_storage_health(Arc::clone(&broker))); + + let (completed_task, terminal_result) = tokio::select! { + result = &mut tcp_task => ( + Some(RuntimeTask::Tcp), + unexpected_task_exit("TCP", result), + ), + result = &mut http_task => ( + Some(RuntimeTask::Http), + unexpected_task_exit("HTTP", result), + ), + result = &mut kodo_http_task => ( + Some(RuntimeTask::KodoHttp), + unexpected_task_exit("Kodo HTTP", result), + ), + result = &mut scrub_task => ( + Some(RuntimeTask::Scrub), + unexpected_task_exit("storage scrub", result), + ), + result = &mut disk_task => ( + Some(RuntimeTask::DiskGuard), + unexpected_task_exit("disk guard", result), + ), + result = &mut storage_health_task => ( + Some(RuntimeTask::StorageHealth), + unexpected_task_exit("storage health", result), + ), _ = shutdown_signal() => { info!("shutdown signal received"); - accepting.store(false, Ordering::Release); - delivering.store(false, Ordering::Release); - let deadline = tokio::time::Instant::now() - + std::time::Duration::from_secs(config.shutdown.grace_seconds); - while metrics.tcp_connections.load(Ordering::Acquire) > 0 - && tokio::time::Instant::now() < deadline - { - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - } - let released = broker.release_all_in_flight(); - broker.flush().await.context("flush queue state")?; - info!(released, "share-nothing broker stopped cleanly"); + (None, Ok(())) }, + }; + accepting.store(false, Ordering::Release); + delivering.store(false, Ordering::Release); + let _ = shutdown_tx.send(true); + let deadline = + tokio::time::Instant::now() + std::time::Duration::from_secs(config.shutdown.grace_seconds); + while shutdown_work_in_progress( + &metrics, + tcp_task.is_finished() && http_task.is_finished() && kodo_http_task.is_finished(), + ) && tokio::time::Instant::now() < deadline + { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + if shutdown_work_in_progress( + &metrics, + tcp_task.is_finished() && http_task.is_finished() && kodo_http_task.is_finished(), + ) { + warn!( + tcp_connections = metrics.tcp_connections.load(Ordering::Acquire), + publish_inflight_bytes = metrics.publish_inflight_bytes.load(Ordering::Acquire), + "broker shutdown grace expired with client work in progress" + ); + } + let mut cleanup_error = None; + if completed_task != Some(RuntimeTask::Tcp) { + record_cleanup_error( + &mut cleanup_error, + stop_service_task(&mut tcp_task, "TCP").await, + ); + } + if completed_task != Some(RuntimeTask::Http) { + record_cleanup_error( + &mut cleanup_error, + stop_service_task(&mut http_task, "HTTP").await, + ); + } + if completed_task != Some(RuntimeTask::KodoHttp) { + record_cleanup_error( + &mut cleanup_error, + stop_service_task(&mut kodo_http_task, "Kodo HTTP").await, + ); + } + if completed_task != Some(RuntimeTask::Scrub) { + record_cleanup_error( + &mut cleanup_error, + stop_background_task(&mut scrub_task, "storage scrub").await, + ); + } + if completed_task != Some(RuntimeTask::DiskGuard) { + record_cleanup_error( + &mut cleanup_error, + stop_background_task(&mut disk_task, "disk guard").await, + ); + } + if completed_task != Some(RuntimeTask::StorageHealth) { + record_cleanup_error( + &mut cleanup_error, + stop_background_task(&mut storage_health_task, "storage health").await, + ); + } + let released = broker.release_all_in_flight(); + record_cleanup_error( + &mut cleanup_error, + broker.flush().await.context("flush queue state"), + ); + info!(released, "share-nothing broker stopped cleanly"); + combine_shutdown_results(terminal_result, cleanup_error) +} + +fn shutdown_work_in_progress(metrics: &Metrics, services_finished: bool) -> bool { + !services_finished || metrics.publish_inflight_bytes.load(Ordering::Acquire) > 0 +} + +#[derive(Clone, Copy, Eq, PartialEq)] +enum RuntimeTask { + Tcp, + Http, + KodoHttp, + Scrub, + DiskGuard, + StorageHealth, +} + +fn unexpected_task_exit( + name: &str, + result: Result, tokio::task::JoinError>, +) -> anyhow::Result<()> { + match result { + Ok(Ok(())) => anyhow::bail!("{name} task stopped unexpectedly"), + Ok(Err(error)) => Err(error).with_context(|| format!("{name} task stopped with an error")), + Err(error) => Err(error).with_context(|| format!("{name} task join")), + } +} + +fn record_cleanup_error(first: &mut Option, result: anyhow::Result<()>) { + if let Err(error) = result { + warn!(%error, "broker shutdown cleanup step failed"); + if first.is_none() { + *first = Some(error); + } + } +} + +fn combine_shutdown_results( + terminal: anyhow::Result<()>, + cleanup: Option, +) -> anyhow::Result<()> { + match (terminal, cleanup) { + (Ok(()), None) => Ok(()), + (Err(error), None) | (Ok(()), Some(error)) => Err(error), + (Err(terminal), Some(cleanup)) => Err(anyhow::anyhow!( + "{terminal:#}; shutdown cleanup also failed: {cleanup:#}" + )), + } +} + +async fn stop_service_task( + task: &mut tokio::task::JoinHandle>, + name: &str, +) -> anyhow::Result<()> { + if !task.is_finished() { + warn!( + service = name, + "service did not stop before the broker deadline; aborting it" + ); + task.abort(); + } + match task.await { + Ok(result) => result.with_context(|| format!("{name} service stopped with an error")), + Err(error) if error.is_cancelled() => Ok(()), + Err(error) => Err(error).with_context(|| format!("{name} service task join")), + } +} + +async fn stop_background_task( + task: &mut tokio::task::JoinHandle>, + name: &str, +) -> anyhow::Result<()> { + task.abort(); + match task.await { + Ok(result) => result.with_context(|| format!("{name} task stopped with an error")), + Err(error) if error.is_cancelled() => Ok(()), + Err(error) => Err(error).with_context(|| format!("{name} task join")), } - Ok(()) } fn write_binary_capabilities(path: &std::path::Path) -> anyhow::Result<()> { - let bytes = serde_json::to_vec(&rustqueue_storage::binary_capabilities())?; + let bytes = serde_json::to_vec(&config::runtime_capabilities())?; std::fs::write(path, bytes).with_context(|| format!("write capabilities to {}", path.display())) } @@ -214,3 +384,24 @@ async fn shutdown_signal() { let terminate = std::future::pending::<()>(); tokio::select! { _ = ctrl_c => {}, _ = terminate => {} } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn shutdown_waits_for_service_completion_and_publish_reservations() { + let metrics = Metrics::default(); + assert!(!shutdown_work_in_progress(&metrics, true)); + assert!(shutdown_work_in_progress(&metrics, false)); + + metrics.tcp_connections.store(1, Ordering::Release); + assert!(!shutdown_work_in_progress(&metrics, true)); + metrics.tcp_connections.store(0, Ordering::Release); + + metrics + .publish_inflight_bytes + .store(1024, Ordering::Release); + assert!(shutdown_work_in_progress(&metrics, true)); + } +} diff --git a/crates/server/src/metrics.rs b/crates/server/src/metrics.rs index db7307b..cced357 100644 --- a/crates/server/src/metrics.rs +++ b/crates/server/src/metrics.rs @@ -381,14 +381,18 @@ mod tests { topics: vec![TopicStats { name: "events".into(), paused: false, + published_count: 7, message_count: 7, segment_count: 1, segment_bytes: 100, channels: vec![ChannelStats { name: "workers".into(), depth: 3, + message_count: 7, in_flight_count: 2, deferred_count: 1, + requeue_count: 0, + timeout_count: 0, paused: false, ephemeral: false, ack_cursor: 4, diff --git a/crates/server/src/subscriptions.rs b/crates/server/src/subscriptions.rs new file mode 100644 index 0000000..56f1a41 --- /dev/null +++ b/crates/server/src/subscriptions.rs @@ -0,0 +1,276 @@ +use parking_lot::Mutex; +use serde::Serialize; +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +#[derive(Clone, Debug, Default)] +pub(crate) struct ClientIdentity { + pub client_id: String, + pub hostname: String, + pub remote_address: String, + pub user_agent: String, + pub sample_rate: u8, + pub tls: bool, + pub deflate: bool, + pub snappy: bool, + pub authed: bool, +} + +#[derive(Clone, Debug, Serialize)] +pub(crate) struct ClientSnapshot { + pub client_id: String, + pub hostname: String, + pub remote_address: String, + pub user_agent: String, + pub state: i32, + pub ready_count: u64, + pub in_flight_count: u64, + pub message_count: u64, + pub finish_count: u64, + pub requeue_count: u64, + pub connect_ts: i64, + pub sample_rate: u8, + pub tls: bool, + pub deflate: bool, + pub snappy: bool, + pub authed: bool, +} + +#[derive(Clone, Default)] +pub(crate) struct SubscriptionRegistry { + inner: Arc>, +} + +#[derive(Default)] +struct RegistryState { + next_id: u64, + channels: BTreeMap>>, + deleting: BTreeSet, +} + +#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)] +struct ChannelKey { + topic: String, + channel: String, +} + +struct ClientRuntime { + identity: ClientIdentity, + connected_at: i64, + ready_count: AtomicU64, + in_flight_count: AtomicU64, + message_count: AtomicU64, + finish_count: AtomicU64, + requeue_count: AtomicU64, +} + +pub(crate) struct SubscriptionLease { + registry: SubscriptionRegistry, + key: ChannelKey, + id: u64, + client: Arc, +} + +pub(crate) struct DeletePermit { + registry: SubscriptionRegistry, + key: ChannelKey, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum DeleteBlocked { + ActiveClients(usize), + InProgress, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum RegisterBlocked { + DeleteInProgress, +} + +impl SubscriptionRegistry { + pub fn register( + &self, + topic: &str, + channel: &str, + identity: ClientIdentity, + ) -> Result { + let key = ChannelKey { + topic: topic.into(), + channel: channel.into(), + }; + let mut state = self.inner.lock(); + if state.deleting.contains(&key) { + return Err(RegisterBlocked::DeleteInProgress); + } + state.next_id = state.next_id.wrapping_add(1).max(1); + let id = state.next_id; + let client = Arc::new(ClientRuntime { + identity, + connected_at: unix_seconds(), + ready_count: AtomicU64::new(0), + in_flight_count: AtomicU64::new(0), + message_count: AtomicU64::new(0), + finish_count: AtomicU64::new(0), + requeue_count: AtomicU64::new(0), + }); + state + .channels + .entry(key.clone()) + .or_default() + .insert(id, Arc::clone(&client)); + Ok(SubscriptionLease { + registry: self.clone(), + key, + id, + client, + }) + } + + pub fn clients(&self, topic: &str, channel: &str) -> Vec { + let key = ChannelKey { + topic: topic.into(), + channel: channel.into(), + }; + self.inner + .lock() + .channels + .get(&key) + .into_iter() + .flat_map(BTreeMap::values) + .map(|client| client.snapshot()) + .collect() + } + + pub fn client_count(&self, topic: &str, channel: &str) -> usize { + let key = ChannelKey { + topic: topic.into(), + channel: channel.into(), + }; + self.inner + .lock() + .channels + .get(&key) + .map_or(0, BTreeMap::len) + } + + pub fn begin_delete(&self, topic: &str, channel: &str) -> Result { + let key = ChannelKey { + topic: topic.into(), + channel: channel.into(), + }; + let mut state = self.inner.lock(); + if state.deleting.contains(&key) { + return Err(DeleteBlocked::InProgress); + } + let clients = state.channels.get(&key).map_or(0, BTreeMap::len); + if clients > 0 { + return Err(DeleteBlocked::ActiveClients(clients)); + } + state.deleting.insert(key.clone()); + Ok(DeletePermit { + registry: self.clone(), + key, + }) + } +} + +impl SubscriptionLease { + pub fn update_flow(&self, ready: u64, in_flight: usize) { + self.client.ready_count.store(ready, Ordering::Relaxed); + self.client + .in_flight_count + .store(in_flight as u64, Ordering::Relaxed); + } + + pub fn observe_delivery(&self) { + self.client.message_count.fetch_add(1, Ordering::Relaxed); + } + + pub fn observe_finish(&self) { + self.client.finish_count.fetch_add(1, Ordering::Relaxed); + } + + pub fn observe_requeue(&self) { + self.client.requeue_count.fetch_add(1, Ordering::Relaxed); + } +} + +impl Drop for SubscriptionLease { + fn drop(&mut self) { + let mut state = self.registry.inner.lock(); + let remove_channel = state.channels.get_mut(&self.key).is_some_and(|clients| { + clients.remove(&self.id); + clients.is_empty() + }); + if remove_channel { + state.channels.remove(&self.key); + } + } +} + +impl Drop for DeletePermit { + fn drop(&mut self) { + self.registry.inner.lock().deleting.remove(&self.key); + } +} + +impl ClientRuntime { + fn snapshot(&self) -> ClientSnapshot { + ClientSnapshot { + client_id: self.identity.client_id.clone(), + hostname: self.identity.hostname.clone(), + remote_address: self.identity.remote_address.clone(), + user_agent: self.identity.user_agent.clone(), + state: 3, + ready_count: self.ready_count.load(Ordering::Relaxed), + in_flight_count: self.in_flight_count.load(Ordering::Relaxed), + message_count: self.message_count.load(Ordering::Relaxed), + finish_count: self.finish_count.load(Ordering::Relaxed), + requeue_count: self.requeue_count.load(Ordering::Relaxed), + connect_ts: self.connected_at, + sample_rate: self.identity.sample_rate, + tls: self.identity.tls, + deflate: self.identity.deflate, + snappy: self.identity.snappy, + authed: self.identity.authed, + } + } +} + +fn unix_seconds() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + .min(i64::MAX as u64) as i64 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn delete_barrier_blocks_new_subscriptions_and_counts_live_clients() { + let registry = SubscriptionRegistry::default(); + let lease = registry + .register("events", "workers", ClientIdentity::default()) + .unwrap(); + assert_eq!(registry.clients("events", "workers").len(), 1); + assert!(matches!( + registry.begin_delete("events", "workers"), + Err(DeleteBlocked::ActiveClients(1)) + )); + drop(lease); + + let permit = registry.begin_delete("events", "workers").unwrap(); + assert!(matches!( + registry.register("events", "workers", ClientIdentity::default()), + Err(RegisterBlocked::DeleteInProgress) + )); + drop(permit); + assert!(registry + .register("events", "workers", ClientIdentity::default()) + .is_ok()); + } +} diff --git a/crates/server/src/tcp.rs b/crates/server/src/tcp.rs index 2ad738d..469d259 100644 --- a/crates/server/src/tcp.rs +++ b/crates/server/src/tcp.rs @@ -21,6 +21,7 @@ use crate::auth::{AuthError, AuthSession, Authenticator}; use crate::compression::{self, BoxIo}; use crate::config::Config; use crate::metrics::Metrics; +use crate::subscriptions::{ClientIdentity, SubscriptionLease, SubscriptionRegistry}; use crate::tls; use anyhow::Context; use bytes::Bytes; @@ -42,11 +43,22 @@ use tokio::io::{ AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, ReadHalf, WriteHalf, }; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::Semaphore; +use tokio::sync::{watch, Semaphore}; +use tokio::task::JoinSet; use tokio::time::{interval, MissedTickBehavior}; use tokio_rustls::TlsAcceptor; use tracing::{debug, info, warn}; +pub(crate) fn broker_storage_error(error: &BrokerError) -> bool { + matches!( + error, + BrokerError::StorageUnavailable + | BrokerError::Storage(_) + | BrokerError::Io(_) + | BrokerError::InvalidRecord(_) + ) +} + #[derive(Debug)] struct ParsedCommand { command: Command, @@ -63,6 +75,7 @@ enum Compression { struct Subscription { topic: String, channel: String, + lease: SubscriptionLease, } const MAX_FETCH_MESSAGES: u16 = 64; @@ -107,6 +120,7 @@ struct SessionState { rdy: u64, in_flight: HashMap, closing: bool, + client_identity: ClientIdentity, } impl SessionState { @@ -118,8 +132,17 @@ impl SessionState { self.sample_cursor = (self.sample_cursor + 1) % 100; selected } + + fn update_subscription_flow(&self) { + if let Some(subscription) = &self.subscription { + subscription + .lease + .update_flow(self.rdy, self.in_flight.len()); + } + } } +#[allow(clippy::too_many_arguments)] pub async fn serve( config: Arc, broker: Arc, @@ -127,6 +150,9 @@ pub async fn serve( accepting: Arc, delivering: Arc, publish_admission: Arc, + subscriptions: SubscriptionRegistry, + mut shutdown: watch::Receiver, + shutdown_grace: Duration, ) -> anyhow::Result<()> { let listener = TcpListener::bind(config.network.tcp_address).await?; let tls_acceptor = tls::acceptor(config.security.tls.as_ref())?; @@ -137,8 +163,22 @@ pub async fn serve( let ephemeral_consumers = EphemeralConsumers::default(); info!(address = %config.network.tcp_address, "NSQ TCP listener ready"); + let mut sessions = JoinSet::new(); loop { - let (stream, peer) = listener.accept().await?; + let accepted = tokio::select! { + result = listener.accept() => Some(result?), + result = sessions.join_next(), if !sessions.is_empty() => { + log_session_result(result); + None + } + _ = shutdown.changed() => break, + }; + let Some((stream, peer)) = accepted else { + continue; + }; + if *shutdown.borrow() { + break; + } let permit = match Arc::clone(&permits).try_acquire_owned() { Ok(permit) => permit, Err(_) => { @@ -155,7 +195,8 @@ pub async fn serve( let accepting = Arc::clone(&accepting); let delivering = Arc::clone(&delivering); let publish_admission = Arc::clone(&publish_admission); - tokio::spawn(async move { + let subscriptions = subscriptions.clone(); + sessions.spawn(async move { let _permit = permit; metrics.tcp_connections.fetch_add(1, Ordering::Relaxed); if let Err(error) = handle_connection( @@ -170,6 +211,7 @@ pub async fn serve( accepting, delivering, publish_admission, + subscriptions, ) .await { @@ -178,6 +220,31 @@ pub async fn serve( metrics.tcp_connections.fetch_sub(1, Ordering::Relaxed); }); } + info!("NSQ TCP listener stopped accepting new connections"); + if tokio::time::timeout(shutdown_grace, drain_sessions(&mut sessions)) + .await + .is_err() + { + warn!( + active_sessions = sessions.len(), + "NSQ TCP shutdown grace expired" + ); + sessions.abort_all(); + drain_sessions(&mut sessions).await; + } + Ok(()) +} + +async fn drain_sessions(sessions: &mut JoinSet<()>) { + while let Some(result) = sessions.join_next().await { + log_session_result(Some(result)); + } +} + +fn log_session_result(result: Option>) { + if let Some(Err(error)) = result { + warn!(%error, "NSQ TCP session task failed"); + } } #[allow(clippy::too_many_arguments)] @@ -193,6 +260,7 @@ async fn handle_connection( accepting: Arc, delivering: Arc, publish_admission: Arc, + subscriptions: SubscriptionRegistry, ) -> anyhow::Result<()> { stream.set_nodelay(true)?; let handshake_timeout = Duration::from_millis(config.limits.client_handshake_timeout_ms); @@ -207,20 +275,20 @@ async fn handle_connection( let connection_budget = Arc::new(ConnectionBudget::new( config.limits.connection_publish_inflight_bytes, )); - let first = match tokio::time::timeout( - handshake_timeout, - read_initial_command(&mut stream, config, &publish_admission, &connection_budget), - ) - .await - .context("initial command timeout")? - { - Ok(command) => command, - Err(CommandReadError::Io(error)) => return Err(error.into()), - Err(CommandReadError::Protocol { code, detail }) => { - write_error(&mut stream, code, &detail).await?; - return Ok(()); - } - }; + let first = + match read_initial_command(&mut stream, config, &publish_admission, &connection_budget) + .await + { + Ok(command) => command, + Err(CommandReadError::Io(error)) => return Err(error.into()), + Err(CommandReadError::Protocol { code, detail }) => { + if disconnect_on_retriable_protocol_error(config, code) { + return Ok(()); + } + write_error(&mut stream, code, &detail).await?; + return Ok(()); + } + }; let mut state = SessionState { identified: false, encrypted: false, @@ -239,6 +307,10 @@ async fn handle_connection( rdy: 0, in_flight: HashMap::new(), closing: false, + client_identity: ClientIdentity { + remote_address: peer.to_string(), + ..ClientIdentity::default() + }, }; let (mut io, pending, negotiated): (BoxIo, Option, Option) = @@ -271,6 +343,10 @@ async fn handle_connection( state.output_buffer_size = output.size; state.output_buffer_timeout = output.timeout; state.sample_rate = sample_rate; + state.client_identity.client_id = identify.client_id.clone(); + state.client_identity.hostname = identify.hostname.clone(); + state.client_identity.user_agent = identify.user_agent.clone(); + state.client_identity.sample_rate = sample_rate; let use_tls = identify.feature_negotiation && identify.tls_v1 && tls_acceptor.is_some(); let negotiated = match negotiate_compression(&identify, config) { Ok(compression) => compression, @@ -300,6 +376,7 @@ async fn handle_connection( let mut io: BoxIo = Box::new(tls_stream); write_frame(&mut io, FrameType::Response, OK).await?; state.encrypted = true; + state.client_identity.tls = true; (io, None, negotiated) } else { (Box::new(stream), None, negotiated) @@ -314,6 +391,8 @@ async fn handle_connection( } if let Some(compression) = negotiated { + state.client_identity.snappy = matches!(compression, Compression::Snappy); + state.client_identity.deflate = matches!(compression, Compression::Deflate(_)); io = match compression { Compression::Snappy => compression::snappy(io), Compression::Deflate(level) => compression::deflate(io, level), @@ -334,7 +413,12 @@ async fn handle_connection( delivering, publish_admission, connection_budget, + subscriptions, state, ) .await } + +fn disconnect_on_retriable_protocol_error(config: &Config, code: &str) -> bool { + config.limits.disconnect_on_retriable_publish_error && code == "E_THROTTLED" +} diff --git a/crates/server/src/tcp/codec.rs b/crates/server/src/tcp/codec.rs index cc9d38d..79572c9 100644 --- a/crates/server/src/tcp/codec.rs +++ b/crates/server/src/tcp/codec.rs @@ -35,29 +35,37 @@ impl From for CommandReadError { } } -pub(super) async fn read_initial_command( - stream: &mut TcpStream, +pub(super) async fn read_initial_command( + stream: &mut R, config: &Config, admission: &PublishAdmission, connection_budget: &ConnectionBudget, -) -> Result { - let mut line = Vec::with_capacity(64); - loop { - let byte = stream.read_u8().await?; - line.push(byte); - if byte == b'\n' { - break; - } - if line.len() > 1024 { - return Err(CommandReadError::protocol( - "E_INVALID", - "command line exceeds limit", - )); - } - } +) -> Result +where + R: AsyncRead + Unpin, +{ + let handshake_timeout = Duration::from_millis(config.limits.client_handshake_timeout_ms); + let first = tokio::time::timeout(handshake_timeout, stream.read_u8()) + .await + .map_err(|_| CommandReadError::protocol("E_INVALID", "initial command timed out"))??; + let line = tokio::time::timeout(handshake_timeout, read_command_line(stream, first)) + .await + .map_err(|_| CommandReadError::protocol("E_INVALID", "initial command timed out"))??; let command = parse_command(&line)?; - let (body, reservation) = - read_command_body(stream, &command, config, admission, connection_budget).await?; + let body_timeout = if matches!( + command, + Command::Publish { .. } | Command::MultiPublish { .. } | Command::DeferredPublish { .. } + ) { + Duration::from_millis(config.limits.tcp_command_timeout_ms) + } else { + handshake_timeout + }; + let (body, reservation) = tokio::time::timeout( + body_timeout, + read_command_body(stream, &command, config, admission, connection_budget), + ) + .await + .map_err(|_| CommandReadError::protocol("E_INVALID", "initial command body timed out"))??; Ok(ParsedCommand { command, body, @@ -91,6 +99,21 @@ async fn read_started_command( admission: &PublishAdmission, connection_budget: &ConnectionBudget, ) -> Result +where + R: AsyncRead + Unpin, +{ + let line = read_command_line(reader, first).await?; + let command = parse_command(&line)?; + let (body, reservation) = + read_command_body(reader, &command, config, admission, connection_budget).await?; + Ok(ParsedCommand { + command, + body, + publish_reservation: reservation, + }) +} + +async fn read_command_line(reader: &mut R, first: u8) -> Result, CommandReadError> where R: AsyncRead + Unpin, { @@ -105,14 +128,7 @@ where )); } } - let command = parse_command(&line)?; - let (body, reservation) = - read_command_body(reader, &command, config, admission, connection_budget).await?; - Ok(ParsedCommand { - command, - body, - publish_reservation: reservation, - }) + Ok(line) } async fn read_command_body( @@ -492,4 +508,26 @@ mod tests { other => panic!("unexpected error: {other}"), } } + + #[tokio::test] + async fn first_publish_body_uses_command_timeout_not_handshake_timeout() { + let (mut peer, mut server) = tokio::io::duplex(1024); + let writer = tokio::spawn(async move { + peer.write_all(b"PUB events\n\0\0\0\x04").await.unwrap(); + tokio::time::sleep(Duration::from_millis(30)).await; + peer.write_all(b"body").await.unwrap(); + }); + let mut config = Config::default(); + config.limits.client_handshake_timeout_ms = 10; + config.limits.tcp_command_timeout_ms = 100; + let metrics = Arc::new(Metrics::default()); + let admission = PublishAdmission::new(config.limits.node_publish_inflight_bytes, metrics); + let connection = ConnectionBudget::new(config.limits.connection_publish_inflight_bytes); + let command = read_initial_command(&mut server, &config, &admission, &connection) + .await + .unwrap(); + assert!(matches!(command.command, Command::Publish { .. })); + assert_eq!(command.body.as_deref(), Some(b"body".as_slice())); + writer.await.unwrap(); + } } diff --git a/crates/server/src/tcp/commands.rs b/crates/server/src/tcp/commands.rs index 0207040..7511020 100644 --- a/crates/server/src/tcp/commands.rs +++ b/crates/server/src/tcp/commands.rs @@ -9,6 +9,7 @@ pub(super) async fn process_command( metrics: &Metrics, authenticator: Option<&Authenticator>, ephemeral_consumers: &EphemeralConsumers, + subscriptions: &SubscriptionRegistry, state: &mut SessionState, writer: &mut ClientWriter, ) -> anyhow::Result { @@ -53,6 +54,7 @@ pub(super) async fn process_command( }); state.auth = Some(session); state.auth_secret = Some(secret.to_vec()); + state.client_identity.authed = true; write_frame(writer, FrameType::Response, &serde_json::to_vec(&response)?) .await?; } @@ -91,6 +93,15 @@ pub(super) async fn process_command( { return Ok(false); } + let lease = + match subscriptions.register(&topic, &channel, state.client_identity.clone()) { + Ok(lease) => lease, + Err(_) => { + write_error(writer, "E_SUB_FAILED", "channel deletion is in progress") + .await?; + return Ok(false); + } + }; let create_result = if channel.ends_with("#ephemeral") { ephemeral_consumers.register(broker, &topic, &channel).await } else { @@ -98,7 +109,11 @@ pub(super) async fn process_command( }; match create_result { Ok(()) => { - state.subscription = Some(Subscription { topic, channel }); + state.subscription = Some(Subscription { + topic, + channel, + lease, + }); write_frame(writer, FrameType::Response, OK).await?; } Err(error) => { @@ -120,6 +135,7 @@ pub(super) async fn process_command( vec![body.unwrap_or_default()], Duration::ZERO, publish_reservation, + config.limits.disconnect_on_retriable_publish_error, ) .await? { @@ -147,6 +163,7 @@ pub(super) async fn process_command( messages, Duration::ZERO, publish_reservation, + config.limits.disconnect_on_retriable_publish_error, ) .await? { @@ -170,6 +187,7 @@ pub(super) async fn process_command( vec![body.unwrap_or_default()], Duration::from_millis(delay_ms), publish_reservation, + config.limits.disconnect_on_retriable_publish_error, ) .await? { @@ -185,6 +203,7 @@ pub(super) async fn process_command( return Ok(false); } else { state.rdy = count; + state.update_subscription_flow(); } } Command::Finish(id) => { @@ -201,6 +220,10 @@ pub(super) async fn process_command( match finish_result { Ok(()) => { state.in_flight.remove(&id); + if let Some(subscription) = &state.subscription { + subscription.lease.observe_finish(); + } + state.update_subscription_flow(); metrics.finished_messages.fetch_add(1, Ordering::Relaxed); } Err(error) => write_broker_error(writer, "E_FIN_FAILED", error).await?, @@ -227,6 +250,10 @@ pub(super) async fn process_command( match requeue_result { Ok(()) => { state.in_flight.remove(&id); + if let Some(subscription) = &state.subscription { + subscription.lease.observe_requeue(); + } + state.update_subscription_flow(); metrics.requeued_messages.fetch_add(1, Ordering::Relaxed); } Err(error) => write_broker_error(writer, "E_REQ_FAILED", error).await?, @@ -263,6 +290,7 @@ pub(super) async fn process_command( } state.closing = true; state.rdy = 0; + state.update_subscription_flow(); write_frame(writer, FrameType::Response, CLOSE_WAIT).await?; } Command::Noop => {} @@ -280,6 +308,7 @@ async fn publish_tcp( messages: Vec, delay: Duration, reservation: Option, + disconnect_on_retriable_error: bool, ) -> anyhow::Result { let message_count = messages.len() as u64; let byte_count = messages.iter().map(Bytes::len).sum::() as u64; @@ -296,13 +325,35 @@ async fn publish_tcp( Ok(true) } Err(error) => { - metrics.storage_errors.fetch_add(1, Ordering::Relaxed); - write_broker_error(writer, error_code, error).await?; + if broker_storage_error(&error) { + metrics.storage_errors.fetch_add(1, Ordering::Relaxed); + } + let retryable = precommit_retryable_publish_error(&error); + if disconnect_on_retriable_error && retryable { + return Ok(false); + } + write_broker_error( + writer, + if retryable { "E_PUB_RETRY" } else { error_code }, + error, + ) + .await?; Ok(false) } } } +fn precommit_retryable_publish_error(error: &BrokerError) -> bool { + matches!( + error, + BrokerError::TopicRetiring + | BrokerError::ManagementUnavailable + | BrokerError::TopicLimit + | BrokerError::PublishWorkerLimit + | BrokerError::SequenceExhausted + ) +} + pub(super) async fn publish_messages( broker: &Broker, topic: &str, @@ -328,3 +379,44 @@ pub(super) async fn finish_message( ) -> Result<(), BrokerError> { broker.finish(topic, channel, id).await } + +#[cfg(test)] +mod tests { + use super::{broker_storage_error, precommit_retryable_publish_error}; + use rustqueue_queue::BrokerError; + + #[test] + fn only_guaranteed_precommit_failures_are_retryable() { + assert!(precommit_retryable_publish_error( + &BrokerError::TopicRetiring + )); + assert!(precommit_retryable_publish_error( + &BrokerError::PublishWorkerLimit + )); + assert!(precommit_retryable_publish_error( + &BrokerError::SequenceExhausted + )); + assert!(!precommit_retryable_publish_error( + &BrokerError::StorageUnavailable + )); + assert!(!precommit_retryable_publish_error( + &BrokerError::InvalidRecord("corrupt local state".into()) + )); + assert!(!precommit_retryable_publish_error( + &BrokerError::MessageTooLarge + )); + assert!(!precommit_retryable_publish_error( + &BrokerError::TopicTombstoned + )); + } + + #[test] + fn only_storage_failures_increment_the_storage_error_metric() { + assert!(broker_storage_error(&BrokerError::StorageUnavailable)); + assert!(broker_storage_error(&BrokerError::InvalidRecord( + "corrupt local state".into() + ))); + assert!(!broker_storage_error(&BrokerError::TopicLimit)); + assert!(!broker_storage_error(&BrokerError::TopicRetiring)); + } +} diff --git a/crates/server/src/tcp/session.rs b/crates/server/src/tcp/session.rs index f6fb33f..e71bba5 100644 --- a/crates/server/src/tcp/session.rs +++ b/crates/server/src/tcp/session.rs @@ -1,6 +1,7 @@ use super::*; -type FetchFuture<'a> = Pin> + Send + 'a>>; +type FetchFuture<'a> = + Pin> + Send + 'a>>; struct PendingFetch<'a> { request: FetchRequest, @@ -9,7 +10,7 @@ struct PendingFetch<'a> { async fn poll_pending_fetch( pending: &mut Option>, -) -> Result { +) -> Result { pending .as_mut() .expect("disabled fetch future is never polled") @@ -32,6 +33,7 @@ pub(super) async fn run_session( delivering: Arc, publish_admission: Arc, connection_budget: Arc, + subscriptions: SubscriptionRegistry, mut state: SessionState, ) -> anyhow::Result<()> { let (read_half, write_half) = tokio::io::split(io); @@ -76,11 +78,25 @@ pub(super) async fn run_session( let session_result: anyhow::Result<()> = async { loop { + let expired = expire_client_deadlines(&mut state.in_flight, Instant::now()); + if expired { + if let Some(subscription) = state.subscription.as_ref() { + broker.expire_channel_in_flight( + &subscription.topic, + &subscription.channel, + ) + .await?; + } + state.update_subscription_flow(); + } if state.closing && state.in_flight.is_empty() { break; } if let Some(command) = pending.take() { if !accepting.load(Ordering::Acquire) && publish_command(&command.command) { + if config.limits.disconnect_on_retriable_publish_error { + break; + } write_error_timed( &mut writer, state.heartbeat, @@ -91,6 +107,9 @@ pub(super) async fn run_session( continue; } if !publish_admission.storage_ready() && publish_command(&command.command) { + if config.limits.disconnect_on_retriable_publish_error { + break; + } write_error_timed( &mut writer, state.heartbeat, @@ -100,6 +119,17 @@ pub(super) async fn run_session( .await?; continue; } + if !broker.storage_healthy() && publish_command(&command.command) { + metrics.storage_errors.fetch_add(1, Ordering::Relaxed); + write_error_timed( + &mut writer, + state.heartbeat, + "E_PUB_RETRY", + "local storage was unavailable before publish", + ) + .await?; + continue; + } if state.closing && !shutdown_command_allowed(&command.command) { write_error_timed( &mut writer, @@ -122,6 +152,7 @@ pub(super) async fn run_session( metrics, authenticator.as_deref(), &ephemeral_consumers, + &subscriptions, &mut state, &mut writer, ), @@ -134,8 +165,7 @@ pub(super) async fn run_session( continue; } - let now = Instant::now(); - state.in_flight.retain(|_, deadline| *deadline > now); + state.update_subscription_flow(); let available = state.rdy.saturating_sub(state.in_flight.len() as u64); if pending_fetch.is_none() && !state.closing @@ -159,6 +189,7 @@ pub(super) async fn run_session( pending_fetch = Some(PendingFetch { request, future }); } } + let in_flight_deadline = state.in_flight.values().copied().min(); tokio::select! { command = command_rx.recv() => { @@ -169,18 +200,38 @@ pub(super) async fn run_session( Ok(command) => command, Err(CommandReadError::Io(_)) => break, Err(CommandReadError::Protocol { code, detail }) => { + if disconnect_on_retriable_protocol_error(config, code) { + break; + } write_error_timed(&mut writer, state.heartbeat, code, &detail).await?; break; } }; if !accepting.load(Ordering::Acquire) && publish_command(&command.command) { + if config.limits.disconnect_on_retriable_publish_error { + break; + } write_error_timed(&mut writer, state.heartbeat, "E_DRAINING", "broker is draining").await?; continue; } if !publish_admission.storage_ready() && publish_command(&command.command) { + if config.limits.disconnect_on_retriable_publish_error { + break; + } write_error_timed(&mut writer, state.heartbeat, "E_THROTTLED", "local disk is above its publish watermark").await?; continue; } + if !broker.storage_healthy() && publish_command(&command.command) { + metrics.storage_errors.fetch_add(1, Ordering::Relaxed); + write_error_timed( + &mut writer, + state.heartbeat, + "E_PUB_RETRY", + "local storage was unavailable before publish", + ) + .await?; + continue; + } if state.closing && !shutdown_command_allowed(&command.command) { write_error_timed(&mut writer, state.heartbeat, "E_CLOSING", "node is shutting down").await?; continue; @@ -197,6 +248,7 @@ pub(super) async fn run_session( metrics, authenticator.as_deref(), &ephemeral_consumers, + &subscriptions, &mut state, &mut writer, ), @@ -288,13 +340,25 @@ pub(super) async fn run_session( state .in_flight .insert(delivery.id, Instant::now() + state.message_timeout); + if let Some(subscription) = &state.subscription { + subscription.lease.observe_delivery(); + } delivery_guard.accept(delivery.id); metrics.delivered_messages.fetch_add(1, Ordering::Relaxed); } + state.update_subscription_flow(); } Err(error) => { - metrics.storage_errors.fetch_add(1, Ordering::Relaxed); - write_error_timed(&mut writer, state.heartbeat, "E_DELIVERY_FAILED", &error).await?; + if broker_storage_error(&error) { + metrics.storage_errors.fetch_add(1, Ordering::Relaxed); + } + write_error_timed( + &mut writer, + state.heartbeat, + "E_DELIVERY_FAILED", + &error.to_string(), + ) + .await?; break; } } @@ -314,6 +378,7 @@ pub(super) async fn run_session( _ = output_buffer_tick.tick(), if state.output_buffer_timeout.is_some() && writer.has_pending() => { flush_timed(&mut writer, state.heartbeat).await?; } + _ = wait_for_in_flight_deadline(in_flight_deadline) => {} } } Ok(()) @@ -344,7 +409,10 @@ pub(super) async fn run_session( session_result } -async fn fetch_deliveries(broker: &Broker, request: FetchRequest) -> Result { +async fn fetch_deliveries( + broker: &Broker, + request: FetchRequest, +) -> Result { let batch = broker .fetch_batch_retained( &request.topic, @@ -354,8 +422,7 @@ async fn fetch_deliveries(broker: &Broker, request: FetchRequest) -> Result, id: u64) -> bool { in_flight.contains_key(&id) } +fn expire_client_deadlines(in_flight: &mut HashMap, now: Instant) -> bool { + let before = in_flight.len(); + in_flight.retain(|_, deadline| *deadline > now); + in_flight.len() != before +} + +async fn wait_for_in_flight_deadline(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline.into()).await, + None => std::future::pending().await, + } +} + #[cfg(test)] mod tests { use super::*; @@ -446,4 +526,34 @@ mod tests { assert!(!delivery_is_outstanding(&in_flight, 7)); assert!(!delivery_is_outstanding(&in_flight, 8)); } + + #[test] + fn client_deadlines_expire_without_waiting_for_another_fetch() { + let now = Instant::now(); + let mut in_flight = HashMap::from([ + (7, now - Duration::from_millis(1)), + (8, now + Duration::from_secs(1)), + ]); + assert!(expire_client_deadlines(&mut in_flight, now)); + assert_eq!(in_flight.keys().copied().collect::>(), vec![8]); + assert!(!expire_client_deadlines(&mut in_flight, now)); + } + + #[test] + fn kodo_publish_admission_pressure_forces_producer_failover() { + let mut config = Config::default(); + assert!(!disconnect_on_retriable_protocol_error( + &config, + "E_THROTTLED" + )); + config.limits.disconnect_on_retriable_publish_error = true; + assert!(disconnect_on_retriable_protocol_error( + &config, + "E_THROTTLED" + )); + assert!(!disconnect_on_retriable_protocol_error( + &config, + "E_BAD_MESSAGE" + )); + } } diff --git a/crates/storage/build.rs b/crates/storage/build.rs index f6e152c..85f3b3f 100644 --- a/crates/storage/build.rs +++ b/crates/storage/build.rs @@ -1,7 +1,12 @@ fn main() { println!("cargo:rerun-if-env-changed=RUSTQUEUE_MAX_STORAGE_FEATURE_LEVEL"); println!("cargo:rustc-check-cfg=cfg(rustqueue_storage_feature_level_2)"); - if std::env::var("RUSTQUEUE_MAX_STORAGE_FEATURE_LEVEL").as_deref() == Ok("2") { - println!("cargo:rustc-cfg=rustqueue_storage_feature_level_2"); + match std::env::var("RUSTQUEUE_MAX_STORAGE_FEATURE_LEVEL") + .unwrap_or_default() + .trim() + { + "" | "2" => println!("cargo:rustc-cfg=rustqueue_storage_feature_level_2"), + "1" => {} + value => panic!("RUSTQUEUE_MAX_STORAGE_FEATURE_LEVEL must be 1 or 2, got {value}"), } } diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs index bf9520b..dc47e3b 100644 --- a/crates/storage/src/lib.rs +++ b/crates/storage/src/lib.rs @@ -18,7 +18,9 @@ pub use disk::{disk_space, DiskSpace}; pub use failpoint::crash_failpoint; pub use format::{ensure_data_format, read_data_format, DataFormat, DATA_FORMAT_VERSION}; pub use payload::PayloadRef; -pub use record::{Record, RecordHeader, RecordKind, HEADER_LEN, MAX_RECORD_BYTES}; +pub use record::{ + Record, RecordHeader, RecordKind, HEADER_LEN, LEGACY_MAX_RECORD_BYTES, MAX_RECORD_BYTES, +}; pub use segment::{ RecordLocation, RecoveryMetadataRef, RecoveryReport, ScrubKind, ScrubTarget, SegmentLog, StorageError, diff --git a/crates/storage/src/record.rs b/crates/storage/src/record.rs index b535081..10b8f1f 100644 --- a/crates/storage/src/record.rs +++ b/crates/storage/src/record.rs @@ -1,9 +1,10 @@ use std::io; pub const HEADER_LEN: usize = 48; -// Public publish payload is capped at 64 MiB. The durable envelope also needs -// room for per-message IDs, timestamps, lengths and checksums. -pub const MAX_RECORD_BYTES: usize = 72 * 1024 * 1024; +pub const LEGACY_MAX_RECORD_BYTES: usize = 72 * 1024 * 1024; +// A 128 MiB command body can grow by 20 bytes per message in the durable +// envelope. Keep enough bounded headroom for the maximum supported batch. +pub const MAX_RECORD_BYTES: usize = 160 * 1024 * 1024; const MAGIC: &[u8; 4] = b"RQV7"; const VERSION: u8 = 7; @@ -32,8 +33,9 @@ impl TryFrom for RecordKind { } impl RecordKind { - pub const fn required_writer_feature_level(self) -> u32 { + pub const fn required_writer_feature_level(self, payload_len: usize) -> u32 { match self { + Self::PublishBatch if payload_len > LEGACY_MAX_RECORD_BYTES => 2, Self::PublishBatch | Self::EvictionGap | Self::Noop => 1, } } @@ -85,12 +87,7 @@ impl Record { } pub fn decode(header: &[u8; HEADER_LEN], payload: Vec) -> io::Result { - if &header[0..4] != MAGIC || header[4] != VERSION { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "invalid record magic or version", - )); - } + let decoded = RecordHeader::decode(header)?; let expected_len = u32::from_be_bytes(header[40..44].try_into().unwrap()) as usize; if expected_len != payload.len() || expected_len > MAX_RECORD_BYTES { return Err(io::Error::new( @@ -107,12 +104,12 @@ impl Record { )); } Ok(Self { - kind: RecordKind::try_from(header[5])?, - flags: u16::from_be_bytes(header[6..8].try_into().unwrap()), - index: u64::from_be_bytes(header[8..16].try_into().unwrap()), - timestamp_ns: i64::from_be_bytes(header[16..24].try_into().unwrap()), - message_id: u64::from_be_bytes(header[24..32].try_into().unwrap()), - available_at_ms: i64::from_be_bytes(header[32..40].try_into().unwrap()), + kind: decoded.kind, + flags: decoded.flags, + index: decoded.index, + timestamp_ns: decoded.timestamp_ns, + message_id: decoded.message_id, + available_at_ms: decoded.available_at_ms, payload, }) } @@ -136,6 +133,23 @@ impl Record { } impl RecordHeader { + pub fn decode(header: &[u8; HEADER_LEN]) -> io::Result { + if &header[0..4] != MAGIC || header[4] != VERSION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid record magic or version", + )); + } + Ok(Self { + kind: RecordKind::try_from(header[5])?, + flags: u16::from_be_bytes(header[6..8].try_into().unwrap()), + index: u64::from_be_bytes(header[8..16].try_into().unwrap()), + timestamp_ns: i64::from_be_bytes(header[16..24].try_into().unwrap()), + message_id: u64::from_be_bytes(header[24..32].try_into().unwrap()), + available_at_ms: i64::from_be_bytes(header[32..40].try_into().unwrap()), + }) + } + pub fn encode(&self, payload: &[&[u8]]) -> io::Result<[u8; HEADER_LEN]> { let payload_len = payload .iter() @@ -168,3 +182,25 @@ impl RecordHeader { Ok(header) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn oversized_v7_publish_records_require_feature_level_two() { + assert_eq!( + RecordKind::PublishBatch.required_writer_feature_level(LEGACY_MAX_RECORD_BYTES), + 1 + ); + assert_eq!( + RecordKind::PublishBatch + .required_writer_feature_level(LEGACY_MAX_RECORD_BYTES.saturating_add(1)), + 2 + ); + assert_eq!( + RecordKind::EvictionGap.required_writer_feature_level(MAX_RECORD_BYTES), + 1 + ); + } +} diff --git a/crates/storage/src/segment.rs b/crates/storage/src/segment.rs index 6448451..28468f9 100644 --- a/crates/storage/src/segment.rs +++ b/crates/storage/src/segment.rs @@ -13,7 +13,9 @@ mod scrub; use crate::{crash_failpoint, Record, RecordHeader, BASE_STORAGE_FEATURE_LEVEL, HEADER_LEN}; use append::write_parts; -use files::{read_record, scan_segment, segment_base_index, segment_path, segment_paths}; +use files::{ + read_record, scan_segment, segment_base_index, segment_path, segment_paths, visit_record, +}; use std::collections::{BTreeMap, BTreeSet}; use std::fs::{self, File, OpenOptions}; use std::io; @@ -279,13 +281,23 @@ impl SegmentLog { durable: bool, ) -> Result { self.ensure_available()?; - let required = record.kind.required_writer_feature_level(); + let payload_len = payload + .iter() + .try_fold(0usize, |total, part| total.checked_add(part.len())) + .ok_or_else(|| { + StorageError::Io(io::Error::new( + io::ErrorKind::InvalidInput, + "record payload length overflow", + )) + })?; + let required = record.kind.required_writer_feature_level(payload_len); if required > self.active_writer_feature_level { return Err(StorageError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!( - "record kind {:?} requires writer feature level {required}, active level is {}", - record.kind, self.active_writer_feature_level + "record kind {:?} with {payload_len} payload bytes requires writer feature level {required}, active level is {}", + record.kind, + self.active_writer_feature_level ), ))); } @@ -373,19 +385,18 @@ impl SegmentLog { } pub fn read_location(&self, location: &RecordLocation) -> Result { - let result = read_record(location).and_then(|record| { - if record.index != location.index { - return Err(StorageError::Corrupt { - path: location.segment.as_ref().clone(), - offset: location.offset, - reason: "recovery index points at a different record".into(), - }); - } - Ok(record) - }); + let result = read_record(location); self.observe_read(result) } + pub fn read_location_with( + &self, + location: &RecordLocation, + visitor: impl FnOnce(RecordHeader, usize, &mut dyn io::Read) -> io::Result, + ) -> Result { + self.observe_read(visit_record(location, visitor)) + } + pub fn read_all(&self) -> Result, StorageError> { self.all_locations()? .iter() diff --git a/crates/storage/src/segment/files.rs b/crates/storage/src/segment/files.rs index 6b5a82f..43322de 100644 --- a/crates/storage/src/segment/files.rs +++ b/crates/storage/src/segment/files.rs @@ -1,5 +1,5 @@ use super::{RecordLocation, StorageError}; -use crate::{Record, HEADER_LEN}; +use crate::{Record, RecordHeader, HEADER_LEN}; use std::fs::{self, File, OpenOptions}; use std::io::{self, Read, Seek, SeekFrom}; use std::path::{Path, PathBuf}; @@ -9,14 +9,77 @@ const SEGMENT_PREFIX: &str = "segment-"; const SEGMENT_SUFFIX: &str = ".rqlog"; pub(super) fn read_record(location: &RecordLocation) -> Result { + visit_record(location, |header, payload_len, reader| { + let mut payload = Vec::with_capacity(payload_len); + reader.read_to_end(&mut payload)?; + Ok(Record { + kind: header.kind, + flags: header.flags, + index: header.index, + timestamp_ns: header.timestamp_ns, + message_id: header.message_id, + available_at_ms: header.available_at_ms, + payload, + }) + }) +} + +pub(super) fn visit_record( + location: &RecordLocation, + visitor: impl FnOnce(RecordHeader, usize, &mut dyn Read) -> io::Result, +) -> Result { let mut file = File::open(location.segment.as_ref())?; file.seek(SeekFrom::Start(location.offset))?; let mut header = [0u8; HEADER_LEN]; file.read_exact(&mut header)?; let payload_len = Record::payload_len(&header)?; - let mut payload = vec![0; payload_len]; - file.read_exact(&mut payload)?; - Record::decode(&header, payload).map_err(StorageError::Io) + if location.encoded_len != HEADER_LEN as u64 + payload_len as u64 { + return Err(corrupt( + location.segment.as_ref(), + location.offset, + "record location length does not match its header", + )); + } + let decoded = RecordHeader::decode(&header)?; + if decoded.index != location.index { + return Err(corrupt( + location.segment.as_ref(), + location.offset, + "recovery index points at a different record", + )); + } + let expected_crc = u32::from_be_bytes(header[44..48].try_into().unwrap()); + let mut reader = ChecksummedReader { + inner: file.take(payload_len as u64), + crc32c: crc32c::crc32c(&header[..44]), + bytes_read: 0, + }; + let value = visitor(decoded, payload_len, &mut reader).map_err(|error| { + corrupt( + location.segment.as_ref(), + location.offset, + error.to_string(), + ) + })?; + let mut buffer = [0u8; 64 * 1024]; + while reader.bytes_read < payload_len { + let read = reader.read(&mut buffer).map_err(StorageError::Io)?; + if read == 0 { + return Err(corrupt( + location.segment.as_ref(), + location.offset, + "partial record payload", + )); + } + } + if reader.crc32c != expected_crc { + return Err(corrupt( + location.segment.as_ref(), + location.offset, + "record checksum mismatch", + )); + } + Ok(value) } pub(super) fn scan_segment( @@ -58,27 +121,51 @@ pub(super) fn scan_segment( } return Err(corrupt(path, offset, "partial record payload")); } - let mut payload = vec![0; payload_len]; - file.read_exact(&mut payload)?; - let record_checksum = - crc32c::crc32c_append(crc32c::crc32c_append(checksum, &header), &payload); - match Record::decode(&header, payload) { - Ok(record) => { - checksum = record_checksum; - locations.push(RecordLocation { - index: record.index, - segment: Arc::clone(&segment), - offset, - encoded_len: record_len, - }); - } - Err(error) => return Err(corrupt(path, offset, error.to_string())), + let decoded = RecordHeader::decode(&header) + .map_err(|error| corrupt(path, offset, error.to_string()))?; + checksum = crc32c::crc32c_append(checksum, &header); + let mut record_crc = crc32c::crc32c(&header[..44]); + let mut remaining = payload_len; + let mut buffer = [0u8; 64 * 1024]; + while remaining > 0 { + let wanted = remaining.min(buffer.len()); + file.read_exact(&mut buffer[..wanted])?; + checksum = crc32c::crc32c_append(checksum, &buffer[..wanted]); + record_crc = crc32c::crc32c_append(record_crc, &buffer[..wanted]); + remaining -= wanted; + } + let expected_crc = u32::from_be_bytes(header[44..48].try_into().unwrap()); + if record_crc != expected_crc { + return Err(corrupt(path, offset, "record checksum mismatch")); } + locations.push(RecordLocation { + index: decoded.index, + segment: Arc::clone(&segment), + offset, + encoded_len: record_len, + }); offset += record_len; } Ok((locations, 0, offset, checksum)) } +struct ChecksummedReader { + inner: R, + crc32c: u32, + bytes_read: usize, +} + +impl Read for ChecksummedReader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + let read = self.inner.read(buffer)?; + if read > 0 { + self.crc32c = crc32c::crc32c_append(self.crc32c, &buffer[..read]); + self.bytes_read = self.bytes_read.saturating_add(read); + } + Ok(read) + } +} + fn corrupt(path: &Path, offset: u64, reason: impl Into) -> StorageError { StorageError::Corrupt { path: path.to_path_buf(), diff --git a/deploy/helm/rustqueue/Chart.yaml b/deploy/helm/rustqueue/Chart.yaml index 963c42d..38eaa7c 100644 --- a/deploy/helm/rustqueue/Chart.yaml +++ b/deploy/helm/rustqueue/Chart.yaml @@ -2,8 +2,8 @@ apiVersion: v2 name: rustqueue description: Kubernetes-native share-nothing NSQ-compatible message queue type: application -version: 0.7.2 -appVersion: "0.7.2" +version: 0.8.0 +appVersion: "0.8.0" kubeVersion: ">=1.28.0-0" keywords: - messaging diff --git a/deploy/helm/rustqueue/crds/rustqueue.io_rustqueues.yaml b/deploy/helm/rustqueue/crds/rustqueue.io_rustqueues.yaml index f7dd586..c34dfde 100644 --- a/deploy/helm/rustqueue/crds/rustqueue.io_rustqueues.yaml +++ b/deploy/helm/rustqueue/crds/rustqueue.io_rustqueues.yaml @@ -46,10 +46,10 @@ spec: minFreeBytes: {type: integer, format: int64, minimum: 0, default: 10737418240} diskHighWatermarkPercent: {type: integer, format: int32, minimum: 1, maximum: 100, default: 85} diskLowWatermarkPercent: {type: integer, format: int32, minimum: 0, maximum: 99, default: 75} - protectiveEvictionEnabled: {type: boolean, default: true} + protectiveEvictionEnabled: {type: boolean, default: false} diskPressureGraceSeconds: {type: integer, format: int64, minimum: 0, default: 60} bootstrapRetentionSeconds: {type: integer, format: int64, minimum: 1, default: 90} - maxMessageBytes: {type: integer, format: int64, minimum: 1, maximum: 33554432, default: 20971520} + maxMessageBytes: {type: integer, format: int64, minimum: 1, maximum: 104857600, default: 20971520} maxTopics: {type: integer, format: int64, minimum: 1, default: 10000} maxPublishWorkers: {type: integer, format: int64, minimum: 1, default: 1024} publishWorkerIdleSeconds: {type: integer, format: int64, minimum: 1, default: 60} @@ -64,6 +64,32 @@ spec: default: {} proxyTcpMaxConnectionAgeSeconds: {type: integer, format: int64, minimum: 0, default: 300} discoveryReplicas: {type: integer, format: int32, minimum: 2, default: 2} + kodoCompatibility: + type: object + default: + enabled: false + decommissionConfirmed: false + producerRestartNonce: "" + cleanupEnabled: false + cutoverGraceSeconds: 630 + allowedPodSelector: + app.kubernetes.io/name: kodo + allowedNamespaceSelector: {} + properties: + enabled: {type: boolean, default: false} + decommissionConfirmed: {type: boolean, default: false} + producerRestartNonce: {type: string, default: ""} + cleanupEnabled: {type: boolean, enum: [false], default: false} + cutoverGraceSeconds: {type: integer, format: int64, minimum: 630, maximum: 86400, default: 630} + allowedPodSelector: + type: object + additionalProperties: {type: string} + default: + app.kubernetes.io/name: kodo + allowedNamespaceSelector: + type: object + additionalProperties: {type: string} + default: {} maintenance: type: object nullable: true @@ -171,3 +197,6 @@ spec: default: [] items: {type: string} desiredStorageSize: {type: string, default: ""} + kodoProducerRestartBaselineNonce: + type: string + nullable: true diff --git a/deploy/helm/rustqueue/templates/monitoring.yaml b/deploy/helm/rustqueue/templates/monitoring.yaml index 6274e72..d4719fa 100644 --- a/deploy/helm/rustqueue/templates/monitoring.yaml +++ b/deploy/helm/rustqueue/templates/monitoring.yaml @@ -24,6 +24,26 @@ spec: --- apiVersion: monitoring.coreos.com/v1 kind: ServiceMonitor +metadata: + name: {{ include "rustqueue.fullname" . }}-kodo-gateway + namespace: {{ .Release.Namespace }} + labels: + {{- include "rustqueue.labels" . | nindent 4 }} +spec: + namespaceSelector: + matchNames: [{{ .Release.Namespace | quote }}] + selector: + matchLabels: + app.kubernetes.io/instance: {{ .Values.queue.name | quote }} + app.kubernetes.io/component: kodo-gateway + rustqueue.io/metrics: "true" + endpoints: + - port: metrics + path: /metrics + interval: {{ .Values.monitoring.serviceMonitor.interval }} +--- +apiVersion: monitoring.coreos.com/v1 +kind: ServiceMonitor metadata: name: {{ include "rustqueue.fullname" . }}-operator namespace: {{ .Release.Namespace }} @@ -55,35 +75,84 @@ spec: - name: rustqueue.production rules: - alert: RustQueueOperatorHasNoLeader - expr: (sum(rustqueue_operator_leader) or vector(0)) < 1 + expr: (sum(rustqueue_operator_leader{namespace="{{ .Release.Namespace }}",service="{{ include "rustqueue.fullname" . }}-operator"}) or vector(0)) < 1 for: 1m labels: {severity: critical} annotations: summary: RustQueue Operator has no active leader + - alert: RustQueueBrokerMetricsMissing + expr: absent(rustqueue_disk_pressure{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-brokers"}) + for: 1m + labels: {severity: critical} + annotations: + summary: RustQueue has no Broker metrics - alert: RustQueueBrokerDiskPressure - expr: rustqueue_disk_pressure == 1 + expr: rustqueue_disk_pressure{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-brokers"} == 1 for: 1m labels: {severity: critical} annotations: summary: RustQueue Broker is rejecting publishes due to disk pressure - alert: RustQueueStorageErrors - expr: increase(rustqueue_storage_errors_total[5m]) > 0 + expr: increase(rustqueue_storage_errors_total{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-brokers"}[5m]) > 0 labels: {severity: critical} annotations: summary: RustQueue Broker reported a storage error - alert: RustQueueProtectiveEviction - expr: increase(rustqueue_protective_evicted_messages_total[5m]) > 0 + expr: increase(rustqueue_protective_evicted_messages_total{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-brokers"}[5m]) > 0 labels: {severity: critical} annotations: summary: RustQueue intentionally evicted messages to protect its disk - alert: RustQueuePublishThrottling - expr: rate(rustqueue_publish_throttled_requests_total[5m]) > 0 + expr: rate(rustqueue_publish_throttled_requests_total{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-brokers"}[5m]) > 0 for: 5m labels: {severity: warning} annotations: summary: RustQueue publish admission has been throttling continuously + - alert: RustQueueAmbiguousPublish + expr: increase(rustqueue_proxy_producer_ambiguous_failures_total{namespace="{{ .Release.Namespace }}",service=~"{{ .Values.queue.name }}-(proxy|kodo-publish)"}[5m]) > 0 + labels: {severity: critical} + annotations: + summary: RustQueue Proxy or Kodo Gateway observed a publish with an unknown commit outcome + - alert: RustQueueDiscoverySourceUnavailable + expr: (min(rustqueue_discovery_source_ready{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-discovery"}) or vector(0)) == 0 + for: 1m + labels: {severity: critical} + annotations: + summary: RustQueue Discovery has no fresh Broker registry source + - alert: RustQueueDiscoveryLookupIncomplete + expr: (min(rustqueue_discovery_source_ready{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-discovery"}) == 1) and ((min(rustqueue_discovery_lookup_ready{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-discovery"}) or vector(0)) == 0) + for: 1m + labels: {severity: warning} + annotations: + summary: RustQueue Discovery cannot serve a complete compatible lookup inventory + - alert: RustQueueProxyHasNoPublishBackend + expr: (min(rustqueue_proxy_publish_backends{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-proxy"}) or vector(0)) < 1 + for: 1m + labels: {severity: critical} + annotations: + summary: RustQueue Proxy has no publish-ready Broker + {{- if .Values.queue.kodoCompatibility.enabled }} + - alert: RustQueueKodoGatewayMetricsMissing + expr: (count(rustqueue_proxy_kodo_gateway{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-kodo-publish"} == 1) or vector(0)) < 3 + for: 1m + labels: {severity: critical} + annotations: + summary: RustQueue cannot scrape all three Kodo Gateways + - alert: RustQueueKodoGatewayBelowStartupQuorum + expr: (min(rustqueue_proxy_publish_backends{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-kodo-publish"}) or vector(0)) < 2 + for: 1m + labels: {severity: warning} + annotations: + summary: RustQueue Kodo Gateway has fewer than two publish-ready Brokers + - alert: RustQueueKodoStatsInventoryIncomplete + expr: (min(rustqueue_proxy_broker_backends{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-kodo-publish"}) or vector(-1)) != 3 + for: 1m + labels: {severity: critical} + annotations: + summary: RustQueue Kodo Gateway cannot serve all three Stats shards + {{- end }} - alert: RustQueueFsyncP99High - expr: histogram_quantile(0.99, sum by (le) (rate(rustqueue_storage_fsync_duration_seconds_bucket[5m]))) > 1 + expr: histogram_quantile(0.99, sum by (le) (rate(rustqueue_storage_fsync_duration_seconds_bucket{namespace="{{ .Release.Namespace }}",service="{{ .Values.queue.name }}-brokers"}[5m]))) > 1 for: 10m labels: {severity: warning} annotations: diff --git a/deploy/helm/rustqueue/templates/operator-rbac.yaml b/deploy/helm/rustqueue/templates/operator-rbac.yaml index 002b1cb..f17c0e2 100644 --- a/deploy/helm/rustqueue/templates/operator-rbac.yaml +++ b/deploy/helm/rustqueue/templates/operator-rbac.yaml @@ -1,7 +1,8 @@ apiVersion: rbac.authorization.k8s.io/v1 -kind: ClusterRole +kind: Role metadata: name: {{ include "rustqueue.fullname" . }}-operator + namespace: {{ .Release.Namespace }} labels: {{- include "rustqueue.labels" . | nindent 4 }} rules: @@ -9,35 +10,35 @@ rules: resources: ["rustqueues", "rustqueues/status", "rustqueues/finalizers"] verbs: ["get", "list", "watch", "patch", "update"] - apiGroups: [""] - resources: ["nodes"] - verbs: ["get", "list", "watch"] - - apiGroups: [""] - resources: ["services", "serviceaccounts", "configmaps", "secrets"] + resources: ["serviceaccounts", "configmaps", "secrets"] verbs: ["get", "list", "watch", "create", "patch", "update"] + - apiGroups: [""] + resources: ["services"] + verbs: ["get", "list", "watch", "create", "patch", "update", "delete"] - apiGroups: [""] resources: ["pods"] verbs: ["get", "list", "watch", "create", "delete"] - apiGroups: [""] resources: ["persistentvolumeclaims"] verbs: ["get", "list", "watch", "patch", "update"] - - apiGroups: ["storage.k8s.io"] - resources: ["storageclasses"] - verbs: ["get", "list", "watch"] - apiGroups: ["policy"] resources: ["poddisruptionbudgets"] - verbs: ["get", "list", "watch", "create", "patch", "update"] + verbs: ["get", "list", "watch", "create", "patch", "update", "delete"] - apiGroups: ["coordination.k8s.io"] resources: ["leases"] verbs: ["get", "list", "watch", "create", "patch", "update"] - apiGroups: ["apps"] - resources: ["statefulsets", "deployments", "daemonsets"] + resources: ["deployments", "daemonsets"] verbs: ["get", "list", "watch", "create", "patch", "update"] + - apiGroups: ["apps"] + resources: ["statefulsets"] + verbs: ["get", "list", "watch", "create", "patch", "update", "delete"] - apiGroups: ["discovery.k8s.io"] resources: ["endpointslices"] verbs: ["get", "list", "watch"] - apiGroups: ["networking.k8s.io"] resources: ["networkpolicies"] - verbs: ["get", "list", "watch", "create", "patch", "update"] + verbs: ["get", "list", "watch", "create", "patch", "update", "delete"] - apiGroups: ["rbac.authorization.k8s.io"] resources: ["roles", "rolebindings"] verbs: ["get", "list", "watch", "create", "patch", "update"] @@ -45,14 +46,44 @@ rules: apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRoleBinding metadata: - name: {{ include "rustqueue.fullname" . }}-operator + name: {{ include "rustqueue.fullname" . }}-operator-cluster-read labels: {{- include "rustqueue.labels" . | nindent 4 }} roleRef: apiGroup: rbac.authorization.k8s.io kind: ClusterRole + name: {{ include "rustqueue.fullname" . }}-operator-cluster-read +subjects: + - kind: ServiceAccount + name: {{ include "rustqueue.operatorServiceAccount" . }} + namespace: {{ .Release.Namespace }} +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: {{ include "rustqueue.fullname" . }}-operator + namespace: {{ .Release.Namespace }} + labels: + {{- include "rustqueue.labels" . | nindent 4 }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role name: {{ include "rustqueue.fullname" . }}-operator subjects: - kind: ServiceAccount name: {{ include "rustqueue.operatorServiceAccount" . }} namespace: {{ .Release.Namespace }} +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: {{ include "rustqueue.fullname" . }}-operator-cluster-read + labels: + {{- include "rustqueue.labels" . | nindent 4 }} +rules: + - apiGroups: [""] + resources: ["nodes"] + verbs: ["get", "list", "watch"] + - apiGroups: ["storage.k8s.io"] + resources: ["storageclasses"] + verbs: ["get", "list", "watch"] diff --git a/deploy/helm/rustqueue/templates/queue.yaml b/deploy/helm/rustqueue/templates/queue.yaml index ff68d2b..59558e2 100644 --- a/deploy/helm/rustqueue/templates/queue.yaml +++ b/deploy/helm/rustqueue/templates/queue.yaml @@ -1,4 +1,10 @@ {{- if .Values.queue.create }} +{{- if .Values.queue.kodoCompatibility.cleanupEnabled }} +{{- fail "queue.kodoCompatibility.cleanupEnabled is disabled until cluster-wide atomic deletion is available" }} +{{- end }} +{{- if and .Values.queue.kodoCompatibility.enabled (ne .Values.queue.imagePullPolicy "Never") (not (regexMatch "@sha256:[a-f0-9]{64}$" .Values.queue.image)) }} +{{- fail "queue.kodoCompatibility.enabled requires queue.image pinned by @sha256, or queue.imagePullPolicy=Never for a preloaded local image" }} +{{- end }} apiVersion: rustqueue.io/v1alpha1 kind: RustQueue metadata: @@ -9,23 +15,23 @@ metadata: spec: image: {{ .Values.queue.image | quote }} imagePullPolicy: {{ .Values.queue.imagePullPolicy | quote }} - minBrokers: {{ .Values.queue.minBrokers }} - maxBrokers: {{ .Values.queue.maxBrokers }} + minBrokers: {{ ternary 3 .Values.queue.minBrokers .Values.queue.kodoCompatibility.enabled }} + maxBrokers: {{ ternary 3 .Values.queue.maxBrokers .Values.queue.kodoCompatibility.enabled }} eligibleNodeSelector: {{ .Values.queue.eligibleNodeSelector | quote }} storageClassName: {{ .Values.queue.storageClassName | quote }} storageSize: {{ .Values.queue.storageSize | quote }} - storageFeatureLevel: {{ .Values.queue.storageFeatureLevel }} + storageFeatureLevel: {{ ternary 2 .Values.queue.storageFeatureLevel .Values.queue.kodoCompatibility.enabled }} messageIndexCacheBytes: {{ .Values.queue.messageIndexCacheBytes | int64 }} maintenanceStartupDelaySeconds: {{ .Values.queue.maintenanceStartupDelaySeconds | int64 }} - nodeDeliveryInflightBytes: {{ .Values.queue.nodeDeliveryInflightBytes | int64 }} - connectionDeliveryInflightBytes: {{ .Values.queue.connectionDeliveryInflightBytes | int64 }} + nodeDeliveryInflightBytes: {{ ternary 536870912 .Values.queue.nodeDeliveryInflightBytes .Values.queue.kodoCompatibility.enabled | int64 }} + connectionDeliveryInflightBytes: {{ ternary 134217728 .Values.queue.connectionDeliveryInflightBytes .Values.queue.kodoCompatibility.enabled | int64 }} minFreeBytes: {{ .Values.queue.minFreeBytes | int64 }} diskHighWatermarkPercent: {{ .Values.queue.diskHighWatermarkPercent }} diskLowWatermarkPercent: {{ .Values.queue.diskLowWatermarkPercent }} protectiveEvictionEnabled: {{ .Values.queue.protectiveEvictionEnabled }} diskPressureGraceSeconds: {{ .Values.queue.diskPressureGraceSeconds }} - bootstrapRetentionSeconds: {{ .Values.queue.bootstrapRetentionSeconds }} - maxMessageBytes: {{ .Values.queue.maxMessageBytes | int64 }} + bootstrapRetentionSeconds: {{ ternary 180 .Values.queue.bootstrapRetentionSeconds .Values.queue.kodoCompatibility.enabled }} + maxMessageBytes: {{ ternary 104857600 .Values.queue.maxMessageBytes .Values.queue.kodoCompatibility.enabled | int64 }} maxTopics: {{ .Values.queue.maxTopics | int64 }} maxPublishWorkers: {{ .Values.queue.maxPublishWorkers | int64 }} publishWorkerIdleSeconds: {{ .Values.queue.publishWorkerIdleSeconds | int64 }} @@ -36,6 +42,16 @@ spec: proxyNodeSelector: {{- toYaml .Values.queue.proxyNodeSelector | nindent 4 }} proxyTcpMaxConnectionAgeSeconds: {{ .Values.queue.proxyTcpMaxConnectionAgeSeconds | int64 }} + kodoCompatibility: + enabled: {{ .Values.queue.kodoCompatibility.enabled }} + decommissionConfirmed: {{ .Values.queue.kodoCompatibility.decommissionConfirmed }} + producerRestartNonce: {{ .Values.queue.kodoCompatibility.producerRestartNonce | quote }} + cleanupEnabled: {{ and .Values.queue.kodoCompatibility.enabled .Values.queue.kodoCompatibility.cleanupEnabled }} + cutoverGraceSeconds: {{ .Values.queue.kodoCompatibility.cutoverGraceSeconds | int64 }} + allowedPodSelector: + {{- toYaml .Values.queue.kodoCompatibility.allowedPodSelector | nindent 6 }} + allowedNamespaceSelector: + {{- toYaml .Values.queue.kodoCompatibility.allowedNamespaceSelector | nindent 6 }} rollout: paused: {{ .Values.queue.rollout.paused }} requireCanaryApproval: {{ .Values.queue.rollout.requireCanaryApproval }} @@ -55,8 +71,8 @@ spec: tolerations: {{- toYaml .Values.queue.brokerScheduling.tolerations | nindent 6 }} brokerResources: - cpuRequest: {{ .Values.queue.brokerResources.cpuRequest | quote }} - memoryRequest: {{ .Values.queue.brokerResources.memoryRequest | quote }} + cpuRequest: {{ ternary .Values.queue.kodoCompatibility.brokerCpuRequest .Values.queue.brokerResources.cpuRequest .Values.queue.kodoCompatibility.enabled | quote }} + memoryRequest: {{ ternary .Values.queue.kodoCompatibility.brokerMemoryRequest .Values.queue.brokerResources.memoryRequest .Values.queue.kodoCompatibility.enabled | quote }} {{- if .Values.queue.brokerResources.cpuLimit }} cpuLimit: {{ .Values.queue.brokerResources.cpuLimit | quote }} {{- end }} diff --git a/deploy/helm/rustqueue/values.yaml b/deploy/helm/rustqueue/values.yaml index eef45be..faf1650 100644 --- a/deploy/helm/rustqueue/values.yaml +++ b/deploy/helm/rustqueue/values.yaml @@ -5,7 +5,7 @@ operator: replicas: 2 image: repository: rustqueue-operator - tag: "0.7.2" + tag: "0.8.0" pullPolicy: IfNotPresent log: rustqueue_operator=info resources: @@ -20,7 +20,7 @@ operator: queue: create: true name: rustqueue - image: rustqueue:0.7.2 + image: rustqueue:0.8.0 imagePullPolicy: IfNotPresent minBrokers: 1 maxBrokers: 500 @@ -35,8 +35,12 @@ queue: minFreeBytes: 10737418240 diskHighWatermarkPercent: 85 diskLowWatermarkPercent: 75 - protectiveEvictionEnabled: true + # Explicit opt-in only: eviction permanently removes the oldest durable + # messages after sustained disk pressure. + protectiveEvictionEnabled: false diskPressureGraceSeconds: 60 + # Kodo compatibility overrides this to 180 seconds so one failed default + # go-nsq Lookupd poll cannot age out a newly discovered owner's messages. bootstrapRetentionSeconds: 90 maxMessageBytes: 20971520 maxTopics: 10000 @@ -51,6 +55,32 @@ queue: # Rotate opaque TCP tunnels so a small producer pool is redistributed as the # Broker fleet changes. Set to 0 to keep connections pinned indefinitely. proxyTcpMaxConnectionAgeSeconds: 300 + # Kodo compatibility is opt-in. Enabling it pins the managed profile to + # three Brokers and three publish Gateway Pods behind one highly available + # Service with a 100 MiB message cap. + kodoCompatibility: + enabled: false + # Set true only after every Kodo workload using this Discovery service has + # stopped. Required to remove an existing Gateway deployment. + decommissionConfirmed: false + # After KodoGatewaysAdvertised=True, restart every Kodo publisher and then + # change this value. Broker disruption remains blocked until it changes. + producerRestartNonce: "" + # Kodo refreshes /nodes every 300s by default. Keep this at least two + # nsq.refresh_at intervals plus 30s so one failed refresh remains safe. + cutoverGraceSeconds: 630 + # Hard-disabled until RustQueue has a cluster-wide atomic, authenticated + # deletion transaction. Setting this to true makes chart rendering fail. + cleanupEnabled: false + brokerMemoryRequest: 2Gi + brokerCpuRequest: "1" + # Gateway access is fail-closed to Kodo Pods. Match these labels on the + # Kodo workload, or override the selector for your deployment. + allowedPodSelector: + app.kubernetes.io/name: kodo + # Empty means the RustQueue namespace. Set labels here when Kodo runs in a + # separately labelled namespace. + allowedNamespaceSelector: {} rollout: paused: false requireCanaryApproval: false diff --git a/docs/architecture/share-nothing-v7.md b/docs/architecture/share-nothing-v7.md index 1822d15..2633035 100644 --- a/docs/architecture/share-nothing-v7.md +++ b/docs/architecture/share-nothing-v7.md @@ -1,7 +1,7 @@ # RustQueue format v7 share-nothing architecture Status: accepted implementation contract -Target release: 0.7.2 +Target release: 0.8.0 Data format: v7, clean directories only ## 1. Goal @@ -60,14 +60,18 @@ There is no broker-to-broker message path and no cluster consensus path. ### 2.3 Bootstrap retention Every broker retains each topic message for at least 90 seconds even when no -channel exists. When a local channel is created, its initial cursor starts at -the oldest message still inside this bootstrap window. +channel exists. The Kodo compatibility profile requires 180 seconds so the +default Go consumer can survive one failed 60-second Lookupd poll, including +its 30% initial jitter. When a local channel is created, its initial cursor +starts at the oldest message still inside this bootstrap window. This deliberately prefers duplicates over misses during the normal discovery window. The guarantee is bounded: - If discovery and `SUB` complete within 90 seconds, messages accepted by a new fallback owner remain consumable by the newly created local channel. +- In the Kodo profile, the equivalent bound is 180 seconds and includes a + second default lookup poll after one failed request. - A newly created channel may receive a small amount of data published before its `SUB`. - If all consumers are absent for longer than the bootstrap window and a topic @@ -131,7 +135,9 @@ Each replica: 3. expires stale broker observations after 5 seconds; 4. maintains topic, channel, consumer, and publisher indexes incrementally, so lookup requests do not scan every Broker registry; -5. serves `/lookup`, `/topics`, `/channels`, `/nodes`, `/ping`, and `/info`. +5. serves `/lookup`, `/topics`, `/channels`, `/nodes`, `/ping`, and `/info`; + `/lookup` returns the healthy owners during a partial Broker outage while + the complete-inventory metric remains false. Discovery state is derived and is never authoritative message metadata. A restart reconstructs the complete index from ready brokers. Replicas do not @@ -166,7 +172,9 @@ membership reconciler. It manages: - one broker StatefulSet; - one RWO SSD PVC per ordinal with `Retain` semantics; - required one-broker-per-node anti-affinity; -- automatic scale-up to the count of eligible labelled nodes; +- automatic scaling to the count of explicitly labelled member nodes, bounded + by `minBrokers` and `maxBrokers`; cordon and transient Node readiness changes + do not express scale-down intent; - conservative highest-ordinal scale-down; - declarative targeted Broker maintenance; - durable operation status and bounded operation history; diff --git a/docs/operations/kubernetes.md b/docs/operations/kubernetes.md index befa2cf..b42ae7f 100644 --- a/docs/operations/kubernetes.md +++ b/docs/operations/kubernetes.md @@ -8,6 +8,9 @@ predictable; they do not turn a Broker into a replicated shard. - Use Kubernetes 1.28 or newer. - Label only SSD-capable nodes with `rustqueue.io/eligible=true`. +- Treat that label as persistent membership intent. Cordon and Node NotReady + do not reduce the Broker target; remove the label only for an intentional, + drain-aware scale-down. - Use an RWO StorageClass with `allowVolumeExpansion: true` and working `fsGroup` ownership. - Keep at least two eligible nodes if image rolling is required. A one-Broker @@ -73,7 +76,7 @@ Canary approval is optional: ```sh helm upgrade rustqueue deploy/helm/rustqueue \ --namespace rustqueue \ - --set queue.image=registry.example/rustqueue:0.7.2 \ + --set queue.image=registry.example/rustqueue:0.8.0 \ --set queue.rollout.requireCanaryApproval=true rustqueuectl -n rustqueue rollout approve @@ -85,7 +88,7 @@ Useful controls: rustqueuectl -n rustqueue rollout pause rustqueuectl -n rustqueue rollout resume rustqueuectl -n rustqueue rollout retry -rustqueuectl -n rustqueue rollout rollback registry.example/rustqueue:0.7.2 +rustqueuectl -n rustqueue rollout rollback registry.example/rustqueue:0.8.0 rustqueuectl -n rustqueue rollout forward ``` @@ -130,9 +133,18 @@ Operator CRDs are installed. Treat these default alerts as follows: removed old messages to keep the disk operable; - `RustQueueOperatorHasNoLeader`: deployment changes are frozen, but existing Broker traffic remains independent; +- `RustQueueBrokerMetricsMissing`, `RustQueueDiscoverySourceUnavailable`, or + `RustQueueProxyHasNoPublishBackend`: treat missing telemetry as a runtime + outage until the corresponding Service endpoints are verified; +- `RustQueueKodoGatewayMetricsMissing` or + `RustQueueKodoStatsInventoryIncomplete`: restore all three Gateway and Broker + Stats shards before trusting Kodo depth or channel-idleness observations; - sustained throttling or high fsync p99: investigate PVC latency/capacity and producer arrival rate. +The bundled rules select the current Helm namespace and queue Services, so a +healthy RustQueue release cannot mask a failed release in the same Prometheus. + PDBs affect voluntary Eviction API calls only. Node loss, `kubectl delete pod`, and the Operator's already-drained replacement are not prevented by a PDB. Permanent PVC loss permanently loses that Broker's unconsumed messages. diff --git a/fuzz/Cargo.lock b/fuzz/Cargo.lock index fa93dd4..5396328 100644 --- a/fuzz/Cargo.lock +++ b/fuzz/Cargo.lock @@ -2017,7 +2017,7 @@ dependencies = [ [[package]] name = "rustqueue-discovery" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "axum", @@ -2031,6 +2031,7 @@ dependencies = [ "rustqueue-telemetry", "serde", "serde_json", + "sha2", "tokio", "tracing", "tracing-subscriber", @@ -2054,7 +2055,7 @@ dependencies = [ [[package]] name = "rustqueue-protocol" -version = "0.7.2" +version = "0.8.0" dependencies = [ "bytes", "serde", @@ -2064,7 +2065,7 @@ dependencies = [ [[package]] name = "rustqueue-proxy" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "axum", @@ -2073,6 +2074,7 @@ dependencies = [ "rand 0.8.7", "reqwest", "rustls", + "rustqueue-protocol", "rustqueue-telemetry", "serde", "serde_json", @@ -2083,7 +2085,7 @@ dependencies = [ [[package]] name = "rustqueue-queue" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "bytes", @@ -2103,7 +2105,7 @@ dependencies = [ [[package]] name = "rustqueue-storage" -version = "0.7.2" +version = "0.8.0" dependencies = [ "anyhow", "crc32c", @@ -2115,7 +2117,7 @@ dependencies = [ [[package]] name = "rustqueue-telemetry" -version = "0.7.2" +version = "0.8.0" dependencies = [ "serde", ] diff --git a/rustqueue.example.toml b/rustqueue.example.toml index 63e6bb0..aac875c 100644 --- a/rustqueue.example.toml +++ b/rustqueue.example.toml @@ -28,11 +28,15 @@ maintenance_startup_delay_seconds = 30 disk_high_watermark_percent = 85 disk_low_watermark_percent = 75 min_free_bytes = 10737418240 -protective_eviction_enabled = true +# Explicit opt-in: this permanently removes the oldest durable messages after +# sustained disk pressure. +protective_eviction_enabled = false disk_pressure_grace_seconds = 60 [queue] -# The stable v7 wire limit is 32 MiB; 20 MiB is the production default. +# The stable v7 wire limit is 100 MiB; 20 MiB remains the conservative default. +# Messages above the legacy durable-record bound require storage feature level 2 +# plus correspondingly larger body, publish, delivery, and segment limits. max_message_bytes = 20971520 # Backlog is not capped by message count. Disk watermarks are the publish limit. message_timeout_ms = 60000 diff --git a/scripts/acceptance-kodo-gateway.sh b/scripts/acceptance-kodo-gateway.sh new file mode 100755 index 0000000..3d9fd82 --- /dev/null +++ b/scripts/acceptance-kodo-gateway.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PROJECT="${PROJECT:-rustqueue-kodo-gateway-e2e}" +BROKER_IMAGE="${BROKER_IMAGE:-rustqueue:dev}" +COMPAT_IMAGE="${COMPAT_IMAGE:-rustqueue-go-compat:kodo-e2e}" +BUILD_IMAGES="${BUILD_IMAGES:-1}" +COMPOSE="$ROOT/tests/kodo-gateway/compose.yaml" + +cleanup() { + code=$? + if [[ $code -ne 0 ]]; then + docker compose -p "$PROJECT" -f "$COMPOSE" logs --no-color || true + fi + docker compose -p "$PROJECT" -f "$COMPOSE" down --volumes --remove-orphans >/dev/null 2>&1 || true + exit "$code" +} +trap cleanup EXIT + +[[ "$PROJECT" =~ ^[a-z0-9][a-z0-9_-]*$ ]] || { + echo "PROJECT must be a narrow Docker Compose project name" >&2 + exit 2 +} +if [[ "$BUILD_IMAGES" == "1" ]]; then + make -C "$ROOT" image + docker build -t "$COMPAT_IMAGE" "$ROOT/tests/compat/go" +fi + +docker compose -p "$PROJECT" -f "$COMPOSE" down --volumes --remove-orphans >/dev/null 2>&1 || true +BROKER_IMAGE="$BROKER_IMAGE" COMPAT_IMAGE="$COMPAT_IMAGE" \ + docker compose -p "$PROJECT" -f "$COMPOSE" up \ + --abort-on-container-exit --exit-code-from acceptance diff --git a/scripts/acceptance-multi-broker-k8s.sh b/scripts/acceptance-multi-broker-k8s.sh index 578ae8c..b7e6e99 100755 --- a/scripts/acceptance-multi-broker-k8s.sh +++ b/scripts/acceptance-multi-broker-k8s.sh @@ -20,6 +20,7 @@ STORAGE_CLASS_CREATED=0 CONSOLE_FORWARD_PID="" source "$(dirname "$0")/lib/console-multi-owner.sh" +source "$(dirname "$0")/lib/token-rotation.sh" require() { command -v "$1" >/dev/null || { echo "missing required command: $1" >&2; exit 1; } @@ -30,6 +31,9 @@ diagnostics() { kubectl -n "$NAMESPACE" describe rustqueue "$QUEUE" || true kubectl -n "$NAMESPACE" logs -l app.kubernetes.io/component=operator --tail=300 || true kubectl -n "$NAMESPACE" logs -l app.kubernetes.io/component=console --tail=300 || true + kubectl -n "$NAMESPACE" logs -l app.kubernetes.io/component=discovery --tail=300 || true + kubectl -n "$NAMESPACE" logs -l app.kubernetes.io/component=broker --tail=300 --prefix || true + kubectl -n "$NAMESPACE" logs default-lookup-bootstrap --tail=300 || true kubectl -n "$NAMESPACE" logs operational-ledger --tail=300 || true } @@ -71,15 +75,23 @@ wait_namespace_deleted() { } wait_queue_ready() { - local brokers=$1 deadline=$((SECONDS + ${2:-360})) + local brokers=$1 deadline=$((SECONDS + ${2:-360})) minimum_generation=${3:-0} + local active_feature_level=${4:-} while (( SECONDS < deadline )); do - local phase desired ready observed generation - phase=$(kubectl -n "$NAMESPACE" get rustqueue "$QUEUE" -o jsonpath='{.status.phase}' 2>/dev/null || true) - desired=$(kubectl -n "$NAMESPACE" get rustqueue "$QUEUE" -o jsonpath='{.status.desiredBrokers}' 2>/dev/null || true) - ready=$(kubectl -n "$NAMESPACE" get rustqueue "$QUEUE" -o jsonpath='{.status.readyBrokers}' 2>/dev/null || true) - observed=$(kubectl -n "$NAMESPACE" get rustqueue "$QUEUE" -o jsonpath='{.status.observedGeneration}' 2>/dev/null || true) - generation=$(kubectl -n "$NAMESPACE" get rustqueue "$QUEUE" -o jsonpath='{.metadata.generation}' 2>/dev/null || true) - if [[ "$phase" == "Ready" && "$desired" == "$brokers" && "$ready" == "$brokers" && "$observed" == "$generation" ]]; then + local queue + queue=$(kubectl -n "$NAMESPACE" get rustqueue "$QUEUE" -o json 2>/dev/null || true) + if jq -e \ + --argjson brokers "$brokers" \ + --argjson minimum_generation "$minimum_generation" \ + --arg active_feature_level "$active_feature_level" \ + '(.metadata.generation // 0) >= $minimum_generation + and .status.phase == "Ready" + and .status.desiredBrokers == $brokers + and .status.readyBrokers == $brokers + and .status.observedGeneration == .metadata.generation + and ($active_feature_level == "" + or (.status.activeStorageFeatureLevel | tostring) == $active_feature_level)' \ + <<<"$queue" >/dev/null 2>&1; then return 0 fi sleep 2 @@ -331,9 +343,9 @@ require jq } if [[ "$BUILD_IMAGES" == "1" ]]; then - BUILD_VERSION=0.7.2-e2e-a MAX_STORAGE_FEATURE_LEVEL=1 make image + BUILD_VERSION=0.8.0-e2e-a MAX_STORAGE_FEATURE_LEVEL=1 make image docker tag rustqueue:dev "$BROKER_IMAGE_A" - BUILD_VERSION=0.7.2-e2e-b MAX_STORAGE_FEATURE_LEVEL=2 make image-from-dist + BUILD_VERSION=0.8.0-e2e-b MAX_STORAGE_FEATURE_LEVEL=2 make image-from-dist docker tag rustqueue:dev "$BROKER_IMAGE_B" [[ "$(docker image inspect "$BROKER_IMAGE_A" -f '{{.Id}}')" != \ "$(docker image inspect "$BROKER_IMAGE_B" -f '{{.Id}}')" ]] || { @@ -429,6 +441,8 @@ kubectl -n "$NAMESPACE" exec "$QUEUE-0" -c broker -- \ http://127.0.0.1:4151/v1/drain >/dev/null wait_queue_ready 3 120 +ADMIN_TOKEN=$(accept_admin_token_rotation "$NAMESPACE" "$QUEUE" 3) + kubectl -n "$NAMESPACE" patch rustqueue "$QUEUE" --type=merge \ -p "{\"spec\":{\"maintenance\":{\"broker\":\"$QUEUE-2\",\"enabled\":true}}}" >/dev/null wait_queue_phase Maintenance 120 @@ -517,13 +531,10 @@ kubectl -n "$NAMESPACE" get pods -l app.kubernetes.io/instance="$QUEUE",app.kube jq -e --arg image "$BROKER_IMAGE_B" \ '.items | (length == 3) and all(.[]; all(.spec.containers[]; .image == $image))' >/dev/null -kubectl -n "$NAMESPACE" patch rustqueue "$QUEUE" --type=merge \ - -p '{"spec":{"storageFeatureLevel":2,"rollout":{"requireCanaryApproval":false,"approvedRevision":null}}}' >/dev/null -wait_queue_ready 3 420 -[[ "$(kubectl -n "$NAMESPACE" get rustqueue "$QUEUE" -o jsonpath='{.status.activeStorageFeatureLevel}')" == "2" ]] || { - echo "feature level 2 was not activated after the compatible rollout" >&2 - exit 1 -} +feature_generation=$(kubectl -n "$NAMESPACE" patch rustqueue "$QUEUE" --type=merge \ + -p '{"spec":{"storageFeatureLevel":2,"rollout":{"requireCanaryApproval":false,"approvedRevision":null}}}' \ + -o json | jq -r '.metadata.generation') +wait_queue_ready 3 420 "$feature_generation" 2 uids_before=$(kubectl -n "$NAMESPACE" get pods \ -l app.kubernetes.io/instance="$QUEUE",app.kubernetes.io/component=broker \ diff --git a/scripts/kodo-replay.sh b/scripts/kodo-replay.sh new file mode 100755 index 0000000..e7c9858 --- /dev/null +++ b/scripts/kodo-replay.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +KODO_SOURCE="${KODO_SOURCE_DIR:-"$ROOT/../kodo-master"}" +REVIEWED_SOURCES=( + "go.mod:6952906dc0c673292f761febfb024bbdc3fdbe84dffd9804fac204f1c84f9823" + "config/cfg.go:5b6e53e0ee0fe01871d4b715b15ef0d4b7b0b2ce9b125bcc81f20d2754770fb7" + "nsq/nsq.go:731670a0cb1ed223314a23e78b81e6c51b683a75efe4170a39c49f01fe189c8c" + "nsq/lookupd.go:a693864a774949d2225cf57a1eb42e0064596516a6ef97462e449b890e0c9802" + "nsq/nsqadmin.go:e12dde1526d05aa0af2e704340d6ca750a89509cf7f25a953f04c14427c4a39b" + "nsq/metrics.go:51894ffaeba1e198c6ce7d129884d7b700d7d66e7eec7d2232f8e4e06de9928f" + "utils/err.go:299d3b9503ddb11aadb378799927c930caafabf899f92d4dd14b5c0b92a27a5b" +) + +if [[ ! -f "$KODO_SOURCE/go.mod" ]]; then + echo "KODO_SOURCE_DIR must point to a Kodo source checkout" >&2 + exit 1 +fi +for reviewed in "${REVIEWED_SOURCES[@]}"; do + source_path="${reviewed%%:*}" + expected_sha256="${reviewed#*:}" + if [[ ! -f "$KODO_SOURCE/$source_path" ]]; then + echo "Kodo source is missing $source_path" >&2 + exit 1 + fi + actual_sha256="$(shasum -a 256 "$KODO_SOURCE/$source_path" | awk '{print $1}')" + if [[ "$actual_sha256" != "$expected_sha256" ]]; then + echo "Kodo $source_path does not match the reviewed upstream source" >&2 + echo "expected: $expected_sha256" >&2 + echo "actual: $actual_sha256" >&2 + exit 1 + fi +done + +DOCKER_ARGS=( + --rm + --network none + -e GOCACHE=/tmp/go-build + -e GOTOOLCHAIN=local + -e GOMAXPROCS=1 + -e GOFLAGS=-p=1 + -e RUSTQUEUE_REPLAY_FIXTURES=/replay/fixtures + -v "$KODO_SOURCE:/source/kodo:ro" + -v "$ROOT/tests/kodo-replay:/replay:ro" +) + +docker run "${DOCKER_ARGS[@]}" golang:1.25-bookworm sh -ec ' + mkdir -p /tmp/kodo/nsq /tmp/kodo/config + cp /source/kodo/nsq/nsqadmin.go /tmp/kodo/nsq/nsqadmin.go + cp /replay/kodo_replay_test.go /tmp/kodo/nsq/kodo_replay_test.go + cp /replay/logger_stub.go /tmp/kodo/nsq/logger_stub.go + cp /replay/config_stub.go /tmp/kodo/config/config_stub.go + cp /replay/kodo_stub.go /tmp/kodo/kodo_stub.go + cp /replay/go.mod /tmp/kodo/go.mod + cd /tmp/kodo + go test ./nsq -run "^TestRustQueueStatsReplay$" -count=1 +' diff --git a/scripts/lib/token-rotation.sh b/scripts/lib/token-rotation.sh new file mode 100644 index 0000000..c87b7da --- /dev/null +++ b/scripts/lib/token-rotation.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash + +accept_admin_token_rotation() { + local namespace=$1 queue=$2 brokers=$3 + local deployment replicas old_token new_token old_revision deadline code + deployment=$(kubectl -n "$namespace" get deployment \ + -l app.kubernetes.io/component=operator \ + -o jsonpath='{.items[0].metadata.name}') + replicas=$(kubectl -n "$namespace" get deployment "$deployment" -o jsonpath='{.spec.replicas}') + old_token=$(kubectl -n "$namespace" get secret "$queue-auth" \ + -o go-template='{{index .data "admin-token" | base64decode}}') + old_revision=$(kubectl -n "$namespace" get statefulset "$queue" \ + -o jsonpath='{.spec.template.metadata.annotations.rustqueue\.io/revision}') + new_token="admin-rotation-$(date +%s)-$RANDOM" + + kubectl -n "$namespace" scale deployment "$deployment" --replicas=0 >/dev/null + kubectl -n "$namespace" wait --for=delete pod \ + -l app.kubernetes.io/component=operator --timeout=90s >/dev/null + kubectl -n "$namespace" patch secret "$queue-auth" --type=merge \ + -p "$(jq -cn --arg token "$new_token" '{stringData:{"admin-token":$token}}')" >/dev/null + + deadline=$((SECONDS + 120)) + code="" + while (( SECONDS < deadline )); do + code=$(kubectl -n "$namespace" exec "$queue-0" -c broker -- \ + curl -sS -o /dev/null -w '%{http_code}' \ + -X POST -H 'Content-Type: application/json' -d '{"enabled":false}' \ + -H "Authorization: Bearer $new_token" \ + http://127.0.0.1:4151/v1/drain 2>/dev/null || true) + [[ "$code" == "200" ]] && break + sleep 2 + done + if [[ "$code" != "200" ]]; then + kubectl -n "$namespace" scale deployment "$deployment" --replicas="$replicas" >/dev/null + echo "running Broker did not hot-reload the rotated admin token" >&2 + return 1 + fi + + code=$(kubectl -n "$namespace" exec "$queue-0" -c broker -- \ + curl -sS -o /dev/null -w '%{http_code}' \ + -X POST -H 'Content-Type: application/json' -d '{"enabled":false}' \ + -H "Authorization: Bearer $old_token" \ + http://127.0.0.1:4151/v1/drain) + if [[ "$code" != "401" ]]; then + kubectl -n "$namespace" scale deployment "$deployment" --replicas="$replicas" >/dev/null + echo "running Broker still accepted the old admin token: HTTP $code" >&2 + return 1 + fi + + kubectl -n "$namespace" scale deployment "$deployment" --replicas="$replicas" >/dev/null + deadline=$((SECONDS + 420)) + while (( SECONDS < deadline )); do + local phase desired ready observed generation target_revision converged + phase=$(kubectl -n "$namespace" get rustqueue "$queue" -o jsonpath='{.status.phase}' 2>/dev/null || true) + desired=$(kubectl -n "$namespace" get rustqueue "$queue" -o jsonpath='{.status.desiredBrokers}' 2>/dev/null || true) + ready=$(kubectl -n "$namespace" get rustqueue "$queue" -o jsonpath='{.status.readyBrokers}' 2>/dev/null || true) + observed=$(kubectl -n "$namespace" get rustqueue "$queue" -o jsonpath='{.status.observedGeneration}' 2>/dev/null || true) + generation=$(kubectl -n "$namespace" get rustqueue "$queue" -o jsonpath='{.metadata.generation}' 2>/dev/null || true) + target_revision=$(kubectl -n "$namespace" get statefulset "$queue" \ + -o jsonpath='{.spec.template.metadata.annotations.rustqueue\.io/revision}' 2>/dev/null || true) + converged=$(kubectl -n "$namespace" get pods \ + -l app.kubernetes.io/instance="$queue",app.kubernetes.io/component=broker \ + -o json 2>/dev/null | jq -r --arg revision "$target_revision" --argjson brokers "$brokers" \ + '(.items | length) == $brokers and all(.items[]; + .metadata.annotations["rustqueue.io/revision"] == $revision + and any(.status.conditions[]?; .type == "Ready" and .status == "True"))' || true) + if [[ -n "$target_revision" && "$target_revision" != "$old_revision" \ + && "$converged" == "true" \ + && "$phase" == "Ready" && "$desired" == "$brokers" && "$ready" == "$brokers" \ + && "$observed" == "$generation" ]]; then + printf '%s\n' "$new_token" + return 0 + fi + sleep 2 + done + echo "cluster did not recover after admin token rotation" >&2 + return 1 +} diff --git a/scripts/release-gate.sh b/scripts/release-gate.sh index 20c6482..161db55 100755 --- a/scripts/release-gate.sh +++ b/scripts/release-gate.sh @@ -4,6 +4,7 @@ set -euo pipefail FULL="${FULL:-1}" K8S_ACCEPTANCE="${K8S_ACCEPTANCE:-0}" FUZZ_SECONDS="${FUZZ_SECONDS:-1}" +KODO_ACCEPTANCE="${KODO_ACCEPTANCE:-0}" make fmt make check @@ -29,4 +30,9 @@ if [[ "$K8S_ACCEPTANCE" == "1" ]]; then make k8s-multi-acceptance fi +if [[ "$KODO_ACCEPTANCE" == "1" ]]; then + make kodo-replay + BUILD_IMAGES=1 make kodo-gateway-acceptance +fi + echo "RustQueue release gate passed" diff --git a/tests/compat/go/fake_discovery.go b/tests/compat/go/fake_discovery.go new file mode 100644 index 0000000..9e30b31 --- /dev/null +++ b/tests/compat/go/fake_discovery.go @@ -0,0 +1,93 @@ +package main + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "os" + "strconv" + "strings" +) + +type fakeProducer struct { + RemoteAddress string `json:"remote_address"` + Hostname string `json:"hostname"` + BroadcastAddress string `json:"broadcast_address"` + TCPPort int `json:"tcp_port"` + HTTPPort int `json:"http_port"` + Version string `json:"version"` + NodeID uint64 `json:"node_id"` +} + +func runFakeDiscovery() error { + brokers, err := parseFakeProducers(os.Getenv("RUSTQUEUE_FAKE_BROKERS")) + if err != nil { + return fmt.Errorf("parse fake Brokers: %w", err) + } + gateways, err := parseFakeProducers(os.Getenv("RUSTQUEUE_FAKE_GATEWAYS")) + if err != nil { + return fmt.Errorf("parse fake Gateways: %w", err) + } + if len(brokers) != 3 || len(gateways) != 3 { + return fmt.Errorf("fake Discovery requires exactly three Brokers and Gateways") + } + publishers := brokers + if os.Getenv("RUSTQUEUE_FAKE_EMPTY_PUBLISHERS") == "1" { + publishers = []fakeProducer{} + } + mux := http.NewServeMux() + write := func(value any) http.HandlerFunc { + return func(response http.ResponseWriter, _ *http.Request) { + response.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(response).Encode(value); err != nil { + panic(err) + } + } + } + mux.HandleFunc("/v1/publishers/head", write(map[string]any{ + "revision": 1, "broker_count": len(publishers), + })) + mux.HandleFunc("/v1/publishers", write(map[string]any{ + "revision": 1, "producers": publishers, + })) + mux.HandleFunc("/v1/brokers", write(map[string]any{ + "revision": 1, "producers": brokers, + })) + mux.HandleFunc("/nodes", write(map[string]any{"producers": gateways})) + mux.HandleFunc("/topics", write(map[string]any{"topics": []string{}})) + mux.HandleFunc("/channels", write(map[string]any{"channels": []string{}})) + return http.ListenAndServe(":4161", mux) +} + +func parseFakeProducers(value string) ([]fakeProducer, error) { + var producers []fakeProducer + for _, item := range strings.Split(value, ",") { + parts := strings.Split(strings.TrimSpace(item), "|") + if len(parts) != 4 { + return nil, fmt.Errorf("%q must be host|tcp|http|node-id", item) + } + tcpPort, err := strconv.Atoi(parts[1]) + if err != nil { + return nil, err + } + httpPort, err := strconv.Atoi(parts[2]) + if err != nil { + return nil, err + } + nodeID, err := strconv.ParseUint(parts[3], 10, 64) + if err != nil { + return nil, err + } + producers = append(producers, fakeProducer{ + RemoteAddress: net.JoinHostPort(parts[0], parts[1]), + Hostname: parts[0], + BroadcastAddress: parts[0], + TCPPort: tcpPort, + HTTPPort: httpPort, + Version: "0.8.0", + NodeID: nodeID, + }) + } + return producers, nil +} diff --git a/tests/compat/go/kodo_gateway.go b/tests/compat/go/kodo_gateway.go new file mode 100644 index 0000000..fd8346f --- /dev/null +++ b/tests/compat/go/kodo_gateway.go @@ -0,0 +1,470 @@ +package main + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/url" + "os" + "strconv" + "strings" + "time" + + "github.com/nsqio/go-nsq" +) + +// config.C.NSQ.MaxMessageSize in the reviewed Kodo source. RustQueue keeps a +// 100 MiB protocol ceiling, leaving exactly 100 bytes of compatibility headroom. +const kodoMessageBytes = 104857500 + +type kodoNode struct { + BroadcastAddress string `json:"broadcast_address"` + TCPPort int `json:"tcp_port"` + HTTPPort int `json:"http_port"` + NodeID uint64 `json:"node_id"` +} + +type kodoTopicStats struct { + TopicName string `json:"topic_name"` + MessageCount uint64 `json:"message_count"` +} + +func runKodoGatewayAcceptance(lookupHTTP, gatewayMetricsHTTP string) error { + nodes, err := waitKodoNodes(lookupHTTP) + if err != nil { + return err + } + for _, node := range nodes { + if err := requireHTTPPublishDisabled(node.httpAddress()); err != nil { + return err + } + } + for _, node := range nodes[:2] { + if err := waitGatewayBackendCount(net.JoinHostPort(node.BroadcastAddress, "4160"), 3); err != nil { + return err + } + } + brokers, err := fakeBrokerNodes() + if err != nil { + return err + } + topic, channel := uniqueTopic("kodo_max_message"), "workers" + for _, broker := range brokers { + if err := createChannel(broker.httpAddress(), topic, channel); err != nil { + return fmt.Errorf("create channel on Broker %d: %w", broker.NodeID, err) + } + } + + body := bytes.Repeat([]byte{0x5a}, kodoMessageBytes) + budgetHold, err := holdGatewayIngressBudget(nodes[0], 40*1024*1024) + if err != nil { + return err + } + defer budgetHold.Close() + started := time.Now() + producer, publishingGateway, failovers, err := publishLikeKodo(nodes, topic, body) + if err != nil { + return fmt.Errorf( + "publish Kodo's 104857500-byte maximum with the reviewed go-nsq timeouts: %w", + err, + ) + } + budgetHold.Close() + defer producer.Stop() + if failovers != 1 || publishingGateway.BroadcastAddress != nodes[1].BroadcastAddress { + return fmt.Errorf( + "Kodo connection-error failover used %d retries and selected %s", + failovers, publishingGateway.BroadcastAddress, + ) + } + fmt.Printf( + "Kodo maximum-size publish completed in %s after %d Gateway failover\n", + time.Since(started), failovers, + ) + + counts, err := waitMessageCount(brokers, topic, 1) + if err != nil { + return err + } + if _, err := waitMessageCount(nodes, topic, 1); err != nil { + return fmt.Errorf("aggregate the three Kodo Gateway Stats shards: %w", err) + } + owner, err := singleOwner(brokers, counts) + if err != nil { + return err + } + if err := consumeLargeMessage(owner, topic, channel, body); err != nil { + return err + } + body = nil + + if _, _, _, err := publishLikeKodo( + nodes[2:], uniqueTopic("kodo_empty_gateway"), []byte("must-fail-before-commit"), + ); !kodoConnectionError(err) { + return fmt.Errorf("Gateway without a Broker did not request Kodo connection failover: %v", err) + } + if err := requireReconnectMetric(gatewayMetricsHTTP); err != nil { + return err + } + + if err := setDrain(owner, true); err != nil { + return err + } + defer setDrain(owner, false) //nolint:errcheck + if err := producer.Publish(topic, []byte("after-drain")); err != nil { + return fmt.Errorf("publish through a draining current Broker: %w", err) + } + counts, err = waitMessageCount(brokers, topic, 2) + if err != nil { + return err + } + if _, err := waitMessageCount(nodes, topic, 2); err != nil { + return fmt.Errorf("aggregate the three Kodo Gateway Stats shards after failover: %w", err) + } + if counts[owner.NodeID] != 1 { + return fmt.Errorf("draining Broker %d accepted another publish", owner.NodeID) + } + if len(nonzeroOwners(counts)) != 2 { + return fmt.Errorf("publish did not fail over to a second Broker: %v", counts) + } + if err := producer.DeferredPublish( + topic, 10*time.Millisecond, []byte("deferred-after-drain"), + ); err != nil { + return fmt.Errorf("deferred publish through a draining current Broker: %w", err) + } + if _, err := waitMessageCount(brokers, topic, 3); err != nil { + return err + } + if _, err := waitMessageCount(nodes, topic, 3); err != nil { + return fmt.Errorf( + "aggregate the three Kodo Gateway Stats shards after deferred publish: %w", + err, + ) + } + if err := requireRetryMetric(net.JoinHostPort(publishingGateway.BroadcastAddress, "4160")); err != nil { + return err + } + return nil +} + +func holdGatewayIngressBudget(node kodoNode, bodyBytes uint32) (net.Conn, error) { + conn, err := net.DialTimeout("tcp", node.tcpAddress(), 3*time.Second) + if err != nil { + return nil, fmt.Errorf("connect Gateway budget hold: %w", err) + } + closeOnError := func(err error) (net.Conn, error) { + conn.Close() + return nil, err + } + if err := conn.SetWriteDeadline(time.Now().Add(3 * time.Second)); err != nil { + return closeOnError(err) + } + request := append([]byte(" V2PUB ingress_budget_hold\n"), make([]byte, 4)...) + binary.BigEndian.PutUint32(request[len(request)-4:], bodyBytes) + request = append(request, 0x5a) + if _, err := io.Copy(conn, bytes.NewReader(request)); err != nil { + return closeOnError(fmt.Errorf("reserve Gateway ingress budget: %w", err)) + } + if err := conn.SetWriteDeadline(time.Time{}); err != nil { + return closeOnError(err) + } + // The Gateway reserves the declared body before waiting for the remainder. + time.Sleep(250 * time.Millisecond) + return conn, nil +} + +func publishLikeKodo( + nodes []kodoNode, topic string, body []byte, +) (*nsq.Producer, kodoNode, int, error) { + var lastErr error + for index, node := range nodes { + config := nsq.NewConfig() + if config.WriteTimeout != time.Second { + return nil, kodoNode{}, index, fmt.Errorf( + "go-nsq default WriteTimeout changed to %s", config.WriteTimeout, + ) + } + if config.ReadTimeout != 60*time.Second { + return nil, kodoNode{}, index, fmt.Errorf( + "go-nsq default ReadTimeout changed to %s", config.ReadTimeout, + ) + } + config.DialTimeout = 3 * time.Second + config.WriteTimeout = 3 * time.Second + producer, err := quietProducer(node.tcpAddress(), config) + if err == nil { + err = producer.Publish(topic, body) + } + if err == nil { + return producer, node, index, nil + } + if producer != nil { + producer.Stop() + } + lastErr = err + if !kodoConnectionError(err) { + return nil, kodoNode{}, index, err + } + } + return nil, kodoNode{}, len(nodes), lastErr +} + +func kodoConnectionError(err error) bool { + if err == nil { + return false + } + detail := err.Error() + for _, fragment := range []string{ + "not connected", + "connection reset by peer", + "broken pipe", + "EOF", + "use of closed network connection", + } { + if strings.Contains(detail, fragment) { + return true + } + } + return false +} + +func requireHTTPPublishDisabled(address string) error { + request, err := http.NewRequest( + http.MethodPost, + "http://"+address+"/pub?topic=must_not_publish", + strings.NewReader("blocked"), + ) + if err != nil { + return err + } + response, err := http.DefaultClient.Do(request) + if err != nil { + return fmt.Errorf("probe disabled Gateway HTTP publishing: %w", err) + } + defer response.Body.Close() + if response.StatusCode != http.StatusNotFound { + return fmt.Errorf( + "Gateway Stats port exposed HTTP publishing: expected 404, got %s", + response.Status, + ) + } + return nil +} + +func waitKodoNodes(address string) ([]kodoNode, error) { + deadline := time.Now().Add(30 * time.Second) + for time.Now().Before(deadline) { + response, err := http.Get("http://" + address + "/nodes") + if err == nil { + var document struct { + Producers []kodoNode `json:"producers"` + } + err = json.NewDecoder(response.Body).Decode(&document) + response.Body.Close() + if err == nil && len(document.Producers) == 3 { + return document.Producers, nil + } + } + time.Sleep(100 * time.Millisecond) + } + return nil, fmt.Errorf("Kodo Discovery did not return exactly three Gateways") +} + +func fakeBrokerNodes() ([]kodoNode, error) { + producers, err := parseFakeProducers(os.Getenv("RUSTQUEUE_FAKE_BROKERS")) + if err != nil { + return nil, err + } + nodes := make([]kodoNode, 0, len(producers)) + for _, producer := range producers { + nodes = append(nodes, kodoNode{ + BroadcastAddress: producer.BroadcastAddress, + TCPPort: producer.TCPPort, + HTTPPort: producer.HTTPPort, + NodeID: producer.NodeID, + }) + } + return nodes, nil +} + +func waitMessageCount(brokers []kodoNode, topic string, expected uint64) (map[uint64]uint64, error) { + deadline := time.Now().Add(30 * time.Second) + for time.Now().Before(deadline) { + counts := make(map[uint64]uint64, len(brokers)) + var total uint64 + valid := true + for _, broker := range brokers { + count, err := topicMessageCount(broker.httpAddress(), topic) + if err != nil { + valid = false + break + } + counts[broker.NodeID] = count + total += count + } + if valid && total == expected { + return counts, nil + } + time.Sleep(100 * time.Millisecond) + } + return nil, fmt.Errorf("topic %s did not reach cumulative message count %d", topic, expected) +} + +func topicMessageCount(address, topic string) (uint64, error) { + endpoint := fmt.Sprintf( + "http://%s/stats?format=json&topic=%s", + address, url.QueryEscape(topic), + ) + response, err := http.Get(endpoint) + if err != nil { + return 0, err + } + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + return 0, fmt.Errorf("stats returned %s", response.Status) + } + var document struct { + Topics []kodoTopicStats `json:"topics"` + } + if err := json.NewDecoder(response.Body).Decode(&document); err != nil { + return 0, err + } + for _, stats := range document.Topics { + if stats.TopicName == topic { + return stats.MessageCount, nil + } + } + return 0, nil +} + +func singleOwner(brokers []kodoNode, counts map[uint64]uint64) (kodoNode, error) { + owners := nonzeroOwners(counts) + if len(owners) != 1 { + return kodoNode{}, fmt.Errorf("maximum-size publish has unexpected owners: %v", counts) + } + for _, broker := range brokers { + if broker.NodeID == owners[0] { + return broker, nil + } + } + return kodoNode{}, fmt.Errorf("owner %d is not in the Broker inventory", owners[0]) +} + +func nonzeroOwners(counts map[uint64]uint64) []uint64 { + var owners []uint64 + for nodeID, count := range counts { + if count > 0 { + owners = append(owners, nodeID) + } + } + return owners +} + +func consumeLargeMessage(broker kodoNode, topic, channel string, expected []byte) error { + config := nsq.NewConfig() + config.MaxInFlight = 1 + consumer, err := quietConsumer(topic, channel, config) + if err != nil { + return err + } + received := make(chan error, 1) + consumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error { + if !bytes.Equal(message.Body, expected) { + received <- fmt.Errorf("maximum-size delivery body mismatch") + } else { + received <- nil + } + return nil + })) + if err := consumer.ConnectToNSQD(broker.tcpAddress()); err != nil { + return err + } + defer stopConsumer(consumer) + select { + case err := <-received: + return err + case <-time.After(90 * time.Second): + return fmt.Errorf("maximum-size delivery timed out") + } +} + +func setDrain(broker kodoNode, enabled bool) error { + body := strings.NewReader(fmt.Sprintf(`{"enabled":%t}`, enabled)) + request, err := http.NewRequest( + http.MethodPost, "http://"+broker.httpAddress()+"/v1/drain", body, + ) + if err != nil { + return err + } + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + if err != nil { + return err + } + defer response.Body.Close() + if response.StatusCode/100 != 2 { + detail, _ := io.ReadAll(io.LimitReader(response.Body, 4096)) + return fmt.Errorf("drain Broker %d returned %s: %s", broker.NodeID, response.Status, detail) + } + return nil +} + +func requireRetryMetric(address string) error { + value, err := proxyMetricValue(address, "rustqueue_proxy_producer_retries_total") + if err == nil && value > 0 { + return nil + } + return fmt.Errorf("Gateway did not record an explicit pre-commit retry") +} + +func requireReconnectMetric(address string) error { + value, err := proxyMetricValue(address, "rustqueue_proxy_publish_backends") + if err == nil && value == 0 { + return nil + } + return fmt.Errorf("isolated Gateway unexpectedly had a publish backend") +} + +func waitGatewayBackendCount(address string, expected uint64) error { + deadline := time.Now().Add(15 * time.Second) + for time.Now().Before(deadline) { + value, err := proxyMetricValue(address, "rustqueue_proxy_publish_backends") + if err == nil && value == expected { + return nil + } + time.Sleep(100 * time.Millisecond) + } + return fmt.Errorf("Gateway %s did not discover %d publish backends", address, expected) +} + +func proxyMetricValue(address, name string) (uint64, error) { + response, err := http.Get("http://" + address + "/metrics") + if err != nil { + return 0, err + } + defer response.Body.Close() + body, err := io.ReadAll(io.LimitReader(response.Body, 1024*1024)) + if err != nil { + return 0, err + } + for _, line := range strings.Split(string(body), "\n") { + fields := strings.Fields(line) + if len(fields) == 2 && fields[0] == name { + value, err := strconv.ParseUint(fields[1], 10, 64) + return value, err + } + } + return 0, fmt.Errorf("Gateway metric %s is absent", name) +} + +func (node kodoNode) tcpAddress() string { + return net.JoinHostPort(node.BroadcastAddress, strconv.Itoa(node.TCPPort)) +} + +func (node kodoNode) httpAddress() string { + return net.JoinHostPort(node.BroadcastAddress, strconv.Itoa(node.HTTPPort)) +} diff --git a/tests/compat/go/lookup_default.go b/tests/compat/go/lookup_default.go index 7c36170..63eced1 100644 --- a/tests/compat/go/lookup_default.go +++ b/tests/compat/go/lookup_default.go @@ -2,7 +2,9 @@ package main import ( "fmt" + "log" "net/http" + "os" "time" "github.com/nsqio/go-nsq" @@ -32,6 +34,7 @@ func runDefaultLookupBootstrap(lookupHTTP, seedHTTP, newOwnerHTTP string) error if err != nil { return err } + consumer.SetLogger(log.New(os.Stderr, "nsq: ", log.LstdFlags|log.Lmicroseconds), nsq.LogLevelInfo) consumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error { delivery <- append([]byte(nil), message.Body...) return nil @@ -40,7 +43,10 @@ func runDefaultLookupBootstrap(lookupHTTP, seedHTTP, newOwnerHTTP string) error return err } defer stopConsumer(consumer) - if err := waitForConnections(consumer, 1, 15*time.Second); err != nil { + // Either Discovery replica may answer the first request while it is still + // one refresh behind. The unchanged client deliberately recovers on its next + // jittered 60-second poll, which the bootstrap-retention contract covers. + if err := waitForConnections(consumer, 1, 90*time.Second); err != nil { return err } diff --git a/tests/compat/go/main.go b/tests/compat/go/main.go index 8f0363c..bc3a402 100644 --- a/tests/compat/go/main.go +++ b/tests/compat/go/main.go @@ -12,6 +12,22 @@ import ( func main() { mode := argument(1, "core") + if mode == "fake-discovery" { + if err := runFakeDiscovery(); err != nil { + log.Fatal(err) + } + return + } + if mode == "kodo-gateway" { + if len(os.Args) != 4 { + log.Fatal("kodo-gateway mode requires lookup HTTP and Gateway metrics addresses") + } + if err := runKodoGatewayAcceptance(os.Args[2], os.Args[3]); err != nil { + log.Fatal(err) + } + fmt.Println("Kodo Gateway maximum-size PUB/DPUB and Broker failover acceptance: ok") + return + } if mode == "operational-ledger" { if len(os.Args) != 6 { log.Fatal("operational-ledger mode requires proxy HTTP, lookup HTTP, duration seconds, and minimum brokers") diff --git a/tests/kodo-gateway/compose.yaml b/tests/kodo-gateway/compose.yaml new file mode 100644 index 0000000..287994a --- /dev/null +++ b/tests/kodo-gateway/compose.yaml @@ -0,0 +1,85 @@ +services: + broker-0: &broker + image: ${BROKER_IMAGE:-rustqueue:dev} + entrypoint: ["rustqueued", "--config", "/etc/rustqueue/rustqueue.toml"] + environment: + POD_NAME: queue-0 + RUSTQUEUE_BROADCAST_ADDRESS: broker-0 + volumes: + - ./rustqueue.toml:/etc/rustqueue/rustqueue.toml:ro + tmpfs: + - /data:size=512m,uid=65532,gid=65532,mode=0700 + broker-1: + <<: *broker + environment: + POD_NAME: queue-1 + RUSTQUEUE_BROADCAST_ADDRESS: broker-1 + broker-2: + <<: *broker + environment: + POD_NAME: queue-2 + RUSTQUEUE_BROADCAST_ADDRESS: broker-2 + + discovery: + image: ${COMPAT_IMAGE:-rustqueue-go-compat:kodo-e2e} + command: ["fake-discovery"] + environment: &inventory + RUSTQUEUE_FAKE_BROKERS: broker-0|4150|4151|1,broker-1|4150|4151|2,broker-2|4150|4151|3 + RUSTQUEUE_FAKE_GATEWAYS: gateway-0|4150|4151|1000000,gateway-1|4150|4154|1000001,gateway-2|4150|4155|1000002 + + empty-discovery: + image: ${COMPAT_IMAGE:-rustqueue-go-compat:kodo-e2e} + command: ["fake-discovery"] + environment: + <<: *inventory + RUSTQUEUE_FAKE_EMPTY_PUBLISHERS: "1" + + gateway-0: &gateway + image: ${BROKER_IMAGE:-rustqueue:dev} + entrypoint: ["rustqueue-proxy"] + environment: + RUSTQUEUE_DISCOVERY_URLS: http://discovery:4161 + RUSTQUEUE_KODO_COMPATIBILITY_ENABLED: "true" + RUSTQUEUE_KODO_GATEWAY_ORDINAL: "0" + RUSTQUEUE_PROXY_MAX_MESSAGE_BYTES: "104857600" + RUSTQUEUE_PROXY_MAX_BODY_BYTES: "134217728" + RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES: "134217728" + RUSTQUEUE_PROXY_TCP_COMMAND_TIMEOUT_MS: "120000" + RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS: "0" + depends_on: + - empty-discovery + - broker-0 + - broker-1 + - broker-2 + gateway-1: + <<: *gateway + environment: + RUSTQUEUE_DISCOVERY_URLS: http://discovery:4161 + RUSTQUEUE_KODO_COMPATIBILITY_ENABLED: "true" + RUSTQUEUE_KODO_GATEWAY_ORDINAL: "1" + RUSTQUEUE_PROXY_MAX_MESSAGE_BYTES: "104857600" + RUSTQUEUE_PROXY_MAX_BODY_BYTES: "134217728" + RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES: "536870912" + RUSTQUEUE_PROXY_TCP_COMMAND_TIMEOUT_MS: "120000" + RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS: "0" + gateway-2: + <<: *gateway + environment: + RUSTQUEUE_DISCOVERY_URLS: http://empty-discovery:4161 + RUSTQUEUE_KODO_COMPATIBILITY_ENABLED: "true" + RUSTQUEUE_KODO_GATEWAY_ORDINAL: "2" + RUSTQUEUE_PROXY_MAX_MESSAGE_BYTES: "104857600" + RUSTQUEUE_PROXY_MAX_BODY_BYTES: "134217728" + RUSTQUEUE_PROXY_MAX_INFLIGHT_BYTES: "536870912" + RUSTQUEUE_PROXY_TCP_COMMAND_TIMEOUT_MS: "120000" + RUSTQUEUE_PROXY_TCP_MAX_CONNECTION_AGE_SECONDS: "0" + + acceptance: + image: ${COMPAT_IMAGE:-rustqueue-go-compat:kodo-e2e} + command: ["kodo-gateway", "discovery:4161", "gateway-2:4160"] + environment: + <<: *inventory + depends_on: + - gateway-0 + - gateway-1 + - gateway-2 diff --git a/tests/kodo-gateway/rustqueue.toml b/tests/kodo-gateway/rustqueue.toml new file mode 100644 index 0000000..68300f6 --- /dev/null +++ b/tests/kodo-gateway/rustqueue.toml @@ -0,0 +1,18 @@ +[storage] +data_path = "/data" +feature_level = 2 +max_segment_bytes = 268435456 +min_free_bytes = 0 +protective_eviction_enabled = false + +[queue] +max_message_bytes = 104857600 + +[limits] +max_body_bytes = 134217728 +node_publish_inflight_bytes = 1073741824 +connection_publish_inflight_bytes = 167772160 +node_delivery_inflight_bytes = 536870912 +connection_delivery_inflight_bytes = 134217728 +tcp_command_timeout_ms = 120000 +http_body_timeout_ms = 120000 diff --git a/tests/kodo-replay/config_stub.go b/tests/kodo-replay/config_stub.go new file mode 100644 index 0000000..bff7456 --- /dev/null +++ b/tests/kodo-replay/config_stub.go @@ -0,0 +1,10 @@ +package config + +var C struct { + Global struct { + NsqLatencyThresholdRatio float64 + } + NSQ struct { + Lookupd string + } +} diff --git a/tests/kodo-replay/fixtures/stats-1.json b/tests/kodo-replay/fixtures/stats-1.json new file mode 100644 index 0000000..94b47a9 --- /dev/null +++ b/tests/kodo-replay/fixtures/stats-1.json @@ -0,0 +1,29 @@ +{ + "version": "0.8.0", + "health": "OK", + "start_time": 1700000001, + "topics": [{ + "topic_name": "events", + "name": "events", + "depth": 0, + "memory_depth": 0, + "backend_depth": 0, + "message_count": 10, + "paused": false, + "channels": [{ + "channel_name": "workers", + "name": "workers", + "depth": 4, + "memory_depth": 0, + "backend_depth": 4, + "message_count": 10, + "in_flight_count": 1, + "deferred_count": 2, + "requeue_count": 3, + "timeout_count": 4, + "client_count": 0, + "clients": [], + "paused": false + }] + }] +} diff --git a/tests/kodo-replay/fixtures/stats-2.json b/tests/kodo-replay/fixtures/stats-2.json new file mode 100644 index 0000000..92d62ad --- /dev/null +++ b/tests/kodo-replay/fixtures/stats-2.json @@ -0,0 +1,29 @@ +{ + "version": "0.8.0", + "health": "OK", + "start_time": 1700000002, + "topics": [{ + "topic_name": "events", + "name": "events", + "depth": 0, + "memory_depth": 0, + "backend_depth": 0, + "message_count": 20, + "paused": false, + "channels": [{ + "channel_name": "workers", + "name": "workers", + "depth": 5, + "memory_depth": 0, + "backend_depth": 5, + "message_count": 20, + "in_flight_count": 2, + "deferred_count": 1, + "requeue_count": 5, + "timeout_count": 6, + "client_count": 0, + "clients": [], + "paused": false + }] + }] +} diff --git a/tests/kodo-replay/fixtures/stats-3.json b/tests/kodo-replay/fixtures/stats-3.json new file mode 100644 index 0000000..7170028 --- /dev/null +++ b/tests/kodo-replay/fixtures/stats-3.json @@ -0,0 +1,29 @@ +{ + "version": "0.8.0", + "health": "OK", + "start_time": 1700000003, + "topics": [{ + "topic_name": "events", + "name": "events", + "depth": 0, + "memory_depth": 0, + "backend_depth": 0, + "message_count": 30, + "paused": false, + "channels": [{ + "channel_name": "workers", + "name": "workers", + "depth": 6, + "memory_depth": 0, + "backend_depth": 6, + "message_count": 30, + "in_flight_count": 0, + "deferred_count": 3, + "requeue_count": 7, + "timeout_count": 8, + "client_count": 0, + "clients": [], + "paused": false + }] + }] +} diff --git a/tests/kodo-replay/fixtures/stats-missing.json b/tests/kodo-replay/fixtures/stats-missing.json new file mode 100644 index 0000000..5a0ee03 --- /dev/null +++ b/tests/kodo-replay/fixtures/stats-missing.json @@ -0,0 +1,15 @@ +{ + "version": "0.8.0", + "health": "OK", + "start_time": 1700000001, + "topics": [{ + "topic_name": "events", + "name": "events", + "depth": 0, + "memory_depth": 0, + "backend_depth": 0, + "message_count": 10, + "paused": false, + "channels": [] + }] +} diff --git a/tests/kodo-replay/go.mod b/tests/kodo-replay/go.mod new file mode 100644 index 0000000..f311f63 --- /dev/null +++ b/tests/kodo-replay/go.mod @@ -0,0 +1,3 @@ +module gitlab.jiagouyun.com/cloudcare-tools/kodo + +go 1.25 diff --git a/tests/kodo-replay/kodo_replay_test.go b/tests/kodo-replay/kodo_replay_test.go new file mode 100644 index 0000000..63c6d96 --- /dev/null +++ b/tests/kodo-replay/kodo_replay_test.go @@ -0,0 +1,179 @@ +package nsq + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestRustQueueStatsReplay(t *testing.T) { + fixtureDir := os.Getenv("RUSTQUEUE_REPLAY_FIXTURES") + if fixtureDir == "" { + t.Fatal("RUSTQUEUE_REPLAY_FIXTURES is required") + } + + var nodeDeletes atomic.Int32 + nodes := make([]*httptest.Server, 0, 3) + for index := 1; index <= 3; index++ { + stats := mustReadFixture(t, filepath.Join(fixtureDir, fmt.Sprintf("stats-%d.json", index))) + missingStats := mustReadFixture(t, filepath.Join(fixtureDir, "stats-missing.json")) + node := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodGet && r.URL.Path == "/stats": + if r.URL.Query().Get("format") != "json" || + r.URL.Query().Get("include_clients") != "false" || + r.URL.Query().Get("topic") != "events" { + http.Error(w, "unexpected stats query", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + switch r.URL.Query().Get("channel") { + case "", "workers": + _, _ = w.Write(stats) + case "missing": + _, _ = w.Write(missingStats) + default: + http.Error(w, "unexpected channel", http.StatusBadRequest) + } + case r.Method == http.MethodPost && r.URL.Path == "/channel/delete": + nodeDeletes.Add(1) + http.Error(w, "E_NOT_FOUND Kodo cleanup compatibility is disabled", http.StatusNotFound) + default: + http.NotFound(w, r) + } + })) + _ = node.Listener.Close() + listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.%d:0", index)) + if err != nil { + t.Fatal(err) + } + node.Listener = listener + node.Start() + nodes = append(nodes, node) + t.Cleanup(node.Close) + } + + producers := make([]map[string]any, 0, len(nodes)) + for index, node := range nodes { + endpoint, err := url.Parse(node.URL) + if err != nil { + t.Fatal(err) + } + port, err := strconv.Atoi(endpoint.Port()) + if err != nil { + t.Fatal(err) + } + producers = append(producers, map[string]any{ + "remote_address": fmt.Sprintf("%s:4150", endpoint.Hostname()), + "hostname": fmt.Sprintf("rustqueue-broker-%d", index), + "broadcast_address": endpoint.Hostname(), + "tcp_port": 4150, + "http_port": port, + "version": "0.8.0", + "topics": []string{"events"}, + "tombstones": []bool{false}, + }) + } + + var lookupDelete atomic.Int32 + lookup := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodGet && r.URL.Path == "/nodes": + writeReplayJSON(w, map[string]any{"producers": producers}) + case r.Method == http.MethodGet && r.URL.Path == "/topics": + writeReplayJSON(w, map[string]any{"topics": []string{"events"}}) + case r.Method == http.MethodGet && r.URL.Path == "/channels": + writeReplayJSON(w, map[string]any{"channels": []string{"workers"}}) + case r.Method == http.MethodGet && r.URL.Path == "/lookup": + writeReplayJSON(w, map[string]any{ + "channels": []string{"workers"}, + "producers": producers, + }) + case r.Method == http.MethodPost && r.URL.Path == "/channel/delete": + lookupDelete.Add(1) + http.Error(w, "E_NOT_FOUND Kodo cleanup compatibility is disabled", http.StatusNotFound) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(lookup.Close) + + admin := NewNsqAdmin(lookup.URL, 5*time.Second) + topicNodes, err := admin.GetTopicNodes("events") + if err != nil { + t.Fatalf("Kodo lookup replay failed: %v", err) + } + if len(topicNodes.Producers) != 3 { + t.Fatalf("Kodo discovered %d topic owners, want 3", len(topicNodes.Producers)) + } + for _, producer := range topicNodes.Producers { + if !strings.HasPrefix(producer.Hostname, "rustqueue-broker-") { + t.Fatalf("Kodo discovered an unexpected topic owner: %q", producer.Hostname) + } + } + + stats, err := admin.GetChannelStats("events", "workers") + if err != nil { + t.Fatalf("Kodo channel stats replay failed: %v", err) + } + assertReplayChannelStats(t, stats) + + byChannel, err := admin.GetTopicStats("events") + if err != nil { + t.Fatalf("Kodo topic stats replay failed: %v", err) + } + assertReplayChannelStats(t, byChannel["workers"]) + + deleted, err := admin.DeleteChannels("events", []string{"workers"}) + if err == nil || !strings.Contains(err.Error(), "cleanup compatibility is disabled") { + t.Fatalf("Kodo disabled-cleanup error = %v", err) + } + if len(deleted) != 0 || lookupDelete.Load() != 1 || nodeDeletes.Load() != 0 { + t.Fatalf( + "Kodo disabled cleanup returned deleted=%v lookup=%d nodes=%d", + deleted, + lookupDelete.Load(), + nodeDeletes.Load(), + ) + } +} + +func assertReplayChannelStats(t *testing.T, stats *ChannelStats) { + t.Helper() + if stats == nil { + t.Fatal("Kodo did not match channel_name=workers") + } + if stats.TopicName != "events" || stats.ChannelName != "workers" || + stats.Depth != 15 || stats.MemoryDepth != 0 || stats.BackendDepth != 15 || + stats.MessageCount != 60 || stats.InFlightCount != 3 || + stats.DeferredCount != 6 || stats.RequeueCount != 15 || + stats.TimeoutCount != 18 || stats.ClientCount != 0 { + t.Fatalf("unexpected Kodo aggregate: %+v", stats) + } +} + +func mustReadFixture(t *testing.T, path string) []byte { + t.Helper() + body, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + return body +} + +func writeReplayJSON(w http.ResponseWriter, value any) { + if err := json.NewEncoder(w).Encode(value); err != nil { + panic(err) + } +} diff --git a/tests/kodo-replay/kodo_stub.go b/tests/kodo-replay/kodo_stub.go new file mode 100644 index 0000000..e812997 --- /dev/null +++ b/tests/kodo-replay/kodo_stub.go @@ -0,0 +1,18 @@ +package kodo + +import "time" + +type BoolCommand struct{} + +func (*BoolCommand) Result() (bool, error) { + return false, nil +} + +type IntCommand struct{} + +type RedisStore interface { + SetNX(string, any, time.Duration) *BoolCommand + Del(...string) *IntCommand +} + +var Redis RedisStore diff --git a/tests/kodo-replay/logger_stub.go b/tests/kodo-replay/logger_stub.go new file mode 100644 index 0000000..c2e5a37 --- /dev/null +++ b/tests/kodo-replay/logger_stub.go @@ -0,0 +1,9 @@ +package nsq + +type replayLogger struct{} + +func (replayLogger) Warnf(string, ...any) {} +func (replayLogger) Infof(string, ...any) {} +func (replayLogger) Errorf(string, ...any) {} + +var l replayLogger