feat(pkt-meta): Add initial flow table implementation

This is the first cut for a flow table.
There is more work to be done, but this is a good checkpoint.
Interfaces and organization may change as we start using this
with the stateful NAT code.

Signed-off-by: Manish Vachharajani <[email protected]>

Author: mvachhar

Cargo.lock

@@ -6,11 +6,15 @@ publish = false
 license = "Apache-2.0"
 
 [features]
+shuttle = ["concurrency/shuttle"]
+std = []
 testing = []
+bolero = ["dep:bolero", "net/bolero"]
 
 [dependencies]
 ahash = { workspace = true }
 atomic-instant-full = { workspace = true }
+bolero = { workspace = true, optional = true }
 lpm = { workspace = true }
 concurrency = { workspace = true }
 config = { workspace = true }
@@ -27,5 +31,6 @@ tracing = { workspace = true }
 [dev-dependencies]
 bolero = { workspace = true, default-features = false }
 lpm = { workspace = true, features = ["testing"] }
+net = { workspace = true, features = ["bolero"] }
 tracing-test = { workspace = true, features = [] }
-
+shuttle = { workspace = true }

pkt-meta/Cargo.toml

@@ -0,0 +1,28 @@
+# Flow Table
+
+The current implementation of flow table uses `dash_map` and per-thread priority queue's (for timeouts) along with `Arc` and `Weak` to get a reasonable flow table with timeouts.
+However, it leaves a lot of room for optimizations.
+
+## Flow Table Implementation
+
+The main `DashMap` holds `Weak` references to all the flow entries so that the memory gets automatically deallocated when the entry times out.  
+
+The priority queue's hold `Arc` references to the flow entries to keep them alive when they are not in any packet meta data.
+When the entry times-out and is removed from the priority queue and the last packet referencing that flow is dropped, the memory for the entry is freed.
+
+Note that in the current implementation, a flow is not removed from the flow table until the last Arc to the flow_info is dropped or the flow entry is replaced.  This can be changed if needed, or even have it be an option on the flow as to whether timeout removes the flow or not.
+
+## Optimizations
+
+In the current implementation, there has to be periodic or on-timeout reaping the Weak reference in the hash table.  
+This is better done by having a version of `DashMap` that can reap the dead `Weak` reference as it walks the table on lookups, instead of waiting for key collisions.
+The hope, for now, is that the entries in the hash table array will contain a small pointer and not take up too much extra memory.
+Those dead `Weak` pointers will prevent shrinking of the hash table though, if the implementation supports that.
+
+Second, the `priority_queue` crate uses a `HashMap` inside the queue in order to allow fast removal and re-insertion.  
+However, this wastes space and requires extra hashes.  
+The better way to do this is to have a custom priority queue integrated with the custom weak-reaping hash map so that the same hash table can be used for both operations.
+This improves cache locality, reduces memory utlization, and avoids multiple hash table lookups in many cases.
+
+However, in the interest of time to completion for the code, this module currently uses existing data structures instead of full custom implementations of everything.
+However, the interface should be able to hide a change from the current to the optimized implementation.

