Add bounded shard clients and atomic membership fencing

crates/transport/src/client.rs

@@ -17,6 +17,17 @@ use crate::clock::Clock;
 use crate::frame::{read_envelope, write_envelope};
 use crate::pool::Pool;
 
+/// Counts complete application envelopes, including the hello exchange.
+/// A framed envelope adds 12 bytes. TCP/IP headers, retransmission, and partial
+/// failed writes/reads are excluded; these counters are not NIC wire telemetry.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct ClientIoStats {
+    pub rpc_body_bytes_sent: u64,
+    pub rpc_body_bytes_received: u64,
+    pub framed_bytes_sent: u64,
+    pub framed_bytes_received: u64,
+}
+
 pub struct TcpClient {
     stream: TcpStream,
     auth: AuthContext,
@@ -25,6 +36,7 @@ pub struct TcpClient {
     epoch: u64,
     clock: Clock,
     usable: bool,
+    io_stats: ClientIoStats,
 }
 
 impl TcpClient {
@@ -73,9 +85,27 @@ impl TcpClient {
             epoch,
             clock: Clock::new(envelope.server_now),
             usable: true,
+            io_stats: ClientIoStats {
+                rpc_body_bytes_sent: hello.len() as u64,
+                rpc_body_bytes_received: envelope.body().len() as u64,
+                framed_bytes_sent: hello.len() as u64 + 12,
+                framed_bytes_received: envelope.body().len() as u64 + 12,
+            },
         })
     }
 
+    pub fn io_stats(&self) -> ClientIoStats {
+        self.io_stats
+    }
+
+    pub(crate) fn is_usable(&self) -> bool {
+        self.usable
+    }
+
+    pub(crate) fn server_now(&self) -> u64 {
+        self.clock.server_now()
+    }
+
     pub fn epoch(&self) -> u64 {
         self.epoch
     }
@@ -158,7 +188,23 @@ impl TcpClient {
         // incomplete exchange, a late reply must never answer a later call.
         self.usable = false;
         write_envelope(&mut self.stream, 0, &bytes, self.max_frame)?;
+        self.io_stats.rpc_body_bytes_sent = self
+            .io_stats
+            .rpc_body_bytes_sent
+            .saturating_add(bytes.len() as u64);
+        self.io_stats.framed_bytes_sent = self
+            .io_stats
+            .framed_bytes_sent
+            .saturating_add(bytes.len() as u64 + 12);
         let envelope = read_envelope(&mut self.stream, self.pool.checkout()?, self.max_frame)?;
+        self.io_stats.rpc_body_bytes_received = self
+            .io_stats
+            .rpc_body_bytes_received
+            .saturating_add(envelope.body().len() as u64);
+        self.io_stats.framed_bytes_received = self
+            .io_stats
+            .framed_bytes_received
+            .saturating_add(envelope.body().len() as u64 + 12);
         match decode_reply(envelope.body())? {
             ReplyBody::Failed {
                 op,
@@ -497,4 +543,25 @@ mod tests {
             })
         ));
     }
+
+    #[test]
+    fn io_counters_separate_complete_rpc_body_and_framed_bytes() {
+        let (mut client, server) = fake_server(vec![encode_reply(&lookup_reply(1, &[1])).unwrap()]);
+        let before = client.io_stats();
+        assert_eq!(before.framed_bytes_sent - before.rpc_body_bytes_sent, 12);
+        assert_eq!(
+            before.framed_bytes_received - before.rpc_body_bytes_received,
+            12
+        );
+        client.lookup(vec![lookup_item(1)]).unwrap();
+        let after = client.io_stats();
+        assert_eq!(after.framed_bytes_sent - after.rpc_body_bytes_sent, 24);
+        assert_eq!(
+            after.framed_bytes_received - after.rpc_body_bytes_received,
+            24
+        );
+        assert!(after.rpc_body_bytes_sent > before.rpc_body_bytes_sent);
+        assert!(after.rpc_body_bytes_received > before.rpc_body_bytes_received);
+        server.join().unwrap();
+    }
 }

crates/transport/src/directory.rs

