feat(cell): add token bucket cell (#96)

Resolves: #85
Co-authored-by: BobAnkh <[email protected]>

Author: EwecaBcD

config.example.toml

@@ -279,6 +279,32 @@ interface_id = 1
 prefix = "10.2.0.0/16"
 # no interface_id, meaning drop
 
+# Token Bucket Cell Example
+[cells.token_bucket]
+type = "TokenBucket"
+
+# Token Bucket configuration
+# We have 5 fields for TokenBucket: 
+# "limit" : Size of the queue of TokenBucket cell in bytes.
+# Default: Unlimited
+limit = 1024
+
+# "burst" : Size of the bucket in bytes.
+# Default: Unlimited
+burst = "5KB"
+
+# "rate" : The rate limit of the bucket in bytes per second.
+# Default: Unlimited
+rate = "500Kbps"
+
+# - "minburst" : Specifies the size of the peakrate bucket.
+# Default: Unlimited
+minburst = "1540B"
+
+# - "peakrate" : Maximum depletion rate of the bucket.
+# Default: Unlimited
+peakrate = "1Mbps"
+
 # ----- Cells Section End -----
 
 

rattan-core/Cargo.toml

@@ -13,6 +13,7 @@ async-trait = "0.1.68"
 axum = { version = "0.7.5", optional = true }
 backon = "1.2.0"
 bandwidth = "0.3.0"
+bytesize = "1.3.0"
 figment = { workspace = true }
 camellia = { git = "https://github.com/minhuw/camellia.git", optional = true }
 etherparse = "0.16.0"
@@ -74,6 +75,7 @@ serde = [
     "dep:serde",
     "dep:serde_json",
     "bandwidth/serde",
+    "bytesize/serde",
 ]
 camellia = ["dep:camellia"]
 

rattan-core/src/cells/bandwidth/queue.rs

@@ -28,6 +28,19 @@ where
 
     // If the queue is empty, return `None`
     fn dequeue(&mut self) -> Option<P>;
+
+    fn is_empty(&self) -> bool;
+
+    // If the queue is empty, return `None`
+    fn get_front_size(&self) -> Option<usize>;
+
+    fn length(&self) -> usize;
+
+    fn retain<F>(&mut self, _f: F)
+    where
+        F: FnMut(&P) -> bool,
+    {
+    }
 }
 
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
@@ -81,6 +94,25 @@ where
     fn dequeue(&mut self) -> Option<P> {
         self.queue.pop_front()
     }
+
+    fn is_empty(&self) -> bool {
+        self.queue.is_empty()
+    }
+
+    fn get_front_size(&self) -> Option<usize> {
+        self.queue.front().map(|packet| packet.l3_length())
+    }
+
+    fn length(&self) -> usize {
+        self.queue.len()
+    }
+
+    fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&P) -> bool,
+    {
+        self.queue.retain(|packet| f(packet));
+    }
 }
 
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
@@ -137,6 +169,10 @@ impl<P> DropTailQueue<P> {
             now_bytes: 0,
         }
     }
+
+    pub fn get_extra_length(&self) -> usize {
+        self.bw_type.extra_length()
+    }
 }
 
 impl<P> Default for DropTailQueue<P> {
@@ -186,6 +222,27 @@ where
             None => None,
         }
     }
+
+    fn is_empty(&self) -> bool {
+        self.queue.is_empty()
+    }
+
+    fn get_front_size(&self) -> Option<usize> {
+        self.queue
+            .front()
+            .map(|packet| packet.l3_length() + self.bw_type.extra_length())
+    }
+
+    fn length(&self) -> usize {
+        self.queue.len()
+    }
+
+    fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&P) -> bool,
+    {
+        self.queue.retain(|packet| f(packet));
+    }
 }
 
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
@@ -291,6 +348,27 @@ where
             None => None,
         }
     }
+
+    fn is_empty(&self) -> bool {
+        self.queue.is_empty()
+    }
+
+    fn get_front_size(&self) -> Option<usize> {
+        self.queue
+            .front()
+            .map(|packet| packet.l3_length() + self.bw_type.extra_length())
+    }
+
+    fn length(&self) -> usize {
+        self.queue.len()
+    }
+
+    fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&P) -> bool,
+    {
+        self.queue.retain(|packet| f(packet));
+    }
 }
 
 // CoDel Queue Implementation Reference:
@@ -558,4 +636,25 @@ where
             }
         }
     }
+
+    fn is_empty(&self) -> bool {
+        self.queue.is_empty()
+    }
+
+    fn get_front_size(&self) -> Option<usize> {
+        self.queue
+            .front()
+            .map(|packet| packet.l3_length() + self.config.bw_type.extra_length())
+    }
+
+    fn length(&self) -> usize {
+        self.queue.len()
+    }
+
+    fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&P) -> bool,
+    {
+        self.queue.retain(|packet| f(packet));
+    }
 }

rattan-core/src/cells/mod.rs

@@ -13,6 +13,7 @@ pub mod external;
 pub mod loss;
 pub mod router;
 pub mod shadow;
+pub mod token_bucket;
 
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[derive(Debug, Clone, Hash, Eq, PartialEq)]