pkt-meta/src/flow_table/README.md

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Open Network Fabric Authors
+
+use atomic_instant_full;
+use std::fmt::{Debug, Formatter};
+use std::ops::Deref;
+use std::sync::atomic::Ordering;
+use std::time::Instant;
+
+#[repr(transparent)]
+pub struct AtomicInstant(atomic_instant_full::AtomicInstant);
+
+impl Debug for AtomicInstant {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "({:?})", self.0.load(Ordering::Relaxed))
+    }
+}
+
+impl Deref for AtomicInstant {
+    type Target = atomic_instant_full::AtomicInstant;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl AtomicInstant {
+    #[must_use]
+    pub fn new(instant: Instant) -> Self {
+        Self(atomic_instant_full::AtomicInstant::new(instant))
+    }
+}

pkt-meta/src/flow_table/atomic_instant.rs

@@ -0,0 +1,180 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Open Network Fabric Authors
+
+use std::fmt::Debug;
+use std::time::{Duration, Instant};
+
+use concurrency::sync::RwLock;
+use net::packet::VpcDiscriminant;
+
+use crate::flow_table::AtomicInstant;
+
+use std::sync::atomic::{AtomicU8, Ordering};
+
+#[derive(Debug, thiserror::Error)]
+pub enum FlowInfoError {
+    #[error("flow expired")]
+    FlowExpired(Instant),
+    #[error("no such status")]
+    NoSuchStatus(u8),
+}
+
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum FlowStatus {
+    Active = 0,
+    Expired = 1,
+    Removed = 2,
+}
+
+impl TryFrom<u8> for FlowStatus {
+    type Error = FlowInfoError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(FlowStatus::Active),
+            1 => Ok(FlowStatus::Expired),
+            2 => Ok(FlowStatus::Removed),
+            v => Err(FlowInfoError::NoSuchStatus(v)),
+        }
+    }
+}
+
+impl From<FlowStatus> for u8 {
+    fn from(status: FlowStatus) -> Self {
+        status as u8
+    }
+}
+
+pub struct AtomicFlowStatus(AtomicU8);
+
+impl Debug for AtomicFlowStatus {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:?}", self.load(std::sync::atomic::Ordering::Relaxed))
+    }
+}
+
+impl AtomicFlowStatus {
+    /// Load the flow status.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the the stored flow status is invalid, which should never happen.
+    ///
+    #[must_use]
+    pub fn load(&self, ordering: Ordering) -> FlowStatus {
+        let value = self.0.load(ordering);
+        FlowStatus::try_from(value).expect("Invalid enum state")
+    }
+
+    pub fn store(&self, state: FlowStatus, ordering: Ordering) {
+        self.0.store(u8::from(state), ordering);
+    }
+
+    /// Atomic compare and exchange of the flow status.
+    ///
+    /// # Errors
+    ///
+    /// Returns previous `FlowStatus` if the compare and exchange fails.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the the stored flow status is invalid, which should never happen.
+    ///
+    pub fn compare_exchange(
+        &self,
+        current: FlowStatus,
+        new: FlowStatus,
+        success: Ordering,
+        failure: Ordering,
+    ) -> Result<FlowStatus, FlowStatus> {
+        match self
+            .0
+            .compare_exchange(current as u8, new as u8, success, failure)
+        {
+            Ok(prev) => Ok(FlowStatus::try_from(prev).expect("Invalid enum state")),
+            Err(prev) => Err(FlowStatus::try_from(prev).expect("Invalid enum state")),
+        }
+    }
+}
+
+impl From<FlowStatus> for AtomicFlowStatus {
+    fn from(status: FlowStatus) -> Self {
+        Self(AtomicU8::new(status as u8))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct FlowInfoLocked {
+    pub dst_vpcd: Option<VpcDiscriminant>,
+}
+
+#[derive(Debug)]
+pub struct FlowInfo {
+    expires_at: AtomicInstant,
+    status: AtomicFlowStatus,
+    pub locked: RwLock<FlowInfoLocked>,
+}
+
+// TODO: We need a way to stuff an Arc<FlowInfo> into the packet
+// meta data.  That means this has to move to net or we need a generic
+// meta data extension method.
+impl FlowInfo {
+    #[must_use]
+    pub fn new(expires_at: Instant) -> Self {
+        Self {
+            expires_at: AtomicInstant::new(expires_at),
+            status: AtomicFlowStatus::from(FlowStatus::Active),
+            locked: RwLock::new(FlowInfoLocked { dst_vpcd: None }),
+        }
+    }
+
+    pub fn expires_at(&self) -> Instant {
+        self.expires_at.load(std::sync::atomic::Ordering::Relaxed)
+    }
+
+    /// Extend the expiry of the flow if it is not expired.
+    ///
+    /// # Errors
+    ///
+    /// Returns `FlowInfoError::FlowExpired` if the flow is expired with the expiry `Instant`
+    ///
+    pub fn extend_expiry(&self, duration: Duration) -> Result<(), FlowInfoError> {
+        if self.status.load(std::sync::atomic::Ordering::Relaxed) == FlowStatus::Expired {
+            return Err(FlowInfoError::FlowExpired(self.expires_at()));
+        }
+        self.extend_expiry_unchecked(duration);
+        Ok(())
+    }
+
+    /// Extend the expiry of the flow without checking if it is already expired.
+    ///
+    /// # Thread Safety
+    ///
+    /// This method is thread-safe.
+    ///
+    pub fn extend_expiry_unchecked(&self, duration: Duration) {
+        self.expires_at
+            .fetch_add(duration, std::sync::atomic::Ordering::Relaxed);
+    }
+
+    pub fn status(&self) -> FlowStatus {
+        self.status.load(std::sync::atomic::Ordering::Relaxed)
+    }
+
+    /// Update the flow status.
+    ///
+    /// # Thread Safety
+    ///
+    /// This method is thread-safe.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the status transition is invalid.
+    ///
+    pub fn update_status(&self, status: FlowStatus) -> Result<(), FlowInfoError> {
+        self.status
+            .store(status, std::sync::atomic::Ordering::Relaxed);
+        Ok(())
+    }
+}