@@ -97,6 +97,17 @@ struct Flight {
     waiters: BTreeSet<[u8; 16]>,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct DirectoryStats {
+    pub flights: usize,
+    pub request_entries: usize,
+    pub notices: usize,
+    pub tenants: usize,
+    pub members: usize,
+    pub fence_count: u64,
+    pub fenced_fills: u64,
+}
+
 pub struct ShardDirectory {
     shard_id: u64,
     epoch: u64,
@@ -115,6 +126,8 @@ pub struct ShardDirectory {
     notices: BTreeMap<(Vec<u8>, [u8; 16]), FillNotice>,
     notice_order: VecDeque<(Vec<u8>, [u8; 16])>,
     max_notices: usize,
+    fence_count: u64,
+    fenced_fills: u64,
 }
 
 impl ShardDirectory {
@@ -145,13 +158,63 @@ impl ShardDirectory {
             next_flight: 1,
             notices: BTreeMap::new(),
             notice_order: VecDeque::new(),
+            fence_count: 0,
+            fenced_fills: 0,
             max_notices: params
                 .max_flights
                 .saturating_mul(params.max_waiters.saturating_add(1))
                 .max(1),
         }
     }
 
+    pub fn stats(&self) -> DirectoryStats {
+        DirectoryStats {
+            flights: self.flights.len(),
+            request_entries: self.by_request.len(),
+            notices: self.notices.len(),
+            tenants: self.tenant_used.len(),
+            members: self.members.len(),
+            fence_count: self.fence_count,
+            fenced_fills: self.fenced_fills,
+        }
+    }
+
+    /// Apply one bounded membership snapshot atomically. Validation and epoch
+    /// capacity checks finish before membership, credits, or notices change.
+    /// There is no update queue: callers retry rejected snapshots explicitly.
+    pub fn apply_membership(&mut self, members: &[MemberView]) -> Result<u64, CacheError> {
+        if members.is_empty() || members.len() > self.max_members {
+            return Err(CacheError::Malformed {
+                reason: "membership",
+            });
+        }
+        let mut seen = BTreeSet::new();
+        for member in members {
+            if member.incarnation == 0 || !seen.insert(member.node_id) {
+                return Err(CacheError::Malformed {
+                    reason: "membership",
+                });
+            }
+            if self
+                .members
+                .iter()
+                .any(|old| old.node_id == member.node_id && old.incarnation > member.incarnation)
+            {
+                return Err(CacheError::FencingRejected { reason: "member" });
+            }
+        }
+        if members.len() == self.members.len()
+            && members.iter().all(|member| self.members.contains(member))
+        {
+            return Ok(self.epoch);
+        }
+        self.ensure_epoch_can_advance()?;
+        let epoch = self.fence_inflight()?;
+        self.members = members.to_vec();
+        self.members.sort_by_key(|member| member.node_id);
+        Ok(epoch)
+    }
+
     pub fn shard_id(&self) -> u64 {
         self.shard_id
     }
@@ -386,6 +449,8 @@ impl ShardDirectory {
             kind: QuotaKind::InFlightRequests,
         })?;
         self.epoch = next;
+        self.fence_count = self.fence_count.saturating_add(1);
+        self.fenced_fills = self.fenced_fills.saturating_add(self.flights.len() as u64);
         self.by_request.clear();
         self.by_id.clear();
         let flights = std::mem::take(&mut self.flights);
@@ -873,4 +938,114 @@ mod tests {
             FillNotice::LeaderFinished
         );
     }
