diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b0bb790..10a813d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,6 +4,7 @@ on: pull_request: push: branches: [main] + tags: ["v*.*.*"] permissions: contents: read @@ -17,5 +18,169 @@ jobs: - uses: azure/setup-helm@v4 with: version: v3.17.3 + - name: Verify release metadata + if: startsWith(github.ref, 'refs/tags/') + run: ./scripts/verify-release-version.sh "${GITHUB_REF_NAME#v}" - name: Core, compatibility, fuzz and packaging gates run: FULL=1 K8S_ACCEPTANCE=0 FUZZ_SECONDS=1 ./scripts/release-gate.sh + + release-context: + needs: release-gate + runs-on: ubuntu-24.04 + outputs: + publish: ${{ steps.release.outputs.publish }} + tag: ${{ steps.release.outputs.tag }} + version: ${{ steps.release.outputs.version }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - id: release + name: Resolve unpublished release + run: | + if [[ "$GITHUB_EVENT_NAME" == "pull_request" ]]; then + echo "publish=false" >> "$GITHUB_OUTPUT" + exit 0 + fi + + if [[ "$GITHUB_REF" == refs/tags/v* ]]; then + version="${GITHUB_REF_NAME#v}" + else + version="$(awk ' + /^\[workspace\.package\]$/ { workspace = 1; next } + /^\[/ { workspace = 0 } + workspace && $1 == "version" { + gsub(/"/, "", $3) + print $3 + exit + } + ' Cargo.toml)" + fi + ./scripts/verify-release-version.sh "$version" + + tag="v$version" + echo "tag=$tag" >> "$GITHUB_OUTPUT" + echo "version=$version" >> "$GITHUB_OUTPUT" + if [[ "$GITHUB_REF" == refs/tags/* ]] || + ! git ls-remote --exit-code --tags origin "refs/tags/$tag" >/dev/null 2>&1 + then + echo "publish=true" >> "$GITHUB_OUTPUT" + else + echo "publish=false" >> "$GITHUB_OUTPUT" + fi + + release-binaries: + if: needs.release-context.outputs.publish == 'true' + needs: release-context + strategy: + fail-fast: false + matrix: + include: + - runner: ubuntu-24.04 + arch: x86_64 + - runner: ubuntu-24.04-arm + arch: aarch64 + runs-on: ${{ matrix.runner }} + timeout-minutes: 60 + steps: + - uses: actions/checkout@v4 + - name: Verify release metadata + env: + RELEASE_VERSION: ${{ needs.release-context.outputs.version }} + run: ./scripts/verify-release-version.sh "$RELEASE_VERSION" + - name: Build native binaries and Console UI + run: make release-bin operator-release-bin console-ui-build + - name: Package native bundle + env: + RELEASE_VERSION: ${{ needs.release-context.outputs.version }} + run: ./scripts/package-release.sh binaries "$RELEASE_VERSION" release "${{ matrix.arch }}" + - uses: actions/upload-artifact@v4 + with: + name: rustqueue-${{ matrix.arch }} + path: release/rustqueue-*-linux-${{ matrix.arch }}.tar.gz + if-no-files-found: error + retention-days: 1 + + publish-release: + if: needs.release-context.outputs.publish == 'true' + needs: [release-context, release-binaries] + runs-on: ubuntu-24.04 + timeout-minutes: 15 + permissions: + contents: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: azure/setup-helm@v4 + with: + version: v3.17.3 + - name: Require the release commit on main + run: | + git fetch --no-tags origin +refs/heads/main:refs/remotes/origin/main + git merge-base --is-ancestor "$GITHUB_SHA" refs/remotes/origin/main + - uses: actions/download-artifact@v4 + with: + path: release + merge-multiple: true + - name: Package source and Helm Chart + env: + RELEASE_VERSION: ${{ needs.release-context.outputs.version }} + run: ./scripts/package-release.sh common "$RELEASE_VERSION" release + - name: Create and verify checksums + env: + RELEASE_VERSION: ${{ needs.release-context.outputs.version }} + run: | + version="$RELEASE_VERSION" + cd release + sha256sum \ + "rustqueue-$version-linux-aarch64.tar.gz" \ + "rustqueue-$version-linux-x86_64.tar.gz" \ + "rustqueue-$version-source.tar.gz" \ + "rustqueue-$version.tgz" \ + > "SHA256SUMS-$version" + sha256sum --check "SHA256SUMS-$version" + - name: Ensure release tag + env: + RELEASE_TAG: ${{ needs.release-context.outputs.tag }} + RELEASE_VERSION: ${{ needs.release-context.outputs.version }} + run: | + git fetch --tags origin + if git rev-parse --verify --quiet "refs/tags/$RELEASE_TAG" >/dev/null; then + test "$(git rev-list -n 1 "$RELEASE_TAG")" = "$GITHUB_SHA" + exit 0 + fi + git config user.name "github-actions[bot]" + git config user.email "41898282+github-actions[bot]@users.noreply.github.com" + git tag -a "$RELEASE_TAG" -m "RustQueue $RELEASE_VERSION" "$GITHUB_SHA" + git push origin "refs/tags/$RELEASE_TAG" + - name: Publish GitHub Release + env: + GH_TOKEN: ${{ github.token }} + RELEASE_TAG: ${{ needs.release-context.outputs.tag }} + RELEASE_VERSION: ${{ needs.release-context.outputs.version }} + run: | + version="$RELEASE_VERSION" + notes="docs/releases/$RELEASE_TAG.md" + assets=( + "release/rustqueue-$version-linux-aarch64.tar.gz" + "release/rustqueue-$version-linux-x86_64.tar.gz" + "release/rustqueue-$version-source.tar.gz" + "release/rustqueue-$version.tgz" + "release/SHA256SUMS-$version" + ) + if gh release view "$RELEASE_TAG" >/dev/null 2>&1; then + gh release upload "$RELEASE_TAG" "${assets[@]}" --clobber + else + gh release create "$RELEASE_TAG" \ + --draft \ + --verify-tag \ + --title "RustQueue $version" \ + --notes-file "$notes" + gh release upload "$RELEASE_TAG" "${assets[@]}" + fi + gh release edit "$RELEASE_TAG" \ + --title "RustQueue $version" \ + --notes-file "$notes" \ + --draft=false \ + --latest diff --git a/Cargo.lock b/Cargo.lock index 8e26b4e..0a8fadc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2283,7 +2283,7 @@ dependencies = [ [[package]] name = "rustqueue-bench" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "clap", @@ -2295,7 +2295,7 @@ dependencies = [ [[package]] name = "rustqueue-console" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "axum", @@ -2323,7 +2323,7 @@ dependencies = [ [[package]] name = "rustqueue-discovery" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "axum", @@ -2346,7 +2346,7 @@ dependencies = [ [[package]] name = "rustqueue-operator" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "axum", @@ -2373,7 +2373,7 @@ dependencies = [ [[package]] name = "rustqueue-protocol" -version = "0.8.0" +version = "0.8.1" dependencies = [ "bytes", "serde", @@ -2383,7 +2383,7 @@ dependencies = [ [[package]] name = "rustqueue-proxy" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "axum", @@ -2404,7 +2404,7 @@ dependencies = [ [[package]] name = "rustqueue-queue" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "bytes", @@ -2426,7 +2426,7 @@ dependencies = [ [[package]] name = "rustqueue-server" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "async-compression", @@ -2435,8 +2435,9 @@ dependencies = [ "clap", "crc32c", "flate2", + "futures", "parking_lot", - "regex", + "regex-automata", "reqwest", "rustls", "rustls-pemfile", @@ -2461,7 +2462,7 @@ dependencies = [ [[package]] name = "rustqueue-storage" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "crc32c", @@ -2474,14 +2475,14 @@ dependencies = [ [[package]] name = "rustqueue-telemetry" -version = "0.8.0" +version = "0.8.1" dependencies = [ "serde", ] [[package]] name = "rustqueuectl" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "clap", diff --git a/Cargo.toml b/Cargo.toml index abe764a..2f08dd4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ members = [ exclude = ["fuzz"] [workspace.package] -version = "0.8.0" +version = "0.8.1" edition = "2021" license = "Apache-2.0" rust-version = "1.88" @@ -37,7 +37,7 @@ libc = "0.2.186" hdrhistogram = "7.5.4" parking_lot = "0.12.3" rand = "0.8.5" -regex = "1.11.1" +regex-automata = "0.4.16" reqwest = { version = "0.12.12", default-features = false, features = ["http2", "json", "rustls-tls"] } rustls = "0.23.21" rustls-pemfile = "2.2.0" diff --git a/README.md b/README.md index 9a2ccf7..a520e34 100644 --- a/README.md +++ b/README.md @@ -1,16 +1,23 @@ # RustQueue +**Durable, NSQ-compatible messaging for Kubernetes — built in Rust.** + [![CI](https://github.com/SamuelSupe/rustqueue/actions/workflows/ci.yml/badge.svg)](https://github.com/SamuelSupe/rustqueue/actions/workflows/ci.yml) [![Release](https://img.shields.io/github/v/release/SamuelSupe/rustqueue)](https://github.com/SamuelSupe/rustqueue/releases/latest) [![Rust](https://img.shields.io/badge/rust-1.88%2B-orange.svg)](https://www.rust-lang.org/) [![Kubernetes](https://img.shields.io/badge/kubernetes-1.28%2B-326CE5.svg)](https://kubernetes.io/) -RustQueue 0.8 is a Kubernetes-native, NSQ V2-compatible message queue for +[Architecture](docs/architecture/share-nothing-v7.md) · +[Kubernetes operations](docs/operations/kubernetes.md) · +[Console operations](docs/operations/console.md) · +[v0.8.1 release](https://github.com/SamuelSupe/rustqueue/releases/tag/v0.8.1) + +RustQueue 0.8.1 is a Kubernetes-native, NSQ V2-compatible message queue for trusted internal networks. It is written in Rust and uses a deliberately simple share-nothing model: each Broker owns one durable RWO PVC, while Kubernetes provides scheduling, rollout and discovery. -> Current release: [v0.8.0](https://github.com/SamuelSupe/rustqueue/releases/tag/v0.8.0). +> Current release: [v0.8.1](https://github.com/SamuelSupe/rustqueue/releases/tag/v0.8.1). > RustQueue is a production candidate for workloads that accept single-PVC > durability and at-least-once delivery. It does not replicate messages between > Brokers and is not an HA replacement for a replicated log. @@ -36,31 +43,54 @@ messages stored on that Broker are lost. Configure disk pressure protection, monitor the exported metrics, and choose PVC/storage failure policies that fit your workload before deploying to production. -## What's new in 0.8 - -- **Kodo compatibility lives in RustQueue.** Discovery `/nodes` advertises - three stable Gateway identities for publishing, while `/lookup` continues to - return the real Broker owners used by consumers. Existing Kodo producer and - consumer behavior is preserved. -- **Safe producer failover.** Gateways route only to publish-ready Brokers and - retry another Broker only after an explicit pre-commit rejection. A failure - after a full body write is reported as ambiguous and is never replayed - automatically. -- **Accurate NSQ Stats.** `/stats` exposes standard `topic_name`, - `channel_name`, depth, client and cumulative-count fields. Topic and Channel - message, requeue and timeout counters remain durable and monotonic across - restart, empty and eviction. -- **100 MiB large-message support.** The v7 storage format now has a stable - 100 MiB protocol/storage ceiling. The Kodo profile raises the runtime limits - and release acceptance publishes and consumes Kodo's exact - 104,857,500-byte maximum. -- **Fail-closed administration and rollout.** Unauthenticated Kodo channel - cleanup stays disabled. Native management uses scoped tokens, and the - Operator blocks Broker maintenance until Gateway cutover and the explicit - Kodo producer-restart fence are complete. - -See the [v0.8.0 release notes](https://github.com/SamuelSupe/rustqueue/releases/tag/v0.8.0) -for the complete upgrade and validation record. +## What's new in 0.8.1 + +- **Truthful end-to-end benchmarks.** `rustqueue-bench` now starts durable, + isolated consumers before publishing, counts unique deliveries and + duplicates, reports publish and receive throughput separately, and fails if + the requested messages do not arrive before the drain deadline. +- **Delivery-state correctness.** Generation tokens reject stale `FIN`, `REQ` + and `TOUCH` commands after redelivery. Initial leases cover buffered writes, + and disconnects no longer release messages while their durable channel + operation is still pending. +- **Crash-safe DLQ and management operations.** Dead-letter transfers are + serialized as one durable transaction, recover without replaying a completed + copy, and respect Topic/Channel fences during concurrent administrative + changes. +- **Cancellation-safe storage.** Payload and recovery-index workers retain + their guards and byte budgets until blocking I/O actually finishes. + Corruption marks storage unhealthy before a response can escape, while + retired Topics are reclaimed after their final reader drains. +- **Bounded control planes.** AUTH responses, compiled authorization regexes, + Broker management bodies, Kodo Stats aggregation and proxy error bodies all + have explicit node-wide limits and timeouts. Invalid semaphore or timer + configurations fail during startup instead of panicking later. + +The patch keeps disk format v7 and the NSQ/Kodo compatibility contract from +0.8.0. See the +[v0.8.1 release notes](https://github.com/SamuelSupe/rustqueue/releases/tag/v0.8.1) +for the complete fix and validation record. + +## Download 0.8.1 + +Every release contains native Linux binaries, the Console UI, source, the Helm +Chart and a checksum manifest: + +| Asset | Contents | +| --- | --- | +| `rustqueue-0.8.1-linux-x86_64.tar.gz` | Linux x86_64 binaries, Console UI and example configuration | +| `rustqueue-0.8.1-linux-aarch64.tar.gz` | Linux ARM64 binaries, Console UI and example configuration | +| `rustqueue-0.8.1-source.tar.gz` | Source archive for the tagged commit | +| `rustqueue-0.8.1.tgz` | Helm Chart | +| `SHA256SUMS-0.8.1` | SHA-256 checksums for every downloadable artifact | + +```sh +arch="$(uname -m)" +curl -LO "https://github.com/SamuelSupe/rustqueue/releases/download/v0.8.1/rustqueue-0.8.1-linux-${arch}.tar.gz" +curl -LO "https://github.com/SamuelSupe/rustqueue/releases/download/v0.8.1/SHA256SUMS-0.8.1" +sha256sum --check --ignore-missing SHA256SUMS-0.8.1 +tar -xzf "rustqueue-0.8.1-linux-${arch}.tar.gz" +``` ## Architecture @@ -193,7 +223,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.8.0 \ + --set queue.image=registry.example/rustqueue:0.8.1 \ --set queue.storageClassName=ssd-rwo ``` @@ -430,10 +460,12 @@ test-only direct Pod placement; production anti-affinity is unchanged. A unit fixture covers discovery indexing for 500 brokers. No 500-broker deployment or load test is part of the functional gate. -The v0.8.0 release additionally passed the unmodified Kodo source replay, -an exact 104,857,500-byte `PUB`/`DPUB` with one Gateway failover, and a -three-Broker operational ledger with 3,239 expected and unique messages, zero -missing, zero duplicates and zero publish errors. +The v0.8.1 CI/CD workflow publishes a Release only after the non-Kubernetes +release gate, both native Linux builds, packaging and checksum verification +succeed. The v0.8.0 Kodo compatibility baseline additionally passed the +unmodified Kodo source replay, an exact 104,857,500-byte `PUB`/`DPUB` with one +Gateway failover, and a three-Broker operational ledger with 3,239 expected and +unique messages, zero missing, zero duplicates and zero publish errors. ## Disk pressure diff --git a/console-ui/package.json b/console-ui/package.json index f0f63ac..425c50f 100644 --- a/console-ui/package.json +++ b/console-ui/package.json @@ -1,7 +1,7 @@ { "name": "rustqueue-console-ui", "private": true, - "version": "0.8.0", + "version": "0.8.1", "packageManager": "pnpm@11.9.0", "type": "module", "scripts": { diff --git a/crates/bench/src/consumer.rs b/crates/bench/src/consumer.rs new file mode 100644 index 0000000..6570312 --- /dev/null +++ b/crates/bench/src/consumer.rs @@ -0,0 +1,395 @@ +use anyhow::Context; +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio::sync::{mpsc, watch, Notify}; +use tokio::task::JoinSet; + +const RDY_COUNT: u64 = 2_500; +const RDY_REFILL_AT: u64 = RDY_COUNT / 4; +const STOP_TIMEOUT: Duration = Duration::from_secs(30); + +#[derive(Default)] +struct IdRanges { + ranges: BTreeMap, +} + +impl IdRanges { + fn insert(&mut self, id: u64) -> bool { + let previous = self + .ranges + .range(..=id) + .next_back() + .map(|(&start, &end)| (start, end)); + if previous.is_some_and(|(_, end)| id <= end) { + return false; + } + let next = self + .ranges + .range(id..) + .next() + .map(|(&start, &end)| (start, end)); + let joins_previous = previous.is_some_and(|(_, end)| end.checked_add(1) == Some(id)); + let joins_next = next.is_some_and(|(start, _)| id.checked_add(1) == Some(start)); + + match (previous, next, joins_previous, joins_next) { + (Some((start, _)), Some((next_start, next_end)), true, true) => { + *self.ranges.get_mut(&start).expect("previous range exists") = next_end; + self.ranges.remove(&next_start); + } + (Some((start, _)), _, true, false) => { + *self.ranges.get_mut(&start).expect("previous range exists") = id; + } + (_, Some((next_start, next_end)), false, true) => { + self.ranges.remove(&next_start); + self.ranges.insert(id, next_end); + } + _ => { + self.ranges.insert(id, id); + } + } + true + } +} + +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct DeliverySnapshot { + pub(crate) unique: u64, + pub(crate) total: u64, +} + +impl DeliverySnapshot { + pub(crate) fn duplicates(self) -> u64 { + self.total.saturating_sub(self.unique) + } +} + +pub(crate) struct DeliveryWait { + pub(crate) snapshot: DeliverySnapshot, + pub(crate) complete: bool, +} + +pub(crate) struct ConsumerProgress { + ids: Mutex, + unique: AtomicU64, + total: AtomicU64, + changed: Notify, +} + +impl Default for ConsumerProgress { + fn default() -> Self { + Self { + ids: Mutex::new(IdRanges::default()), + unique: AtomicU64::new(0), + total: AtomicU64::new(0), + changed: Notify::new(), + } + } +} + +impl ConsumerProgress { + fn observe(&self, id: u64) { + self.total.fetch_add(1, Ordering::Relaxed); + let inserted = self + .ids + .lock() + .expect("consumer progress lock poisoned") + .insert(id); + if inserted { + self.unique.fetch_add(1, Ordering::Release); + self.changed.notify_one(); + } + } + + pub(crate) fn snapshot(&self) -> DeliverySnapshot { + DeliverySnapshot { + unique: self.unique.load(Ordering::Acquire), + total: self.total.load(Ordering::Relaxed), + } + } + + pub(crate) async fn wait_for(&self, target: u64, timeout: Duration) -> DeliveryWait { + let deadline = tokio::time::Instant::now() + timeout; + loop { + let changed = self.changed.notified(); + let snapshot = self.snapshot(); + if snapshot.unique >= target { + return DeliveryWait { + snapshot, + complete: true, + }; + } + if tokio::time::timeout_at(deadline, changed).await.is_err() { + return DeliveryWait { + snapshot: self.snapshot(), + complete: false, + }; + } + } + } +} + +pub(crate) struct ConsumerGroup { + stop: watch::Sender, + tasks: JoinSet>, + failures: mpsc::UnboundedReceiver, +} + +impl ConsumerGroup { + pub(crate) async fn failure(&mut self) -> String { + self.failures + .recv() + .await + .unwrap_or_else(|| "all consumers exited unexpectedly".into()) + } + + pub(crate) async fn stop(mut self) -> anyhow::Result<()> { + let _ = self.stop.send(true); + match tokio::time::timeout(STOP_TIMEOUT, join_consumers(&mut self.tasks)).await { + Ok(result) => result, + Err(_) => { + self.tasks.abort_all(); + while self.tasks.join_next().await.is_some() {} + anyhow::bail!("benchmark consumers did not stop within {STOP_TIMEOUT:?}"); + } + } + } +} + +async fn join_consumers(tasks: &mut JoinSet>) -> anyhow::Result<()> { + let mut first_error = None; + while let Some(result) = tasks.join_next().await { + let result = result + .context("benchmark consumer panicked") + .and_then(|result| result); + if let Err(error) = result { + if first_error.is_none() { + first_error = Some(error); + } + } + } + match first_error { + Some(error) => Err(error), + None => Ok(()), + } +} + +pub(crate) async fn start_consumers( + address: &str, + topic: &str, + channel: &str, + count: usize, + progress: Arc, +) -> anyhow::Result { + let (stop, stop_rx) = watch::channel(false); + let (ready, mut ready_rx) = mpsc::channel(count); + let (failure_tx, mut failures) = mpsc::unbounded_channel(); + let mut tasks = JoinSet::new(); + for _ in 0..count { + let failure_tx = failure_tx.clone(); + let worker = consume_worker( + address.to_owned(), + topic.to_owned(), + channel.to_owned(), + ready.clone(), + stop_rx.clone(), + Arc::clone(&progress), + ); + tasks.spawn(async move { + let result = worker.await; + if let Err(error) = &result { + let _ = failure_tx.send(format!("{error:#}")); + } + result + }); + } + drop(ready); + drop(failure_tx); + for _ in 0..count { + tokio::select! { + ready = ready_rx.recv() => { + if ready.is_none() { + let group = ConsumerGroup { stop, tasks, failures }; + let _ = group.stop().await; + anyhow::bail!("consumer exited before subscription became ready"); + } + } + failure = failures.recv() => { + let failure = + failure.unwrap_or_else(|| "consumer exited before subscription became ready".into()); + let group = ConsumerGroup { stop, tasks, failures }; + let _ = group.stop().await; + anyhow::bail!("consumer failed before subscription became ready: {failure}"); + } + } + } + Ok(ConsumerGroup { + stop, + tasks, + failures, + }) +} + +async fn consume_worker( + address: String, + topic: String, + channel: String, + ready: mpsc::Sender<()>, + mut stop: watch::Receiver, + progress: Arc, +) -> anyhow::Result<()> { + let mut stream = TcpStream::connect(&address) + .await + .with_context(|| format!("connect consumer to {address}"))?; + stream.set_nodelay(true)?; + stream.write_all(b" V2").await?; + stream + .write_all(format!("SUB {topic} {channel}\n").as_bytes()) + .await?; + wait_for_ok(&mut stream).await?; + stream + .write_all(format!("RDY {RDY_COUNT}\n").as_bytes()) + .await?; + ready.send(()).await.ok(); + drop(ready); + + let mut remaining_rdy = RDY_COUNT; + loop { + tokio::select! { + changed = stop.changed() => { + if changed.is_err() || *stop.borrow() { + return close_consumer(&mut stream).await; + } + } + frame = read_frame(&mut stream) => { + let (frame_type, response) = match frame { + Ok(frame) => frame, + Err(_) if *stop.borrow() => return Ok(()), + Err(error) => return Err(error), + }; + match frame_type { + 0 if response == b"_heartbeat_" => stream.write_all(b"NOP\n").await?, + 2 if response.len() >= 26 => { + let id: [u8; 16] = response[10..26] + .try_into() + .expect("message frame ID length was checked"); + let numeric_id = parse_message_id(&id)?; + stream.write_all(b"FIN ").await?; + stream.write_all(&id).await?; + stream.write_all(b"\n").await?; + remaining_rdy = remaining_rdy.saturating_sub(1); + if remaining_rdy <= RDY_REFILL_AT { + stream + .write_all(format!("RDY {RDY_COUNT}\n").as_bytes()) + .await?; + remaining_rdy = RDY_COUNT; + } + progress.observe(numeric_id); + } + 1 => anyhow::bail!( + "consumer error: {}", + String::from_utf8_lossy(&response) + ), + _ => {} + } + } + } + } +} + +async fn close_consumer(stream: &mut TcpStream) -> anyhow::Result<()> { + stream.write_all(b"CLS\n").await?; + stream.flush().await?; + loop { + let (frame_type, response) = read_frame(stream).await?; + match frame_type { + 0 if response == b"CLOSE_WAIT" => return Ok(()), + 0 if response == b"_heartbeat_" => stream.write_all(b"NOP\n").await?, + 2 if response.len() >= 26 => { + stream.write_all(b"FIN ").await?; + stream.write_all(&response[10..26]).await?; + stream.write_all(b"\n").await?; + } + 1 => anyhow::bail!( + "consumer close failed: {}", + String::from_utf8_lossy(&response) + ), + _ => {} + } + } +} + +fn parse_message_id(id: &[u8; 16]) -> anyhow::Result { + let id = std::str::from_utf8(id).context("message ID is not ASCII")?; + u64::from_str_radix(id, 16).context("message ID is not hexadecimal") +} + +async fn wait_for_ok(stream: &mut TcpStream) -> anyhow::Result<()> { + loop { + let (frame_type, response) = read_frame(stream).await?; + if frame_type == 0 && response == b"_heartbeat_" { + stream.write_all(b"NOP\n").await?; + } else if frame_type == 0 && response == b"OK" { + return Ok(()); + } else if frame_type == 1 { + anyhow::bail!("subscribe failed: {}", String::from_utf8_lossy(&response)); + } + } +} + +async fn read_frame(stream: &mut TcpStream) -> anyhow::Result<(i32, Vec)> { + let size = stream.read_u32().await? as usize; + if !(4..=16 * 1024 * 1024).contains(&size) { + anyhow::bail!("server returned invalid frame size {size}"); + } + let frame_type = stream.read_i32().await?; + let mut response = vec![0; size - 4]; + stream.read_exact(&mut response).await?; + Ok((frame_type, response)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn progress_counts_unique_deliveries_and_duplicates() { + let progress = ConsumerProgress::default(); + progress.observe(1); + progress.observe(1); + progress.observe(2); + + let snapshot = progress.snapshot(); + assert_eq!(snapshot.unique, 2); + assert_eq!(snapshot.total, 3); + assert_eq!(snapshot.duplicates(), 1); + } + + #[test] + fn id_ranges_merge_out_of_order_ids_without_per_message_storage() { + let mut ids = IdRanges::default(); + assert!(ids.insert(7)); + assert!(ids.insert(9)); + assert!(ids.insert(8)); + assert!(!ids.insert(8)); + assert_eq!(ids.ranges, BTreeMap::from([(7, 9)])); + } + + #[test] + fn parses_rustqueue_and_nsq_hex_ids() { + assert_eq!(parse_message_id(b"000000000123abcd").unwrap(), 0x0123_abcd); + assert!(parse_message_id(b"not-a-message-id").is_err()); + } + + #[tokio::test] + async fn wait_reports_an_incomplete_drain() { + let progress = ConsumerProgress::default(); + let result = progress.wait_for(1, Duration::from_millis(1)).await; + + assert!(!result.complete); + assert_eq!(result.snapshot.unique, 0); + } +} diff --git a/crates/bench/src/main.rs b/crates/bench/src/main.rs index 1dd1559..f7a964a 100644 --- a/crates/bench/src/main.rs +++ b/crates/bench/src/main.rs @@ -1,15 +1,20 @@ -use anyhow::Context; +mod consumer; +mod producer; + use clap::Parser; +use consumer::{start_consumers, ConsumerProgress, DeliverySnapshot}; use hdrhistogram::Histogram; +use producer::run_workers; use serde::Serialize; use std::sync::Arc; -use std::time::{Duration, Instant}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpStream; -use tokio::sync::{mpsc, watch, Mutex}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use tokio::sync::Mutex; #[derive(Parser, Debug)] -#[command(name = "rustqueue-bench", about = "NSQ V2 durable publish benchmark")] +#[command( + name = "rustqueue-bench", + about = "NSQ V2 publish and delivery benchmark" +)] struct Args { #[arg(long, default_value = "127.0.0.1:4150")] address: String, @@ -40,13 +45,34 @@ struct Args { rate: Option, #[arg(long, default_value_t = false)] json: bool, + #[arg( + long, + default_value_t = false, + help = "Use --topic verbatim; existing retained messages can affect receive counts" + )] + reuse_topic: bool, + #[arg( + long, + default_value_t = 300, + help = "Maximum seconds to wait for consumers after publishing" + )] + drain_timeout_seconds: u64, } #[derive(Serialize)] struct Report { address: String, + requested_topic: String, topic: String, + channel: Option, messages: u64, + received_unique_messages: u64, + received_total_messages: u64, + duplicate_messages: u64, + missing_messages: u64, + delivery_verified: bool, + delivery_complete: bool, + drain_timed_out: bool, message_bytes: usize, producers: usize, consumers: usize, @@ -55,6 +81,12 @@ struct Report { elapsed_seconds: f64, messages_per_second: f64, mebibytes_per_second: f64, + publish_elapsed_seconds: f64, + publish_messages_per_second: f64, + drain_elapsed_seconds: f64, + end_to_end_elapsed_seconds: Option, + receive_messages_per_second: Option, + receive_mebibytes_per_second: Option, latency_us_p50: u64, latency_us_p95: u64, latency_us_p99: u64, @@ -62,12 +94,6 @@ struct Report { latency_us_max: u64, } -#[derive(Clone, Copy)] -struct MessageShape { - bytes: usize, - batch_size: usize, -} - #[tokio::main] async fn main() -> anyhow::Result<()> { let args = Args::parse(); @@ -80,45 +106,180 @@ async fn main() -> anyhow::Result<()> { if args.duration_seconds == Some(0) { anyhow::bail!("duration-seconds must be greater than zero when specified"); } + if args.drain_timeout_seconds == 0 { + anyhow::bail!("drain-timeout-seconds must be greater than zero"); + } if args.topic.len() > 64 { anyhow::bail!("topic exceeds NSQ's 64-byte limit"); } - let (stop_consumers, consumer_tasks) = if args.consumers == 0 { - let (stop, _) = watch::channel(false); - (stop, Vec::new()) - } else { - start_consumers(&args).await? - }; + if !args.topic.is_ascii() { + anyhow::bail!("topic must use NSQ's ASCII name character set"); + } if args.warmup_seconds > 0 { + let warmup_topic = isolated_topic(&args.topic, "warmup"); + let mut warmup_group = if args.consumers == 0 { + None + } else { + let progress = Arc::new(ConsumerProgress::default()); + let channel = benchmark_name("warmup-channel"); + Some( + start_consumers( + &args.address, + &warmup_topic, + &channel, + args.consumers, + progress, + ) + .await?, + ) + }; let warmup = Arc::new(Mutex::new(Histogram::::new_with_max(60_000_000, 3)?)); + let warmup_result = if let Some(group) = warmup_group.as_mut() { + tokio::select! { + result = run_workers( + &args, + &warmup_topic, + None, + Some(Duration::from_secs(args.warmup_seconds)), + warmup, + ) => result, + failure = group.failure() => { + Err(anyhow::anyhow!("warmup consumer failed: {failure}")) + } + } + } else { + run_workers( + &args, + &warmup_topic, + None, + Some(Duration::from_secs(args.warmup_seconds)), + warmup, + ) + .await + }; + if let Some(group) = warmup_group.take() { + let stop_result = group.stop().await; + warmup_result?; + stop_result?; + } else { + warmup_result?; + } + } + + let delivery_progress = (args.consumers > 0).then(|| Arc::new(ConsumerProgress::default())); + let measured_topic = if args.reuse_topic { + args.topic.clone() + } else { + isolated_topic(&args.topic, "run") + }; + let channel = delivery_progress + .as_ref() + .map(|_| benchmark_name("measured")); + let mut consumer_group = match (&delivery_progress, &channel) { + (Some(progress), Some(channel)) => Some( + start_consumers( + &args.address, + &measured_topic, + channel, + args.consumers, + Arc::clone(progress), + ) + .await?, + ), + _ => None, + }; + let histogram = Arc::new(Mutex::new(Histogram::::new_with_max(60_000_000, 3)?)); + let start = Instant::now(); + let measured_result = if let Some(group) = consumer_group.as_mut() { + tokio::select! { + result = run_workers( + &args, + &measured_topic, + args.duration_seconds.is_none().then_some(args.messages), + args.duration_seconds.map(Duration::from_secs), + Arc::clone(&histogram), + ) => result, + failure = group.failure() => { + Err(anyhow::anyhow!("benchmark consumer failed: {failure}")) + } + } + } else { run_workers( &args, - None, - Some(Duration::from_secs(args.warmup_seconds)), - warmup, + &measured_topic, + args.duration_seconds.is_none().then_some(args.messages), + args.duration_seconds.map(Duration::from_secs), + Arc::clone(&histogram), ) - .await?; + .await + }; + let publish_elapsed = start.elapsed(); + let measured_messages = match measured_result { + Ok(messages) => messages, + Err(error) => { + if let Some(group) = consumer_group.take() { + let _ = group.stop().await; + } + return Err(error); + } + }; + if measured_messages == 0 { + if let Some(group) = consumer_group.take() { + let _ = group.stop().await; + } + anyhow::bail!("benchmark published no messages during the measurement window"); } - let histogram = Arc::new(Mutex::new(Histogram::::new_with_max(60_000_000, 3)?)); - let start = Instant::now(); - let measured_messages = run_workers( - &args, - args.duration_seconds.is_none().then_some(args.messages), - args.duration_seconds.map(Duration::from_secs), - Arc::clone(&histogram), - ) - .await?; - let elapsed = start.elapsed(); - let _ = stop_consumers.send(true); - for task in consumer_tasks { - task.await.context("benchmark consumer panicked")??; + + let drain_start = Instant::now(); + let delivery_result = match (&delivery_progress, consumer_group.as_mut()) { + (Some(progress), Some(group)) => { + tokio::select! { + waited = progress.wait_for( + measured_messages, + Duration::from_secs(args.drain_timeout_seconds), + ) => Ok((waited.snapshot, waited.complete)), + failure = group.failure() => { + Err(anyhow::anyhow!("benchmark consumer failed: {failure}")) + } + } + } + _ => Ok((DeliverySnapshot::default(), false)), + }; + let (delivery, delivery_complete) = match delivery_result { + Ok(delivery) => delivery, + Err(error) => { + if let Some(group) = consumer_group.take() { + let _ = group.stop().await; + } + return Err(error); + } + }; + let receive_elapsed = delivery_progress.as_ref().map(|_| start.elapsed()); + let drain_elapsed = delivery_progress + .as_ref() + .map_or(Duration::ZERO, |_| drain_start.elapsed()); + if let Some(group) = consumer_group.take() { + group.stop().await?; } + let histogram = histogram.lock().await; - let seconds = elapsed.as_secs_f64(); + let publish_seconds = publish_elapsed.as_secs_f64(); + let received_unique = delivery.unique; + let missing_messages = measured_messages.saturating_sub(received_unique); + let receive_seconds = receive_elapsed.map(|elapsed| elapsed.as_secs_f64()); let report = Report { address: args.address, - topic: args.topic, + requested_topic: args.topic, + topic: measured_topic, + channel, messages: measured_messages, + received_unique_messages: received_unique, + received_total_messages: delivery.total, + duplicate_messages: delivery.duplicates(), + missing_messages, + delivery_verified: delivery_progress.is_some(), + delivery_complete, + drain_timed_out: delivery_progress.is_some() && !delivery_complete, message_bytes: args.message_bytes, producers: args.producers, consumers: args.consumers, @@ -128,11 +289,20 @@ async fn main() -> anyhow::Result<()> { } else { "saturation" }, - elapsed_seconds: seconds, - messages_per_second: measured_messages as f64 / seconds, + elapsed_seconds: publish_seconds, + messages_per_second: measured_messages as f64 / publish_seconds, mebibytes_per_second: measured_messages as f64 * args.message_bytes as f64 - / seconds + / publish_seconds / (1024.0 * 1024.0), + publish_elapsed_seconds: publish_seconds, + publish_messages_per_second: measured_messages as f64 / publish_seconds, + drain_elapsed_seconds: drain_elapsed.as_secs_f64(), + end_to_end_elapsed_seconds: receive_seconds, + receive_messages_per_second: receive_seconds + .map(|seconds| received_unique as f64 / seconds), + receive_mebibytes_per_second: receive_seconds.map(|seconds| { + received_unique as f64 * args.message_bytes as f64 / seconds / (1024.0 * 1024.0) + }), latency_us_p50: histogram.value_at_quantile(0.50), latency_us_p95: histogram.value_at_quantile(0.95), latency_us_p99: histogram.value_at_quantile(0.99), @@ -143,7 +313,7 @@ async fn main() -> anyhow::Result<()> { println!("{}", serde_json::to_string_pretty(&report)?); } else { println!( - "{} messages in {:.3}s: {:.0} msg/s, {:.2} MiB/s", + "publish ACK: {} messages in {:.3}s, {:.0} msg/s, {:.2} MiB/s", report.messages, report.elapsed_seconds, report.messages_per_second, @@ -157,256 +327,85 @@ async fn main() -> anyhow::Result<()> { report.latency_us_p999, report.latency_us_max ); + if report.delivery_verified { + println!( + "receive: {} unique / {} published in {:.3}s, {:.0} msg/s; duplicates={}, missing={}, drain={:.3}s{}", + report.received_unique_messages, + report.messages, + report.end_to_end_elapsed_seconds.unwrap_or_default(), + report.receive_messages_per_second.unwrap_or_default(), + report.duplicate_messages, + report.missing_messages, + report.drain_elapsed_seconds, + if report.delivery_complete { + "" + } else { + " (drain timeout)" + } + ); + } else { + println!("receive: not measured (--consumers=0)"); + } } + require_complete_delivery( + report.delivery_verified, + report.delivery_complete, + report.received_unique_messages, + report.messages, + )?; Ok(()) } -async fn start_consumers( - args: &Args, -) -> anyhow::Result<( - watch::Sender, - Vec>>, -)> { - let (stop, stop_rx) = watch::channel(false); - let (ready, mut ready_rx) = mpsc::channel(args.consumers); - let mut tasks = Vec::with_capacity(args.consumers); - for _ in 0..args.consumers { - tasks.push(tokio::spawn(consume_worker( - args.address.clone(), - args.topic.clone(), - ready.clone(), - stop_rx.clone(), - ))); - } - drop(ready); - for _ in 0..args.consumers { - ready_rx - .recv() - .await - .context("consumer exited before subscription became ready")?; - } - Ok((stop, tasks)) -} - -async fn consume_worker( - address: String, - topic: String, - ready: mpsc::Sender<()>, - mut stop: watch::Receiver, +fn require_complete_delivery( + verified: bool, + complete: bool, + received: u64, + published: u64, ) -> anyhow::Result<()> { - let mut stream = TcpStream::connect(&address) - .await - .with_context(|| format!("connect consumer to {address}"))?; - stream.set_nodelay(true)?; - stream.write_all(b" V2").await?; - stream - .write_all(format!("SUB {topic} benchmark\n").as_bytes()) - .await?; - loop { - let (frame_type, response) = read_frame(&mut stream).await?; - if frame_type == 0 && response == b"_heartbeat_" { - stream.write_all(b"NOP\n").await?; - } else if frame_type == 0 && response == b"OK" { - break; - } else if frame_type == 1 { - anyhow::bail!("subscribe failed: {}", String::from_utf8_lossy(&response)); - } + if verified && !complete { + anyhow::bail!( + "delivery verification failed: received {received} unique messages out of {published} before the drain timeout" + ); } - stream.write_all(b"RDY 2500\n").await?; - ready.send(()).await.ok(); + Ok(()) +} - loop { - tokio::select! { - changed = stop.changed() => { - if changed.is_err() || *stop.borrow() { - return Ok(()); - } - } - frame = read_frame(&mut stream) => { - let (frame_type, response) = match frame { - Ok(frame) => frame, - Err(_) if *stop.borrow() => return Ok(()), - Err(error) => return Err(error), - }; - match frame_type { - 0 if response == b"_heartbeat_" => stream.write_all(b"NOP\n").await?, - 2 if response.len() >= 26 => { - let id = &response[10..26]; - stream.write_all(b"FIN ").await?; - stream.write_all(id).await?; - stream.write_all(b"\n").await?; - } - 1 => anyhow::bail!( - "consumer error: {}", - String::from_utf8_lossy(&response) - ), - _ => {} - } - } - } - } +fn benchmark_name(label: &str) -> String { + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + format!("bench-{label}-{}-{nonce}", std::process::id()) } -async fn read_frame(stream: &mut TcpStream) -> anyhow::Result<(i32, Vec)> { - let size = stream.read_u32().await? as usize; - if !(4..=16 * 1024 * 1024).contains(&size) { - anyhow::bail!("server returned invalid frame size {size}"); - } - let frame_type = stream.read_i32().await?; - let mut response = vec![0; size - 4]; - stream.read_exact(&mut response).await?; - Ok((frame_type, response)) +fn isolated_topic(base: &str, label: &str) -> String { + let suffix = format!("-{label}-{}", benchmark_name("topic")); + let keep = 64usize.saturating_sub(suffix.len()); + format!("{}{}", &base[..base.len().min(keep)], suffix) } -async fn run_workers( - args: &Args, - total_messages: Option, - duration: Option, - histogram: Arc>>, -) -> anyhow::Result { - let base = total_messages.map(|messages| messages / args.producers as u64); - let remainder = total_messages.map(|messages| messages % args.producers as u64); - let deadline = duration.map(|duration| Instant::now() + duration); - let mut tasks = Vec::with_capacity(args.producers); - for producer in 0..args.producers { - let count = base.map(|base| { - base + u64::from((producer as u64) < remainder.expect("count has remainder")) - }); - let address = args.address.clone(); - let topic = args.topic.clone(); - let message_bytes = args.message_bytes; - let batch_size = args.batch_size; - let histogram = Arc::clone(&histogram); - let producer_rate = args.rate.map(|rate| { - let base = rate / args.producers as u64; - base + u64::from((producer as u64) < rate % args.producers as u64) - }); - tasks.push(tokio::spawn(async move { - publish_worker( - &address, - &topic, - count, - deadline, - MessageShape { - bytes: message_bytes, - batch_size, - }, - producer_rate, - histogram, - ) - .await - })); - } - let mut messages = 0u64; - for task in tasks { - messages = messages.saturating_add(task.await.context("benchmark worker panicked")??); +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn benchmark_channels_are_durable_isolated_and_nsq_compatible() { + let channel = benchmark_name("measured"); + assert!(!channel.ends_with("#ephemeral")); + assert!(channel.len() <= 64); } - Ok(messages) -} -async fn publish_worker( - address: &str, - topic: &str, - count: Option, - deadline: Option, - shape: MessageShape, - rate: Option, - histogram: Arc>>, -) -> anyhow::Result { - let MessageShape { - bytes: message_bytes, - batch_size, - } = shape; - let mut stream = TcpStream::connect(address) - .await - .with_context(|| format!("connect to {address}"))?; - stream.set_nodelay(true)?; - stream.write_all(b" V2").await?; - let command = format!("{} {topic}\n", if batch_size == 1 { "PUB" } else { "MPUB" }); - let body = vec![b'x'; message_bytes]; - let max_batch_body = 4usize - .checked_add( - batch_size - .checked_mul(4usize.saturating_add(message_bytes)) - .context("batch body length overflow")?, - ) - .context("batch body length overflow")?; - if max_batch_body > u32::MAX as usize { - anyhow::bail!("batch body exceeds the NSQ 32-bit frame limit"); + #[test] + fn isolated_topics_are_nsq_bounded() { + let topic = isolated_topic(&"x".repeat(64), "run"); + assert_eq!(topic.len(), 64); + assert!(topic.contains("-run-bench-topic-")); } - let mut local = Histogram::::new_with_max(60_000_000, 3)?; - let period = match rate { - Some(0) => anyhow::bail!("fixed rate is lower than producer count"), - Some(rate) => { - let period_ns = 1_000_000_000u64 / rate.max(1); - Some(Duration::from_nanos(period_ns.max(1))) - } - None => None, - }; - let mut scheduled = Instant::now(); - let mut sent = 0u64; - loop { - if count.is_some_and(|count| sent >= count) - || deadline.is_some_and(|deadline| Instant::now() >= deadline) - { - break; - } - let send_messages = count - .map(|count| count.saturating_sub(sent).min(batch_size as u64)) - .unwrap_or(batch_size as u64) as usize; - let started = if let Some(period) = period { - scheduled += period - .checked_mul(u32::try_from(send_messages).context("batch-size is too large")?) - .context("fixed-rate schedule overflow")?; - tokio::time::sleep_until(tokio::time::Instant::from_std(scheduled)).await; - if deadline.is_some_and(|deadline| Instant::now() >= deadline) { - break; - } - scheduled - } else { - Instant::now() - }; - stream.write_all(command.as_bytes()).await?; - if batch_size == 1 { - stream - .write_all(&(message_bytes as u32).to_be_bytes()) - .await?; - stream.write_all(&body).await?; - } else { - let batch_body = 4usize + send_messages * (4 + message_bytes); - stream.write_all(&(batch_body as u32).to_be_bytes()).await?; - stream - .write_all(&(send_messages as u32).to_be_bytes()) - .await?; - for _ in 0..send_messages { - stream - .write_all(&(message_bytes as u32).to_be_bytes()) - .await?; - stream.write_all(&body).await?; - } - } - loop { - let size = stream.read_u32().await? as usize; - if !(4..=1024).contains(&size) { - anyhow::bail!("server returned invalid frame size {size}"); - } - let frame_type = stream.read_i32().await?; - let mut response = vec![0; size - 4]; - stream.read_exact(&mut response).await?; - if frame_type == 0 && response == b"_heartbeat_" { - stream.write_all(b"NOP\n").await?; - continue; - } - if frame_type != 0 || response != b"OK" { - anyhow::bail!("publish failed: {}", String::from_utf8_lossy(&response)); - } - break; - } - let latency = started.elapsed().as_micros().min(u64::MAX as u128) as u64; - local.record(latency.max(1))?; - sent += send_messages as u64; + #[test] + fn incomplete_verified_delivery_fails_the_benchmark() { + assert!(require_complete_delivery(true, false, 99, 100).is_err()); + assert!(require_complete_delivery(true, true, 100, 100).is_ok()); + assert!(require_complete_delivery(false, false, 0, 100).is_ok()); } - histogram.lock().await.add(&local)?; - Ok(sent) } diff --git a/crates/bench/src/producer.rs b/crates/bench/src/producer.rs new file mode 100644 index 0000000..022b06f --- /dev/null +++ b/crates/bench/src/producer.rs @@ -0,0 +1,167 @@ +use super::Args; +use anyhow::Context; +use hdrhistogram::Histogram; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio::sync::Mutex; + +#[derive(Clone, Copy)] +struct MessageShape { + bytes: usize, + batch_size: usize, +} + +pub(crate) async fn run_workers( + args: &Args, + topic: &str, + total_messages: Option, + duration: Option, + histogram: Arc>>, +) -> anyhow::Result { + let base = total_messages.map(|messages| messages / args.producers as u64); + let remainder = total_messages.map(|messages| messages % args.producers as u64); + let deadline = duration.map(|duration| Instant::now() + duration); + let mut tasks = tokio::task::JoinSet::new(); + for producer in 0..args.producers { + let count = base.map(|base| { + base + u64::from((producer as u64) < remainder.expect("count has remainder")) + }); + let address = args.address.clone(); + let topic = topic.to_owned(); + let histogram = Arc::clone(&histogram); + let producer_rate = args.rate.map(|rate| { + let base = rate / args.producers as u64; + base + u64::from((producer as u64) < rate % args.producers as u64) + }); + let shape = MessageShape { + bytes: args.message_bytes, + batch_size: args.batch_size, + }; + tasks.spawn(async move { + publish_worker( + &address, + &topic, + count, + deadline, + shape, + producer_rate, + histogram, + ) + .await + }); + } + let mut messages = 0u64; + while let Some(result) = tasks.join_next().await { + messages = messages.saturating_add(result.context("benchmark worker panicked")??); + } + Ok(messages) +} + +async fn publish_worker( + address: &str, + topic: &str, + count: Option, + deadline: Option, + shape: MessageShape, + rate: Option, + histogram: Arc>>, +) -> anyhow::Result { + let MessageShape { + bytes: message_bytes, + batch_size, + } = shape; + let mut stream = TcpStream::connect(address) + .await + .with_context(|| format!("connect to {address}"))?; + stream.set_nodelay(true)?; + stream.write_all(b" V2").await?; + let command = format!("{} {topic}\n", if batch_size == 1 { "PUB" } else { "MPUB" }); + let body = vec![b'x'; message_bytes]; + let max_batch_body = 4usize + .checked_add( + batch_size + .checked_mul(4usize.saturating_add(message_bytes)) + .context("batch body length overflow")?, + ) + .context("batch body length overflow")?; + if max_batch_body > u32::MAX as usize { + anyhow::bail!("batch body exceeds the NSQ 32-bit frame limit"); + } + let mut local = Histogram::::new_with_max(60_000_000, 3)?; + let period = match rate { + Some(0) => anyhow::bail!("fixed rate is lower than producer count"), + Some(rate) => { + let period_ns = 1_000_000_000u64 / rate.max(1); + Some(Duration::from_nanos(period_ns.max(1))) + } + None => None, + }; + let mut scheduled = Instant::now(); + + let mut sent = 0u64; + loop { + if count.is_some_and(|count| sent >= count) + || deadline.is_some_and(|deadline| Instant::now() >= deadline) + { + break; + } + let send_messages = count + .map(|count| count.saturating_sub(sent).min(batch_size as u64)) + .unwrap_or(batch_size as u64) as usize; + let started = if let Some(period) = period { + scheduled += period + .checked_mul(u32::try_from(send_messages).context("batch-size is too large")?) + .context("fixed-rate schedule overflow")?; + tokio::time::sleep_until(tokio::time::Instant::from_std(scheduled)).await; + if deadline.is_some_and(|deadline| Instant::now() >= deadline) { + break; + } + scheduled + } else { + Instant::now() + }; + stream.write_all(command.as_bytes()).await?; + if batch_size == 1 { + stream + .write_all(&(message_bytes as u32).to_be_bytes()) + .await?; + stream.write_all(&body).await?; + } else { + let batch_body = 4usize + send_messages * (4 + message_bytes); + stream.write_all(&(batch_body as u32).to_be_bytes()).await?; + stream + .write_all(&(send_messages as u32).to_be_bytes()) + .await?; + for _ in 0..send_messages { + stream + .write_all(&(message_bytes as u32).to_be_bytes()) + .await?; + stream.write_all(&body).await?; + } + } + loop { + let size = stream.read_u32().await? as usize; + if !(4..=1024).contains(&size) { + anyhow::bail!("server returned invalid frame size {size}"); + } + let frame_type = stream.read_i32().await?; + let mut response = vec![0; size - 4]; + stream.read_exact(&mut response).await?; + if frame_type == 0 && response == b"_heartbeat_" { + stream.write_all(b"NOP\n").await?; + continue; + } + if frame_type != 0 || response != b"OK" { + anyhow::bail!("publish failed: {}", String::from_utf8_lossy(&response)); + } + break; + } + let latency = started.elapsed().as_micros().min(u64::MAX as u128) as u64; + local.record(latency.max(1))?; + sent += send_messages as u64; + } + histogram.lock().await.add(&local)?; + Ok(sent) +} diff --git a/crates/operator/src/controller/drain.rs b/crates/operator/src/controller/drain.rs index 0de9e4d..37e395c 100644 --- a/crates/operator/src/controller/drain.rs +++ b/crates/operator/src/controller/drain.rs @@ -568,15 +568,14 @@ async fn drain_status( ip: &str, auth: &AuthSecret, ) -> anyhow::Result { - Ok(context + let response = context .http .get(format!("{}/v1/drain", origin(ip))) .bearer_auth(&auth.registry_token) .send() .await? - .error_for_status()? - .json::() - .await?) + .error_for_status()?; + super::read_broker_json(response).await } fn pod_revision(pod: &Pod) -> Option<&str> { diff --git a/crates/operator/src/controller/mod.rs b/crates/operator/src/controller/mod.rs index 213b0db..f880993 100644 --- a/crates/operator/src/controller/mod.rs +++ b/crates/operator/src/controller/mod.rs @@ -26,6 +26,7 @@ use std::time::Duration; use tokio::sync::watch; const KODO_BOOTSTRAP_RETENTION_SECONDS: u64 = 180; +const MAX_BROKER_CONTROL_RESPONSE_BYTES: usize = 64 * 1024; pub(super) struct ContextData { pub client: Client, @@ -34,6 +35,25 @@ pub(super) struct ContextData { pub leadership: watch::Receiver, } +async fn read_broker_json( + mut response: reqwest::Response, +) -> anyhow::Result { + if response + .content_length() + .is_some_and(|length| length > MAX_BROKER_CONTROL_RESPONSE_BYTES as u64) + { + anyhow::bail!("broker control response exceeds its byte limit"); + } + let mut body = Vec::new(); + while let Some(chunk) = response.chunk().await? { + if body.len().saturating_add(chunk.len()) > MAX_BROKER_CONTROL_RESPONSE_BYTES { + anyhow::bail!("broker control response exceeds its byte limit"); + } + body.extend_from_slice(&chunk); + } + Ok(serde_json::from_slice(&body)?) +} + #[derive(Debug, thiserror::Error)] #[error(transparent)] pub struct ReconcileError(#[from] anyhow::Error); @@ -1544,6 +1564,34 @@ fn watch_namespace() -> String { mod tests { use super::*; + #[tokio::test] + async fn broker_control_json_is_bounded() { + 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, + axum::Router::new().route( + "/", + axum::routing::get(|| async { + "x".repeat(MAX_BROKER_CONTROL_RESPONSE_BYTES + 1) + }), + ), + ) + .await + .unwrap(); + }); + let response = reqwest::get(format!("http://{address}")).await.unwrap(); + + let error = read_broker_json::(response) + .await + .unwrap_err(); + assert!(error.to_string().contains("byte limit")); + task.abort(); + } + #[test] fn existing_gateway_replicas_preserve_activation_without_status() { let gateway: StatefulSet = serde_json::from_value(serde_json::json!({ diff --git a/crates/operator/src/controller/preflight.rs b/crates/operator/src/controller/preflight.rs index f6689ad..68d6d8e 100644 --- a/crates/operator/src/controller/preflight.rs +++ b/crates/operator/src/controller/preflight.rs @@ -191,14 +191,13 @@ pub(super) async fn current_brokers( let Some(ip) = ip else { return Ok::<_, anyhow::Error>((name, None)); }; - let report = http + let response = http .get(format!("{}/v1/capabilities", origin(&ip))) .bearer_auth(&token) .send() .await? - .error_for_status()? - .json::() - .await?; + .error_for_status()?; + let report = super::read_broker_json(response).await?; Ok((name, Some(report))) } })) diff --git a/crates/protocol/src/identify.rs b/crates/protocol/src/identify.rs index e600255..8ede636 100644 --- a/crates/protocol/src/identify.rs +++ b/crates/protocol/src/identify.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -#[derive(Clone, Debug, Default, Deserialize)] +#[derive(Clone, Debug, Default, Deserialize, Serialize)] pub struct IdentifyRequest { #[serde(default)] pub client_id: String, diff --git a/crates/proxy/src/backend.rs b/crates/proxy/src/backend.rs index fc63c71..da483cb 100644 --- a/crates/proxy/src/backend.rs +++ b/crates/proxy/src/backend.rs @@ -166,6 +166,22 @@ impl BackendPool { self.inner.read().backends.clone() } + pub fn matching_bounded( + &self, + maximum: usize, + mut predicate: impl FnMut(&Backend) -> bool, + ) -> Option> { + let state = self.inner.read(); + let mut matching = Vec::new(); + for backend in state.backends.iter().filter(|backend| predicate(backend)) { + if matching.len() == maximum { + return None; + } + matching.push(backend.clone()); + } + Some(matching) + } + pub fn len(&self) -> usize { self.inner.read().backends.len() } diff --git a/crates/proxy/src/http.rs b/crates/proxy/src/http.rs index d3e3b94..016ba71 100644 --- a/crates/proxy/src/http.rs +++ b/crates/proxy/src/http.rs @@ -16,6 +16,7 @@ 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; +const KODO_STATS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); #[derive(Clone)] struct ProxyState { @@ -143,13 +144,26 @@ async fn kodo_stats(State(state): State, request: Request) -> let Some(config) = &state.kodo else { return forward(State(state), request).await; }; - kodo::stats( - config, - &state.broker_pool, - &state.client, - &request.uri().to_string(), + let Ok(permits) = u32::try_from(kodo::STATS_WORKING_SET_BYTES) else { + return stats_throttled(); + }; + let Ok(_permit) = Arc::clone(&state.inflight_bytes).try_acquire_many_owned(permits) else { + return stats_throttled(); + }; + match tokio::time::timeout( + KODO_STATS_TIMEOUT, + kodo::stats( + config, + &state.broker_pool, + &state.client, + &request.uri().to_string(), + ), ) .await + { + Ok(response) => response, + Err(_) => stats_timeout(), + } } #[derive(serde::Deserialize)] @@ -407,6 +421,27 @@ fn throttled() -> Response { response } +fn stats_throttled() -> Response { + let mut response = ( + StatusCode::TOO_MANY_REQUESTS, + "E_THROTTLED proxy stats byte budget is exhausted", + ) + .into_response(); + response.headers_mut().insert( + header::RETRY_AFTER, + axum::http::HeaderValue::from_static("1"), + ); + response +} + +fn stats_timeout() -> Response { + ( + StatusCode::GATEWAY_TIMEOUT, + "E_STATS_TIMEOUT proxy stats collection timed out", + ) + .into_response() +} + #[cfg(test)] mod tests { use super::*; @@ -524,6 +559,39 @@ mod tests { assert_eq!(response.status(), StatusCode::NOT_FOUND); } + #[tokio::test] + async fn kodo_stats_respects_the_node_byte_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: Some(KodoConfig { + ordinal: 0, + cleanup_enabled: false, + cleanup_token: None, + registry_token: None, + }), + }; + let response = Router::new() + .route("/stats", get(kodo_stats)) + .with_state(state) + .oneshot( + Request::builder() + .uri("/stats") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(response.headers().get(header::RETRY_AFTER).unwrap(), "1"); + } + #[tokio::test] async fn rejects_a_request_that_exceeds_the_node_inflight_budget() { let state = ProxyState { diff --git a/crates/proxy/src/kodo.rs b/crates/proxy/src/kodo.rs index 81f7c10..cd1c2c6 100644 --- a/crates/proxy/src/kodo.rs +++ b/crates/proxy/src/kodo.rs @@ -10,7 +10,11 @@ use std::time::{SystemTime, UNIX_EPOCH}; const EXPECTED_BROKERS: usize = 3; const MAX_STATS_BYTES: usize = 16 * 1024 * 1024; +const MAX_STATS_AGGREGATE_BYTES: usize = 32 * 1024 * 1024; +const MAX_STATS_BACKENDS_PER_SHARD: usize = 16; const MAX_REGISTRY_HEAD_BYTES: usize = 64 * 1024; +const MAX_BACKEND_ERROR_BYTES: usize = 64 * 1024; +pub(crate) const STATS_WORKING_SET_BYTES: usize = 3 * MAX_STATS_AGGREGATE_BYTES; #[derive(Clone)] pub(crate) struct KodoConfig { @@ -94,7 +98,10 @@ pub(crate) async fn stats( client: &reqwest::Client, path_and_query: &str, ) -> Response { - let backends = sharded_backends(pool, config.ordinal); + let Some(backends) = sharded_backends(pool, config.ordinal, MAX_STATS_BACKENDS_PER_SHARD) + else { + return unavailable("gateway stats shard has too many brokers"); + }; if backends.is_empty() { return unavailable("gateway stats shard has no broker"); } @@ -106,6 +113,7 @@ pub(crate) async fn stats( topics: Vec::new(), }; let path = force_json_stats(path_and_query); + let mut aggregate_bytes = 0usize; for backend in backends { let response = match client .get(format!("{}{path}", backend.http_origin())) @@ -119,13 +127,25 @@ pub(crate) async fn stats( return unavailable("broker stats are unavailable"); } }; - let stats: StatsResponse = match read_json_bounded(response, MAX_STATS_BYTES).await { + let (stats, response_bytes): (StatsResponse, usize) = match read_json_bounded_with_size( + 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"); } }; + aggregate_bytes = match aggregate_bytes.checked_add(response_bytes) { + Some(total) if total <= MAX_STATS_AGGREGATE_BYTES => total, + _ => { + tracing::warn!("Kodo stats aggregate exceeded its response budget"); + return unavailable("broker stats aggregate is too large"); + } + }; if stats.start_time > 0 { aggregate.start_time = aggregate.start_time.min(stats.start_time); } @@ -243,10 +263,11 @@ async fn delete_from_backend( return Ok(()); } let status = response.status(); - let detail = response - .text() + let detail = read_bytes_bounded(response, MAX_BACKEND_ERROR_BYTES) .await - .unwrap_or_else(|_| "broker cleanup failed".into()); + .ok() + .and_then(|bytes| String::from_utf8(bytes).ok()) + .unwrap_or_else(|| "broker cleanup failed".into()); let outward = if status == StatusCode::CONFLICT { StatusCode::CONFLICT } else if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { @@ -258,26 +279,21 @@ async fn delete_from_backend( } fn complete_broker_set(pool: &BackendPool) -> Option> { - let backends = pool.all(); + let backends = pool.matching_bounded(EXPECTED_BROKERS, |_| true)?; (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) + sharded_backends(pool, ordinal, 1)?.into_iter().next() } -fn sharded_backends(pool: &BackendPool, ordinal: usize) -> Vec { +fn sharded_backends(pool: &BackendPool, ordinal: usize, maximum: usize) -> Option> { if ordinal >= EXPECTED_BROKERS { - return Vec::new(); + return Some(Vec::new()); } - pool.all() - .into_iter() - .filter(|backend| { - backend.node_id.saturating_sub(1) % EXPECTED_BROKERS as u64 == ordinal as u64 - }) - .collect() + pool.matching_bounded(maximum, |backend| { + backend.node_id.saturating_sub(1) % EXPECTED_BROKERS as u64 == ordinal as u64 + }) } fn complete_sharded_backend(pool: &BackendPool, ordinal: usize) -> Option { @@ -385,9 +401,27 @@ fn merge_channel(existing: &mut ChannelStats, mut incoming: ChannelStats) { } async fn read_json_bounded( - mut response: reqwest::Response, + response: reqwest::Response, maximum: usize, ) -> anyhow::Result { + read_json_bounded_with_size(response, maximum) + .await + .map(|(value, _)| value) +} + +async fn read_json_bounded_with_size( + response: reqwest::Response, + maximum: usize, +) -> anyhow::Result<(T, usize)> { + let bytes = read_bytes_bounded(response, maximum).await?; + let length = bytes.len(); + Ok((serde_json::from_slice(&bytes)?, length)) +} + +async fn read_bytes_bounded( + mut response: reqwest::Response, + maximum: usize, +) -> anyhow::Result> { if response .content_length() .is_some_and(|length| length > maximum as u64) @@ -401,7 +435,7 @@ async fn read_json_bounded( } bytes.extend_from_slice(&chunk); } - Ok(serde_json::from_slice(&bytes)?) + Ok(bytes) } fn stable_hash(topic: &str, channel: &str) -> u64 { @@ -521,12 +555,14 @@ mod tests { .collect(), ); assert_eq!( - sharded_backends(&pool, 0) + sharded_backends(&pool, 0, 3) + .unwrap() .into_iter() .map(|backend| backend.node_id) .collect::>(), vec![1, 4, 7] ); + assert!(sharded_backends(&pool, 0, 2).is_none()); assert!(sharded_backend(&pool, 0).is_none()); } @@ -568,4 +604,31 @@ mod tests { .is_ok()); task.abort(); } + + #[tokio::test] + async fn backend_error_body_is_bounded() { + 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( + "/error", + get(|| async { vec![b'x'; MAX_BACKEND_ERROR_BYTES + 1] }), + ), + ) + .await + .unwrap(); + }); + let response = reqwest::get(format!("http://{address}/error")) + .await + .unwrap(); + + assert!(read_bytes_bounded(response, MAX_BACKEND_ERROR_BYTES) + .await + .is_err()); + task.abort(); + } } diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index 8c343ca..185de43 100644 --- a/crates/proxy/src/main.rs +++ b/crates/proxy/src/main.rs @@ -206,11 +206,37 @@ fn validate_proxy_limits(cli: &Cli) -> anyhow::Result<()> { "proxy limits must be non-zero, fit the 100 MiB message and 128 MiB batch contract, and fit the inflight byte budget" ); } + if cli.max_connections > tokio::sync::Semaphore::MAX_PERMITS + || cli.max_inflight_bytes > tokio::sync::Semaphore::MAX_PERMITS + { + anyhow::bail!("proxy limits exceed the runtime semaphore capacity"); + } + let now = std::time::Instant::now(); + let shutdown_timeout = + Duration::from_secs(cli.shutdown_grace_seconds).checked_add(Duration::from_secs(2)); + let timers = [ + Duration::from_millis(cli.http_body_timeout_ms), + Duration::from_millis(cli.tcp_command_timeout_ms), + Duration::from_secs(cli.tcp_max_connection_age_seconds), + ]; + if shutdown_timeout + .and_then(|duration| now.checked_add(duration)) + .is_none() + || timers + .into_iter() + .any(|duration| now.checked_add(duration).is_none()) + { + anyhow::bail!("proxy timeouts exceed the platform timer range"); + } if cli.kodo_compatibility_enabled && (cli.max_message_bytes != rustqueue_protocol::MAX_MESSAGE_BYTES - || cli.max_body_bytes != rustqueue_protocol::MAX_BATCH_BYTES) + || cli.max_body_bytes != rustqueue_protocol::MAX_BATCH_BYTES + || cli.max_inflight_bytes + < tcp::maximum_gateway_working_set(cli.max_message_bytes, cli.max_body_bytes)) { - anyhow::bail!("Kodo compatibility requires the 100 MiB message and 128 MiB batch limits"); + anyhow::bail!( + "Kodo compatibility requires the 100 MiB message and 128 MiB batch limits plus their parsing working set" + ); } Ok(()) } @@ -338,6 +364,17 @@ mod tests { assert!(validate_proxy_limits(&cli).is_err()); } + #[test] + fn proxy_limits_reject_runtime_panics() { + let mut cli = default_cli(); + cli.max_connections = tokio::sync::Semaphore::MAX_PERMITS + 1; + assert!(validate_proxy_limits(&cli).is_err()); + + let mut cli = default_cli(); + cli.shutdown_grace_seconds = u64::MAX; + assert!(validate_proxy_limits(&cli).is_err()); + } + #[test] fn kodo_mode_requires_the_full_hundred_mebibyte_profile() { let mut cli = default_cli(); @@ -348,5 +385,12 @@ mod tests { cli.max_message_bytes -= 1; assert!(validate_proxy_limits(&cli).is_err()); + + 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; + cli.max_inflight_bytes = cli.max_body_bytes; + assert!(validate_proxy_limits(&cli).is_err()); } } diff --git a/crates/proxy/src/tcp.rs b/crates/proxy/src/tcp.rs index 6818148..44cb12e 100644 --- a/crates/proxy/src/tcp.rs +++ b/crates/proxy/src/tcp.rs @@ -4,6 +4,7 @@ mod gateway; use crate::backend::BackendPool; use crate::metrics::ProxyMetrics; use rand::Rng; +use rustqueue_protocol::{Command, MAX_MPUB_MESSAGES}; use std::net::SocketAddr; use std::sync::atomic::Ordering; use std::sync::Arc; @@ -11,6 +12,37 @@ use std::time::Duration; use tokio::sync::{watch, Semaphore}; use tokio::task::JoinSet; +pub(crate) const MAX_CONTROL_BODY_BYTES: usize = 1024 * 1024; +const IDENTIFY_WORKING_SET_COPIES: usize = 3; +const IDENTIFY_FIXED_WORKING_BYTES: usize = 4096; +const MPUB_MESSAGE_WORKING_BYTES: usize = 128; + +fn command_working_set(command: &Command, bytes: usize) -> usize { + match command { + Command::Identify => bytes + .saturating_mul(IDENTIFY_WORKING_SET_COPIES) + .saturating_add(IDENTIFY_FIXED_WORKING_BYTES), + Command::MultiPublish { .. } => { + bytes.saturating_add(MAX_MPUB_MESSAGES.saturating_mul(MPUB_MESSAGE_WORKING_BYTES)) + } + _ => bytes, + } +} + +pub(crate) fn maximum_gateway_working_set( + max_message_bytes: usize, + max_body_bytes: usize, +) -> usize { + command_working_set(&Command::Identify, MAX_CONTROL_BODY_BYTES) + .max(command_working_set( + &Command::MultiPublish { + topic: String::new(), + }, + max_body_bytes, + )) + .max(max_message_bytes) +} + pub struct Limits { pub max_connections: usize, pub max_connection_age: Duration, diff --git a/crates/proxy/src/tcp/gateway.rs b/crates/proxy/src/tcp/gateway.rs index cdbd91b..b11eb0d 100644 --- a/crates/proxy/src/tcp/gateway.rs +++ b/crates/proxy/src/tcp/gateway.rs @@ -4,7 +4,7 @@ use axum::body::Bytes; use rustqueue_protocol::{ encode_frame, Command, FrameType, IdentifyRequest, IdentifyResponse, HEARTBEAT, MAGIC_V2, OK, }; -use serde_json::{json, Value}; +use serde_json::json; use std::collections::BTreeSet; use std::io; use std::sync::atomic::Ordering; @@ -16,7 +16,6 @@ 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); @@ -35,7 +34,7 @@ struct ClientCommand { command: Command, line: Vec, body: Option, - _permit: Option, + permit: Option, } struct BrokerSession { @@ -106,6 +105,7 @@ pub(super) async fn run( let mut heartbeat = Box::pin(tokio::time::sleep(DEFAULT_HEARTBEAT)); let mut backend_identify = default_backend_identify(); let mut backend_session = None; + let mut _identify_permit = None; loop { tokio::select! { @@ -115,7 +115,7 @@ pub(super) async fn run( let Some(command) = command else { break; }; - let command = match command { + let mut command = match command { Ok(command) => command, Err(ReadError::RetryByReconnect(detail)) => { tracing::debug!( @@ -165,14 +165,16 @@ pub(super) async fn run( } None => Some(DEFAULT_HEARTBEAT), }; - backend_identify = backend_identify_body(body)?; + let msg_timeout = request.msg_timeout.unwrap_or(60_000).max(1_000); + backend_identify = backend_identify_body(request)?; identified = true; + _identify_permit = command.permit.take(); 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), + msg_timeout, tls_v1: false, deflate: false, deflate_level: 0, @@ -425,17 +427,13 @@ fn retriable_error(body: &[u8]) -> bool { || 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 backend_identify_body(mut request: IdentifyRequest) -> anyhow::Result> { + request.feature_negotiation = true; + request.heartbeat_interval = Some(-1); + request.tls_v1 = false; + request.snappy = false; + request.deflate = false; + Ok(serde_json::to_vec(&request)?) } fn default_backend_identify() -> Vec { @@ -508,7 +506,7 @@ where 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::Identify | Command::Auth => Some(("E_BAD_BODY", super::MAX_CONTROL_BODY_BYTES)), Command::Publish { .. } | Command::DeferredPublish { .. } => { Some(("E_BAD_MESSAGE", max_message_bytes)) } @@ -520,7 +518,7 @@ where command, line, body: None, - _permit: None, + permit: None, }); }; let length = reader.read_u32().await? as usize; @@ -530,8 +528,11 @@ where format!("command body size {length} is outside 1..={maximum}"), )); } + let working_set = super::command_working_set(&command, length); + let permits = u32::try_from(working_set) + .map_err(|_| ReadError::RetryByReconnect("publish Gateway byte budget is exhausted"))?; let permit = Arc::clone(inflight_bytes) - .try_acquire_many_owned(length as u32) + .try_acquire_many_owned(permits) .map_err(|_| ReadError::RetryByReconnect("publish Gateway byte budget is exhausted"))?; let mut body = vec![0; length]; reader.read_exact(&mut body).await?; @@ -544,7 +545,7 @@ where command, line, body: Some(body), - _permit: Some(permit), + permit: Some(permit), }) } @@ -730,7 +731,7 @@ mod tests { }, line: b"PUB events\n".to_vec(), body: Some(Bytes::from_static(b"payload")), - _permit: None, + permit: None, }; let mut current = None; assert_eq!( @@ -765,7 +766,7 @@ mod tests { }, line: b"PUB events\n".to_vec(), body: Some(Bytes::from_static(b"payload")), - _permit: None, + permit: None, }; let mut current = None; assert_eq!( @@ -851,6 +852,48 @@ mod tests { gateway_task.await.unwrap(); } + #[tokio::test] + async fn identify_budget_is_retained_until_the_client_disconnects() { + let identify = b"{}"; + let working_set = super::super::command_working_set(&Command::Identify, identify.len()); + let budget = Arc::new(Semaphore::new(working_set)); + let gateway_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let gateway_address = gateway_listener.local_addr().unwrap(); + let task_budget = Arc::clone(&budget); + let gateway_task = tokio::spawn(async move { + let (client, _) = gateway_listener.accept().await.unwrap(); + let (_shutdown_tx, shutdown) = watch::channel(false); + let _ = run( + client, + BackendPool::default(), + Limits { + max_message_bytes: 1024, + max_body_bytes: 1024, + command_timeout: Duration::from_secs(1), + inflight_bytes: task_budget, + }, + ProxyMetrics::default(), + shutdown, + ) + .await; + }); + + let mut client = TcpStream::connect(gateway_address).await.unwrap(); + client.write_all(MAGIC_V2).await.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); + assert_eq!(budget.available_permits(), 0); + + drop(client); + tokio::time::timeout(Duration::from_secs(1), gateway_task) + .await + .unwrap() + .unwrap(); + assert_eq!(budget.available_permits(), working_set); + } + #[tokio::test] async fn single_message_limit_is_independent_from_the_batch_body_limit() { let (mut client, mut gateway) = tokio::io::duplex(1024); @@ -876,13 +919,41 @@ mod tests { 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 working_set = super::super::command_working_set( + &Command::MultiPublish { + topic: "events".into(), + }, + body.len(), + ); + let budget = Arc::new(Semaphore::new(working_set)); 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); + assert_eq!(budget.available_permits(), working_set); + } + + #[tokio::test] + async fn mpub_metadata_is_admitted_before_parsing_a_malformed_count() { + let body = (rustqueue_protocol::MAX_MPUB_MESSAGES as u32).to_be_bytes(); + let command = Command::MultiPublish { + topic: "events".into(), + }; + let working_set = super::super::command_working_set(&command, body.len()); + 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(working_set - 1)); + + let error = + match read_command(&mut gateway, 1024, 1024, Duration::from_secs(1), &budget).await { + Ok(_) => panic!("malformed MPUB bypassed its metadata budget"), + Err(error) => error, + }; + assert!(matches!(error, ReadError::RetryByReconnect(_))); + assert_eq!(budget.available_permits(), working_set - 1); } #[tokio::test] @@ -894,7 +965,7 @@ mod tests { }, line: b"PUB events\n".to_vec(), body: Some(Bytes::from(vec![0x5a; 4096])), - _permit: None, + permit: None, }; let error = send_before_commit(&mut gateway, &command, Duration::from_millis(10)) .await @@ -926,7 +997,7 @@ mod tests { }, line: b"PUB events\n".to_vec(), body: Some(Bytes::from_static(b"payload")), - _permit: None, + permit: None, }; let mut current = Some(first); publish(&pool, &identify, &command, &mut current, &metrics) @@ -989,7 +1060,7 @@ mod tests { max_message_bytes: 1024, max_body_bytes: 1024, command_timeout: Duration::from_secs(1), - inflight_bytes: Arc::new(Semaphore::new(1024)), + inflight_bytes: Arc::new(Semaphore::new(8192)), }, ProxyMetrics::default(), shutdown, diff --git a/crates/queue/src/broker.rs b/crates/queue/src/broker.rs index ccee8bd..951f970 100644 --- a/crates/queue/src/broker.rs +++ b/crates/queue/src/broker.rs @@ -35,7 +35,7 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use thiserror::Error; const FEATURE_LEVEL_2_MAX_MESSAGE_BYTES: usize = 100 * 1024 * 1024; @@ -115,6 +115,8 @@ pub enum BrokerError { RevisionConflict { expected: u64, actual: u64 }, #[error("management operation conflicts with an existing operation ID")] OperationConflict, + #[error("active tombstone deadline is required")] + InvalidTombstone, #[error("message does not exist")] MessageNotFound, #[error("message is not in flight")] @@ -165,9 +167,10 @@ struct BrokerInner { fences: Mutex, management_ops_path: PathBuf, management_ops: Mutex, + outbox_moves: tokio::sync::Mutex<()>, management_fences_ready: AtomicBool, registry_revision: AtomicU64, - storage_healthy: AtomicBool, + storage_healthy: Arc, gc_cursor: AtomicUsize, publish_groups: group_commit::PublishGroups, channel_groups: channel_commit::ChannelGroups, @@ -195,6 +198,14 @@ impl Broker { "topic and publish worker limits must be greater than zero".into(), )); } + let now = Instant::now(); + if now.checked_add(config.message_timeout).is_none() + || now.checked_add(config.publish_worker_idle).is_none() + { + return Err(BrokerError::InvalidRecord( + "broker timeouts exceed the platform timer range".into(), + )); + } let readable_message_bytes = if config.storage_feature_level >= 2 { FEATURE_LEVEL_2_MAX_MESSAGE_BYTES } else { @@ -239,6 +250,7 @@ impl Broker { } }; let metrics = QueueMetrics::default(); + let storage_healthy = Arc::new(AtomicBool::new(true)); let (fences_path, fences) = FenceCatalog::load(&config.data_path)?; let (management_ops_path, management_ops) = OperationCatalog::load(&config.data_path)?; let require_fence_sync = config.require_management_fence_sync; @@ -247,11 +259,13 @@ impl Broker { config.payload_read_workers, config.payload_read_queue, Arc::clone(&metrics.payload_read), + Arc::clone(&storage_healthy), ); let message_index_cache = MessageIndexCache::new( config.message_index_cache_bytes, config.payload_read_workers, config.payload_read_queue, + Arc::clone(&storage_healthy), ); let mut topics = HashMap::new(); for entry in std::fs::read_dir(&topics_root)? { @@ -326,9 +340,10 @@ impl Broker { fences: Mutex::new(fences), management_ops_path, management_ops: Mutex::new(management_ops), - management_fences_ready: AtomicBool::new(true), + outbox_moves: tokio::sync::Mutex::new(()), + management_fences_ready: AtomicBool::new(!require_fence_sync), registry_revision: AtomicU64::new(revision), - storage_healthy: AtomicBool::new(true), + storage_healthy, gc_cursor: AtomicUsize::new(0), publish_groups, channel_groups, @@ -338,12 +353,6 @@ impl Broker { }; broker.reserve_message_metadata(0)?; broker.recover_outbox()?; - if require_fence_sync { - broker - .inner - .management_fences_ready - .store(false, Ordering::Release); - } Ok(broker) } @@ -352,10 +361,15 @@ impl Broker { } pub async fn create_topic(&self, name: &str) -> Result<(), BrokerError> { + validate_name(name).map_err(|_| BrokerError::InvalidTopic)?; let broker = self.clone(); let name = name.to_owned(); - self.storage_task(move || broker.get_or_create_topic(&name).map(|_| ())) - .await + self.storage_task(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&name, None)?; + broker.get_or_create_topic_locked(&name).map(|_| ()) + }) + .await } pub async fn delete_topic(&self, name: &str) -> Result<(), BrokerError> { @@ -365,6 +379,7 @@ impl Broker { let name = name.to_owned(); self.storage_task(move || { let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&name, None)?; if broker.delete_topic_locked(&name)? { broker.bump_registry()?; } @@ -374,13 +389,16 @@ impl Broker { } pub async fn create_channel(&self, topic: &str, channel: &str) -> Result<(), BrokerError> { + validate_name(topic).map_err(|_| BrokerError::InvalidTopic)?; validate_channel(channel)?; self.ensure_management_access(topic, Some(channel))?; let broker = self.clone(); let topic = topic.to_owned(); let channel = channel.to_owned(); self.storage_task(move || { - let handle = broker.get_or_create_topic(&topic)?; + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&topic, Some(&channel))?; + let handle = broker.get_or_create_topic_locked(&topic)?; if handle .state .lock() @@ -400,6 +418,8 @@ impl Broker { let topic = topic.to_owned(); let channel = channel.to_owned(); self.storage_task(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&topic, Some(&channel))?; broker .topic(&topic)? .state @@ -415,8 +435,12 @@ impl Broker { self.ensure_management_access(topic, None)?; let broker = self.clone(); let topic = topic.to_owned(); - self.storage_task(move || broker.topic(&topic)?.state.lock().set_paused(paused)) - .await + self.storage_task(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&topic, None)?; + broker.topic(&topic)?.state.lock().set_paused(paused) + }) + .await } pub async fn set_channel_paused( @@ -430,6 +454,8 @@ impl Broker { let topic = topic.to_owned(); let channel = channel.to_owned(); self.storage_task(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&topic, Some(&channel))?; broker .topic(&topic)? .state @@ -443,8 +469,12 @@ impl Broker { self.ensure_management_access(topic, None)?; let broker = self.clone(); let topic = topic.to_owned(); - self.storage_task(move || broker.topic(&topic)?.state.lock().empty_topic()) - .await + self.storage_task(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&topic, None)?; + broker.topic(&topic)?.state.lock().empty_topic() + }) + .await } pub async fn empty_channel(&self, topic: &str, channel: &str) -> Result<(), BrokerError> { @@ -452,8 +482,12 @@ impl Broker { let broker = self.clone(); let topic = topic.to_owned(); let channel = channel.to_owned(); - self.storage_task(move || broker.topic(&topic)?.state.lock().empty_channel(&channel)) - .await + self.storage_task(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.ensure_management_access(&topic, Some(&channel))?; + broker.topic(&topic)?.state.lock().empty_channel(&channel) + }) + .await } pub fn topic_names(&self) -> Vec { @@ -509,7 +543,13 @@ impl Broker { task: impl FnOnce() -> Result + Send + 'static, ) -> Result { self.ensure_storage_healthy()?; - let result = blocking(task).await; + let observer = self.clone(); + let result = blocking(move || { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(task)) + .unwrap_or(Err(BrokerError::StorageUnavailable)); + observer.observe_storage_result(result) + }) + .await; self.observe_storage_result(result) } diff --git a/crates/queue/src/broker/channel_commit.rs b/crates/queue/src/broker/channel_commit.rs index e340d5d..2fa4e47 100644 --- a/crates/queue/src/broker/channel_commit.rs +++ b/crates/queue/src/broker/channel_commit.rs @@ -13,8 +13,16 @@ const MAX_GROUP_REQUESTS: usize = 64; const COALESCE_DELAY: Duration = Duration::from_millis(1); pub(super) enum ChannelOperation { - Finish { id: u64, require_in_flight: bool }, - Requeue { id: u64, available_at_ms: i64 }, + Finish { + id: u64, + require_in_flight: bool, + token: Option, + }, + Requeue { + id: u64, + available_at_ms: i64, + token: Option, + }, } pub(super) struct ChannelGroups { @@ -187,11 +195,13 @@ impl Broker { ChannelOperation::Finish { id, require_in_flight, - } => topic_state.finish_buffered(&channel, id, require_in_flight), + token, + } => topic_state.finish_buffered(&channel, id, require_in_flight, token), ChannelOperation::Requeue { id, available_at_ms, - } => topic_state.requeue_buffered(&channel, id, available_at_ms), + token, + } => topic_state.requeue_buffered(&channel, id, available_at_ms, token), }; match result { Ok(()) => { @@ -333,6 +343,7 @@ mod tests { operation: ChannelOperation::Finish { id: 1, require_in_flight: true, + token: None, }, enqueued_at: Instant::now(), reply, diff --git a/crates/queue/src/broker/delivery.rs b/crates/queue/src/broker/delivery.rs index fdb91d8..1491d75 100644 --- a/crates/queue/src/broker/delivery.rs +++ b/crates/queue/src/broker/delivery.rs @@ -40,6 +40,7 @@ impl Broker { .min(self.inner.delivery_budget.max_payload_bytes()); let mut batch = self .reserve_deliveries( + topic, Arc::clone(&handle), channel, max_messages.clamp(1, 64), @@ -51,6 +52,7 @@ impl Broker { let _ = tokio::time::timeout(wait.min(Duration::from_secs(1)), wake.changed()).await; batch = self .reserve_deliveries( + topic, Arc::clone(&handle), channel, max_messages.clamp(1, 64), @@ -65,13 +67,19 @@ impl Broker { let bytes = batch.payload_bytes(); let hold = self.inner.delivery_budget.acquire(bytes).await?; let lease = self.inner.payload_reader.retain(batch.payloads()); - let bodies = match self.inner.payload_reader.read_retained(lease).await { - Ok(bodies) => bodies, + let (bodies, hold) = match self + .inner + .payload_reader + .read_retained(lease, Some(hold)) + .await + { + Ok(result) => result, Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { return Ok(DeliveryBatch::new(Vec::new(), DeliveryGuard::empty())); } Err(error) => return self.observe_storage_result(Err(error.into())), }; + let hold = hold.expect("delivery payload read returns its byte-budget hold"); let handle = Arc::clone(&batch.handle); let channel = batch.channel.clone(); let reservations = batch.disarm(); @@ -91,6 +99,7 @@ impl Broker { async fn reserve_deliveries( &self, + topic: &str, handle: Arc, channel: &str, max_messages: usize, @@ -105,8 +114,9 @@ impl Broker { } let handle = Arc::clone(&batch.handle); let action = { - let mut topic = handle.state.lock(); - let action = topic.reserve_batch( + let mut topic_state = handle.state.lock(); + self.ensure_management_access(topic, Some(channel))?; + let action = topic_state.reserve_batch( channel, remaining_messages, max_bytes.saturating_sub(batch.payload_bytes()).max(1), @@ -153,7 +163,16 @@ impl Broker { if batch.payload_bytes() >= max_bytes || batch.len() >= max_messages { break; } - if let Err(error) = self.inner.message_index_cache.load(request).await { + let read_lease = self + .inner + .payload_reader + .retain_paths(vec![request.segment_path().to_path_buf()]); + if let Err(error) = self + .inner + .message_index_cache + .load(request, read_lease) + .await + { if matches!(&error, BrokerError::Io(io) if io.kind() == std::io::ErrorKind::WouldBlock) { return Ok(ReservedBatch::new(Arc::clone(&batch.handle), channel)); diff --git a/crates/queue/src/broker/group_commit.rs b/crates/queue/src/broker/group_commit.rs index 1a7bb1d..981941e 100644 --- a/crates/queue/src/broker/group_commit.rs +++ b/crates/queue/src/broker/group_commit.rs @@ -430,6 +430,7 @@ pub(super) fn copy_error(error: &BrokerError) -> BrokerError { actual: *actual, }, BrokerError::OperationConflict => BrokerError::OperationConflict, + BrokerError::InvalidTombstone => BrokerError::InvalidTombstone, BrokerError::MessageNotFound => BrokerError::MessageNotFound, BrokerError::MessageNotInFlight => BrokerError::MessageNotInFlight, BrokerError::MessageTooLarge => BrokerError::MessageTooLarge, diff --git a/crates/queue/src/broker/io.rs b/crates/queue/src/broker/io.rs index 6ab8a32..0ec08b2 100644 --- a/crates/queue/src/broker/io.rs +++ b/crates/queue/src/broker/io.rs @@ -11,6 +11,27 @@ pub(super) const SEQUENCE_RESERVATION: u64 = 1 << 20; impl Broker { pub async fn finish(&self, topic: &str, channel: &str, id: u64) -> Result<(), BrokerError> { + self.finish_with_token(topic, channel, id, None).await + } + + pub async fn finish_delivery( + &self, + topic: &str, + channel: &str, + id: u64, + token: u64, + ) -> Result<(), BrokerError> { + self.finish_with_token(topic, channel, id, Some(token)) + .await + } + + async fn finish_with_token( + &self, + topic: &str, + channel: &str, + id: u64, + token: Option, + ) -> Result<(), BrokerError> { let _timer = self.inner.metrics.channel_ack.timer(); self.ensure_storage_healthy()?; self.ensure_management_access(topic, Some(channel))?; @@ -23,30 +44,52 @@ impl Broker { super::channel_commit::ChannelOperation::Finish { id, require_in_flight: true, + token, }, ) .await } - async fn finish_inner( + async fn finish_outbox_source( &self, topic: &str, channel: &str, id: u64, - require_in_flight: bool, ) -> Result<(), BrokerError> { - self.ensure_management_access(topic, Some(channel))?; let broker = self.clone(); let topic = topic.to_owned(); let channel = channel.to_owned(); - self.storage_task(move || { - broker - .topic(&topic)? - .state - .lock() - .finish(&channel, id, require_in_flight) - }) - .await + let result = self + .storage_task(move || { + broker + .topic(&topic)? + .state + .lock() + .finish(&channel, id, false) + }) + .await; + match result { + Ok(()) + | Err( + BrokerError::TopicNotFound + | BrokerError::ChannelNotFound + | BrokerError::MessageNotFound, + ) => Ok(()), + Err(error) => Err(error), + } + } + + fn source_message_is_unacknowledged( + &self, + topic: &str, + channel: &str, + id: u64, + ) -> Result, BrokerError> { + match self.topic(topic) { + Ok(topic) => topic.state.lock().message_is_unacknowledged(channel, id), + Err(BrokerError::TopicNotFound) => Ok(None), + Err(error) => Err(error), + } } pub async fn requeue( @@ -55,6 +98,30 @@ impl Broker { channel: &str, id: u64, delay: Duration, + ) -> Result<(), BrokerError> { + self.requeue_with_token(topic, channel, id, None, delay) + .await + } + + pub async fn requeue_delivery( + &self, + topic: &str, + channel: &str, + id: u64, + token: u64, + delay: Duration, + ) -> Result<(), BrokerError> { + self.requeue_with_token(topic, channel, id, Some(token), delay) + .await + } + + async fn requeue_with_token( + &self, + topic: &str, + channel: &str, + id: u64, + token: Option, + delay: Duration, ) -> Result<(), BrokerError> { let _timer = self.inner.metrics.channel_ack.timer(); self.ensure_storage_healthy()?; @@ -70,6 +137,7 @@ impl Broker { super::channel_commit::ChannelOperation::Requeue { id, available_at_ms: available, + token, }, ) .await; @@ -85,13 +153,56 @@ impl Broker { ) -> Result<(), BrokerError> { self.ensure_storage_healthy()?; self.ensure_management_access(topic, Some(channel))?; - self.topic(topic)?.state.lock().touch( + let handle = self.topic(topic)?; + let mut topic_state = handle.state.lock(); + self.ensure_management_access(topic, Some(channel))?; + topic_state.touch( channel, id, timeout.unwrap_or(self.inner.config.message_timeout), ) } + pub fn touch_delivery( + &self, + topic: &str, + channel: &str, + id: u64, + token: u64, + timeout: Option, + ) -> Result<(), BrokerError> { + self.ensure_storage_healthy()?; + self.ensure_management_access(topic, Some(channel))?; + let handle = self.topic(topic)?; + let mut topic_state = handle.state.lock(); + self.ensure_management_access(topic, Some(channel))?; + topic_state.touch_with_token( + channel, + id, + Some(token), + timeout.unwrap_or(self.inner.config.message_timeout), + ) + } + + pub fn touch_deliveries( + &self, + topic: &str, + channel: &str, + deliveries: &[(u64, u64)], + timeout: Option, + ) -> Result<(), BrokerError> { + self.ensure_storage_healthy()?; + self.ensure_management_access(topic, Some(channel))?; + let handle = self.topic(topic)?; + let mut topic_state = handle.state.lock(); + self.ensure_management_access(topic, Some(channel))?; + topic_state.touch_deliveries_with_tokens( + channel, + deliveries, + timeout.unwrap_or(self.inner.config.message_timeout), + ) + } + pub fn release(&self, topic: &str, channel: &str, ids: &[u64]) { if let Ok(handle) = self.topic(topic) { handle.state.lock().release(channel, ids); @@ -99,6 +210,13 @@ impl Broker { } } + pub fn release_deliveries(&self, topic: &str, channel: &str, deliveries: &[(u64, u64)]) { + if let Ok(handle) = self.topic(topic) { + handle.state.lock().release_with_tokens(channel, deliveries); + handle.signal(); + } + } + pub async fn move_to_dead_letter( &self, source_topic: &str, @@ -106,8 +224,46 @@ impl Broker { message_id: u64, target_topic: &str, body: Bytes, - ) -> Result<(), BrokerError> { + ) -> Result { self.ensure_storage_healthy()?; + let broker = self.clone(); + let source_topic = source_topic.to_owned(); + let source_channel = source_channel.to_owned(); + let target_topic = target_topic.to_owned(); + let task = tokio::spawn(async move { + let move_guard = broker.inner.outbox_moves.lock().await; + let result = broker + .move_to_dead_letter_transaction( + &source_topic, + &source_channel, + message_id, + &target_topic, + body, + ) + .await; + drop(move_guard); + result + }); + match task.await { + Ok(result) => result, + Err(_) => self.observe_storage_result(Err(BrokerError::StorageUnavailable)), + } + } + + async fn move_to_dead_letter_transaction( + &self, + source_topic: &str, + source_channel: &str, + message_id: u64, + target_topic: &str, + body: Bytes, + ) -> Result { + if !self + .source_message_is_unacknowledged(source_topic, source_channel, message_id)? + .unwrap_or(false) + { + return Ok(false); + } let entry = OutboxEntry { source_topic: source_topic.into(), source_channel: source_channel.into(), @@ -127,15 +283,11 @@ impl Broker { broker.publish_durable_body_sync(&target_topic, &[body], Duration::ZERO) }) .await?; - self.finish_inner( - &entry.source_topic, - &entry.source_channel, - entry.message_id, - false, - ) - .await?; + self.finish_outbox_source(&entry.source_topic, &entry.source_channel, entry.message_id) + .await?; self.storage_task(move || crate::outbox::remove(&path)) - .await + .await?; + Ok(true) } fn publish_durable_body_sync( @@ -148,6 +300,7 @@ impl Broker { let mut metadata = self.reserve_message_metadata(bodies.len())?; let handle = self.get_or_create_topic(topic)?; let mut state = handle.state.lock(); + self.ensure_management_access(topic, None)?; let ids = self.append_publish_to_topic(&mut state, bodies, delay, true, &mut metadata)?; if self.inner.message_index_cache.over_budget() { state.spill_message_metadata()?; @@ -240,7 +393,32 @@ 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_durable_body_sync(&entry.target_topic, &[entry.body], Duration::ZERO)?; + if self.source_message_is_unacknowledged( + &entry.source_topic, + &entry.source_channel, + entry.message_id, + )? == Some(false) + { + crate::outbox::remove(&path)?; + continue; + } + if let Err(error) = + self.publish_durable_body_sync(&entry.target_topic, &[entry.body], Duration::ZERO) + { + if matches!( + error, + BrokerError::ManagementUnavailable + | BrokerError::TopicRetiring + | BrokerError::TopicTombstoned + | BrokerError::TopicLimit + ) { + // Keep both the durable intent and its source payload. A + // consumer retry can complete the move after the + // management fence or capacity constraint is cleared. + continue; + } + return Err(error); + } let finish = self.topic(&entry.source_topic).and_then(|topic| { topic .state @@ -261,3 +439,64 @@ impl Broker { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[tokio::test] + async fn durable_internal_publish_rechecks_fences_after_acquiring_topic() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_topic("dead-letter").await.unwrap(); + + let topic = broker.topic("dead-letter").unwrap(); + let (locked_tx, locked_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let lock_holder = { + let topic = Arc::clone(&topic); + std::thread::spawn(move || { + let _topic_state = topic.state.lock(); + locked_tx.send(()).unwrap(); + let _ = release_rx.recv(); + }) + }; + locked_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let baseline_references = Arc::strong_count(&topic); + let publisher = { + let broker = broker.clone(); + std::thread::spawn(move || { + broker.publish_durable_body_sync( + "dead-letter", + &[Bytes::from_static(b"poison")], + Duration::ZERO, + ) + }) + }; + tokio::time::timeout(Duration::from_secs(1), async { + while Arc::strong_count(&topic) <= baseline_references { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + { + let mut fences = broker.inner.fences.lock(); + fences.set_topic("dead-letter", now_ms() + 60_000); + fences.store(&broker.inner.fences_path).unwrap(); + } + release_tx.send(()).unwrap(); + lock_holder.join().unwrap(); + + assert!(matches!( + publisher.join().unwrap(), + Err(BrokerError::TopicTombstoned) + )); + assert_eq!(broker.stats().topics[0].message_count, 0); + } +} diff --git a/crates/queue/src/broker/maintenance.rs b/crates/queue/src/broker/maintenance.rs index 2719660..d911e8e 100644 --- a/crates/queue/src/broker/maintenance.rs +++ b/crates/queue/src/broker/maintenance.rs @@ -15,6 +15,10 @@ impl Broker { let _timer = self.inner.metrics.gc.timer(); let broker = self.clone(); self.storage_task(move || { + { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + broker.cleanup_drained_retired_topics()?; + } let mut outbox_ids = HashMap::>::new(); for (source_topic, message_id) in crate::outbox::retained_sources(&broker.inner.config.data_path.join("dlq-outbox"))? diff --git a/crates/queue/src/broker/management.rs b/crates/queue/src/broker/management.rs index 420cd22..f890ebe 100644 --- a/crates/queue/src/broker/management.rs +++ b/crates/queue/src/broker/management.rs @@ -16,6 +16,7 @@ impl Broker { ) -> Result<(), BrokerError> { let broker = self.clone(); self.storage_task(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); let mut fences = broker.inner.fences.lock(); fences.replace(snapshot); fences.store(&broker.inner.fences_path)?; @@ -44,15 +45,47 @@ 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, operation_id) = 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 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), + } + }; + if pending_operation.is_none() + && matches!( + action, + TopicManagementAction::Delete | TopicManagementAction::Tombstone + ) + { + valid_tombstone(tombstone_until_ms, false)?; + } + if pending_operation.is_none() + && matches!( + action, + TopicManagementAction::Pause + | TopicManagementAction::Unpause + | TopicManagementAction::Empty + ) + { + broker.topic(&topic)?; + } + 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 { @@ -140,6 +173,33 @@ impl Broker { OperationLookup::New => operations.pending_id(&topic, &fingerprint), } }; + if pending_operation.is_none() + && matches!( + action, + ChannelManagementAction::Delete | ChannelManagementAction::Tombstone + ) + { + valid_tombstone(tombstone_until_ms, false)?; + } + if pending_operation.is_none() + && matches!( + action, + ChannelManagementAction::Pause + | ChannelManagementAction::Unpause + | ChannelManagementAction::Empty + ) + { + let handle = broker.topic(&topic)?; + if !handle + .state + .lock() + .channel_names() + .iter() + .any(|name| name == &channel) + { + return Err(BrokerError::ChannelNotFound); + } + } let idle_handle = if action == ChannelManagementAction::Delete && require_idle { match broker.topic(&topic) { Ok(handle) => Some(handle), @@ -315,6 +375,9 @@ impl Broker { if let Some(pending_id) = operations.pending_id(topic, fingerprint) { return Ok(PreparedOperation::Pending(pending_id)); } + if operations.blocks_topic(topic) { + return Err(BrokerError::OperationConflict); + } self.check_revision(expected_revision)?; operations .prepare( @@ -352,7 +415,7 @@ enum PreparedOperation { fn valid_tombstone(value: Option, replayed: bool) -> Result { value .filter(|until| replayed || *until > now_ms()) - .ok_or_else(|| BrokerError::InvalidRecord("active tombstone deadline is required".into())) + .ok_or(BrokerError::InvalidTombstone) } fn set_channel_fence( diff --git a/crates/queue/src/broker/topics.rs b/crates/queue/src/broker/topics.rs index eff6844..8f3b82e 100644 --- a/crates/queue/src/broker/topics.rs +++ b/crates/queue/src/broker/topics.rs @@ -61,6 +61,25 @@ impl Broker { Ok(()) } + pub(super) fn cleanup_drained_retired_topics(&self) -> Result { + let ready: Vec<_> = self + .inner + .retired_topics + .lock() + .iter() + .filter_map(|(name, handle)| { + let directory = topic_directory(&self.inner.config.data_path, name); + (Arc::strong_count(handle) == 1 + && !self.inner.payload_reader.has_active_under(&directory)) + .then(|| name.clone()) + }) + .collect(); + for name in &ready { + self.cleanup_retired_topic(name)?; + } + Ok(ready.len()) + } + pub(super) fn delete_topic_locked(&self, name: &str) -> Result { let Some(handle) = self.inner.topics.read().get(name).cloned() else { return Ok(false); diff --git a/crates/queue/src/broker_tests.rs b/crates/queue/src/broker_tests.rs index 5a5629e..68caeaf 100644 --- a/crates/queue/src/broker_tests.rs +++ b/crates/queue/src/broker_tests.rs @@ -130,6 +130,335 @@ async fn startup_replays_dlq_outbox_before_finishing_the_source() { assert_eq!(&*dlq.body, b"poison"); } +#[tokio::test] +async fn blocked_dlq_target_does_not_prevent_broker_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 id = broker + .publish("events", vec![b"poison".to_vec()], Duration::ZERO) + .await + .unwrap()[0]; + let target = "events.workers.DLQ"; + crate::outbox::store( + &root.path().join("dlq-outbox"), + &OutboxEntry { + source_topic: "events".into(), + source_channel: "workers".into(), + message_id: id, + target_topic: target.into(), + body: bytes::Bytes::from_static(b"poison"), + }, + ) + .unwrap(); + broker + .sync_management_fences(ManagementFenceSnapshot { + revision: "blocked-dlq-target".into(), + topics: BTreeMap::from([(target.into(), now_ms() + 60_000)]), + channels: Vec::new(), + }) + .await + .unwrap(); + drop(broker); + + let broker = Broker::open(config).unwrap(); + assert_eq!( + std::fs::read_dir(root.path().join("dlq-outbox")) + .unwrap() + .count(), + 1 + ); + let stats = broker.stats(); + let source = stats + .topics + .iter() + .find(|topic| topic.name == "events") + .unwrap(); + assert_eq!(source.channels[0].depth, 1); + assert!(!stats.topics.iter().any(|topic| topic.name == target)); +} + +#[tokio::test] +async fn startup_does_not_replay_outbox_before_required_fence_sync() { + let root = tempdir().unwrap(); + let config = BrokerConfig { + data_path: root.path().into(), + require_management_fence_sync: true, + ..BrokerConfig::default() + }; + let broker = Broker::open(config.clone()).unwrap(); + broker + .sync_management_fences(ManagementFenceSnapshot::default()) + .await + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + let id = broker + .publish("events", vec![b"poison".to_vec()], Duration::ZERO) + .await + .unwrap()[0]; + let target = "events.workers.DLQ"; + crate::outbox::store( + &root.path().join("dlq-outbox"), + &OutboxEntry { + source_topic: "events".into(), + source_channel: "workers".into(), + message_id: id, + target_topic: target.into(), + body: bytes::Bytes::from_static(b"poison"), + }, + ) + .unwrap(); + drop(broker); + + let broker = Broker::open(config).unwrap(); + assert!(!broker.management_fences_ready()); + assert_eq!( + std::fs::read_dir(root.path().join("dlq-outbox")) + .unwrap() + .count(), + 1 + ); + let stats = broker.stats(); + let source = stats + .topics + .iter() + .find(|topic| topic.name == "events") + .unwrap(); + assert_eq!(source.channels[0].depth, 1); + assert!(!stats.topics.iter().any(|topic| topic.name == target)); +} + +#[tokio::test] +async fn concurrent_dlq_moves_publish_only_one_copy() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + let id = broker + .publish("events", vec![b"poison".to_vec()], Duration::ZERO) + .await + .unwrap()[0]; + + let first = broker.move_to_dead_letter( + "events", + "workers", + id, + "events.workers.DLQ", + bytes::Bytes::from_static(b"poison"), + ); + let second = broker.move_to_dead_letter( + "events", + "workers", + id, + "events.workers.DLQ", + bytes::Bytes::from_static(b"poison"), + ); + let (first, second) = tokio::join!(first, second); + assert_eq!( + [first.unwrap(), second.unwrap()] + .into_iter() + .filter(|moved| *moved) + .count(), + 1 + ); + + let stats = broker.stats(); + let source = stats + .topics + .iter() + .find(|topic| topic.name == "events") + .unwrap(); + assert_eq!(source.channels[0].depth, 0); + let target = stats + .topics + .iter() + .find(|topic| topic.name == "events.workers.DLQ") + .unwrap(); + assert_eq!(target.message_count, 1); +} + +#[tokio::test] +async fn cancelled_dlq_move_keeps_the_transaction_serialized_until_completion() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + let id = broker + .publish("events", vec![b"poison".to_vec()], Duration::ZERO) + .await + .unwrap()[0]; + + let blocker = broker.inner.outbox_moves.lock().await; + let mut moving = Box::pin(broker.move_to_dead_letter( + "events", + "workers", + id, + "events.workers.DLQ", + bytes::Bytes::from_static(b"poison"), + )); + std::future::poll_fn(|context| { + assert!( + std::future::Future::poll(moving.as_mut(), context).is_pending(), + "the blocked move must not complete" + ); + std::task::Poll::Ready(()) + }) + .await; + drop(moving); + drop(blocker); + + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let stats = broker.stats(); + let source_depth = stats + .topics + .iter() + .find(|topic| topic.name == "events") + .map(|topic| topic.channels[0].depth); + let target_count = stats + .topics + .iter() + .find(|topic| topic.name == "events.workers.DLQ") + .map(|topic| topic.message_count); + let outbox_empty = std::fs::read_dir(root.path().join("dlq-outbox")) + .is_ok_and(|entries| entries.count() == 0); + if source_depth == Some(0) && target_count == Some(1) && outbox_empty { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + assert!(!broker + .move_to_dead_letter( + "events", + "workers", + id, + "events.workers.DLQ", + bytes::Bytes::from_static(b"poison"), + ) + .await + .unwrap()); + assert_eq!( + broker + .stats() + .topics + .iter() + .find(|topic| topic.name == "events.workers.DLQ") + .unwrap() + .message_count, + 1 + ); +} + +#[tokio::test] +async fn completed_dlq_outbox_is_not_published_again_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 id = broker + .publish("events", vec![b"poison".to_vec()], Duration::ZERO) + .await + .unwrap()[0]; + let target = "events.workers.DLQ"; + let outbox_path = crate::outbox::store( + &root.path().join("dlq-outbox"), + &OutboxEntry { + source_topic: "events".into(), + source_channel: "workers".into(), + message_id: id, + target_topic: target.into(), + body: bytes::Bytes::from_static(b"poison"), + }, + ) + .unwrap(); + broker + .publish(target, vec![b"poison".to_vec()], Duration::ZERO) + .await + .unwrap(); + let delivery = broker + .next_message("events", "workers", None) + .await + .unwrap() + .unwrap(); + broker + .finish("events", "workers", delivery.id) + .await + .unwrap(); + assert!(outbox_path.exists()); + drop(broker); + + let broker = Broker::open(config).unwrap(); + assert!(!outbox_path.exists()); + let target = broker + .stats() + .topics + .into_iter() + .find(|topic| topic.name == target) + .unwrap(); + assert_eq!(target.message_count, 1); +} + +#[tokio::test] +async fn dlq_outbox_recovers_when_the_source_was_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 id = broker + .publish("events", vec![b"poison".to_vec()], Duration::ZERO) + .await + .unwrap()[0]; + let target = "events.workers.DLQ"; + crate::outbox::store( + &root.path().join("dlq-outbox"), + &OutboxEntry { + source_topic: "events".into(), + source_channel: "workers".into(), + message_id: id, + target_topic: target.into(), + body: bytes::Bytes::from_static(b"poison"), + }, + ) + .unwrap(); + broker.delete_topic("events").await.unwrap(); + drop(broker); + + let broker = Broker::open(config).unwrap(); + assert_eq!( + std::fs::read_dir(root.path().join("dlq-outbox")) + .unwrap() + .count(), + 0 + ); + let target = broker + .stats() + .topics + .into_iter() + .find(|topic| topic.name == target) + .unwrap(); + assert_eq!(target.message_count, 1); +} + #[tokio::test] async fn stats_settle_expired_in_flight_messages_without_another_fetch() { let root = tempdir().unwrap(); @@ -157,6 +486,128 @@ async fn stats_settle_expired_in_flight_messages_without_another_fetch() { assert_eq!(channel.depth, 1); } +#[tokio::test] +async fn stale_delivery_token_cannot_mutate_a_redelivery() { + 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(); + + let first = broker + .fetch_batch_retained( + "events", + "workers", + 1, + usize::MAX, + Duration::ZERO, + Some(Duration::ZERO), + ) + .await + .unwrap(); + let (first, mut first_guard) = first.into_parts(); + let id = first[0].id; + let stale_token = first_guard.accept_with_token(id).unwrap(); + broker + .expire_channel_in_flight("events", "workers") + .await + .unwrap(); + + let redelivery = broker + .fetch_batch_retained("events", "workers", 1, usize::MAX, Duration::ZERO, None) + .await + .unwrap(); + let (redelivery, mut redelivery_guard) = redelivery.into_parts(); + assert_eq!(redelivery[0].id, id); + let current_token = redelivery_guard.accept_with_token(id).unwrap(); + assert_ne!(stale_token, current_token); + + assert!(matches!( + broker + .finish_delivery("events", "workers", id, stale_token) + .await, + Err(BrokerError::MessageNotInFlight) + )); + assert!(matches!( + broker.touch_delivery("events", "workers", id, stale_token, None), + Err(BrokerError::MessageNotInFlight) + )); + assert!(matches!( + broker + .requeue_delivery("events", "workers", id, stale_token, Duration::ZERO) + .await, + Err(BrokerError::MessageNotInFlight) + )); + broker.release_deliveries("events", "workers", &[(id, stale_token)]); + assert_eq!(broker.stats().topics[0].channels[0].in_flight_count, 1); + + broker + .finish_delivery("events", "workers", id, current_token) + .await + .unwrap(); + assert_eq!(broker.stats().topics[0].channels[0].depth, 0); +} + +#[tokio::test] +async fn current_delivery_tokens_can_renew_an_expiring_batch() { + 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"first".to_vec(), b"second".to_vec()], + Duration::ZERO, + ) + .await + .unwrap(); + + let batch = broker + .fetch_batch_retained( + "events", + "workers", + 2, + usize::MAX, + Duration::ZERO, + Some(Duration::ZERO), + ) + .await + .unwrap(); + let (deliveries, mut guard) = batch.into_parts(); + let tokens: Vec<_> = deliveries + .iter() + .map(|delivery| (delivery.id, guard.token(delivery.id).unwrap())) + .collect(); + broker + .touch_deliveries("events", "workers", &tokens, Some(Duration::from_secs(30))) + .unwrap(); + for (id, token) in &tokens { + assert_eq!(guard.accept_with_token(*id), Some(*token)); + } + + broker + .expire_channel_in_flight("events", "workers") + .await + .unwrap(); + assert_eq!(broker.stats().topics[0].channels[0].in_flight_count, 2); + for (id, token) in tokens { + broker + .finish_delivery("events", "workers", id, token) + .await + .unwrap(); + } +} + #[tokio::test] async fn kodo_channel_counters_are_monotonic_across_empty_and_restart() { if rustqueue_storage::MAX_WRITER_FEATURE_LEVEL < 2 { @@ -281,21 +732,35 @@ async fn startup_replays_a_large_dlq_outbox_after_lowering_the_publish_limit() { } #[test] -fn feature_level_two_rejects_a_delivery_budget_that_cannot_read_retained_messages() { +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")); +} + +#[test] +fn broker_rejects_timeouts_that_would_overflow_instant() { 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, + message_timeout: Duration::MAX, ..BrokerConfig::default() }) { - Ok(_) => panic!("undersized delivery budget was accepted"), + Ok(_) => panic!("unrepresentable broker timeout was accepted"), Err(error) => error, }; - assert!(error - .to_string() - .contains("every message readable at the active storage feature level")); + assert!(error.to_string().contains("platform timer range")); } #[tokio::test] @@ -807,6 +1272,29 @@ async fn dropping_an_unhanded_delivery_batch_does_not_consume_an_attempt() { assert_eq!(messages[0].attempts, 1); } +#[tokio::test] +async fn maintenance_reclaims_retired_topics_after_readers_drain() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_topic("events").await.unwrap(); + let handle = broker.topic("events").unwrap(); + let directory = crate::metadata::topic_directory(root.path(), "events"); + + broker.delete_topic("events").await.unwrap(); + assert!(directory.exists()); + assert!(broker.inner.retired_topics.lock().contains_key("events")); + + drop(handle); + broker.compact().await.unwrap(); + + assert!(!directory.exists()); + assert!(!broker.inner.retired_topics.lock().contains_key("events")); +} + #[tokio::test] async fn broker_metadata_budget_spills_active_tails_across_topics() { let root = tempdir().unwrap(); @@ -898,6 +1386,327 @@ async fn concurrent_publishes_wait_for_metadata_spill_instead_of_rejecting() { assert!(broker.inner.message_index_cache.resident_bytes() <= 64 * 1024); } +#[tokio::test] +async fn channel_creation_rechecks_fences_after_waiting_for_management() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_topic("events").await.unwrap(); + + let (locked_tx, locked_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let lock_holder = { + let broker = broker.clone(); + std::thread::spawn(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + locked_tx.send(()).unwrap(); + let _ = release_rx.recv(); + }) + }; + locked_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let mut creation = Box::pin(broker.create_channel("events", "workers")); + assert!( + tokio::time::timeout(Duration::from_millis(100), creation.as_mut()) + .await + .is_err(), + "ordinary channel creation bypassed the management lifecycle barrier" + ); + { + let mut fences = broker.inner.fences.lock(); + fences.set_channel("events", "workers", now_ms() + 60_000); + fences.store(&broker.inner.fences_path).unwrap(); + } + release_tx.send(()).unwrap(); + lock_holder.join().unwrap(); + + assert!(matches!( + tokio::time::timeout(Duration::from_secs(1), creation) + .await + .unwrap(), + Err(BrokerError::ChannelTombstoned) + )); + assert!(broker.channel_names("events").unwrap().is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn waiting_fetch_rechecks_channel_fence_before_reserving() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_channel("events", "workers").await.unwrap(); + + let topic = broker.topic("events").unwrap(); + let baseline_receivers = topic.wake.receiver_count(); + let fetch_broker = broker.clone(); + let fetch = tokio::spawn(async move { + fetch_broker + .fetch_batch_retained("events", "workers", 1, 1024, Duration::from_secs(1), None) + .await + }); + tokio::time::timeout(Duration::from_secs(1), async { + while topic.wake.receiver_count() <= baseline_receivers { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let topic_state = topic.state.lock(); + { + let mut fences = broker.inner.fences.lock(); + fences.set_channel("events", "workers", now_ms() + 60_000); + fences.store(&broker.inner.fences_path).unwrap(); + } + drop(topic_state); + broker + .publish("events", vec![b"body".to_vec()], Duration::ZERO) + .await + .unwrap(); + + assert!(matches!( + tokio::time::timeout(Duration::from_secs(1), fetch) + .await + .unwrap() + .unwrap(), + Err(BrokerError::ChannelTombstoned) + )); + let channel = &broker.stats().topics[0].channels[0]; + assert_eq!(channel.depth, 1); + assert_eq!(channel.in_flight_count, 0); +} + +#[tokio::test] +async fn touch_rechecks_channel_fence_after_waiting_for_topic() { + 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(); + let batch = broker + .fetch_batch_retained("events", "workers", 1, 1024, Duration::ZERO, None) + .await + .unwrap(); + let (deliveries, mut guard) = batch.into_parts(); + let id = deliveries[0].id; + let token = guard.accept_with_token(id).unwrap(); + + let topic = broker.topic("events").unwrap(); + let (locked_tx, locked_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let lock_holder = { + let topic = Arc::clone(&topic); + std::thread::spawn(move || { + let _topic_state = topic.state.lock(); + locked_tx.send(()).unwrap(); + let _ = release_rx.recv(); + }) + }; + locked_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let baseline_references = Arc::strong_count(&topic); + let toucher = { + let broker = broker.clone(); + std::thread::spawn(move || { + broker.touch_delivery( + "events", + "workers", + id, + token, + Some(Duration::from_secs(30)), + ) + }) + }; + tokio::time::timeout(Duration::from_secs(1), async { + while Arc::strong_count(&topic) <= baseline_references { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + { + let mut fences = broker.inner.fences.lock(); + fences.set_channel("events", "workers", now_ms() + 60_000); + fences.store(&broker.inner.fences_path).unwrap(); + } + release_tx.send(()).unwrap(); + lock_holder.join().unwrap(); + + assert!(matches!( + toucher.join().unwrap(), + Err(BrokerError::ChannelTombstoned) + )); +} + +#[tokio::test] +async fn fence_sync_waits_for_active_management_mutation() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + let deadline = now_ms() + 60_000; + + let (locked_tx, locked_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let lock_holder = { + let broker = broker.clone(); + std::thread::spawn(move || { + let _lifecycle = broker.inner.topic_lifecycle.lock(); + locked_tx.send(()).unwrap(); + let _ = release_rx.recv(); + }) + }; + locked_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let mut sync = Box::pin(broker.sync_management_fences(ManagementFenceSnapshot { + revision: "resource-version-1".into(), + topics: BTreeMap::from([("events".into(), deadline)]), + channels: Vec::new(), + })); + assert!( + tokio::time::timeout(Duration::from_millis(100), sync.as_mut()) + .await + .is_err(), + "fence replacement bypassed the management lifecycle barrier" + ); + release_tx.send(()).unwrap(); + lock_holder.join().unwrap(); + tokio::time::timeout(Duration::from_secs(1), sync) + .await + .unwrap() + .unwrap(); + + assert!(matches!( + broker.create_topic("events").await, + Err(BrokerError::TopicTombstoned) + )); +} + +#[tokio::test] +async fn failed_management_preconditions_do_not_leave_durable_blockers() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + + assert!(matches!( + broker + .manage_topic( + "pause-missing-0001", + "missing", + TopicManagementAction::Pause, + broker.registry_revision(), + None, + ) + .await, + Err(BrokerError::TopicNotFound) + )); + broker.create_topic("missing").await.unwrap(); + + broker.create_topic("events").await.unwrap(); + assert!(matches!( + broker + .manage_channel(ChannelManagementCommand { + operation_id: "pause-channel-0001", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Pause, + expected_revision: broker.registry_revision(), + tombstone_until_ms: None, + require_idle: false, + }) + .await, + Err(BrokerError::ChannelNotFound) + )); + broker.create_channel("events", "workers").await.unwrap(); + + assert!(matches!( + broker + .manage_topic( + "delete-no-deadline-0001", + "missing", + TopicManagementAction::Delete, + broker.registry_revision(), + None, + ) + .await, + Err(BrokerError::InvalidTombstone) + )); + broker + .publish("missing", vec![b"still-open".to_vec()], Duration::ZERO) + .await + .unwrap(); + + assert!(matches!( + broker + .manage_channel(ChannelManagementCommand { + operation_id: "delete-expired-0001", + topic: "events", + channel: "workers", + action: ChannelManagementAction::Delete, + expected_revision: broker.registry_revision(), + tombstone_until_ms: Some(now_ms().saturating_sub(1)), + require_idle: false, + }) + .await, + Err(BrokerError::InvalidTombstone) + )); + broker + .publish("events", vec![b"still-open".to_vec()], Duration::ZERO) + .await + .unwrap(); +} + +#[tokio::test] +async fn a_different_management_action_cannot_cross_a_pending_operation() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + broker.create_topic("events").await.unwrap(); + let fingerprint = + serde_json::to_string(&("topic", "events", TopicManagementAction::Pause)).unwrap(); + broker + .inner + .management_ops + .lock() + .prepare( + &broker.inner.management_ops_path, + "pending-pause-0001", + fingerprint, + "events".into(), + ) + .unwrap(); + + assert!(matches!( + broker + .manage_topic( + "delete-events-0002", + "events", + TopicManagementAction::Delete, + broker.registry_revision(), + Some(now_ms() + 60_000), + ) + .await, + Err(BrokerError::OperationConflict) + )); + assert!(broker.topic_names().iter().any(|topic| topic == "events")); +} + #[tokio::test] async fn management_fences_fail_closed_and_survive_restart() { let root = tempdir().unwrap(); @@ -1367,10 +2176,56 @@ async fn concurrent_registry_updates_persist_the_latest_revision() { #[tokio::test] async fn panicked_storage_tasks_fail_the_broker_closed() { - let error = super::blocking::<()>(|| panic!("injected storage task panic")) + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + let error = broker + .storage_task::<()>(|| panic!("injected storage task panic")) .await .unwrap_err(); assert!(matches!(error, BrokerError::StorageUnavailable)); + assert!(!broker.storage_healthy()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_storage_task_still_records_its_failure() { + let root = tempdir().unwrap(); + let broker = Broker::open(BrokerConfig { + data_path: root.path().into(), + ..BrokerConfig::default() + }) + .unwrap(); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let task = { + let broker = broker.clone(); + tokio::spawn(async move { + broker + .storage_task(move || { + started_tx.send(()).unwrap(); + release_rx.recv().unwrap(); + Err::<(), _>(BrokerError::InvalidRecord( + "injected cancelled storage failure".into(), + )) + }) + .await + }) + }; + started_rx.recv().unwrap(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + release_tx.send(()).unwrap(); + + tokio::time::timeout(Duration::from_secs(1), async { + while broker.storage_healthy() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); } #[test] diff --git a/crates/queue/src/channel.rs b/crates/queue/src/channel.rs index bbc0b2b..8637b7c 100644 --- a/crates/queue/src/channel.rs +++ b/crates/queue/src/channel.rs @@ -342,6 +342,18 @@ impl ChannelState { self.in_flight_ids.get(&id).copied() } + pub fn in_flight_position_with_token(&self, id: u64, token: u64) -> Option { + let position = self.in_flight_position(id)?; + self.in_flight + .get(&position) + .is_some_and(|flight| flight.token == token) + .then_some(position) + } + + pub fn is_unacknowledged(&self, position: u64) -> bool { + position > self.ack_floor_position && !self.acknowledged.contains(&position) + } + pub fn delivery_attempts(&self, position: u64) -> u16 { self.attempts.get(&position).copied().unwrap_or_default() } @@ -351,10 +363,14 @@ impl ChannelState { } pub fn touch(&mut self, position: u64, timeout: Duration) -> bool { + self.touch_until(position, Instant::now() + timeout) + } + + pub fn touch_until(&mut self, position: u64, deadline: Instant) -> bool { let Some(flight) = self.in_flight.get_mut(&position) else { return false; }; - flight.deadline = Instant::now() + timeout; + flight.deadline = deadline; true } diff --git a/crates/queue/src/delivery_guard.rs b/crates/queue/src/delivery_guard.rs index d712b88..3bfd84f 100644 --- a/crates/queue/src/delivery_guard.rs +++ b/crates/queue/src/delivery_guard.rs @@ -35,13 +35,25 @@ impl DeliveryGuard { } pub fn accept(&mut self, id: u64) { + let _ = self.accept_with_token(id); + } + + pub fn accept_with_token(&mut self, id: u64) -> Option { if let Some(index) = self .reservations .iter() .position(|reservation| reservation.id == id) { - self.reservations.swap_remove(index); + return Some(self.reservations.swap_remove(index).token); } + None + } + + pub fn token(&self, id: u64) -> Option { + self.reservations + .iter() + .find(|reservation| reservation.id == id) + .map(|reservation| reservation.token) } pub fn accept_all(&mut self) { diff --git a/crates/queue/src/payload_reader.rs b/crates/queue/src/payload_reader.rs index c09ab20..64417b6 100644 --- a/crates/queue/src/payload_reader.rs +++ b/crates/queue/src/payload_reader.rs @@ -1,3 +1,4 @@ +use crate::delivery_budget::DeliveryHold; use parking_lot::Mutex; use rustqueue_storage::PayloadRef; use rustqueue_telemetry::LatencyHistogram; @@ -7,17 +8,20 @@ use std::io; #[cfg(unix)] use std::os::unix::fs::FileExt; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TrySendError}; use std::sync::Arc; use tokio::sync::oneshot; type ReadFailure = (io::ErrorKind, String); type ReadResult = Result>, ReadFailure>; +type ReadResponse = (ReadResult, Option); struct ReadJob { payloads: Vec, _lease: PayloadLease, - response: oneshot::Sender, + hold: Option, + response: oneshot::Sender, } pub(crate) struct PayloadLease { @@ -60,6 +64,7 @@ impl PayloadReader { workers: usize, queue_depth: usize, latency: Arc, + storage_healthy: Arc, ) -> Arc { let (sender, receiver) = sync_channel(queue_depth.max(1)); let receiver = Arc::new(Mutex::new(receiver)); @@ -80,9 +85,10 @@ impl PayloadReader { for index in 0..workers { let receiver = Arc::clone(&receiver); let files = Arc::clone(&files); + let storage_healthy = Arc::clone(&storage_healthy); std::thread::Builder::new() .name(format!("rustqueue-payload-{index}")) - .spawn(move || worker(receiver, files)) + .spawn(move || worker(receiver, files, storage_healthy)) .expect("payload reader worker must start"); } Arc::new(Self { @@ -102,7 +108,9 @@ impl PayloadReader { #[cfg(test)] pub async fn read_many(&self, payloads: &[PayloadRef]) -> io::Result>> { let lease = self.retain(payloads.to_vec()); - self.read_retained(lease).await + self.read_retained(lease, None) + .await + .map(|(bodies, _)| bodies) } pub fn retain(&self, payloads: Vec) -> PayloadLease { @@ -131,7 +139,11 @@ impl PayloadReader { } } - pub async fn read_retained(&self, lease: PayloadLease) -> io::Result>> { + pub async fn read_retained( + &self, + lease: PayloadLease, + hold: Option, + ) -> io::Result<(Vec>, Option)> { let _timer = self.latency.timer(); let payloads = lease.payloads.clone(); let mut output = vec![None; payloads.len()]; @@ -147,7 +159,7 @@ impl PayloadReader { } } if missing.is_empty() { - return Ok(output.into_iter().flatten().collect()); + return Ok((output.into_iter().flatten().collect(), hold)); } let (sender, receiver) = oneshot::channel(); let job = ReadJob { @@ -156,6 +168,7 @@ impl PayloadReader { .map(|index| payloads[*index].clone()) .collect(), _lease: lease, + hold, response: sender, }; match self.sender.try_send(job) { @@ -173,10 +186,10 @@ impl PayloadReader { )); } } - let bodies = receiver + let (bodies, hold) = receiver .await - .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "payload reader stopped"))? - .map_err(|(kind, message)| io::Error::new(kind, message))?; + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "payload reader stopped"))?; + let bodies = bodies.map_err(|(kind, message)| io::Error::new(kind, message))?; let mut cache = self.cache.lock(); for ((index, payload), body) in missing .into_iter() @@ -187,10 +200,13 @@ impl PayloadReader { cache.insert(payload, Arc::clone(&body)); output[index] = Some(body); } - Ok(output - .into_iter() - .map(|body| body.expect("every payload read produced a body")) - .collect()) + Ok(( + output + .into_iter() + .map(|body| body.expect("every payload read produced a body")) + .collect(), + hold, + )) } pub fn retained_paths(&self) -> BTreeSet { @@ -275,13 +291,27 @@ impl PayloadCache { } } -fn worker(receiver: Arc>>, files: Arc>) { +fn worker( + receiver: Arc>>, + files: Arc>, + storage_healthy: Arc, +) { loop { let job = receiver.lock().recv(); let Ok(job) = job else { return }; + let ReadJob { + payloads, + _lease: lease, + hold, + response, + } = job; let result = - read_payloads(&files, &job.payloads).map_err(|error| (error.kind(), error.to_string())); - let _ = job.response.send(result); + read_payloads(&files, &payloads).map_err(|error| (error.kind(), error.to_string())); + if result.is_err() { + storage_healthy.store(false, Ordering::Release); + } + drop(lease); + let _ = response.send((result, hold)); } } @@ -416,9 +446,24 @@ fn release_paths( #[cfg(test)] mod tests { use super::*; + use crate::delivery_budget::DeliveryBudget; use rustqueue_telemetry::LatencyHistogram; use tempfile::tempdir; + fn reader(cache_bytes: usize, queue_depth: usize) -> (Arc, Arc) { + let storage_healthy = Arc::new(AtomicBool::new(true)); + ( + PayloadReader::new( + cache_bytes, + 1, + queue_depth, + Arc::new(LatencyHistogram::default()), + Arc::clone(&storage_healthy), + ), + storage_healthy, + ) + } + #[test] fn payload_lease_keeps_segment_visible_to_gc_until_drop() { let root = tempdir().unwrap(); @@ -429,7 +474,7 @@ mod tests { len: 1, crc32c: 0, }; - let reader = PayloadReader::new(1, 1, 1, Arc::new(LatencyHistogram::default())); + let (reader, _) = reader(1, 1); let lease = reader.retain(vec![payload]); assert!(reader.retained_paths().contains(&path)); assert!(reader.has_active_under(root.path())); @@ -449,7 +494,7 @@ mod tests { len: 7, crc32c: crc32c::crc32c(b"payload"), }; - let reader = PayloadReader::new(1, 1, 4, Arc::new(LatencyHistogram::default())); + let (reader, _) = reader(1, 4); assert_eq!( &*reader .read_many(std::slice::from_ref(&payload)) @@ -478,7 +523,7 @@ mod tests { len: 3, crc32c: crc32c::crc32c(b"old"), }; - let reader = PayloadReader::new(1024, 1, 4, Arc::new(LatencyHistogram::default())); + let (reader, _) = reader(1024, 4); assert_eq!(&*reader.read_many(&[old]).await.unwrap()[0], b"old"); std::fs::remove_file(&path).unwrap(); @@ -492,4 +537,65 @@ mod tests { }; assert_eq!(&*reader.read_many(&[new]).await.unwrap()[0], b"new"); } + + #[tokio::test] + async fn corrupt_payload_marks_the_reader_unhealthy_before_replying() { + let root = tempdir().unwrap(); + let path = root.path().join("payload"); + std::fs::write(&path, b"bad").unwrap(); + let payload = PayloadRef { + path: Arc::new(path), + offset: 0, + len: 3, + crc32c: crc32c::crc32c(b"good"), + }; + let (reader, storage_healthy) = reader(1024, 4); + + assert!(reader.read_many(&[payload]).await.is_err()); + assert!(!storage_healthy.load(Ordering::Acquire)); + } + + #[cfg(unix)] + #[tokio::test] + async fn cancelled_payload_read_keeps_its_byte_budget_until_the_worker_finishes() { + use std::ffi::CString; + use std::os::unix::ffi::OsStrExt; + + let root = tempdir().unwrap(); + let path = root.path().join("payload"); + let fifo = CString::new(path.as_os_str().as_bytes()).unwrap(); + assert_eq!(unsafe { libc::mkfifo(fifo.as_ptr(), 0o600) }, 0); + let payload = PayloadRef { + path: Arc::new(path.clone()), + offset: 0, + len: 1, + crc32c: 0, + }; + let (reader, _) = reader(1, 1); + let lease = reader.retain(vec![payload]); + let budget = DeliveryBudget::new(2); + let hold = budget.acquire(1).await.unwrap(); + assert_eq!(budget.snapshot().in_flight_bytes, 2); + + let mut reading = Box::pin(reader.read_retained(lease, Some(hold))); + std::future::poll_fn(|context| { + assert!( + std::future::Future::poll(reading.as_mut(), context).is_pending(), + "the FIFO-backed payload read must remain pending" + ); + std::task::Poll::Ready(()) + }) + .await; + drop(reading); + assert_eq!(budget.snapshot().in_flight_bytes, 2); + + drop(std::fs::OpenOptions::new().write(true).open(path).unwrap()); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while budget.snapshot().in_flight_bytes != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } } diff --git a/crates/queue/src/topic/delivery.rs b/crates/queue/src/topic/delivery.rs index e719116..0b097fb 100644 --- a/crates/queue/src/topic/delivery.rs +++ b/crates/queue/src/topic/delivery.rs @@ -2,6 +2,7 @@ use super::*; use crate::channel::{MessageAvailability, NextCandidate}; use crate::model::ReservedDelivery; use crate::topic::index::{Lookup, PageRequest}; +use std::time::Instant; pub(crate) enum ReserveBatch { Ready(Vec), @@ -89,7 +90,7 @@ impl Topic { id: u64, require_in_flight: bool, ) -> Result<(), BrokerError> { - let command = self.finish_command(channel, id, require_in_flight)?; + let command = self.finish_command(channel, id, require_in_flight, None)?; self.persist_channel(channel, command) } @@ -98,8 +99,9 @@ impl Topic { channel: &str, id: u64, require_in_flight: bool, + token: Option, ) -> Result<(), BrokerError> { - let command = self.finish_command(channel, id, require_in_flight)?; + let command = self.finish_command(channel, id, require_in_flight, token)?; self.persist_channel_buffered(channel, command) } @@ -108,12 +110,19 @@ impl Topic { channel: &str, id: u64, require_in_flight: bool, + token: Option, ) -> Result { let position = if require_in_flight { - self.channels + let state = &self + .channels .get(channel) - .and_then(|channel| channel.state.in_flight_position(id)) - .ok_or(BrokerError::MessageNotInFlight)? + .ok_or(BrokerError::ChannelNotFound)? + .state; + match token { + Some(token) => state.in_flight_position_with_token(id, token), + None => state.in_flight_position(id), + } + .ok_or(BrokerError::MessageNotInFlight)? } else { self.position_by_id(id)? .ok_or(BrokerError::MessageNotFound)? @@ -129,8 +138,9 @@ impl Topic { channel: &str, id: u64, available_at_ms: i64, + token: Option, ) -> Result<(), BrokerError> { - let command = self.requeue_command(channel, id, available_at_ms)?; + let command = self.requeue_command(channel, id, available_at_ms, token)?; self.persist_channel_buffered(channel, command) } @@ -139,15 +149,17 @@ impl Topic { channel: &str, id: u64, available_at_ms: i64, + token: Option, ) -> Result { let runtime = self .channels .get(channel) .ok_or(BrokerError::ChannelNotFound)?; - let position = runtime - .state - .in_flight_position(id) - .ok_or(BrokerError::MessageNotInFlight)?; + let position = match token { + Some(token) => runtime.state.in_flight_position_with_token(id, token), + None => runtime.state.in_flight_position(id), + } + .ok_or(BrokerError::MessageNotInFlight)?; let attempts = runtime.state.delivery_attempts(position); Ok(ChannelCommand::Requeue { position, @@ -161,14 +173,25 @@ impl Topic { } pub fn touch(&mut self, channel: &str, id: u64, timeout: Duration) -> Result<(), BrokerError> { + self.touch_with_token(channel, id, None, timeout) + } + + pub fn touch_with_token( + &mut self, + channel: &str, + id: u64, + token: Option, + timeout: Duration, + ) -> Result<(), BrokerError> { let channel = self .channels .get_mut(channel) .ok_or(BrokerError::ChannelNotFound)?; - let position = channel - .state - .in_flight_position(id) - .ok_or(BrokerError::MessageNotInFlight)?; + let position = match token { + Some(token) => channel.state.in_flight_position_with_token(id, token), + None => channel.state.in_flight_position(id), + } + .ok_or(BrokerError::MessageNotInFlight)?; if channel.state.touch(position, timeout) { Ok(()) } else { @@ -176,6 +199,32 @@ impl Topic { } } + pub fn touch_deliveries_with_tokens( + &mut self, + channel: &str, + deliveries: &[(u64, u64)], + timeout: Duration, + ) -> Result<(), BrokerError> { + let channel = self + .channels + .get_mut(channel) + .ok_or(BrokerError::ChannelNotFound)?; + let positions: Vec<_> = deliveries + .iter() + .map(|(id, token)| { + channel + .state + .in_flight_position_with_token(*id, *token) + .ok_or(BrokerError::MessageNotInFlight) + }) + .collect::>()?; + let deadline = Instant::now() + timeout; + for position in positions { + debug_assert!(channel.state.touch_until(position, deadline)); + } + Ok(()) + } + pub fn release(&mut self, channel: &str, ids: &[u64]) { let Some(channel) = self.channels.get_mut(channel) else { return; @@ -187,6 +236,17 @@ impl Topic { } } + pub fn release_with_tokens(&mut self, channel: &str, deliveries: &[(u64, u64)]) { + let Some(channel) = self.channels.get_mut(channel) else { + return; + }; + for (id, token) in deliveries { + if let Some(position) = channel.state.in_flight_position_with_token(*id, *token) { + channel.state.release(position); + } + } + } + pub fn release_all(&mut self) -> usize { self.channels .values_mut() @@ -194,6 +254,20 @@ impl Topic { .sum() } + pub fn message_is_unacknowledged( + &self, + channel: &str, + id: u64, + ) -> Result, BrokerError> { + let Some(position) = self.position_by_id(id)? else { + return Ok(None); + }; + Ok(self + .channels + .get(channel) + .map(|runtime| runtime.state.is_unacknowledged(position))) + } + fn position_by_id(&self, id: u64) -> Result, BrokerError> { self.messages.position_by_id(id) } diff --git a/crates/queue/src/topic/index_cache.rs b/crates/queue/src/topic/index_cache.rs index 3792e64..f29633c 100644 --- a/crates/queue/src/topic/index_cache.rs +++ b/crates/queue/src/topic/index_cache.rs @@ -6,6 +6,7 @@ use rustqueue_storage::RecoveryMetadataRef; use std::collections::{HashMap, VecDeque}; use std::io; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TrySendError}; use std::sync::Arc; use tokio::sync::oneshot; @@ -35,6 +36,7 @@ impl PageRequest { struct ReadJob { request: PageRequest, response: oneshot::Sender, BrokerError>>, + _guard: Box, } struct CacheState { @@ -59,7 +61,12 @@ pub(crate) struct MetadataReservation { } impl MessageIndexCache { - pub(crate) fn new(cache_bytes: usize, workers: usize, queue_depth: usize) -> Arc { + pub(crate) fn new( + cache_bytes: usize, + workers: usize, + queue_depth: usize, + storage_healthy: Arc, + ) -> Arc { let (sender, receiver) = sync_channel(queue_depth.max(1)); let receiver = Arc::new(Mutex::new(receiver)); let workers = if workers == 0 { @@ -72,9 +79,10 @@ impl MessageIndexCache { .max(1); for index in 0..workers { let receiver = Arc::clone(&receiver); + let storage_healthy = Arc::clone(&storage_healthy); std::thread::Builder::new() .name(format!("rustqueue-index-{index}")) - .spawn(move || index_worker(receiver)) + .spawn(move || index_worker(receiver, storage_healthy)) .expect("message index reader worker must start"); } Arc::new(Self { @@ -92,7 +100,11 @@ impl MessageIndexCache { }) } - pub(crate) async fn load(&self, request: PageRequest) -> Result<(), BrokerError> { + pub(crate) async fn load( + &self, + request: PageRequest, + guard: impl Send + 'static, + ) -> Result<(), BrokerError> { if self.state.lock().pages.contains_key(&request.key) { return Ok(()); } @@ -100,6 +112,7 @@ impl MessageIndexCache { match self.sender.try_send(ReadJob { request: request.clone(), response: sender, + _guard: Box::new(guard), }) { Ok(()) => {} Err(TrySendError::Full(_)) => { @@ -301,7 +314,7 @@ fn mark_changed(state: &mut CacheState, changed: &Condvar) { changed.notify_all(); } -fn index_worker(receiver: Arc>>) { +fn index_worker(receiver: Arc>>, storage_healthy: Arc) { loop { let job = receiver.lock().recv(); let Ok(job) = job else { return }; @@ -310,6 +323,101 @@ fn index_worker(receiver: Arc>>) { job.request.first_ordinal, job.request.count, ); + if result.as_ref().is_err_and(|error| { + matches!( + error, + BrokerError::StorageUnavailable + | BrokerError::Storage(_) + | BrokerError::Io(_) + | BrokerError::InvalidRecord(_) + ) + }) { + storage_healthy.store(false, Ordering::Release); + } let _ = job.response.send(result); } } + +#[cfg(all(test, unix))] +mod tests { + use super::*; + use rustqueue_storage::{Record, RecordKind, SegmentLog, HEADER_LEN}; + use std::ffi::CString; + use std::os::unix::ffi::OsStrExt; + use tempfile::tempdir; + + struct DropProbe(Arc); + + impl Drop for DropProbe { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + + #[tokio::test] + async fn cancelled_index_load_keeps_its_path_guard_until_the_worker_finishes() { + let root = tempdir().unwrap(); + let mut log = SegmentLog::open(root.path(), HEADER_LEN as u64 + 1).unwrap(); + let record = || Record { + kind: RecordKind::Noop, + flags: 0, + index: 0, + timestamp_ns: 0, + message_id: 0, + available_at_ms: 0, + payload: vec![0], + }; + log.append(record(), true).unwrap(); + let sealed = log.current_segment_path().to_path_buf(); + log.append(record(), true).unwrap(); + log.persist_recovery_index( + &sealed, + vec![0; recovery::HEADER_LEN + recovery::MESSAGE_LEN], + ) + .unwrap(); + let metadata = log.recovery_metadata_ref(&sealed).unwrap(); + let index_path = sealed.with_extension("rqidx"); + drop(log); + std::fs::remove_file(&index_path).unwrap(); + let path = CString::new(index_path.as_os_str().as_bytes()).unwrap(); + assert_eq!(unsafe { libc::mkfifo(path.as_ptr(), 0o600) }, 0); + + let healthy = Arc::new(AtomicBool::new(true)); + let cache = MessageIndexCache::new(MIN_CACHE_BYTES * 2, 1, 1, healthy); + let request = PageRequest { + key: PageKey { + segment: sealed, + page: 0, + }, + metadata, + first_ordinal: 0, + count: 1, + }; + let dropped = Arc::new(AtomicBool::new(false)); + let mut loading = Box::pin(cache.load(request, DropProbe(Arc::clone(&dropped)))); + std::future::poll_fn(|context| { + assert!( + std::future::Future::poll(loading.as_mut(), context).is_pending(), + "the FIFO-backed index read must remain pending" + ); + std::task::Poll::Ready(()) + }) + .await; + drop(loading); + assert!(!dropped.load(Ordering::Acquire)); + + drop( + std::fs::OpenOptions::new() + .write(true) + .open(index_path) + .unwrap(), + ); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while !dropped.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } +} diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml index f03451e..ebd6d80 100644 --- a/crates/server/Cargo.toml +++ b/crates/server/Cargo.toml @@ -16,8 +16,9 @@ axum.workspace = true bytes.workspace = true clap.workspace = true crc32c.workspace = true +futures.workspace = true parking_lot.workspace = true -regex.workspace = true +regex-automata.workspace = true reqwest.workspace = true rustls.workspace = true rustls-pemfile.workspace = true diff --git a/crates/server/src/admission.rs b/crates/server/src/admission.rs index f236443..5d96d8f 100644 --- a/crates/server/src/admission.rs +++ b/crates/server/src/admission.rs @@ -25,6 +25,7 @@ pub struct ConnectionBudget { pub struct PublishReservation { bytes: usize, + publish_metrics: bool, metrics: Arc, _node: OwnedSemaphorePermit, _connection: Option, @@ -57,7 +58,7 @@ impl PublishAdmission { } pub fn try_reserve(&self, bytes: usize) -> Option { - self.try_reserve_inner(bytes, None) + self.try_reserve_inner(bytes, None, true) } pub fn try_reserve_publish( @@ -80,7 +81,7 @@ impl PublishAdmission { return None; } }; - self.try_reserve_inner(bytes, connection) + self.try_reserve_inner(bytes, connection, true) } pub fn try_reserve_connection_publish( @@ -92,12 +93,23 @@ impl PublishAdmission { self.try_reserve_connection(working_set_bytes(bytes, shape), connection) } + pub fn try_reserve_control( + &self, + bytes: usize, + connection: &ConnectionBudget, + ) -> Option { + let bytes = bytes.saturating_mul(2); + let connection = connection.try_acquire(bytes)?; + self.try_reserve_inner(bytes, Some(connection), false) + } + fn try_reserve_inner( &self, bytes: usize, connection: Option, + publish_metrics: bool, ) -> Option { - if !self.storage_ready() { + if publish_metrics && !self.storage_ready() { self.record_rejected(bytes); return None; } @@ -105,15 +117,20 @@ impl PublishAdmission { let node = match Arc::clone(&self.permits).try_acquire_many_owned(count) { Ok(permit) => permit, Err(_) => { - self.record_rejected(bytes); + if publish_metrics { + self.record_rejected(bytes); + } return None; } }; - self.metrics - .publish_inflight_bytes - .fetch_add(bytes as i64, Ordering::Relaxed); + if publish_metrics { + self.metrics + .publish_inflight_bytes + .fetch_add(bytes as i64, Ordering::Relaxed); + } Some(PublishReservation { bytes, + publish_metrics, metrics: Arc::clone(&self.metrics), _node: node, _connection: connection, @@ -143,11 +160,19 @@ impl ConnectionBudget { } } +pub(crate) fn capacity_is_supported(bytes: usize) -> bool { + bytes <= i64::MAX as usize + && units(bytes) <= tokio::sync::Semaphore::MAX_PERMITS + && units(bytes) <= u32::MAX as usize +} + impl Drop for PublishReservation { fn drop(&mut self) { - self.metrics - .publish_inflight_bytes - .fetch_sub(self.bytes as i64, Ordering::Relaxed); + if self.publish_metrics { + self.metrics + .publish_inflight_bytes + .fetch_sub(self.bytes as i64, Ordering::Relaxed); + } } } @@ -158,7 +183,10 @@ fn units(bytes: usize) -> usize { pub(crate) fn working_set_bytes(bytes: usize, shape: PublishShape) -> usize { let messages = match shape { PublishShape::Single => 1, - PublishShape::Multi => rustqueue_protocol::MAX_MPUB_MESSAGES.min(bytes.max(1)), + // A four-byte malformed MPUB body can still declare the maximum count, + // causing the parser to allocate its full message table before it can + // reject the truncated body. + PublishShape::Multi => rustqueue_protocol::MAX_MPUB_MESSAGES, }; bytes .saturating_add(RECORD_FIXED_BYTES) @@ -187,5 +215,37 @@ mod tests { 20 * 1024 * 1024 + RECORD_FIXED_BYTES + MESSAGE_WORKING_BYTES ); assert!(working_set_bytes(64 * 1024 * 1024, PublishShape::Multi) > 64 * 1024 * 1024); + assert_eq!( + working_set_bytes(4, PublishShape::Multi), + 4 + RECORD_FIXED_BYTES + rustqueue_protocol::MAX_MPUB_MESSAGES * MESSAGE_WORKING_BYTES + ); + } + + #[test] + fn admission_capacity_must_fit_runtime_and_metrics_primitives() { + assert!(capacity_is_supported(512 * 1024 * 1024)); + assert!(!capacity_is_supported(usize::MAX)); + } + + #[test] + fn control_bodies_share_the_node_budget_without_publish_metrics() { + let metrics = Arc::new(Metrics::default()); + let admission = PublishAdmission::new(8192, Arc::clone(&metrics)); + let first_connection = ConnectionBudget::new(8192); + let second_connection = ConnectionBudget::new(8192); + admission.set_storage_ready(false); + + let first = admission + .try_reserve_control(4096, &first_connection) + .unwrap(); + assert!(admission + .try_reserve_control(1, &second_connection) + .is_none()); + assert_eq!(metrics.publish_inflight_bytes.load(Ordering::Relaxed), 0); + + drop(first); + assert!(admission + .try_reserve_control(4096, &second_connection) + .is_some()); } } diff --git a/crates/server/src/auth.rs b/crates/server/src/auth.rs index ca9364c..933a319 100644 --- a/crates/server/src/auth.rs +++ b/crates/server/src/auth.rs @@ -1,12 +1,25 @@ use crate::config::Config; use parking_lot::Mutex; -use regex::Regex; +use regex_automata::{meta::Regex, nfa::thompson::WhichCaptures}; use reqwest::Client; use serde::Deserialize; use sha2::{Digest, Sha256}; use std::collections::{HashMap, VecDeque}; +use std::mem::size_of; +use std::sync::Arc; use std::time::{Duration, Instant}; use thiserror::Error; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +const AUTH_MEMORY_UNIT_BYTES: usize = 4 * 1024; +const AUTH_RESPONSE_WORKING_SET_MULTIPLIER: usize = 3; +const AUTH_REGEX_NFA_LIMIT_BYTES: usize = 256 * 1024; +const AUTH_REGEX_ENGINE_LIMIT_BYTES: usize = 64 * 1024; +const AUTH_REGEX_SEARCH_HEADROOM_BYTES: usize = 64 * 1024; +const MAX_AUTHORIZATIONS: usize = 256; +const MAX_AUTH_REGEX_PATTERNS: usize = 256; +const MAX_AUTH_PATTERN_BYTES: usize = 4 * 1024; +const MAX_AUTH_PATTERN_SOURCE_BYTES: usize = 64 * 1024; #[derive(Debug, Error)] pub enum AuthError { @@ -16,6 +29,8 @@ pub enum AuthError { InvalidResponse, #[error("no permissions found")] Unauthorized, + #[error("authorization memory budget is exhausted")] + Overloaded, } pub struct Authenticator { @@ -24,17 +39,22 @@ pub struct Authenticator { max_response_bytes: usize, max_ttl_seconds: u64, cache: Mutex, + memory: Arc, } #[derive(Clone)] pub struct AuthSession { - pub identity: String, - pub identity_url: String, + inner: Arc, +} + +struct AuthSessionInner { + identity: String, + identity_url: String, expires_at: Instant, grants: Vec, + _memory: Vec, } -#[derive(Clone)] struct Grant { publish: bool, subscribe: bool, @@ -48,6 +68,10 @@ struct AuthCache { max_entries: usize, } +struct AuthMemoryReservation { + _permit: OwnedSemaphorePermit, +} + #[derive(Deserialize)] struct AuthResponse { #[serde(default)] @@ -95,6 +119,8 @@ impl Authenticator { format!("{}/auth", address.trim_end_matches('/')) }) .collect(); + let memory_units = + auth_memory_units(config.limits.auth_memory_bytes).ok_or(AuthError::InvalidResponse)?; Ok(Some(Self { client, endpoints, @@ -105,6 +131,7 @@ impl Authenticator { order: VecDeque::new(), max_entries: config.limits.auth_cache_max_entries, }), + memory: Arc::new(Semaphore::new(memory_units as usize)), })) } @@ -120,6 +147,10 @@ impl Authenticator { if let Some(session) = self.cache.lock().get(&cache_key) { return Ok(session); } + let response_working_set = self + .max_response_bytes + .saturating_mul(AUTH_RESPONSE_WORKING_SET_MULTIPLIER); + let response_memory = self.reserve_memory(response_working_set)?; let mut last_error = AuthError::Service; for endpoint in &self.endpoints { match self @@ -127,10 +158,12 @@ impl Authenticator { .await { Ok(session) => { + drop(response_memory); self.cache.lock().insert(cache_key, session.clone()); return Ok(session); } Err(AuthError::Unauthorized) => return Err(AuthError::Unauthorized), + Err(AuthError::Overloaded) => return Err(AuthError::Overloaded), Err(error) => last_error = error, } } @@ -187,7 +220,24 @@ impl Authenticator { } let response: AuthResponse = serde_json::from_slice(&body).map_err(|_| AuthError::InvalidResponse)?; - AuthSession::try_from_response(response, self.max_ttl_seconds) + self.build_session(response) + } + + fn build_session(&self, response: AuthResponse) -> Result { + AuthSession::try_from_response_with_memory(response, self.max_ttl_seconds, |bytes| { + self.reserve_memory(bytes) + }) + } + + fn reserve_memory(&self, bytes: usize) -> Result { + let units = auth_memory_units(bytes).ok_or(AuthError::Overloaded)?; + loop { + match Arc::clone(&self.memory).try_acquire_many_owned(units) { + Ok(permit) => return Ok(AuthMemoryReservation { _permit: permit }), + Err(_) if self.cache.lock().evict_oldest() => {} + Err(_) => return Err(AuthError::Overloaded), + } + } } } @@ -210,13 +260,26 @@ impl AuthCache { self.order.retain(|candidate| candidate != &key); self.order.push_back(key); while self.values.len() > self.max_entries { - if let Some(oldest) = self.order.pop_front() { - self.values.remove(&oldest); + if !self.evict_oldest() { + break; + } + } + } + + fn evict_oldest(&mut self) -> bool { + while let Some(oldest) = self.order.pop_front() { + if self.values.remove(&oldest).is_some() { + return true; } } + false } } +fn auth_memory_units(bytes: usize) -> Option { + u32::try_from(bytes.max(1).div_ceil(AUTH_MEMORY_UNIT_BYTES)).ok() +} + fn auth_cache_key(remote_ip: &str, tls: bool, common_name: &str, secret: &[u8]) -> [u8; 32] { let mut digest = Sha256::new(); for part in [ @@ -232,8 +295,43 @@ fn auth_cache_key(remote_ip: &str, tls: bool, common_name: &str, secret: &[u8]) } impl AuthSession { + #[cfg(test)] fn try_from_response(response: AuthResponse, max_ttl_seconds: u64) -> Result { - let mut grants = Vec::with_capacity(response.authorizations.len()); + Self::build(response, max_ttl_seconds, |_| Ok(None)) + } + + fn try_from_response_with_memory( + response: AuthResponse, + max_ttl_seconds: u64, + mut reserve: impl FnMut(usize) -> Result, + ) -> Result { + Self::build(response, max_ttl_seconds, |bytes| reserve(bytes).map(Some)) + } + + fn build( + response: AuthResponse, + max_ttl_seconds: u64, + mut reserve: impl FnMut(usize) -> Result, AuthError>, + ) -> Result { + let shape = validate_auth_response(&response)?; + let base_bytes = size_of::() + .saturating_add(2 * size_of::()) + .saturating_add(response.identity.capacity()) + .saturating_add(response.identity_url.capacity()) + .saturating_add(shape.grants.saturating_mul(size_of::())) + .saturating_add(shape.channel_patterns.saturating_mul(size_of::())) + .saturating_add( + shape + .patterns + .saturating_add(1) + .saturating_mul(size_of::()), + ); + let mut memory = Vec::with_capacity(shape.patterns.saturating_add(1)); + if let Some(reservation) = reserve(base_bytes)? { + memory.push(reservation); + } + + let mut grants = Vec::with_capacity(shape.grants); for authorization in response.authorizations { let publish = authorization .permissions @@ -246,12 +344,18 @@ impl AuthSession { if !publish && !subscribe { continue; } - let topic = Regex::new(&authorization.topic).map_err(|_| AuthError::InvalidResponse)?; - let channels = authorization - .channels - .iter() - .map(|channel| Regex::new(channel).map_err(|_| AuthError::InvalidResponse)) - .collect::>()?; + let topic = compile_auth_pattern(&authorization.topic)?; + if let Some(reservation) = reserve(auth_regex_memory_bytes(&topic))? { + memory.push(reservation); + } + let mut channels = Vec::with_capacity(authorization.channels.len()); + for channel in authorization.channels { + let pattern = compile_auth_pattern(&channel)?; + if let Some(reservation) = reserve(auth_regex_memory_bytes(&pattern))? { + memory.push(reservation); + } + channels.push(pattern); + } grants.push(Grant { publish, subscribe, @@ -259,36 +363,37 @@ impl AuthSession { channels, }); } - if grants.is_empty() { - return Err(AuthError::Unauthorized); - } let ttl_seconds = response.ttl.min(max_ttl_seconds); Ok(Self { - identity: response.identity, - identity_url: response.identity_url, - expires_at: Instant::now() + Duration::from_secs(ttl_seconds), - grants, + inner: Arc::new(AuthSessionInner { + identity: response.identity, + identity_url: response.identity_url, + expires_at: Instant::now() + Duration::from_secs(ttl_seconds), + grants, + _memory: memory, + }), }) } pub fn can_publish(&self, topic: &str) -> bool { - if Instant::now() >= self.expires_at { + if self.is_expired() { return false; } - self.grants + self.inner + .grants .iter() .any(|grant| grant.publish && grant.topic.is_match(topic)) } pub fn is_expired(&self) -> bool { - Instant::now() >= self.expires_at + Instant::now() >= self.inner.expires_at } pub fn can_subscribe(&self, topic: &str, channel: &str) -> bool { - if Instant::now() >= self.expires_at { + if self.is_expired() { return false; } - self.grants.iter().any(|grant| { + self.inner.grants.iter().any(|grant| { grant.subscribe && grant.topic.is_match(topic) && grant @@ -299,8 +404,90 @@ impl AuthSession { } pub fn permission_count(&self) -> usize { - self.grants.len() + self.inner.grants.len() } + + pub fn identity(&self) -> &str { + &self.inner.identity + } + + pub fn identity_url(&self) -> &str { + &self.inner.identity_url + } +} + +struct AuthResponseShape { + grants: usize, + patterns: usize, + channel_patterns: usize, +} + +fn validate_auth_response(response: &AuthResponse) -> Result { + if response.ttl == 0 || response.authorizations.len() > MAX_AUTHORIZATIONS { + return Err(AuthError::InvalidResponse); + } + let mut grants = 0usize; + let mut patterns = 0usize; + let mut channel_patterns = 0usize; + let mut pattern_bytes = 0usize; + for authorization in &response.authorizations { + let relevant = authorization + .permissions + .iter() + .any(|permission| permission == "publish" || permission == "subscribe"); + if !relevant { + continue; + } + grants = grants.checked_add(1).ok_or(AuthError::InvalidResponse)?; + patterns = patterns + .checked_add(authorization.channels.len().saturating_add(1)) + .ok_or(AuthError::InvalidResponse)?; + channel_patterns = channel_patterns + .checked_add(authorization.channels.len()) + .ok_or(AuthError::InvalidResponse)?; + if patterns > MAX_AUTH_REGEX_PATTERNS { + return Err(AuthError::InvalidResponse); + } + for pattern in std::iter::once(&authorization.topic).chain(&authorization.channels) { + if pattern.len() > MAX_AUTH_PATTERN_BYTES { + return Err(AuthError::InvalidResponse); + } + pattern_bytes = pattern_bytes + .checked_add(pattern.len()) + .ok_or(AuthError::InvalidResponse)?; + if pattern_bytes > MAX_AUTH_PATTERN_SOURCE_BYTES { + return Err(AuthError::InvalidResponse); + } + } + } + if grants == 0 { + return Err(AuthError::Unauthorized); + } + Ok(AuthResponseShape { + grants, + patterns, + channel_patterns, + }) +} + +fn compile_auth_pattern(pattern: &str) -> Result { + Regex::builder() + .configure( + Regex::config() + .which_captures(WhichCaptures::None) + .nfa_size_limit(Some(AUTH_REGEX_NFA_LIMIT_BYTES)) + .onepass_size_limit(Some(AUTH_REGEX_ENGINE_LIMIT_BYTES)) + .hybrid_cache_capacity(AUTH_REGEX_ENGINE_LIMIT_BYTES) + .dfa_size_limit(Some(AUTH_REGEX_ENGINE_LIMIT_BYTES)), + ) + .build(pattern) + .map_err(|_| AuthError::InvalidResponse) +} + +fn auth_regex_memory_bytes(pattern: &Regex) -> usize { + size_of::() + .saturating_add(pattern.memory_usage()) + .saturating_add(AUTH_REGEX_SEARCH_HEADROOM_BYTES) } #[cfg(test)] @@ -329,7 +516,8 @@ mod tests { assert!(session.can_publish("orders")); assert!(session.can_subscribe("orders", "workers-2")); assert!(!session.can_subscribe("orders", "admin")); - assert!(session.expires_at > Instant::now()); + assert!(session.inner.expires_at > Instant::now()); + assert!(Arc::ptr_eq(&session.inner, &session.clone().inner)); } #[test] @@ -361,7 +549,7 @@ mod tests { assert!(cache.get(&[2; 32]).is_some()); let mut expired = session(); - expired.expires_at = Instant::now(); + Arc::get_mut(&mut expired.inner).unwrap().expires_at = Instant::now(); cache.insert([3; 32], expired); assert!(cache.get(&[3; 32]).is_none()); assert_eq!(cache.order.len(), cache.values.len()); @@ -377,14 +565,17 @@ mod tests { for ordinal in 0..1_000u16 { let key = [(ordinal % 251) as u8; 32]; let session = AuthSession { - identity: "worker".into(), - identity_url: String::new(), - expires_at: Instant::now() + Duration::from_millis(1), - grants: Vec::new(), + inner: Arc::new(AuthSessionInner { + identity: "worker".into(), + identity_url: String::new(), + expires_at: Instant::now() + Duration::from_millis(1), + grants: Vec::new(), + _memory: Vec::new(), + }), }; cache.insert(key, session); if let Some(value) = cache.values.get_mut(&key) { - value.expires_at = Instant::now(); + Arc::get_mut(&mut value.inner).unwrap().expires_at = Instant::now(); } assert!(cache.get(&key).is_none()); } @@ -392,6 +583,76 @@ mod tests { assert!(cache.values.is_empty()); } + #[test] + fn rejects_auth_responses_that_amplify_regex_compilation() { + let channels = (0..MAX_AUTH_REGEX_PATTERNS) + .map(|ordinal| format!("^worker-{ordinal}$")) + .collect(); + let response = AuthResponse { + ttl: 60, + identity: "worker".into(), + identity_url: String::new(), + authorizations: vec![Authorization { + permissions: vec!["subscribe".into()], + topic: "^orders$".into(), + channels, + }], + }; + assert!(matches!( + AuthSession::try_from_response(response, 60), + Err(AuthError::InvalidResponse) + )); + + let response = AuthResponse { + ttl: 0, + identity: "worker".into(), + identity_url: String::new(), + authorizations: vec![Authorization { + permissions: vec!["publish".into()], + topic: ".*".into(), + channels: Vec::new(), + }], + }; + assert!(matches!( + AuthSession::try_from_response(response, 60), + Err(AuthError::InvalidResponse) + )); + } + + #[test] + fn auth_memory_budget_bounds_cached_and_live_sessions() { + let budget = 2 * 1024 * 1024; + let mut config = Config::default(); + config.security.auth_http_addresses = vec!["http://127.0.0.1:1".into()]; + config.limits.auth_memory_bytes = budget; + let authenticator = Authenticator::new(&config).unwrap().unwrap(); + + for ordinal in 0..100u8 { + let session = authenticator + .build_session(AuthResponse { + ttl: 60, + identity: format!("worker-{ordinal}"), + identity_url: String::new(), + authorizations: vec![Authorization { + permissions: vec!["publish".into()], + topic: format!("^orders-{ordinal}$"), + channels: Vec::new(), + }], + }) + .unwrap(); + authenticator.cache.lock().insert([ordinal; 32], session); + } + assert!(authenticator.cache.lock().values.len() < 100); + + let held = authenticator.reserve_memory(budget).unwrap(); + assert!(matches!( + authenticator.reserve_memory(1), + Err(AuthError::Overloaded) + )); + drop(held); + assert!(authenticator.reserve_memory(1).is_ok()); + } + #[tokio::test] async fn invalid_auth_replica_fails_over_to_the_next_endpoint() { let (invalid, invalid_task) = auth_server(StatusCode::OK, "not-json").await; diff --git a/crates/server/src/config.rs b/crates/server/src/config.rs index 678c0fd..c4c46aa 100644 --- a/crates/server/src/config.rs +++ b/crates/server/src/config.rs @@ -146,6 +146,7 @@ pub struct LimitsConfig { pub auth_timeout_ms: u64, pub auth_max_ttl_seconds: u64, pub auth_cache_max_entries: usize, + pub auth_memory_bytes: usize, pub http_body_timeout_ms: u64, pub disconnect_on_retriable_publish_error: bool, } @@ -279,6 +280,7 @@ impl Default for LimitsConfig { auth_timeout_ms: 5_000, auth_max_ttl_seconds: 3600, auth_cache_max_entries: 10_000, + auth_memory_bytes: 256 * 1024 * 1024, http_body_timeout_ms: 30_000, disconnect_on_retriable_publish_error: false, } @@ -409,6 +411,7 @@ impl Config { bail!("metrics.max_detailed_series must be greater than zero"); } self.validate_protocol_limits()?; + self.validate_runtime_limits()?; if !matches!(self.log_format.as_str(), "text" | "json") { bail!("log_format must be text or json"); } diff --git a/crates/server/src/config/tests.rs b/crates/server/src/config/tests.rs index d84358e..06f4af9 100644 --- a/crates/server/src/config/tests.rs +++ b/crates/server/src/config/tests.rs @@ -152,6 +152,14 @@ fn rejects_zero_protocol_capacity_and_timeouts() { config.limits.http_body_timeout_ms = 0; assert!(config.validate().is_err()); + let mut config = Config::default(); + config.limits.auth_max_ttl_seconds = 0; + assert!(config.validate().is_err()); + + let mut config = Config::default(); + config.limits.auth_memory_bytes = 0; + assert!(config.validate().is_err()); + let mut config = Config::default(); config.limits.http_body_timeout_ms = 1; config.queue.message_timeout_ms = config.queue.max_message_timeout_ms + 1; @@ -166,6 +174,37 @@ fn rejects_zero_protocol_capacity_and_timeouts() { assert!(config.validate().is_err()); } +#[test] +fn auth_memory_budget_must_cover_response_working_set() { + let mut config = Config::default(); + config.limits.auth_memory_bytes = config.limits.auth_response_bytes * 4 - 1; + assert!(config.validate().is_err()); + + config.limits.auth_memory_bytes = config.limits.auth_response_bytes * 4; + config.validate().unwrap(); +} + +#[test] +fn rejects_limits_that_would_panic_runtime_primitives() { + let mut config = Config::default(); + config.limits.max_connections = tokio::sync::Semaphore::MAX_PERMITS + 1; + assert!(config.validate().is_err()); + + let mut config = Config::default(); + config.shutdown.grace_seconds = u64::MAX; + assert!(config.validate().is_err()); + + let mut config = Config::default(); + config.limits.auth_max_ttl_seconds = u64::MAX; + assert!(config.validate().is_err()); + + let mut config = Config::default(); + config.limits.connection_publish_inflight_bytes = usize::MAX; + config.limits.node_publish_inflight_bytes = usize::MAX; + let error = config.validate().unwrap_err(); + assert!(error.to_string().contains("publish inflight limits")); +} + #[test] fn rejects_an_unbounded_detailed_metric_configuration() { let mut config = Config::default(); diff --git a/crates/server/src/config/validation.rs b/crates/server/src/config/validation.rs index 0b3f294..8443b30 100644 --- a/crates/server/src/config/validation.rs +++ b/crates/server/src/config/validation.rs @@ -1,6 +1,7 @@ use super::{Config, MAX_SUPPORTED_BATCH_BYTES, MAX_SUPPORTED_MESSAGE_BYTES}; -use crate::admission::{working_set_bytes, PublishShape}; +use crate::admission::{capacity_is_supported, working_set_bytes, PublishShape}; use anyhow::bail; +use std::time::{Duration, Instant}; impl Config { pub(super) fn validate_protocol_limits(&self) -> anyhow::Result<()> { @@ -12,6 +13,11 @@ impl Config { working_set_bytes(self.limits.max_body_bytes, PublishShape::Multi).max( working_set_bytes(self.queue.max_message_bytes, PublishShape::Single), ); + if !capacity_is_supported(self.limits.connection_publish_inflight_bytes) + || !capacity_is_supported(self.limits.node_publish_inflight_bytes) + { + bail!("publish inflight limits exceed the runtime semaphore or metrics capacity"); + } if self.limits.connection_publish_inflight_bytes < publish_working_set || self.limits.node_publish_inflight_bytes < self.limits.connection_publish_inflight_bytes @@ -51,14 +57,23 @@ impl Config { if self.limits.client_handshake_timeout_ms == 0 || self.limits.tcp_command_timeout_ms == 0 || self.limits.auth_cache_max_entries == 0 + || self.limits.auth_memory_bytes == 0 || self.limits.max_connections == 0 || self.limits.max_rdy_count == 0 || self.limits.auth_response_bytes == 0 || self.limits.auth_timeout_ms == 0 + || self.limits.auth_max_ttl_seconds == 0 || self.limits.http_body_timeout_ms == 0 { bail!("limits connection, RDY, auth, TCP and HTTP timeouts/sizes must be greater than zero"); } + if !capacity_is_supported(self.limits.auth_memory_bytes) + || self.limits.auth_response_bytes > self.limits.auth_memory_bytes / 4 + { + bail!( + "limits.auth_memory_bytes must fit the runtime semaphore and be at least four times auth_response_bytes" + ); + } if self.limits.heartbeat_interval_ms < 1_000 || self.limits.heartbeat_interval_ms > self.limits.max_heartbeat_interval_ms || self.limits.max_heartbeat_interval_ms > i64::MAX as u64 @@ -80,4 +95,65 @@ impl Config { } Ok(()) } + + pub(super) fn validate_runtime_limits(&self) -> anyhow::Result<()> { + if self.limits.max_connections > tokio::sync::Semaphore::MAX_PERMITS { + bail!("limits.max_connections exceeds the runtime semaphore capacity"); + } + + let maximum_visibility = Duration::from_millis(self.queue.max_message_timeout_ms) + .saturating_add(Duration::from_millis( + self.limits.max_output_buffer_timeout_ms, + )); + let maximum_handoff = Duration::from_millis(self.limits.max_heartbeat_interval_ms) + .saturating_mul(2) + .saturating_mul(64) + .max(maximum_visibility); + let timers = [ + ( + "storage.maintenance_startup_delay_seconds", + Duration::from_secs(self.storage.maintenance_startup_delay_seconds), + ), + ( + "storage.scrub_interval_seconds", + Duration::from_secs(self.storage.scrub_interval_seconds), + ), + ( + "queue.publish_worker_idle_seconds", + Duration::from_secs(self.queue.publish_worker_idle_seconds), + ), + ( + "limits.client_handshake_timeout_ms", + Duration::from_millis(self.limits.client_handshake_timeout_ms), + ), + ( + "limits.tcp_command_timeout_ms", + Duration::from_millis(self.limits.tcp_command_timeout_ms), + ), + ( + "limits.auth_timeout_ms", + Duration::from_millis(self.limits.auth_timeout_ms), + ), + ( + "limits.auth_max_ttl_seconds", + Duration::from_secs(self.limits.auth_max_ttl_seconds), + ), + ( + "limits.http_body_timeout_ms", + Duration::from_millis(self.limits.http_body_timeout_ms), + ), + ( + "shutdown.grace_seconds", + Duration::from_secs(self.shutdown.grace_seconds), + ), + ("maximum delivery handoff timeout", maximum_handoff), + ]; + let now = Instant::now(); + for (name, duration) in timers { + if now.checked_add(duration).is_none() { + bail!("{name} exceeds the platform timer range"); + } + } + Ok(()) + } } diff --git a/crates/server/src/http.rs b/crates/server/src/http.rs index bcf36a7..c0f972b 100644 --- a/crates/server/src/http.rs +++ b/crates/server/src/http.rs @@ -296,6 +296,7 @@ impl From for ApiError { } BrokerError::RevisionConflict { .. } => (StatusCode::CONFLICT, "E_REVISION_CONFLICT"), BrokerError::OperationConflict => (StatusCode::CONFLICT, "E_OPERATION_CONFLICT"), + BrokerError::InvalidTombstone => (StatusCode::BAD_REQUEST, "E_BAD_TOMBSTONE"), BrokerError::InvalidChannel => (StatusCode::BAD_REQUEST, "E_BAD_CHANNEL"), BrokerError::MessageTooLarge | BrokerError::BatchTooLarge => { (StatusCode::BAD_REQUEST, "E_BAD_MESSAGE") diff --git a/crates/server/src/http/manage.rs b/crates/server/src/http/manage.rs index 3a75a01..f909378 100644 --- a/crates/server/src/http/manage.rs +++ b/crates/server/src/http/manage.rs @@ -1,7 +1,8 @@ use super::*; +use crate::subscriptions::DeletePermit; use axum::extract::Path; use rustqueue_queue::{ - ChannelManagementAction, ChannelManagementCommand, ManagementFenceSnapshot, + ChannelManagementAction, ChannelManagementCommand, ManagementFenceSnapshot, ManagementResult, TopicManagementAction, }; @@ -79,7 +80,7 @@ pub(super) async fn delete_idle_channel_compat( ) .into_response(); } - let _permit = match state + let permit = match state .subscriptions .begin_delete(&query.topic, &query.channel) { @@ -94,18 +95,16 @@ pub(super) async fn delete_idle_channel_compat( }; let revision = state.broker.registry_revision(); let operation_id = kodo_compat_operation_id(revision, &query.topic, &query.channel); - let result = state - .broker - .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; + let result = manage_idle_channel( + Arc::clone(&state.broker), + permit, + operation_id, + query.topic.clone(), + query.channel.clone(), + revision, + kodo_cleanup_deadline(), + ) + .await; match result { Ok(result) => { tracing::info!( @@ -136,7 +135,7 @@ async fn apply_channel_management( } else { authorize(&headers, &state.tokens.console, "console")?; } - let _delete_permit = if require_idle { + let delete_permit = if require_idle { Some( state .subscriptions @@ -156,18 +155,31 @@ async fn apply_channel_management( } 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?; + let result = if let Some(permit) = delete_permit { + manage_idle_channel( + Arc::clone(&state.broker), + permit, + request.operation_id.clone(), + request.topic.clone(), + request.channel.clone(), + request.expected_revision, + tombstone_until_ms.expect("idle deletion always has a server deadline"), + ) + .await? + } else { + 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, @@ -178,6 +190,39 @@ async fn apply_channel_management( Ok(Json(json!(result))) } +async fn manage_idle_channel( + broker: Arc, + delete_permit: DeletePermit, + operation_id: String, + topic: String, + channel: String, + expected_revision: u64, + tombstone_until_ms: i64, +) -> Result { + let task = tokio::spawn(async move { + let result = broker + .manage_channel(ChannelManagementCommand { + operation_id: &operation_id, + topic: &topic, + channel: &channel, + action: ChannelManagementAction::Delete, + expected_revision, + tombstone_until_ms: Some(tombstone_until_ms), + require_idle: true, + }) + .await; + drop(delete_permit); + result + }); + match task.await { + Ok(result) => result, + Err(error) => { + tracing::error!(%error, "idle channel management task failed"); + Err(BrokerError::StorageUnavailable) + } + } +} + fn kodo_cleanup_deadline() -> i64 { let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -242,6 +287,9 @@ fn parse_channel_action( #[cfg(test)] mod tests { use super::*; + use crate::subscriptions::{ClientIdentity, RegisterBlocked, SubscriptionRegistry}; + use rustqueue_queue::BrokerConfig; + use tempfile::tempdir; #[test] fn idle_delete_action_is_unavailable_until_kodo_cleanup_is_enabled() { @@ -270,4 +318,58 @@ mod tests { kodo_compat_operation_id(8, "events", "workers") ); } + + #[tokio::test] + async fn cancelled_idle_delete_keeps_new_subscriptions_blocked_until_completion() { + let root = tempdir().unwrap(); + let broker = Arc::new( + Broker::open(BrokerConfig { + data_path: root.path().to_path_buf(), + ..BrokerConfig::default() + }) + .unwrap(), + ); + broker.create_channel("events", "workers").await.unwrap(); + let subscriptions = SubscriptionRegistry::default(); + let permit = subscriptions.begin_delete("events", "workers").unwrap(); + let revision = broker.registry_revision(); + let mut deletion = Box::pin(manage_idle_channel( + Arc::clone(&broker), + permit, + "cancelled-idle-delete-0001".into(), + "events".into(), + "workers".into(), + revision, + kodo_cleanup_deadline(), + )); + std::future::poll_fn(|context| { + assert!( + std::future::Future::poll(deletion.as_mut(), context).is_pending(), + "the detached deletion must not complete in its first poll" + ); + std::task::Poll::Ready(()) + }) + .await; + drop(deletion); + + assert!(matches!( + subscriptions.register("events", "workers", ClientIdentity::default()), + Err(RegisterBlocked::DeleteInProgress) + )); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let channel_removed = broker + .channel_names("events") + .is_ok_and(|channels| channels.is_empty()); + let barrier_released = subscriptions.begin_delete("events", "workers").is_ok(); + if channel_removed && barrier_released { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } } diff --git a/crates/server/src/tcp.rs b/crates/server/src/tcp.rs index 469d259..f7cb83d 100644 --- a/crates/server/src/tcp.rs +++ b/crates/server/src/tcp.rs @@ -1,4 +1,5 @@ mod authorization; +mod channel_ops; mod codec; mod commands; mod dead_letter; @@ -8,6 +9,7 @@ mod time; mod writer; use authorization::*; +use channel_ops::*; use codec::*; use commands::*; use dead_letter::*; @@ -31,7 +33,7 @@ use rustqueue_protocol::{ }; use rustqueue_queue::{Broker, BrokerError, DeliveryGuard}; use serde_json::json; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::net::SocketAddr; use std::pin::Pin; @@ -59,6 +61,22 @@ pub(crate) fn broker_storage_error(error: &BrokerError) -> bool { ) } +fn renew_delivery_lease( + broker: &Broker, + topic: &str, + channel: &str, + id: u64, + token: u64, + timeout: Duration, +) -> Result { + // Keep the session deadline slightly earlier than the broker deadline so a + // client operation cannot pass the local check after its broker lease has + // already expired. + let deadline = Instant::now() + timeout; + broker.touch_delivery(topic, channel, id, token, Some(timeout))?; + Ok(deadline) +} + #[derive(Debug)] struct ParsedCommand { command: Command, @@ -104,6 +122,12 @@ struct RemoteDelivery { body: Bytes, } +#[derive(Clone, Copy, Debug)] +struct InFlightDelivery { + deadline: Instant, + token: u64, +} + struct SessionState { identified: bool, encrypted: bool, @@ -115,10 +139,12 @@ struct SessionState { sample_rate: u8, sample_cursor: u8, auth: Option, - auth_secret: Option>, + auth_secret: Option, + _auth_reservation: Option, subscription: Option, rdy: u64, - in_flight: HashMap, + in_flight: HashMap, + pending_channel_ops: HashSet, closing: bool, client_identity: ClientIdentity, } @@ -140,6 +166,13 @@ impl SessionState { .update_flow(self.rdy, self.in_flight.len()); } } + + fn delivery_for_operation(&self, id: u64) -> Option { + if self.pending_channel_ops.contains(&id) { + return None; + } + self.in_flight.get(&id).copied() + } } #[allow(clippy::too_many_arguments)] @@ -303,9 +336,11 @@ async fn handle_connection( sample_cursor: 0, auth: None, auth_secret: None, + _auth_reservation: None, subscription: None, rdy: 0, in_flight: HashMap::new(), + pending_channel_ops: HashSet::new(), closing: false, client_identity: ClientIdentity { remote_address: peer.to_string(), diff --git a/crates/server/src/tcp/channel_ops.rs b/crates/server/src/tcp/channel_ops.rs new file mode 100644 index 0000000..516b99b --- /dev/null +++ b/crates/server/src/tcp/channel_ops.rs @@ -0,0 +1,252 @@ +use super::*; +use futures::stream::{FuturesUnordered, StreamExt}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +#[derive(Clone, Copy)] +pub(super) enum ChannelOpKind { + Finish, + SampleFinish, + Requeue, +} + +impl ChannelOpKind { + pub(super) fn error_code(self) -> &'static str { + match self { + Self::Finish | Self::SampleFinish => "E_FIN_FAILED", + Self::Requeue => "E_REQ_FAILED", + } + } +} + +enum ChannelOp { + Finish { + topic: String, + channel: String, + id: u64, + token: u64, + sampled: bool, + }, + Requeue { + topic: String, + channel: String, + id: u64, + token: u64, + delay: Duration, + }, +} + +impl ChannelOp { + fn id(&self) -> u64 { + match self { + Self::Finish { id, .. } | Self::Requeue { id, .. } => *id, + } + } + + fn kind(&self) -> ChannelOpKind { + match self { + Self::Finish { sampled: true, .. } => ChannelOpKind::SampleFinish, + Self::Finish { .. } => ChannelOpKind::Finish, + Self::Requeue { .. } => ChannelOpKind::Requeue, + } + } +} + +pub(super) struct ChannelOpCompletion { + pub(super) id: u64, + pub(super) kind: ChannelOpKind, + pub(super) result: Result<(), BrokerError>, +} + +#[derive(Clone)] +pub(super) struct ChannelOpSender { + sender: mpsc::UnboundedSender, +} + +impl ChannelOpSender { + pub(super) fn finish( + &self, + topic: String, + channel: String, + id: u64, + token: u64, + ) -> Result<(), BrokerError> { + self.send(ChannelOp::Finish { + topic, + channel, + id, + token, + sampled: false, + }) + } + + pub(super) fn finish_sampled( + &self, + topic: String, + channel: String, + id: u64, + token: u64, + ) -> Result<(), BrokerError> { + self.send(ChannelOp::Finish { + topic, + channel, + id, + token, + sampled: true, + }) + } + + pub(super) fn requeue( + &self, + topic: String, + channel: String, + id: u64, + token: u64, + delay: Duration, + ) -> Result<(), BrokerError> { + self.send(ChannelOp::Requeue { + topic, + channel, + id, + token, + delay, + }) + } + + fn send(&self, operation: ChannelOp) -> Result<(), BrokerError> { + self.sender + .send(operation) + .map_err(|_| BrokerError::StorageUnavailable) + } +} + +pub(super) fn start_channel_ops( + broker: Broker, +) -> ( + ChannelOpSender, + mpsc::UnboundedReceiver, + JoinHandle<()>, +) { + // The session admits at most one operation per in-flight message, and + // in-flight messages are capped by max_rdy_count. + let (operation_tx, operation_rx) = mpsc::unbounded_channel(); + let (completion_tx, completion_rx) = mpsc::unbounded_channel(); + let task = tokio::spawn(run_channel_ops(broker, operation_rx, completion_tx)); + ( + ChannelOpSender { + sender: operation_tx, + }, + completion_rx, + task, + ) +} + +async fn run_channel_ops( + broker: Broker, + mut operations: mpsc::UnboundedReceiver, + completions: mpsc::UnboundedSender, +) { + type Pending = Pin + Send>>; + + let mut pending = FuturesUnordered::::new(); + let mut receiving = true; + while receiving || !pending.is_empty() { + tokio::select! { + operation = operations.recv(), if receiving => { + match operation { + Some(operation) => { + let broker = broker.clone(); + pending.push(Box::pin(async move { + execute_channel_op(broker, operation).await + })); + } + None => receiving = false, + } + } + completion = pending.next(), if !pending.is_empty() => { + if let Some(completion) = completion { + let _ = completions.send(completion); + } + } + } + } +} + +async fn execute_channel_op(broker: Broker, operation: ChannelOp) -> ChannelOpCompletion { + let id = operation.id(); + let kind = operation.kind(); + let result = match operation { + ChannelOp::Finish { + topic, + channel, + id, + token, + .. + } => broker.finish_delivery(&topic, &channel, id, token).await, + ChannelOp::Requeue { + topic, + channel, + id, + token, + delay, + } => { + broker + .requeue_delivery(&topic, &channel, id, token, delay) + .await + } + }; + ChannelOpCompletion { id, kind, result } +} + +#[cfg(test)] +mod tests { + use super::*; + use rustqueue_queue::BrokerConfig; + use tempfile::tempdir; + + #[tokio::test] + async fn one_session_pipeline_forms_durable_groups() { + 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", + (0..32).map(|_| vec![b'x']).collect(), + Duration::ZERO, + ) + .await + .unwrap(); + let batch = broker + .fetch_batch_retained("events", "workers", 32, usize::MAX, Duration::ZERO, None) + .await + .unwrap(); + let (deliveries, mut guard) = batch.into_parts(); + + let (sender, mut completions, task) = start_channel_ops(broker.clone()); + for delivery in deliveries { + let token = guard.accept_with_token(delivery.id).unwrap(); + sender + .finish("events".into(), "workers".into(), delivery.id, token) + .unwrap(); + } + drop(sender); + + let mut completed = 0; + while let Some(completion) = completions.recv().await { + completion.result.unwrap(); + completed += 1; + } + task.await.unwrap(); + + let stats = broker.stats().channel_group_commit; + assert_eq!(completed, 32); + assert_eq!(stats.requests, 32); + assert!(stats.commits < stats.requests); + assert!(stats.max_batch_requests > 1); + } +} diff --git a/crates/server/src/tcp/codec.rs b/crates/server/src/tcp/codec.rs index 79572c9..4e2f515 100644 --- a/crates/server/src/tcp/codec.rs +++ b/crates/server/src/tcp/codec.rs @@ -176,27 +176,37 @@ where format!("{name} body too big {length} > {maximum}"), )); } - let reservation = if matches!( - command, - Command::Publish { .. } | Command::MultiPublish { .. } | Command::DeferredPublish { .. } - ) { - let shape = if matches!(command, Command::MultiPublish { .. }) { - crate::admission::PublishShape::Multi - } else { - crate::admission::PublishShape::Single - }; - Some( + let reservation = match command { + Command::Identify | Command::Auth => Some( admission - .try_reserve_connection_publish(length, shape, connection_budget) + .try_reserve_control(length, connection_budget) .ok_or_else(|| { CommandReadError::protocol( "E_THROTTLED", - format!("{name} publish byte budget is exhausted; retry later"), + format!("{name} control body byte budget is exhausted; retry later"), ) })?, - ) - } else { - None + ), + Command::Publish { .. } + | Command::MultiPublish { .. } + | Command::DeferredPublish { .. } => { + let shape = if matches!(command, Command::MultiPublish { .. }) { + crate::admission::PublishShape::Multi + } else { + crate::admission::PublishShape::Single + }; + Some( + admission + .try_reserve_connection_publish(length, shape, connection_budget) + .ok_or_else(|| { + CommandReadError::protocol( + "E_THROTTLED", + format!("{name} publish byte budget is exhausted; retry later"), + ) + })?, + ) + } + _ => None, }; let mut body = vec![0; length]; reader.read_exact(&mut body).await.map_err(|error| { @@ -530,4 +540,37 @@ mod tests { assert_eq!(command.body.as_deref(), Some(b"body".as_slice())); writer.await.unwrap(); } + + #[tokio::test] + async fn control_bodies_are_admitted_against_the_node_byte_budget() { + let config = Config::default(); + let metrics = Arc::new(Metrics::default()); + let admission = PublishAdmission::new(8192, metrics); + let first_connection = ConnectionBudget::new(8192); + let (mut first_peer, mut first_server) = tokio::io::duplex(8192); + first_peer.write_all(b"AUTH\n").await.unwrap(); + first_peer.write_u32(4096).await.unwrap(); + first_peer.write_all(&vec![b'x'; 4096]).await.unwrap(); + + let first = read_initial_command(&mut first_server, &config, &admission, &first_connection) + .await + .unwrap(); + assert!(first.publish_reservation.is_some()); + + let second_connection = ConnectionBudget::new(8192); + let (mut second_peer, mut second_server) = tokio::io::duplex(64); + second_peer.write_all(b"IDENTIFY\n").await.unwrap(); + second_peer.write_u32(2).await.unwrap(); + let error = + read_initial_command(&mut second_server, &config, &admission, &second_connection) + .await + .unwrap_err(); + assert!(matches!( + error, + CommandReadError::Protocol { + code: "E_THROTTLED", + .. + } + )); + } } diff --git a/crates/server/src/tcp/commands.rs b/crates/server/src/tcp/commands.rs index 7511020..0298568 100644 --- a/crates/server/src/tcp/commands.rs +++ b/crates/server/src/tcp/commands.rs @@ -10,6 +10,7 @@ pub(super) async fn process_command( authenticator: Option<&Authenticator>, ephemeral_consumers: &EphemeralConsumers, subscriptions: &SubscriptionRegistry, + channel_ops: &ChannelOpSender, state: &mut SessionState, writer: &mut ClientWriter, ) -> anyhow::Result { @@ -32,7 +33,7 @@ pub(super) async fn process_command( write_error(writer, "E_AUTH_DISABLED", "AUTH disabled").await?; return Ok(false); }; - let secret = body.as_deref().unwrap_or_default(); + let secret = body.unwrap_or_default(); if secret.is_empty() { write_error(writer, "E_BAD_BODY", "AUTH invalid body size 0").await?; return Ok(false); @@ -42,18 +43,19 @@ pub(super) async fn process_command( &peer.ip().to_string(), state.encrypted, &state.tls_common_name, - secret, + &secret, ) .await { Ok(session) => { let response = json!({ - "identity": session.identity, - "identity_url": session.identity_url, + "identity": session.identity(), + "identity_url": session.identity_url(), "permission_count": session.permission_count(), }); state.auth = Some(session); - state.auth_secret = Some(secret.to_vec()); + state.auth_secret = Some(secret); + state._auth_reservation = publish_reservation; state.client_identity.authed = true; write_frame(writer, FrameType::Response, &serde_json::to_vec(&response)?) .await?; @@ -211,22 +213,41 @@ pub(super) async fn process_command( write_error(writer, "E_INVALID", "client is not subscribed").await?; return Ok(false); }; - if !state.in_flight.contains_key(&id) { + let Some(delivery) = state.delivery_for_operation(id) else { write_error(writer, "E_FIN_FAILED", "message is not in flight").await?; return Ok(true); - } - let finish_result = - finish_message(broker, &subscription.topic, &subscription.channel, id).await; - 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); + }; + let deadline = match renew_delivery_lease( + broker, + &subscription.topic, + &subscription.channel, + id, + delivery.token, + state.message_timeout, + ) { + Ok(deadline) => deadline, + Err(error) => { + write_broker_error(writer, "E_FIN_FAILED", error).await?; + return Ok(true); } - Err(error) => write_broker_error(writer, "E_FIN_FAILED", error).await?, + }; + let result = channel_ops.finish( + subscription.topic.clone(), + subscription.channel.clone(), + id, + delivery.token, + ); + if let Err(error) = result { + write_broker_error(writer, "E_FIN_FAILED", error).await?; + } else { + state.in_flight.insert( + id, + InFlightDelivery { + deadline, + ..delivery + }, + ); + state.pending_channel_ops.insert(id); } } Command::Requeue { id, delay_ms } => { @@ -235,28 +256,42 @@ pub(super) async fn process_command( write_error(writer, "E_INVALID", "client is not subscribed").await?; return Ok(false); }; - if !state.in_flight.contains_key(&id) { + let Some(delivery) = state.delivery_for_operation(id) else { write_error(writer, "E_REQ_FAILED", "message is not in flight").await?; return Ok(true); - } - let requeue_result = broker - .requeue( - &subscription.topic, - &subscription.channel, - id, - Duration::from_millis(delay_ms), - ) - .await; - 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); + }; + let deadline = match renew_delivery_lease( + broker, + &subscription.topic, + &subscription.channel, + id, + delivery.token, + state.message_timeout, + ) { + Ok(deadline) => deadline, + Err(error) => { + write_broker_error(writer, "E_REQ_FAILED", error).await?; + return Ok(true); } - Err(error) => write_broker_error(writer, "E_REQ_FAILED", error).await?, + }; + let result = channel_ops.requeue( + subscription.topic.clone(), + subscription.channel.clone(), + id, + delivery.token, + Duration::from_millis(delay_ms), + ); + if let Err(error) = result { + write_broker_error(writer, "E_REQ_FAILED", error).await?; + } else { + state.in_flight.insert( + id, + InFlightDelivery { + deadline, + ..delivery + }, + ); + state.pending_channel_ops.insert(id); } } Command::Touch(id) => { @@ -264,21 +299,27 @@ pub(super) async fn process_command( write_error(writer, "E_INVALID", "client is not subscribed").await?; return Ok(false); }; - if !state.in_flight.contains_key(&id) { + let Some(delivery) = state.delivery_for_operation(id) else { write_error(writer, "E_TOUCH_FAILED", "message is not in flight").await?; return Ok(true); - } - let touch_result = broker.touch( + }; + let touch_result = renew_delivery_lease( + broker, &subscription.topic, &subscription.channel, id, - Some(state.message_timeout), + delivery.token, + state.message_timeout, ); match touch_result { - Ok(()) => { - state - .in_flight - .insert(id, Instant::now() + state.message_timeout); + Ok(deadline) => { + state.in_flight.insert( + id, + InFlightDelivery { + deadline, + ..delivery + }, + ); } Err(error) => write_broker_error(writer, "E_TOUCH_FAILED", error).await?, } @@ -371,15 +412,6 @@ pub(super) async fn publish_messages( } } -pub(super) async fn finish_message( - broker: &Broker, - topic: &str, - channel: &str, - id: u64, -) -> Result<(), BrokerError> { - broker.finish(topic, channel, id).await -} - #[cfg(test)] mod tests { use super::{broker_storage_error, precommit_retryable_publish_error}; diff --git a/crates/server/src/tcp/dead_letter.rs b/crates/server/src/tcp/dead_letter.rs index 5740501..a4b932e 100644 --- a/crates/server/src/tcp/dead_letter.rs +++ b/crates/server/src/tcp/dead_letter.rs @@ -20,12 +20,11 @@ pub(super) async fn dead_letter_if_needed( }; let target = dead_letter_topic(topic, channel, &config.queue.dead_letter_suffix) .map_err(BrokerError::InvalidRecord)?; - let result = broker + let moved = broker .move_to_dead_letter(topic, channel, delivery.id, &target, delivery.body.clone()) - .await; - if let Err(error) = result { - broker.release(topic, channel, &[delivery.id]); - return Err(error); + .await?; + if !moved { + return Ok(true); } metrics.dead_letter_messages.fetch_add(1, Ordering::Relaxed); if reason == DeadLetterReason::Retention { diff --git a/crates/server/src/tcp/ephemeral.rs b/crates/server/src/tcp/ephemeral.rs index e9f7502..6f4a0d3 100644 --- a/crates/server/src/tcp/ephemeral.rs +++ b/crates/server/src/tcp/ephemeral.rs @@ -29,17 +29,28 @@ impl EphemeralConsumers { } pub async fn unregister(&self, broker: &Broker, topic: &str, channel: &str) { - let mut counts = self.counts.lock().await; - let key = (topic.to_owned(), channel.to_owned()); - let Some(count) = counts.get_mut(&key) else { - return; - }; - *count = count.saturating_sub(1); - if *count > 0 { - return; + let counts = Arc::clone(&self.counts); + let broker = broker.clone(); + let topic = topic.to_owned(); + let channel = channel.to_owned(); + let task = tokio::spawn(async move { + let mut counts = counts.lock().await; + let key = (topic.clone(), channel.clone()); + let Some(count) = counts.get_mut(&key) else { + return; + }; + *count = count.saturating_sub(1); + if *count > 0 { + return; + } + counts.remove(&key); + let result = broker.delete_channel(&topic, &channel).await; + drop(counts); + let _ = result; + }); + if let Err(error) = task.await { + tracing::error!(%error, "ephemeral consumer unregister task failed"); } - counts.remove(&key); - let _ = broker.delete_channel(topic, channel).await; } } @@ -151,4 +162,47 @@ mod tests { assert!(counts.is_empty()); assert!(broker.channel_names("events").unwrap().is_empty()); } + + #[tokio::test] + async fn cancelled_unregister_still_removes_the_ephemeral_channel() { + let root = tempdir().unwrap(); + let broker = Broker::open(rustqueue_queue::BrokerConfig { + data_path: root.path().to_path_buf(), + ..rustqueue_queue::BrokerConfig::default() + }) + .unwrap(); + let consumers = EphemeralConsumers::default(); + consumers + .register(&broker, "events", "live#ephemeral") + .await + .unwrap(); + + let blocker = consumers.counts.lock().await; + let mut unregister = Box::pin(consumers.unregister(&broker, "events", "live#ephemeral")); + std::future::poll_fn(|context| { + assert!( + std::future::Future::poll(unregister.as_mut(), context).is_pending(), + "the blocked unregister must not complete" + ); + std::task::Poll::Ready(()) + }) + .await; + drop(unregister); + drop(blocker); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let counts_empty = consumers.counts.lock().await.is_empty(); + let channel_removed = broker + .channel_names("events") + .is_ok_and(|channels| channels.is_empty()); + if counts_empty && channel_removed { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } } diff --git a/crates/server/src/tcp/session.rs b/crates/server/src/tcp/session.rs index e71bba5..fd685c8 100644 --- a/crates/server/src/tcp/session.rs +++ b/crates/server/src/tcp/session.rs @@ -8,6 +8,8 @@ struct PendingFetch<'a> { future: FetchFuture<'a>, } +const CHANNEL_OP_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5); + async fn poll_pending_fetch( pending: &mut Option>, ) -> Result { @@ -74,11 +76,15 @@ pub(super) async fn run_session( output_buffer_tick.tick().await; let mut last_command = Instant::now(); let mut pending_fetch: Option> = None; - let mut abandoned_deliveries = Vec::new(); + let (channel_ops, mut channel_op_results, channel_ops_task) = start_channel_ops(broker.clone()); let session_result: anyhow::Result<()> = async { loop { - let expired = expire_client_deadlines(&mut state.in_flight, Instant::now()); + let expired = expire_client_deadlines( + &mut state.in_flight, + &state.pending_channel_ops, + Instant::now(), + ); if expired { if let Some(subscription) = state.subscription.as_ref() { broker.expire_channel_in_flight( @@ -153,6 +159,7 @@ pub(super) async fn run_session( authenticator.as_deref(), &ephemeral_consumers, &subscriptions, + &channel_ops, &mut state, &mut writer, ), @@ -189,7 +196,8 @@ pub(super) async fn run_session( pending_fetch = Some(PendingFetch { request, future }); } } - let in_flight_deadline = state.in_flight.values().copied().min(); + let in_flight_deadline = + next_client_deadline(&state.in_flight, &state.pending_channel_ops); tokio::select! { command = command_rx.recv() => { @@ -249,6 +257,7 @@ pub(super) async fn run_session( authenticator.as_deref(), &ephemeral_consumers, &subscriptions, + &channel_ops, &mut state, &mut writer, ), @@ -259,6 +268,16 @@ pub(super) async fn run_session( break; } } + completion = channel_op_results.recv(), if !state.pending_channel_ops.is_empty() => { + let Some(completion) = completion else { + anyhow::bail!("channel operation pipeline stopped"); + }; + if let Err((code, error)) = + apply_channel_op_completion(completion, &mut state, metrics) + { + write_broker_error(&mut writer, code, error).await?; + } + } delivery_result = poll_pending_fetch(&mut pending_fetch), if pending_fetch.is_some() => { let request = pending_fetch .take() @@ -292,6 +311,33 @@ pub(super) async fn run_session( flush_timed(&mut writer, state.heartbeat).await?; continue; } + let write_timeout = delivery_write_timeout(state.heartbeat); + let visibility_timeout = delivery_visibility_timeout( + state.message_timeout, + state.output_buffer_timeout, + ); + let handoff_timeout = write_timeout + .saturating_mul(deliveries.len() as u32) + .max(visibility_timeout); + let handoff_deadline = Instant::now() + handoff_timeout; + let delivery_tokens = deliveries + .iter() + .map(|delivery| { + delivery_guard + .token(delivery.id) + .map(|token| (delivery.id, token)) + .ok_or_else(|| anyhow::anyhow!("delivery token is missing")) + }) + .collect::>>()?; + broker + .touch_deliveries( + &delivery_topic, + &delivery_channel, + &delivery_tokens, + Some(handoff_timeout), + ) + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + let mut handed_off = Vec::with_capacity(deliveries.len()); for delivery in deliveries { if delivery_is_outstanding( &state.in_flight, @@ -314,14 +360,34 @@ pub(super) async fn run_session( continue; } if !state.accept_sample() { - finish_message( + let token = delivery_guard + .token(delivery.id) + .ok_or_else(|| anyhow::anyhow!("delivery token is missing"))?; + let deadline = renew_delivery_lease( broker, &delivery_topic, &delivery_channel, delivery.id, + token, + state.message_timeout, ) - .await .map_err(|error| anyhow::anyhow!(error.to_string()))?; + channel_ops + .finish_sampled( + delivery_topic.clone(), + delivery_channel.clone(), + delivery.id, + token, + ) + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + state.in_flight.insert( + delivery.id, + InFlightDelivery { + deadline, + token, + }, + ); + state.pending_channel_ops.insert(delivery.id); delivery_guard.accept(delivery.id); continue; } @@ -331,21 +397,49 @@ pub(super) async fn run_session( delivery.id, delivery.body.len(), ); + let token = delivery_guard + .token(delivery.id) + .ok_or_else(|| anyhow::anyhow!("delivery token is missing"))?; tokio::time::timeout( - delivery_write_timeout(state.heartbeat), + write_timeout, writer.write_message_parts(&header, &delivery.body), ) .await .map_err(|_| anyhow::anyhow!("consumer delivery write timed out"))??; - state - .in_flight - .insert(delivery.id, Instant::now() + state.message_timeout); + let accepted_token = delivery_guard + .accept_with_token(delivery.id) + .ok_or_else(|| anyhow::anyhow!("delivery token is missing"))?; + debug_assert_eq!(accepted_token, token); + state.in_flight.insert( + delivery.id, + InFlightDelivery { + deadline: handoff_deadline, + token, + }, + ); + handed_off.push((delivery.id, token)); if let Some(subscription) = &state.subscription { subscription.lease.observe_delivery(); } - delivery_guard.accept(delivery.id); metrics.delivered_messages.fetch_add(1, Ordering::Relaxed); } + if !handed_off.is_empty() { + let deadline = Instant::now() + visibility_timeout; + broker + .touch_deliveries( + &delivery_topic, + &delivery_channel, + &handed_off, + Some(visibility_timeout), + ) + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + for (id, token) in handed_off { + if let Some(delivery) = state.in_flight.get_mut(&id) { + debug_assert_eq!(delivery.token, token); + delivery.deadline = deadline; + } + } + } state.update_subscription_flow(); } Err(error) => { @@ -384,21 +478,29 @@ pub(super) async fn run_session( Ok(()) } .await; - if let Some(mut fetch) = pending_fetch.take() { - if let Ok(Ok(response)) = - tokio::time::timeout(Duration::from_secs(5), fetch.future.as_mut()).await - { - abandoned_deliveries - .extend(response.deliveries.into_iter().map(|delivery| delivery.id)); - } + drop(channel_ops); + let channel_ops_drained = tokio::time::timeout( + CHANNEL_OP_CLEANUP_TIMEOUT, + settle_channel_op_results(&mut channel_op_results, &mut state, metrics), + ) + .await + .unwrap_or(false); + if !channel_ops_drained { + warn!( + pending = state.pending_channel_ops.len(), + "channel operation cleanup did not finish before the connection deadline" + ); + channel_ops_task.abort(); } + let _ = channel_ops_task.await; + // Cancelling the fetch drops its reservation guard, which only releases + // the matching delivery tokens. + drop(pending_fetch.take()); if let Some(subscription) = &state.subscription { - let mut ids: Vec<_> = state.in_flight.into_keys().collect(); - ids.append(&mut abandoned_deliveries); - ids.sort_unstable(); - ids.dedup(); - if !ids.is_empty() { - broker.release(&subscription.topic, &subscription.channel, &ids); + let deliveries = + releaseable_in_flight_deliveries(state.in_flight, &state.pending_channel_ops); + if !deliveries.is_empty() { + broker.release_deliveries(&subscription.topic, &subscription.channel, &deliveries); } if subscription.channel.ends_with("#ephemeral") { ephemeral_consumers @@ -465,16 +567,94 @@ fn publish_command(command: &Command) -> bool { ) } -fn delivery_is_outstanding(in_flight: &HashMap, id: u64) -> bool { +fn delivery_is_outstanding(in_flight: &HashMap, id: u64) -> bool { in_flight.contains_key(&id) } -fn expire_client_deadlines(in_flight: &mut HashMap, now: Instant) -> bool { +fn expire_client_deadlines( + in_flight: &mut HashMap, + pending_channel_ops: &HashSet, + now: Instant, +) -> bool { let before = in_flight.len(); - in_flight.retain(|_, deadline| *deadline > now); + in_flight.retain(|id, delivery| pending_channel_ops.contains(id) || delivery.deadline > now); in_flight.len() != before } +fn next_client_deadline( + in_flight: &HashMap, + pending_channel_ops: &HashSet, +) -> Option { + in_flight + .iter() + .filter_map(|(id, delivery)| { + (!pending_channel_ops.contains(id)).then_some(delivery.deadline) + }) + .min() +} + +fn apply_channel_op_completion( + completion: ChannelOpCompletion, + state: &mut SessionState, + metrics: &Metrics, +) -> Result<(), (&'static str, BrokerError)> { + state.pending_channel_ops.remove(&completion.id); + if let Err(error) = completion.result { + if broker_storage_error(&error) { + metrics.storage_errors.fetch_add(1, Ordering::Relaxed); + } + return Err((completion.kind.error_code(), error)); + } + + state.in_flight.remove(&completion.id); + if let Some(subscription) = &state.subscription { + match completion.kind { + ChannelOpKind::Finish => subscription.lease.observe_finish(), + ChannelOpKind::Requeue => subscription.lease.observe_requeue(), + ChannelOpKind::SampleFinish => {} + } + } + match completion.kind { + ChannelOpKind::Finish => { + metrics.finished_messages.fetch_add(1, Ordering::Relaxed); + } + ChannelOpKind::Requeue => { + metrics.requeued_messages.fetch_add(1, Ordering::Relaxed); + } + ChannelOpKind::SampleFinish => {} + } + state.update_subscription_flow(); + Ok(()) +} + +async fn settle_channel_op_results( + results: &mut tokio::sync::mpsc::UnboundedReceiver, + state: &mut SessionState, + metrics: &Metrics, +) -> bool { + while !state.pending_channel_ops.is_empty() { + let Some(completion) = results.recv().await else { + return false; + }; + let _ = apply_channel_op_completion(completion, state, metrics); + } + true +} + +fn releaseable_in_flight_deliveries( + in_flight: HashMap, + pending_channel_ops: &HashSet, +) -> Vec<(u64, u64)> { + let mut deliveries: Vec<_> = in_flight + .into_iter() + .filter_map(|(id, delivery)| { + (!pending_channel_ops.contains(&id)).then_some((id, delivery.token)) + }) + .collect(); + deliveries.sort_unstable(); + deliveries +} + async fn wait_for_in_flight_deadline(deadline: Option) { match deadline { Some(deadline) => tokio::time::sleep_until(deadline.into()).await, @@ -486,6 +666,10 @@ async fn wait_for_in_flight_deadline(deadline: Option) { mod tests { use super::*; + fn delivery(deadline: Instant, token: u64) -> InFlightDelivery { + InFlightDelivery { deadline, token } + } + #[tokio::test] async fn pending_fetch_survives_an_unrelated_ready_branch() { let request = FetchRequest { @@ -519,7 +703,7 @@ mod tests { #[test] fn duplicate_delivery_is_suppressed_while_in_flight() { let mut in_flight = HashMap::new(); - in_flight.insert(7, Instant::now()); + in_flight.insert(7, delivery(Instant::now(), 70)); assert!(delivery_is_outstanding(&in_flight, 7)); in_flight.clear(); @@ -531,12 +715,43 @@ mod tests { 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)), + (7, delivery(now - Duration::from_millis(1), 70)), + (8, delivery(now + Duration::from_secs(1), 80)), ]); - assert!(expire_client_deadlines(&mut in_flight, now)); + assert!(expire_client_deadlines( + &mut in_flight, + &HashSet::new(), + now + )); assert_eq!(in_flight.keys().copied().collect::>(), vec![8]); - assert!(!expire_client_deadlines(&mut in_flight, now)); + assert!(!expire_client_deadlines( + &mut in_flight, + &HashSet::new(), + now + )); + } + + #[test] + fn pending_channel_operations_do_not_expire() { + let now = Instant::now(); + let mut in_flight = HashMap::from([(7, delivery(now - Duration::from_millis(1), 70))]); + let pending = HashSet::from([7]); + + assert!(!expire_client_deadlines(&mut in_flight, &pending, now)); + assert_eq!(next_client_deadline(&in_flight, &pending), None); + assert!(in_flight.contains_key(&7)); + } + + #[test] + fn unresolved_channel_operations_are_not_released_on_disconnect() { + let now = Instant::now(); + let in_flight = HashMap::from([(7, delivery(now, 70)), (8, delivery(now, 80))]); + let pending = HashSet::from([7]); + + assert_eq!( + releaseable_in_flight_deliveries(in_flight, &pending), + vec![(8, 80)] + ); } #[test] diff --git a/crates/server/src/tcp/writer.rs b/crates/server/src/tcp/writer.rs index eff2ad9..fdb0bec 100644 --- a/crates/server/src/tcp/writer.rs +++ b/crates/server/src/tcp/writer.rs @@ -47,6 +47,13 @@ pub(super) fn delivery_write_timeout(heartbeat: Option) -> Duration { .max(Duration::from_secs(1)) } +pub(super) fn delivery_visibility_timeout( + message_timeout: Duration, + output_buffer_timeout: Option, +) -> Duration { + message_timeout.saturating_add(output_buffer_timeout.unwrap_or_default()) +} + pub(super) fn connection_progress_timeout(heartbeat: Option) -> Duration { heartbeat .map(|interval| interval.saturating_mul(2)) @@ -182,4 +189,16 @@ mod tests { Duration::from_secs(5) ); } + + #[test] + fn initial_delivery_lease_covers_output_buffering() { + assert_eq!( + delivery_visibility_timeout(Duration::from_secs(1), Some(Duration::from_secs(30))), + Duration::from_secs(31) + ); + assert_eq!( + delivery_visibility_timeout(Duration::from_secs(1), None), + Duration::from_secs(1) + ); + } } diff --git a/crates/storage/src/segment/maintenance.rs b/crates/storage/src/segment/maintenance.rs index b88f842..b9a275a 100644 --- a/crates/storage/src/segment/maintenance.rs +++ b/crates/storage/src/segment/maintenance.rs @@ -15,6 +15,7 @@ impl SegmentLog { return Ok(()); }; self.current.sync_all()?; + let directory = File::open(&self.directory)?; let paths = segment_paths(&self.directory)?; let target = paths .iter() @@ -25,20 +26,23 @@ impl SegmentLog { "truncate target segment disappeared", )) })?; - OpenOptions::new() - .write(true) - .open(location.segment.as_ref())? - .set_len(location.offset)?; - recovery_index::remove(location.segment.as_ref())?; - self.sealed_indexes.remove(location.segment.as_ref()); - for path in paths.into_iter().skip(target + 1) { + for path in paths.iter().skip(target + 1).rev() { if path.exists() { - fs::remove_file(&path)?; + fs::remove_file(path)?; } - recovery_index::remove(&path)?; - self.checksums.remove(&path); - self.sealed_indexes.remove(&path); + recovery_index::remove(path)?; + self.checksums.remove(path); + self.sealed_indexes.remove(path); + directory.sync_all()?; } + let target_file = OpenOptions::new() + .write(true) + .open(location.segment.as_ref())?; + target_file.set_len(location.offset)?; + target_file.sync_all()?; + recovery_index::remove(location.segment.as_ref())?; + directory.sync_all()?; + self.sealed_indexes.remove(location.segment.as_ref()); self.resident_records .retain(|record| record.index < from_index); let (locations, _, bytes, crc32c) = scan_segment(location.segment.as_ref(), true)?; @@ -57,7 +61,7 @@ impl SegmentLog { self.refresh_aggregates(); self.start_index = self.first_index().unwrap_or(from_index); self.current.sync_all()?; - File::open(&self.directory)?.sync_all()?; + directory.sync_all()?; Ok(()) } @@ -101,6 +105,7 @@ impl SegmentLog { return Ok(0); } self.current.sync_all()?; + let directory = File::open(&self.directory)?; crash_failpoint("gc_before_segment_delete"); let removed_through = removable .last() @@ -109,6 +114,8 @@ impl SegmentLog { for path in &removable { fs::remove_file(path)?; recovery_index::remove(path)?; + crash_failpoint("gc_after_segment_delete_before_dir_fsync"); + directory.sync_all()?; self.checksums.remove(path); self.sealed_indexes.remove(path); } @@ -118,8 +125,7 @@ impl SegmentLog { self.start_index = self .first_index() .unwrap_or_else(|| through_index.saturating_add(1)); - crash_failpoint("gc_after_segment_delete_before_dir_fsync"); - File::open(&self.directory)?.sync_all()?; + directory.sync_all()?; Ok(removable.len()) } diff --git a/crates/storage/src/segment/tests.rs b/crates/storage/src/segment/tests.rs index 8d90dba..372e970 100644 --- a/crates/storage/src/segment/tests.rs +++ b/crates/storage/src/segment/tests.rs @@ -188,6 +188,26 @@ fn truncates_uncommitted_suffix() { assert_eq!(log.last_index(), Some(2)); } +#[test] +fn multi_segment_suffix_truncation_reopens_and_accepts_replacement() { + let directory = tempdir().unwrap(); + let mut log = SegmentLog::open(directory.path(), 100).unwrap(); + for value in 1..=5 { + log.append(record(0, &[value; 20]), true).unwrap(); + } + assert_eq!(log.segment_paths().unwrap().len(), 5); + + log.truncate_suffix(3).unwrap(); + drop(log); + + let mut log = SegmentLog::open(directory.path(), 100).unwrap(); + assert_eq!((log.first_index(), log.last_index()), (Some(1), Some(2))); + assert_eq!(log.read(2).unwrap().unwrap().payload, vec![2; 20]); + assert!(log.read(3).unwrap().is_none()); + log.append(record(0, b"replacement"), true).unwrap(); + assert_eq!(log.last_index(), Some(3)); +} + #[test] fn refuses_middle_corruption() { let directory = tempdir().unwrap(); diff --git a/deploy/helm/rustqueue/Chart.yaml b/deploy/helm/rustqueue/Chart.yaml index 38eaa7c..29d31e7 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.8.0 -appVersion: "0.8.0" +version: 0.8.1 +appVersion: "0.8.1" kubeVersion: ">=1.28.0-0" keywords: - messaging diff --git a/deploy/helm/rustqueue/values.yaml b/deploy/helm/rustqueue/values.yaml index faf1650..a63a45a 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.8.0" + tag: "0.8.1" pullPolicy: IfNotPresent log: rustqueue_operator=info resources: @@ -20,7 +20,7 @@ operator: queue: create: true name: rustqueue - image: rustqueue:0.8.0 + image: rustqueue:0.8.1 imagePullPolicy: IfNotPresent minBrokers: 1 maxBrokers: 500 diff --git a/docs/architecture/share-nothing-v7.md b/docs/architecture/share-nothing-v7.md index 2633035..3096be6 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.8.0 +Target release: 0.8.1 Data format: v7, clean directories only ## 1. Goal diff --git a/docs/operations/kubernetes.md b/docs/operations/kubernetes.md index b42ae7f..0c0ed85 100644 --- a/docs/operations/kubernetes.md +++ b/docs/operations/kubernetes.md @@ -76,7 +76,7 @@ Canary approval is optional: ```sh helm upgrade rustqueue deploy/helm/rustqueue \ --namespace rustqueue \ - --set queue.image=registry.example/rustqueue:0.8.0 \ + --set queue.image=registry.example/rustqueue:0.8.1 \ --set queue.rollout.requireCanaryApproval=true rustqueuectl -n rustqueue rollout approve @@ -88,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.8.0 +rustqueuectl -n rustqueue rollout rollback registry.example/rustqueue:0.8.1 rustqueuectl -n rustqueue rollout forward ``` diff --git a/docs/releases/v0.8.1.md b/docs/releases/v0.8.1.md new file mode 100644 index 0000000..7aa7e43 --- /dev/null +++ b/docs/releases/v0.8.1.md @@ -0,0 +1,52 @@ +# RustQueue 0.8.1 + +RustQueue 0.8.1 is a reliability patch for the real publish-to-delivery path. +It fixes benchmark accounting that could hide missing deliveries and hardens +the queue, storage, authentication, proxy and operational control planes found +during the follow-up P0/P1 review. + +## Highlights + +- `rustqueue-bench` now uses isolated durable Topics and Channels, waits for + consumers to become ready before publishing, counts unique and duplicate + deliveries, reports publish and receive throughput separately, and exits + unsuccessfully when verified delivery is incomplete. +- Per-delivery generation tokens prevent a stale `FIN`, `REQ` or `TOUCH` from + mutating a later redelivery of the same message. Delivery visibility also + covers output buffering and pending durable channel operations. +- Dead-letter moves are serialized as durable transactions. Cancellation, + restart and concurrent requests cannot publish multiple DLQ copies or + acknowledge the source before the target is durable. +- Payload and recovery-index workers keep path guards and memory reservations + until blocking I/O completes. Storage corruption is marked unhealthy before + a response can escape. +- AUTH response parsing and compiled regex state now share an explicit + node-wide memory budget. Broker management bodies, Kodo Stats aggregation, + proxy control bodies and backend error bodies are bounded as well. +- Invalid semaphore capacities and timer ranges are rejected during startup, + avoiding runtime panics or unbounded waits from malformed configuration. + +## Compatibility + +- The on-disk format remains v7; no data migration is required from 0.8.0. +- NSQ V2 commands, TLS/mTLS, AUTH, compression and the opt-in Kodo profile keep + their 0.8.0 compatibility contract. +- Delivery remains at least once. A message can be redelivered after an + ambiguous disconnect, but a stale acknowledgement can no longer settle that + newer delivery. + +## Release gate + +The `v0.8.1` tag is published only after GitHub Actions completes the full +non-Kubernetes production gate, builds native Linux x86_64 and ARM64 bundles, +packages the tagged source and Helm Chart, and verifies every SHA-256 checksum. + +## Assets + +- `rustqueue-0.8.1-linux-x86_64.tar.gz`: Linux x86_64 binaries, Console UI and + example configuration +- `rustqueue-0.8.1-linux-aarch64.tar.gz`: Linux ARM64 binaries, Console UI and + example configuration +- `rustqueue-0.8.1-source.tar.gz`: source archive for the tagged commit +- `rustqueue-0.8.1.tgz`: Helm Chart +- `SHA256SUMS-0.8.1`: SHA-256 checksums for all assets diff --git a/fuzz/Cargo.lock b/fuzz/Cargo.lock index 5396328..c02910f 100644 --- a/fuzz/Cargo.lock +++ b/fuzz/Cargo.lock @@ -2017,7 +2017,7 @@ dependencies = [ [[package]] name = "rustqueue-discovery" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "axum", @@ -2055,7 +2055,7 @@ dependencies = [ [[package]] name = "rustqueue-protocol" -version = "0.8.0" +version = "0.8.1" dependencies = [ "bytes", "serde", @@ -2065,7 +2065,7 @@ dependencies = [ [[package]] name = "rustqueue-proxy" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "axum", @@ -2085,7 +2085,7 @@ dependencies = [ [[package]] name = "rustqueue-queue" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "bytes", @@ -2105,7 +2105,7 @@ dependencies = [ [[package]] name = "rustqueue-storage" -version = "0.8.0" +version = "0.8.1" dependencies = [ "anyhow", "crc32c", @@ -2117,7 +2117,7 @@ dependencies = [ [[package]] name = "rustqueue-telemetry" -version = "0.8.0" +version = "0.8.1" dependencies = [ "serde", ] diff --git a/rustqueue.example.toml b/rustqueue.example.toml index aac875c..0048cdd 100644 --- a/rustqueue.example.toml +++ b/rustqueue.example.toml @@ -77,6 +77,8 @@ auth_response_bytes = 1048576 auth_timeout_ms = 5000 auth_max_ttl_seconds = 3600 auth_cache_max_entries = 10000 +# Bounds concurrent auth responses plus compiled authorization sessions retained by live clients and the cache. +auth_memory_bytes = 268435456 http_body_timeout_ms = 30000 [metrics] diff --git a/scripts/acceptance-multi-broker-k8s.sh b/scripts/acceptance-multi-broker-k8s.sh index b7e6e99..384ec56 100755 --- a/scripts/acceptance-multi-broker-k8s.sh +++ b/scripts/acceptance-multi-broker-k8s.sh @@ -343,9 +343,9 @@ require jq } if [[ "$BUILD_IMAGES" == "1" ]]; then - BUILD_VERSION=0.8.0-e2e-a MAX_STORAGE_FEATURE_LEVEL=1 make image + BUILD_VERSION=0.8.1-e2e-a MAX_STORAGE_FEATURE_LEVEL=1 make image docker tag rustqueue:dev "$BROKER_IMAGE_A" - BUILD_VERSION=0.8.0-e2e-b MAX_STORAGE_FEATURE_LEVEL=2 make image-from-dist + BUILD_VERSION=0.8.1-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}}')" ]] || { diff --git a/scripts/benchmark-compare.sh b/scripts/benchmark-compare.sh index bbba921..38b498d 100755 --- a/scripts/benchmark-compare.sh +++ b/scripts/benchmark-compare.sh @@ -86,8 +86,24 @@ done for payload in $payloads; do for name in rustqueue-local-fsync nsq-sync-every-1 nsq-sync-every-2500 \ rustqueue-local-fsync-fixed nsq-sync-every-1-fixed; do - jq -s '{runs: length, median: (sort_by(.messages_per_second) | .[length / 2 | floor])}' \ - "$result_dir/$name-$payload-run"*.json > "$result_dir/$name-$payload-median.json" + jq -s ' + def median_by($field): + map(select(.[$field] != null)) + | sort_by(.[$field]) + | if length == 0 then null else .[length / 2 | floor] end; + { + runs: length, + complete_delivery_runs: map(select(.delivery_complete == true)) | length, + incomplete_delivery_runs: map(select(.delivery_complete != true)) | length, + median: median_by("publish_messages_per_second"), + median_publish: median_by("publish_messages_per_second"), + median_receive: ( + map(select(.delivery_complete == true)) + | median_by("receive_messages_per_second") + ) + } + ' "$result_dir/$name-$payload-run"*.json \ + > "$result_dir/$name-$payload-median.json" done done diff --git a/scripts/package-release.sh b/scripts/package-release.sh new file mode 100755 index 0000000..c01bc4a --- /dev/null +++ b/scripts/package-release.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +MODE="${1:-}" +VERSION="${2:-}" +OUTPUT_ARG="${3:-release}" + +"$ROOT/scripts/verify-release-version.sh" "$VERSION" +mkdir -p "$OUTPUT_ARG" +OUTPUT="$(cd "$OUTPUT_ARG" && pwd)" + +case "$MODE" in + binaries) + ARCH="${4:-}" + case "$ARCH" in + x86_64|aarch64) ;; + *) + echo "binary architecture must be x86_64 or aarch64" >&2 + exit 2 + ;; + esac + + BINARIES=( + rustqueued + rustqueue-discovery + rustqueue-proxy + rustqueue-bench + rustqueuectl + rustqueue-console + rustqueue-operator + ) + for binary in "${BINARIES[@]}"; do + [[ -x "$ROOT/.docker-bin/$binary" ]] || { + echo "missing release binary: .docker-bin/$binary" >&2 + exit 1 + } + done + [[ -f "$ROOT/console-ui/dist/index.html" ]] || { + echo "console-ui/dist is missing; run make console-ui-build" >&2 + exit 1 + } + + STAGING="$(mktemp -d)" + trap 'rm -rf "$STAGING"' EXIT + PACKAGE="rustqueue-$VERSION" + mkdir -p "$STAGING/$PACKAGE/bin" "$STAGING/$PACKAGE/console-ui" + for binary in "${BINARIES[@]}"; do + install -m 0755 "$ROOT/.docker-bin/$binary" "$STAGING/$PACKAGE/bin/$binary" + done + cp -R "$ROOT/console-ui/dist/." "$STAGING/$PACKAGE/console-ui/" + install -m 0644 "$ROOT/README.md" "$STAGING/$PACKAGE/README.md" + install -m 0644 \ + "$ROOT/rustqueue.example.toml" \ + "$STAGING/$PACKAGE/rustqueue.example.toml" + tar -C "$STAGING" -czf \ + "$OUTPUT/rustqueue-$VERSION-linux-$ARCH.tar.gz" \ + "$PACKAGE" + ;; + common) + git -C "$ROOT" archive \ + --format=tar.gz \ + --prefix="rustqueue-$VERSION/" \ + --output="$OUTPUT/rustqueue-$VERSION-source.tar.gz" \ + HEAD + helm package "$ROOT/deploy/helm/rustqueue" --destination "$OUTPUT" + ;; + *) + echo "usage: $0 [output-dir] [architecture]" >&2 + exit 2 + ;; +esac diff --git a/scripts/rss-gate.sh b/scripts/rss-gate.sh index d540374..1f1bd88 100755 --- a/scripts/rss-gate.sh +++ b/scripts/rss-gate.sh @@ -60,6 +60,7 @@ docker exec "$container" /usr/local/bin/rustqueue-bench \ --message-bytes "$message_bytes" \ --batch-size "$batch_size" \ --producers 16 \ + --reuse-topic \ --consumers 0 >/tmp/rustqueue-rss-gate-first.out sleep 5 warm=$(rss_bytes) @@ -70,6 +71,7 @@ docker exec "$container" /usr/local/bin/rustqueue-bench \ --message-bytes "$message_bytes" \ --batch-size "$batch_size" \ --producers 16 \ + --reuse-topic \ --consumers 0 >/tmp/rustqueue-rss-gate-second.out sleep 5 after=$(rss_bytes) diff --git a/scripts/verify-release-version.sh b/scripts/verify-release-version.sh new file mode 100755 index 0000000..111ebaa --- /dev/null +++ b/scripts/verify-release-version.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +EXPECTED="${1:-}" + +if [[ ! "$EXPECTED" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then + echo "usage: $0 " >&2 + exit 2 +fi + +read_workspace_version() { + awk ' + /^\[workspace\.package\]$/ { workspace = 1; next } + /^\[/ { workspace = 0 } + workspace && $1 == "version" { + gsub(/"/, "", $3) + print $3 + exit + } + ' "$ROOT/Cargo.toml" +} + +read_chart_field() { + local field="$1" + awk -v field="$field" ' + $1 == field ":" { + gsub(/"/, "", $2) + print $2 + exit + } + ' "$ROOT/deploy/helm/rustqueue/Chart.yaml" +} + +read_ui_version() { + awk -F'"' '$2 == "version" { print $4; exit }' \ + "$ROOT/console-ui/package.json" +} + +require_version() { + local label="$1" + local actual="$2" + if [[ "$actual" != "$EXPECTED" ]]; then + echo "$label version is $actual, expected $EXPECTED" >&2 + exit 1 + fi +} + +verify_lock() { + local lock="$1" + awk -v expected="$EXPECTED" ' + $1 == "name" && $3 ~ /^"rustqueue/ && $3 != "\"rustqueue-fuzz\"" { + package = $3 + getline + wanted = "\"" expected "\"" + if ($1 != "version" || $3 != wanted) { + printf "%s has version %s in %s, expected %s\n", + package, $3, FILENAME, wanted > "/dev/stderr" + failed = 1 + } + checked += 1 + } + END { + if (checked == 0) { + printf "no RustQueue packages found in %s\n", FILENAME > "/dev/stderr" + exit 1 + } + exit failed + } + ' "$lock" +} + +require_version "workspace" "$(read_workspace_version)" +require_version "Helm Chart" "$(read_chart_field version)" +require_version "Helm app" "$(read_chart_field appVersion)" +require_version "Console UI" "$(read_ui_version)" +verify_lock "$ROOT/Cargo.lock" +verify_lock "$ROOT/fuzz/Cargo.lock" + +grep -Fq "tag: \"$EXPECTED\"" "$ROOT/deploy/helm/rustqueue/values.yaml" || { + echo "operator image tag is not $EXPECTED" >&2 + exit 1 +} +grep -Fq "image: rustqueue:$EXPECTED" "$ROOT/deploy/helm/rustqueue/values.yaml" || { + echo "broker image tag is not $EXPECTED" >&2 + exit 1 +} +grep -Fq "Current release: [v$EXPECTED]" "$ROOT/README.md" || { + echo "README current release is not v$EXPECTED" >&2 + exit 1 +} +[[ -f "$ROOT/docs/releases/v$EXPECTED.md" ]] || { + echo "docs/releases/v$EXPECTED.md is missing" >&2 + exit 1 +} + +echo "RustQueue release metadata is consistent at $EXPECTED"