feat(amt-verify): add --packet-count for multi-packet one-shot verification

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

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;
 
@@ -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,
@@ -50,6 +54,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 +75,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 +93,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 {
@@ -97,25 +122,37 @@ struct OneshotReport {
     family: &'static str,
     group: String,
     source: Option<String>,
+    packet_count: u64,
+    byte_count: u64,
     timings_ms: Timings,
     first_packet: FirstPacket,
 }
 
 #[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 {
     let args = Args::parse();
     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 +175,23 @@ 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"
+        )));
+    }
+    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 {
+        ShutdownMode::TeardownOnly
+    } else {
+        ShutdownMode::GracefulLeave
+    };
 
     // ----- Build gateway (explicit relay OR DRIAD path) -----
     let (gw, resolved_relay) = match args.relay {
@@ -211,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
@@ -231,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!(
@@ -252,47 +309,95 @@ 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()
         );
     }
 
     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(())
 }
 
-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<amt_protocol::native::DataEvent>,
     group: IpAddr,
     source: IpAddr,
+    want_count: u64,
     timeout: Duration,
-) -> Result<amt_protocol::native::DataEvent> {
+) -> Result<PacketStats> {
     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())
-            .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(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) => return Err(anyhow!("data broadcast closed before first matching packet")),
+            Err(RecvError::Closed) => {
+                return Err(anyhow!(
+                    "data broadcast closed before first matching packet"
+                ));
+            }
         }
     }
 }
 
 async fn run_watch(
     gw: AsyncAmtGateway,
     mut data_rx: tokio::sync::broadcast::Receiver<amt_protocol::native::DataEvent>,
+    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 +434,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(())
+        }
+    }
+}