+
+    #[test]
+    fn membership_batch_rejects_every_invalid_snapshot_before_fencing() {
+        let mut dir = ShardDirectory::single_node(params());
+        dir.note_member(2, 2).unwrap();
+        dir.begin_or_join(&fill(b"a", b"tenant-a", 10, 1, dir.epoch()))
+            .unwrap();
+        dir.begin_or_join(&fill(b"a", b"tenant-a", 10, 2, dir.epoch()))
+            .unwrap();
+        let before = dir.stats();
+        let members = dir.members().to_vec();
+        for candidate in [
+            vec![],
+            vec![MemberView {
+                node_id: 2,
+                incarnation: 1,
+            }],
+            vec![members[0], members[0]],
+            vec![MemberView {
+                node_id: 3,
+                incarnation: 0,
+            }],
+            (1..=5)
+                .map(|node_id| MemberView {
+                    node_id,
+                    incarnation: 2,
+                })
+                .collect(),
+        ] {
+            assert!(dir.apply_membership(&candidate).is_err());
+            assert_eq!(dir.stats(), before);
+            assert_eq!(dir.members(), members);
+            assert_eq!(dir.credits_used(), 10);
+            assert_eq!(
+                dir.fill_notice(b"tenant-a", 2u128.to_le_bytes()).unwrap(),
+                FillNotice::Pending
+            );
+        }
+        dir.set_epoch_for_test(u64::MAX);
+        assert!(dir.apply_membership(&[members[0]]).is_err());
+        assert_eq!(dir.stats(), before);
+        assert_eq!(dir.members(), members);
+    }
+
+    #[test]
+    fn batch_membership_is_order_independent_and_fences_once() {
+        let mut dir = ShardDirectory::single_node(params());
+        dir.note_member(2, 1).unwrap();
+        dir.begin_or_join(&fill(b"a", b"tenant-a", 10, 1, dir.epoch()))
+            .unwrap();
+        let mut same = dir.members().to_vec();
+        same.reverse();
+        assert_eq!(dir.apply_membership(&same).unwrap(), 2);
+        assert_eq!(dir.inflight(), 1);
+        let snapshot = [
+            MemberView {
+                node_id: 2,
+                incarnation: 2,
+            },
+            MemberView {
+                node_id: 3,
+                incarnation: 1,
+            },
+        ];
+        assert_eq!(dir.apply_membership(&snapshot).unwrap(), 3);
+        assert_eq!(dir.inflight(), 0);
+        assert_eq!(dir.credits_used(), 0);
+        assert_eq!(dir.stats().fenced_fills, 1);
+        assert_eq!(dir.stats().fence_count, 2);
+        assert_eq!(
+            dir.fill_notice(b"tenant-a", 1u128.to_le_bytes()).unwrap(),
+            FillNotice::Fenced
+        );
+    }
+
+    #[test]
+    fn cancel_restart_and_membership_storm_keep_every_table_bounded() {
+        let mut dir = ShardDirectory::single_node(params());
+        for n in 1..=2000u128 {
+            let epoch = dir.epoch();
+            let Admit::Leader { flight_id, .. } = dir
+                .begin_or_join(&fill(b"a", b"tenant-a", 10, n * 2, epoch))
+                .unwrap()
+            else {
+                panic!("leader")
+            };
+            dir.begin_or_join(&fill(b"a", b"tenant-a", 10, n * 2 + 1, epoch))
+                .unwrap();
+            if n % 2 == 0 {
+                dir.fence_epoch().unwrap();
+                assert!(dir
+                    .complete_flight(flight_id, epoch, (n * 2).to_le_bytes())
+                    .is_err());
+            } else {
+                dir.fail_flight(flight_id, epoch, (n * 2).to_le_bytes())
+                    .unwrap();
+                assert_eq!(
+                    dir.complete_flight(flight_id, epoch, (n * 2).to_le_bytes()),
+                    Err(CacheError::UnknownGeneration)
+                );
+            }
+            let stats = dir.stats();
+            assert_eq!(stats.flights, 0);
+            assert_eq!(stats.request_entries, 0);
+            assert_eq!(stats.tenants, 0);
+            assert!(stats.notices <= dir.max_notices);
+            assert!(dir.notice_order.len() <= dir.max_notices);
+            assert_eq!(dir.credits_used(), 0);
+        }
+    }
 }

crates/transport/src/lib.rs

@@ -20,13 +20,16 @@ mod plane;
 mod pool;
 mod rate;
 mod registered;
+mod routing;
 mod server;
 mod shard_demo;
 mod stage;
 
-pub use client::TcpClient;
+pub use client::{ClientIoStats, TcpClient};
 pub use demo::{run_share_demo, ShareReport};
-pub use directory::{Admit, FillNotice, FillRequest, MemberView, ShardDirectory, ShardParams};
+pub use directory::{
+    Admit, DirectoryStats, FillNotice, FillRequest, MemberView, ShardDirectory, ShardParams,
+};
 pub use fault_demo::{run_fault_demo, FaultReport};
 pub use iface::{
     rdma_feature_selected, rdma_probe, submit_preferring_rdma, FallbackReason, HardwarePath,
@@ -35,6 +38,7 @@ pub use iface::{
 pub use plane::{ControlPlane, RebalancePlan, UpgradeStamp};
 pub use rate::RateGate;
 pub use registered::{CheckedBatch, PoolStats, RegisteredPool, SimCompletion, SlotId};
+pub use routing::{ClientLimits, ClientStats, ShardEndpoint, ShardSession, ShardedClient};
 pub use server::{serve, ServerConfig, ServerHandle};
 pub use shard_demo::{run_shard_demo, ShardReport};
 pub use stage::stage_through_tcp;