From 69d5ea3ce30377a745ece4e4c9650d8442a01f9b Mon Sep 17 00:00:00 2001 From: Omar Ramadan Date: Sat, 23 May 2026 09:03:19 +0000 Subject: [PATCH 1/2] Add AMT verifier non-graceful shutdown modes --- docs/runbook-staging-e2e.md | 25 +++++++ src/bin/amt-verify.rs | 127 +++++++++++++++++++++++++++++++----- tests/cli_json.rs | 112 +++++++++++++++++++++++++++++-- 3 files changed, 239 insertions(+), 25 deletions(-) diff --git a/docs/runbook-staging-e2e.md b/docs/runbook-staging-e2e.md index 2a60bff..c761947 100644 --- a/docs/runbook-staging-e2e.md +++ b/docs/runbook-staging-e2e.md @@ -38,6 +38,31 @@ Expected output (one line of JSON): Exit code: 0 on success, 1 on timeout / handshake failure. +## Billing shutdown matrix + +Use the same `--relay`, `--source`, `--group`, and `--timeout` values as the +direct-path check. The final flag selects the AMT shutdown shape the relay must +bill: + +```bash +# Nominal: Membership Update leave, then AMT Teardown. +./target/release/amt-verify --relay "$STAGING_RELAY" \ + --source "$STAGING_SOURCE" --group "$STAGING_GROUP" --timeout 30 --json + +# Non-graceful app leave: skip Membership Update leave, send AMT Teardown. +./target/release/amt-verify --relay "$STAGING_RELAY" \ + --source "$STAGING_SOURCE" --group "$STAGING_GROUP" --timeout 30 --json \ + --no-graceful-leave + +# Hard loss: drop the gateway runtime with no leave and no AMT Teardown. +# Staging sets EBPF_BILLING_STALE_TIMEOUT=45s and +# EBPF_BILLING_STALE_SWEEP_INTERVAL=5s, so billing should finalize within +# roughly 50s after first data. +./target/release/amt-verify --relay "$STAGING_RELAY" \ + --source "$STAGING_SOURCE" --group "$STAGING_GROUP" --timeout 30 --json \ + --drop-without-teardown +``` + ## DRIAD path — from in-cluster pod ```bash diff --git a/src/bin/amt-verify.rs b/src/bin/amt-verify.rs index 5f1f66a..36f93e6 100644 --- a/src/bin/amt-verify.rs +++ b/src/bin/amt-verify.rs @@ -4,7 +4,7 @@ use std::net::IpAddr; use std::process::ExitCode; use std::time::{Duration, Instant}; -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use clap::Parser; use tracing_subscriber::EnvFilter; @@ -50,6 +50,16 @@ struct Args { #[arg(long, default_value_t = false)] watch: bool, + /// On shutdown, skip the Membership Update leave and send only AMT Teardown. + /// Useful for billing-path negative tests. + #[arg(long, default_value_t = false)] + no_graceful_leave: bool, + + /// On shutdown, drop the gateway runtime without Membership Update leave or AMT Teardown. + /// This simulates hard client loss for relay-side stale-expiry billing tests. + #[arg(long, default_value_t = false)] + drop_without_teardown: bool, + /// Machine-readable JSON output (one-shot mode only). /// Rejected with exit 2 if combined with --watch. #[arg(long, default_value_t = false)] @@ -61,7 +71,11 @@ struct Args { } #[derive(Copy, Clone, Debug, clap::ValueEnum)] -enum Family { V4, V6, Auto } +enum Family { + V4, + V6, + Auto, +} /// Exit-code classification per spec: /// 0 → success (one-shot data observed, or watch SIGINT clean teardown) @@ -75,12 +89,19 @@ enum ExitCategory { Fatal(anyhow::Error), } +#[derive(Copy, Clone, Debug)] +enum ShutdownMode { + GracefulLeave, + TeardownOnly, + DropWithoutTeardown, +} + impl ExitCategory { fn code(&self) -> u8 { match self { ExitCategory::HandshakeFail(_) => 1, - ExitCategory::Config(_) => 2, - ExitCategory::Fatal(_) => 3, + ExitCategory::Config(_) => 2, + ExitCategory::Fatal(_) => 3, } } fn err(&self) -> &anyhow::Error { @@ -102,10 +123,16 @@ struct OneshotReport { } #[derive(serde::Serialize)] -struct Timings { first_data: u64 } +struct Timings { + first_data: u64, +} #[derive(serde::Serialize)] -struct FirstPacket { src: String, dst_port: u16, len: usize } +struct FirstPacket { + src: String, + dst_port: u16, + len: usize, +} #[tokio::main(flavor = "current_thread")] async fn main() -> ExitCode { @@ -113,9 +140,13 @@ async fn main() -> ExitCode { let filter = if args.verbose { EnvFilter::new("amt=debug,amt_protocol=debug") } else { - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("amt=info,amt_protocol=info")) + EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("amt=info,amt_protocol=info")) }; - tracing_subscriber::fmt().with_env_filter(filter).with_writer(std::io::stderr).init(); + tracing_subscriber::fmt() + .with_env_filter(filter) + .with_writer(std::io::stderr) + .init(); match run(args).await { Ok(()) => ExitCode::from(0), @@ -138,6 +169,18 @@ async fn run(args: Args) -> std::result::Result<(), ExitCategory> { "--no-driad set but --relay missing" ))); } + if args.no_graceful_leave && args.drop_without_teardown { + return Err(ExitCategory::Config(anyhow!( + "--no-graceful-leave and --drop-without-teardown are mutually exclusive" + ))); + } + let shutdown_mode = if args.drop_without_teardown { + ShutdownMode::DropWithoutTeardown + } else if args.no_graceful_leave { + ShutdownMode::TeardownOnly + } else { + ShutdownMode::GracefulLeave + }; // ----- Build gateway (explicit relay OR DRIAD path) ----- let (gw, resolved_relay) = match args.relay { @@ -259,9 +302,13 @@ async fn run(args: Args) -> std::result::Result<(), ExitCategory> { } if args.watch { - run_watch(gw, data_rx).await.map_err(ExitCategory::Fatal)?; + run_watch(gw, data_rx, args.group, args.source, shutdown_mode) + .await + .map_err(ExitCategory::Fatal)?; } else { - gw.shutdown().await.map_err(ExitCategory::Fatal)?; + finish_gateway(gw, args.group, args.source, shutdown_mode) + .await + .map_err(ExitCategory::Fatal)?; } Ok(()) } @@ -275,17 +322,35 @@ async fn recv_first_matching( use tokio::sync::broadcast::error::RecvError; let deadline = tokio::time::Instant::now() + timeout; loop { - let remaining = deadline.checked_duration_since(tokio::time::Instant::now()) - .ok_or_else(|| anyhow!("timed out after {}s waiting for first data matching ({}, {})", - timeout.as_secs(), group, source))?; - let recv = tokio::time::timeout(remaining, rx.recv()).await - .map_err(|_| anyhow!("timed out after {}s waiting for first data matching ({}, {})", - timeout.as_secs(), group, source))?; + let remaining = deadline + .checked_duration_since(tokio::time::Instant::now()) + .ok_or_else(|| { + anyhow!( + "timed out after {}s waiting for first data matching ({}, {})", + timeout.as_secs(), + group, + source + ) + })?; + let recv = tokio::time::timeout(remaining, rx.recv()) + .await + .map_err(|_| { + anyhow!( + "timed out after {}s waiting for first data matching ({}, {})", + timeout.as_secs(), + group, + source + ) + })?; match recv { Ok(evt) if evt.group == group && evt.src == source => return Ok(evt), Ok(_skip) => continue, Err(RecvError::Lagged(_)) => continue, - Err(RecvError::Closed) => return Err(anyhow!("data broadcast closed before first matching packet")), + Err(RecvError::Closed) => { + return Err(anyhow!( + "data broadcast closed before first matching packet" + )); + } } } } @@ -293,6 +358,9 @@ async fn recv_first_matching( async fn run_watch( gw: AsyncAmtGateway, mut data_rx: tokio::sync::broadcast::Receiver, + group: IpAddr, + source: IpAddr, + shutdown_mode: ShutdownMode, ) -> Result<()> { use tokio::sync::broadcast::error::RecvError; let mut tick = tokio::time::interval(Duration::from_secs(5)); @@ -329,6 +397,29 @@ async fn run_watch( } } } - gw.shutdown().await?; + finish_gateway(gw, group, source, shutdown_mode).await?; Ok(()) } + +async fn finish_gateway( + gw: AsyncAmtGateway, + group: IpAddr, + source: IpAddr, + mode: ShutdownMode, +) -> Result<()> { + match mode { + ShutdownMode::GracefulLeave => { + if let Err(e) = gw.unsubscribe(group, Some(source)).await { + tracing::warn!(target: "amt", error=?e, "unsubscribe before shutdown failed"); + } else { + tokio::time::sleep(Duration::from_millis(100)).await; + } + gw.shutdown().await + } + ShutdownMode::TeardownOnly => gw.shutdown().await, + ShutdownMode::DropWithoutTeardown => { + drop(gw); + Ok(()) + } + } +} diff --git a/tests/cli_json.rs b/tests/cli_json.rs index c2c99d4..e1a5020 100644 --- a/tests/cli_json.rs +++ b/tests/cli_json.rs @@ -2,10 +2,11 @@ mod common; -use common::fake_relay::{synth_v4_udp, FakeRelay}; +use amt_protocol::messages::MessageType; +use common::fake_relay::{FakeRelay, synth_v4_udp}; use std::process::Stdio; -use tokio::process::Command; use tokio::io::AsyncReadExt; +use tokio::process::Command; #[tokio::test(flavor = "current_thread")] async fn json_output_is_parseable() { @@ -16,11 +17,16 @@ async fn json_output_is_parseable() { let bin = env!("CARGO_BIN_EXE_amt-verify"); let mut child = Command::new(bin) .args([ - "--relay", &relay.addr.ip().to_string(), - "--port", &relay.addr.port().to_string(), - "--group", "232.0.0.1", - "--source", "10.0.0.1", - "--timeout", "5", + "--relay", + &relay.addr.ip().to_string(), + "--port", + &relay.addr.port().to_string(), + "--group", + "232.0.0.1", + "--source", + "10.0.0.1", + "--timeout", + "5", "--json", ]) .stdout(Stdio::piped()) @@ -39,3 +45,95 @@ async fn json_output_is_parseable() { assert_eq!(v["group"], "232.0.0.1"); assert_eq!(v["first_packet"]["src"], "10.0.0.1:5004"); } + +#[tokio::test(flavor = "current_thread")] +async fn default_shutdown_sends_leave_then_teardown() { + let relay = run_verify(&[]).await; + let types = captured_types(&relay).await; + + assert_eq!(count(&types, MessageType::MembershipUpdate), 2, "{types:?}"); + assert_eq!(count(&types, MessageType::Teardown), 1, "{types:?}"); +} + +#[tokio::test(flavor = "current_thread")] +async fn no_graceful_leave_sends_teardown_without_leave_update() { + let relay = run_verify(&["--no-graceful-leave"]).await; + let types = captured_types(&relay).await; + + assert_eq!(count(&types, MessageType::MembershipUpdate), 1, "{types:?}"); + assert_eq!(count(&types, MessageType::Teardown), 1, "{types:?}"); +} + +#[tokio::test(flavor = "current_thread")] +async fn drop_without_teardown_sends_no_leave_or_teardown() { + let relay = run_verify(&["--drop-without-teardown"]).await; + let types = captured_types(&relay).await; + + assert_eq!(count(&types, MessageType::MembershipUpdate), 1, "{types:?}"); + assert_eq!(count(&types, MessageType::Teardown), 0, "{types:?}"); +} + +#[tokio::test(flavor = "current_thread")] +async fn non_graceful_modes_are_mutually_exclusive() { + let bin = env!("CARGO_BIN_EXE_amt-verify"); + let output = Command::new(bin) + .args([ + "--relay", + "127.0.0.1", + "--port", + "2268", + "--group", + "232.0.0.1", + "--source", + "10.0.0.1", + "--no-graceful-leave", + "--drop-without-teardown", + ]) + .output() + .await + .expect("spawn amt-verify"); + + assert_eq!(output.status.code(), Some(2)); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("mutually exclusive"), "{stderr}"); +} + +async fn run_verify(extra_args: &[&str]) -> FakeRelay { + let relay = FakeRelay::bind("v4").await; + let inner = synth_v4_udp([10, 0, 0, 1], [232, 0, 0, 1], 5004, 5005, b"x"); + relay.spawn(inner); + + let bin = env!("CARGO_BIN_EXE_amt-verify"); + let mut args = vec![ + "--relay".to_string(), + relay.addr.ip().to_string(), + "--port".to_string(), + relay.addr.port().to_string(), + "--group".to_string(), + "232.0.0.1".to_string(), + "--source".to_string(), + "10.0.0.1".to_string(), + "--timeout".to_string(), + "5".to_string(), + ]; + args.extend(extra_args.iter().map(|arg| arg.to_string())); + + let status = Command::new(bin) + .args(args) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .await + .expect("spawn amt-verify"); + assert!(status.success(), "exit code: {status}"); + relay +} + +async fn captured_types(relay: &FakeRelay) -> Vec { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + relay.captured.lock().await.message_types.clone() +} + +fn count(types: &[u8], msg_type: MessageType) -> usize { + types.iter().filter(|&&t| t == msg_type as u8).count() +} From 6f6408e1fd614a4005f12fb7c3d7508175680d35 Mon Sep 17 00:00:00 2001 From: Omar Ramadan Date: Fri, 26 Jun 2026 03:03:51 +0000 Subject: [PATCH 2/2] feat(amt-verify): add --packet-count for multi-packet one-shot verification Replace recv_first_matching with recv_matching_packets so amt-verify can wait for N matching multicast packets before a successful one-shot shutdown. Report packet_count and byte_count in the JSON OneshotReport; reject --packet-count 0 at config validation. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/bin/amt-verify.rs | 57 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 47 insertions(+), 10 deletions(-) diff --git a/src/bin/amt-verify.rs b/src/bin/amt-verify.rs index 36f93e6..e7e402a 100644 --- a/src/bin/amt-verify.rs +++ b/src/bin/amt-verify.rs @@ -42,6 +42,10 @@ struct Args { #[arg(long, default_value = "30")] timeout: u64, + /// Matching multicast packets to receive before successful one-shot shutdown + #[arg(long, default_value = "1")] + packet_count: u64, + /// Keep-alive interval in seconds #[arg(long, default_value = "60")] keepalive: u64, @@ -118,6 +122,8 @@ struct OneshotReport { family: &'static str, group: String, source: Option, + packet_count: u64, + byte_count: u64, timings_ms: Timings, first_packet: FirstPacket, } @@ -174,6 +180,11 @@ async fn run(args: Args) -> std::result::Result<(), ExitCategory> { "--no-graceful-leave and --drop-without-teardown are mutually exclusive" ))); } + if args.packet_count == 0 { + return Err(ExitCategory::Config(anyhow!( + "--packet-count must be greater than 0" + ))); + } let shutdown_mode = if args.drop_without_teardown { ShutdownMode::DropWithoutTeardown } else if args.no_graceful_leave { @@ -254,10 +265,11 @@ async fn run(args: Args) -> std::result::Result<(), ExitCategory> { .await .map_err(ExitCategory::HandshakeFail)?; - let evt = match recv_first_matching( + let stats = match recv_matching_packets( &mut data_rx, args.group, args.source, + args.packet_count, Duration::from_secs(args.timeout), ) .await @@ -274,13 +286,15 @@ async fn run(args: Args) -> std::result::Result<(), ExitCategory> { family: family_str, group: args.group.to_string(), source: Some(args.source.to_string()), + packet_count: stats.packet_count, + byte_count: stats.byte_count, timings_ms: Timings { first_data: first_data_ms, }, first_packet: FirstPacket { - src: format!("{}:{}", evt.src, evt.src_port), - dst_port: evt.dst_port, - len: evt.payload.len(), + src: format!("{}:{}", stats.first.src, stats.first.src_port), + dst_port: stats.first.dst_port, + len: stats.first.payload.len(), }, }; println!( @@ -295,9 +309,9 @@ async fn run(args: Args) -> std::result::Result<(), ExitCategory> { args.group, args.source, first_data_ms, - evt.src, - evt.src_port, - evt.payload.len() + stats.first.src, + stats.first.src_port, + stats.first.payload.len() ); } @@ -313,14 +327,24 @@ async fn run(args: Args) -> std::result::Result<(), ExitCategory> { Ok(()) } -async fn recv_first_matching( +struct PacketStats { + first: amt_protocol::native::DataEvent, + packet_count: u64, + byte_count: u64, +} + +async fn recv_matching_packets( rx: &mut tokio::sync::broadcast::Receiver, group: IpAddr, source: IpAddr, + want_count: u64, timeout: Duration, -) -> Result { +) -> Result { use tokio::sync::broadcast::error::RecvError; let deadline = tokio::time::Instant::now() + timeout; + let mut first = None; + let mut packet_count = 0; + let mut byte_count = 0; loop { let remaining = deadline .checked_duration_since(tokio::time::Instant::now()) @@ -343,7 +367,20 @@ async fn recv_first_matching( ) })?; match recv { - Ok(evt) if evt.group == group && evt.src == source => return Ok(evt), + Ok(evt) if evt.group == group && evt.src == source => { + if first.is_none() { + first = Some(evt.clone()); + } + packet_count += 1; + byte_count += evt.payload.len() as u64; + if packet_count >= want_count { + return Ok(PacketStats { + first: first.expect("first matching packet recorded"), + packet_count, + byte_count, + }); + } + } Ok(_skip) => continue, Err(RecvError::Lagged(_)) => continue, Err(RecvError::Closed) => {