Repository: cloudflare/pingora (language: rust)

Issue:
TinyUFO unable to run benches
## Describe the bug

When running TinyUFO (benches / tests) I am getting an error due to the Bucket struct not being publicly usable for the `bucket.rs` file.

## Pingora info

Please include the following information about your environment:

**Pingora version**: 0.1.1
**Rust version**: 1.72.1
**Operating system version**: Manjaro, Kernel 6.6

## Steps to reproduce

When I run any tests or benchmarks for TinyUFO I get the observed error.

`cargo bench --bench bench_perf`

## Observed results

error[E0446]: private type `Bucket<T>` in public interface
   --> tinyufo/src/buckets.rs:39:5
    |
39  |     pub fn get(&self, key: &Key) -> Option<Entry<Key, Bucket<T>>> {
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ can't leak private type
    |
   ::: tinyufo/src/lib.rs:124:1
    |
124 | struct Bucket<T> {
    | ---------------- `Bucket<T>` declared as private

## Additional context

These tests / benches were working for me before the TinyUFO compact changes were added where the Bucket was moved. 
 This was fixed by making the Bucket public.  Let me know if you'd like me to create a PR or if I am doing something wrong when trying to run these tests.



Discussion / hints:

Fix patch:
diff --git a/.bleep b/.bleep
index 1d00da409..da156b550 100644
--- a/.bleep
+++ b/.bleep
@@ -1,1 +1,1 @@
-836062d709916982c97b1990f56e446f834e610d
\ No newline at end of file
+fa95ba9cdfb9a347913e0cffbb84ca003833ccaa
diff --git a/.rustfmt.toml b/.rustfmt.toml
index e69de29bb..3a26366d4 100644
--- a/.rustfmt.toml
+++ b/.rustfmt.toml
@@ -0,0 +1,1 @@
+edition = "2021"
diff --git a/pingora-cache/src/eviction/simple_lru.rs b/pingora-cache/src/eviction/simple_lru.rs
index 43009a8dc..caa929330 100644
--- a/pingora-cache/src/eviction/simple_lru.rs
+++ b/pingora-cache/src/eviction/simple_lru.rs
@@ -222,15 +222,18 @@ impl EvictionManager for Manager {
 
     async fn save(&self, dir_path: &str) -> Result<()> {
         let data = self.serialize()?;
-        let dir_path = dir_path.to_owned();
+        let dir_str = dir_path.to_owned();
         tokio::task::spawn_blocking(move || {
-            let dir_path = Path::new(&dir_path);
-            std::fs::create_dir_all(dir_path).or_err(InternalError, "fail to create {dir_path}")?;
+            let dir_path = Path::new(&dir_str);
+            std::fs::create_dir_all(dir_path)
+                .or_err_with(InternalError, || format!("fail to create {dir_str}"))?;
             let file_path = dir_path.join(FILE_NAME);
-            let mut file =
-                File::create(file_path).or_err(InternalError, "fail to create {file_path}")?;
-            file.write_all(&data)
-                .or_err(InternalError, "fail to write to {file_path}")
+            let mut file = File::create(&file_path).or_err_with(InternalError, || {
+                format!("fail to create {}", file_path.display())
+            })?;
+            file.write_all(&data).or_err_with(InternalError, || {
+                format!("fail to write to {}", file_path.display())
+            })
         })
         .await
         .or_err(InternalError, "async blocking IO failure")?
@@ -240,8 +243,9 @@ impl EvictionManager for Manager {
         let dir_path = dir_path.to_owned();
         let data = tokio::task::spawn_blocking(move || {
             let file_path = Path::new(&dir_path).join(FILE_NAME);
-            let mut file =
-                File::open(file_path).or_err(InternalError, "fail to open {file_path}")?;
+            let mut file = File::open(file_path.clone()).or_err_with(InternalError, || {
+                format!("fail to open {}", file_path.display())
+            })?;
             let mut buffer = Vec::with_capacity(8192);
             file.read_to_end(&mut buffer)
                 .or_err(InternalError, "fail to write to {file_path}")?;
diff --git a/pingora-core/src/connectors/tls.rs b/pingora-core/src/connectors/tls.rs
index 7cdbab036..8a6bd6350 100644
--- a/pingora-core/src/connectors/tls.rs
+++ b/pingora-core/src/connectors/tls.rs
@@ -162,8 +162,7 @@ impl Connector {
 */
 fn replace_leftmost_underscore(sni: &str) -> Option<String> {
     // wildcard is only leftmost label
-    let mut s = sni.splitn(2, '.');
-    if let (Some(leftmost), Some(rest)) = (s.next(), s.next()) {
+    if let Some((leftmost, rest)) = sni.split_once('.') {
         // if not a subdomain or leftmost does not contain underscore return
         if !rest.contains('.') || !leftmost.contains('_') {
             return None;
diff --git a/pingora-core/src/listeners/tls.rs b/pingora-core/src/listeners/tls.rs
index c8f0e74a4..655346fd8 100644
--- a/pingora-core/src/listeners/tls.rs
+++ b/pingora-core/src/listeners/tls.rs
@@ -63,10 +63,12 @@ impl TlsSettings {
         )?;
         accept_builder
             .set_private_key_file(key_path, SslFiletype::PEM)
-            .or_err(TLS_CONF_ERR, "fail to read key file {key_path}")?;
+            .or_err_with(TLS_CONF_ERR, || format!("fail to read key file {key_path}"))?;
         accept_builder
             .set_certificate_chain_file(cert_path)
-            .or_err(TLS_CONF_ERR, "fail to read cert file {cert_path}")?;
+            .or_err_with(TLS_CONF_ERR, || {
+                format!("fail to read cert file {cert_path}")
+            })?;
         Ok(TlsSettings {
             accept_builder,
             callbacks: None,
diff --git a/pingora-core/src/protocols/http/server.rs b/pingora-core/src/protocols/http/server.rs
index 6edea75ab..697b75fc9 100644
--- a/pingora-core/src/protocols/http/server.rs
+++ b/pingora-core/src/protocols/http/server.rs
@@ -18,7 +18,7 @@ use super::error_resp;
 use super::v1::server::HttpSession as SessionV1;
 use super::v2::server::HttpSession as SessionV2;
 use super::HttpTask;
-use crate::protocols::{SocketAddr, Stream};
+use crate::protocols::{Digest, SocketAddr, Stream};
 use bytes::Bytes;
 use http::header::AsHeaderName;
 use http::HeaderValue;
@@ -323,7 +323,7 @@ impl Session {
         }
     }
 
-    /// How many response body bytes already sent
+    /// Return how many response body bytes (application, not wire) already sent downstream
     pub fn body_bytes_sent(&self) -> usize {
         match self {
             Self::H1(s) => s.body_bytes_sent(),
@@ -331,6 +331,22 @@ impl Session {
         }
     }
 
+    /// Return how many request body bytes (application, not wire) already read from downstream
+    pub fn body_bytes_read(&self) -> usize {
+        match self {
+            Self::H1(s) => s.body_bytes_read(),
+            Self::H2(s) => s.body_bytes_read(),
+        }
+    }
+
+    /// Return the digest for the session.
+    pub fn digest(&self) -> Option<&Digest> {
+        match self {
+            Self::H1(s) => Some(s.digest()),
+            Self::H2(s) => s.digest(),
+        }
+    }
+
     /// Return the client (peer) address of the connnection.
     pub fn client_addr(&self) -> Option<&SocketAddr> {
         match self {
diff --git a/pingora-core/src/protocols/http/v1/client.rs b/pingora-core/src/protocols/http/v1/client.rs
index 7b125932c..1dfb97ae7 100644
--- a/pingora-core/src/protocols/http/v1/client.rs
+++ b/pingora-core/src/protocols/http/v1/client.rs
@@ -445,12 +445,10 @@ impl HttpSession {
         let mut max = None;
 
         for param in header_value.split(',') {
-            let mut parts = param.splitn(2, '=').map(|s| s.trim());
-            match (parts.next(), parts.next()) {
-                (Some("timeout"), Some(timeout_value)) => {
-                    timeout = timeout_value.trim().parse().ok()
-                }
-                (Some("max"), Some(max_value)) => max = max_value.trim().parse().ok(),
+            let parts = param.split_once('=').map(|(k, v)| (k.trim(), v));
+            match parts {
+                Some(("timeout", timeout_value)) => timeout = timeout_value.trim().parse().ok(),
+                Some(("max", max_value)) => max = max_value.trim().parse().ok(),
                 _ => {}
             }
         }
diff --git a/pingora-core/src/protocols/http/v1/server.rs b/pingora-core/src/protocols/http/v1/server.rs
index f8783db3c..7939ed466 100644
--- a/pingora-core/src/protocols/http/v1/server.rs
+++ b/pingora-core/src/protocols/http/v1/server.rs
@@ -53,6 +53,8 @@ pub struct HttpSession {
     body_write_buf: BytesMut,
     /// Track how many application (not on the wire) body bytes already sent
     body_bytes_sent: usize,
+    /// Track how many application (not on the wire) body bytes already read
+    body_bytes_read: usize,
     /// Whether to update headers like connection, Date
     update_resp_headers: bool,
     /// timeouts:
@@ -61,7 +63,7 @@ pub struct HttpSession {
     write_timeout: Option<Duration>,
     /// A copy of the response that is already written to the client
     response_written: Option<Box<ResponseHeader>>,
-    /// The parse request header
+    /// The parsed request header
     request_header: Option<Box<RequestHeader>>,
     /// An internal buffer that holds a copy of the request body up to a certain size
     retry_buffer: Option<FixedBuffer>,
@@ -100,6 +102,7 @@ impl HttpSession {
             read_timeout: None,
             write_timeout: None,
             body_bytes_sent: 0,
+            body_bytes_read: 0,
             retry_buffer: None,
             upgraded: false,
             digest,
@@ -338,6 +341,7 @@ impl HttpSession {
         let read = self.read_body().await?;
         Ok(read.map(|b| {
             let bytes = Bytes::copy_from_slice(self.get_body(&b));
+            self.body_bytes_read += bytes.len();
             if let Some(buffer) = self.retry_buffer.as_mut() {
                 buffer.write_to_buffer(&bytes);
             }
@@ -627,11 +631,16 @@ impl HttpSession {
         Ok(res)
     }
 
-    /// Return how many (application, not wire) body bytes that have been written
+    /// Return how many response body bytes (application, not wire) already sent downstream
     pub fn body_bytes_sent(&self) -> usize {
         self.body_bytes_sent
     }
 
+    /// Return how many request body bytes (application, not wire) already read from downstream
+    pub fn body_bytes_read(&self) -> usize {
+        self.body_bytes_read
+    }
+
     fn is_chunked_encoding(&self) -> bool {
         is_header_value_chunked_encoding(self.get_header(header::TRANSFER_ENCODING))
     }
@@ -1088,10 +1097,10 @@ mod tests_stream {
             .build();
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_request().await.unwrap();
-        let res = http_stream.read_body().await.unwrap().unwrap();
-        assert_eq!(res, BufRef::new(0, 3));
+        let res = http_stream.read_body_bytes().await.unwrap().unwrap();
+        assert_eq!(res, input3.as_slice());
         assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(3));
-        assert_eq!(input3, http_stream.get_body(&res));
+        assert_eq!(http_stream.body_bytes_read(), 3);
     }
 
     #[tokio::test]
@@ -1110,7 +1119,8 @@ mod tests_stream {
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_timeout = Some(Duration::from_secs(1));
         http_stream.read_request().await.unwrap();
-        let res = http_stream.read_body().await;
+        let res = http_stream.read_body_bytes().await;
+        assert_eq!(http_stream.body_bytes_read(), 0);
         assert_eq!(res.unwrap_err().etype(), &ReadTimedout);
     }
 
@@ -1122,10 +1132,10 @@ mod tests_stream {
         let mock_io = Builder::new().read(&input1[..]).read(&input2[..]).build();
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_request().await.unwrap();
-        let res = http_stream.read_body().await.unwrap().unwrap();
-        assert_eq!(res, BufRef::new(0, 3));
+        let res = http_stream.read_body_bytes().await.unwrap().unwrap();
+        assert_eq!(res, b"abc".as_slice());
         assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(3));
-        assert_eq!(b"abc", http_stream.get_body(&res));
+        assert_eq!(http_stream.body_bytes_read(), 3);
     }
 
     #[tokio::test]
@@ -1143,13 +1153,14 @@ mod tests_stream {
             .build();
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_request().await.unwrap();
-        let res = http_stream.read_body().await.unwrap().unwrap();
-        assert_eq!(res, BufRef::new(0, 1));
+        let res = http_stream.read_body_bytes().await.unwrap().unwrap();
+        assert_eq!(res, input3.as_slice());
         assert_eq!(http_stream.body_reader.body_state, ParseState::HTTP1_0(1));
-        assert_eq!(input3, http_stream.get_body(&res));
-        let res = http_stream.read_body().await.unwrap();
-        assert_eq!(res, None);
+        assert_eq!(http_stream.body_bytes_read(), 1);
+        let res = http_stream.read_body_bytes().await.unwrap();
+        assert!(res.is_none());
         assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(1));
+        assert_eq!(http_stream.body_bytes_read(), 1);
     }
 
     #[tokio::test]
@@ -1167,16 +1178,15 @@ mod tests_stream {
             .build();
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_request().await.unwrap();
-        let res = http_stream.read_body().await.unwrap().unwrap();
-        assert_eq!(res, BufRef::new(0, 1));
+        let res = http_stream.read_body_bytes().await.unwrap().unwrap();
+        assert_eq!(res, b"a".as_slice());
         assert_eq!(http_stream.body_reader.body_state, ParseState::HTTP1_0(1));
-        assert_eq!(b"a", http_stream.get_body(&res));
-        let res = http_stream.read_body().await.unwrap().unwrap();
-        assert_eq!(res, BufRef::new(0, 1));
+        let res = http_stream.read_body_bytes().await.unwrap().unwrap();
+        assert_eq!(res, b"b".as_slice());
         assert_eq!(http_stream.body_reader.body_state, ParseState::HTTP1_0(2));
-        assert_eq!(input3, http_stream.get_body(&res));
-        let res = http_stream.read_body().await.unwrap();
-        assert_eq!(res, None);
+        let res = http_stream.read_body_bytes().await.unwrap();
+        assert_eq!(http_stream.body_bytes_read(), 2);
+        assert!(res.is_none());
         assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(2));
     }
 
@@ -1188,8 +1198,9 @@ mod tests_stream {
         let mock_io = Builder::new().read(&input1[..]).read(&input2[..]).build();
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_request().await.unwrap();
-        let res = http_stream.read_body().await.unwrap();
-        assert_eq!(res, None);
+        let res = http_stream.read_body_bytes().await.unwrap();
+        assert!(res.is_none());
+        assert_eq!(http_stream.body_bytes_read(), 0);
         assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0));
     }
 
@@ -1207,8 +1218,9 @@ mod tests_stream {
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_request().await.unwrap();
         assert!(http_stream.is_chunked_encoding());
-        let res = http_stream.read_body().await.unwrap();
-        assert_eq!(res, None);
+        let res = http_stream.read_body_bytes().await.unwrap();
+        assert!(res.is_none());
+        assert_eq!(http_stream.body_bytes_read(), 0);
         assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0));
     }
 
@@ -1226,14 +1238,15 @@ mod tests_stream {
         let mut http_stream = HttpSession::new(Box::new(mock_io));
         http_stream.read_request().await.unwrap();
         assert!(http_stream.is_chunked_encoding());
-        let res = http_stream.read_body().await.unwrap().unwrap();
-        assert_eq!(res, BufRef::new(3, 1));
+        let res = http_stream.read_body_bytes().await.unwrap().unwrap();
+        assert_eq!(res, b"a".as_slice());
         assert_eq!(
             http_stream.body_reader.body_state,
             ParseState::Chunked(1, 0, 0, 0)
         );
-        let res = http_stream.read_body().await.unwrap();
-        assert_eq!(res, None);
+        let res = http_stream.read_body_bytes().await.unwrap();
+        assert!(res.is_none());
+        assert_eq!(http_stream.body_bytes_read(), 1);
         assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(1));
     }
 
diff --git a/pingora-core/src/protocols/http/v2/server.rs b/pingora-core/src/protocols/http/v2/server.rs
index a81176991..8e3e673f3 100644
--- a/pingora-core/src/protocols/http/v2/server.rs
+++ b/pingora-core/src/protocols/http/v2/server.rs
@@ -90,9 +90,9 @@ pub struct HttpSession {
     // Indicate that whether a END_STREAM is already sent
     // in order to tell whether needs to send one extra FRAME when this response finishes
     ended: bool,
-    // How many request body bytes have been read so far.
+    // How many (application, not wire) request body bytes have been read so far.
     body_read: usize,
-    // How many response body bytes have been sent so far.
+    // How many (application, not wire) response body bytes have been sent so far.
     body_sent: usize,
     // buffered request body for retry logic
     retry_buffer: Option<FixedBuffer>,
@@ -413,11 +413,16 @@ impl HttpSession {
         }
     }
 
-    /// How many response body bytes sent to the client
+    /// Return how many response body bytes (application, not wire) already sent downstream
     pub fn body_bytes_sent(&self) -> usize {
         self.body_sent
     }
 
+    /// Return how many request body bytes (application, not wire) already read from downstream
+    pub fn body_bytes_read(&self) -> usize {
+        self.body_read
+    }
+
     /// Return the [Digest] of the connection.
     pub fn digest(&self) -> Option<&Digest> {
         Some(&self.digest)
@@ -490,6 +495,7 @@ mod test {
                 let body = http.read_body_or_idle(false).await.unwrap().unwrap();
                 assert_eq!(body, client_body);
                 assert!(http.is_body_done());
+                assert_eq!(http.body_bytes_read(), 16);
 
                 let retry_body = http.get_retry_buffer().unwrap();
                 assert_eq!(retry_body, client_body);
@@ -511,6 +517,7 @@ mod test {
 
                 // end: false here to verify finish() closes the stream nicely
                 http.write_body(server_body.into(), false).unwrap();
+                assert_eq!(http.body_bytes_sent(), 16);
 
                 http.finish().unwrap();
             });
diff --git a/pingora-core/src/server/mod.rs b/pingora-core/src/server/mod.rs
index b1e24f2a3..8ebf6fde1 100644
--- a/pingora-core/src/server/mod.rs
+++ b/pingora-core/src/server/mod.rs
@@ -289,7 +289,7 @@ impl Server {
     ///
     /// Note: this function may fork the process for daemonization, so any additional threads created
     /// before this function will be lost to any service logic once this function is called.
-    pub fn run_forever(&mut self) {
+    pub fn run_forever(mut self) -> ! {
         info!("Server starting");
 
         let conf = self.configuration.as_ref();
@@ -364,7 +364,7 @@ impl Server {
             }
         }
         info!("All runtimes exited, exiting now");
-        std::process::exit(0);
+        std::process::exit(0)
     }
 
     fn create_runtime(name: &str, threads: usize, work_steal: bool) -> Runtime {
diff --git a/pingora-limits/src/estimator.rs b/pingora-limits/src/estimator.rs
index 2e98ad451..2abf5a7c1 100644
--- a/pingora-limits/src/estimator.rs
+++ b/pingora-limits/src/estimator.rs
@@ -30,17 +30,12 @@ pub struct Estimator {
 impl Estimator {
     /// Create a new `Estimator` with the given amount of hashes and columns (slots).
     pub fn new(hashes: usize, slots: usize) -> Self {
-        let mut estimator = Vec::with_capacity(hashes);
-        for _ in 0..hashes {
-            let mut slot = Vec::with_capacity(slots);
-            for _ in 0..slots {
-                slot.push(AtomicIsize::new(0));
-            }
-            estimator.push((slot.into_boxed_slice(), RandomState::new()));
-        }
-
-        Estimator {
-            estimator: estimator.into_boxed_slice(),
+        Self {
+            estimator: (0..hashes)
+                .map(|_| (0..slots).map(|_| AtomicIsize::new(0)).collect::<Vec<_>>())
+                .map(|slot| (slot.into_boxed_slice(), RandomState::new()))
+                .collect::<Vec<_>>()
+                .into_boxed_slice(),
         }
     }
 
@@ -48,15 +43,15 @@ impl Estimator {
     /// Note: overflow can happen. When some of the internal counters overflow, a negative number
     /// will be returned. It is up to the caller to catch and handle this case.
     pub fn incr<T: Hash>(&self, key: T, value: isize) -> isize {
-        let mut min = isize::MAX;
-        for (slot, hasher) in self.estimator.iter() {
-            let hash = hash(&key, hasher) as usize;
-            let counter = &slot[hash % slot.len()];
-            // Overflow is allowed for simplicity
-            let current = counter.fetch_add(value, Ordering::Relaxed);
-            min = std::cmp::min(min, current + value);
-        }
-        min
+        self.estimator
+            .iter()
+            .fold(isize::MAX, |min, (slot, hasher)| {
+                let hash = hash(&key, hasher) as usize;
+                let counter = &slot[hash % slot.len()];
+                // Overflow is allowed for simplicity
+                let current = counter.fetch_add(value, Ordering::Relaxed);
+                std::cmp::min(min, current + value)
+            })
     }
 
     /// Decrement `key` by the value given.
@@ -70,23 +65,22 @@ impl Estimator {
 
     /// Get the estimated frequency of `key`.
     pub fn get<T: Hash>(&self, key: T) -> isize {
-        let mut min = isize::MAX;
-        for (slot, hasher) in self.estimator.iter() {
-            let hash = hash(&key, hasher) as usize;
-            let counter = &slot[hash % slot.len()];
-            let current = counter.load(Ordering::Relaxed);
-            min = std::cmp::min(min, current);
-        }
-        min
+        self.estimator
+            .iter()
+            .fold(isize::MAX, |min, (slot, hasher)| {
+                let hash = hash(&key, hasher) as usize;
+                let counter = &slot[hash % slot.len()];
+                let current = counter.load(Ordering::Relaxed);
+                std::cmp::min(min, current)
+            })
     }
 
     /// Reset all values inside this `Estimator`.
     pub fn reset(&self) {
-        for (slot, _) in self.estimator.iter() {
-            for counter in slot.iter() {
-                counter.store(0, Ordering::Relaxed);
-            }
-        }
+        self.estimator.iter().for_each(|(slot, _)| {
+            slot.iter()
+                .for_each(|counter| counter.store(0, Ordering::Relaxed))
+        });
     }
 }
 
@@ -128,4 +122,17 @@ mod tests {
         assert_eq!(est.get("a"), 3);
         assert_eq!(est.get("b"), 3);
     }
+
+    #[test]
+    fn reset() {
+        let est = Estimator::new(8, 8);
+        est.incr("a", 1);
+        est.incr("a", 2);
+        est.incr("b", 1);
+        est.incr("b", 2);
+        est.decr("b", 1);
+        est.reset();
+        assert_eq!(est.get("a"), 0);
+        assert_eq!(est.get("b"), 0);
+    }
 }
diff --git a/pingora-lru/src/lib.rs b/pingora-lru/src/lib.rs
index 15a115bc0..d7c31c7a8 100644
--- a/pingora-lru/src/lib.rs
+++ b/pingora-lru/src/lib.rs
@@ -68,7 +68,7 @@ impl<T, const N: usize> Lru<T, N> {
 
         // Make sure weight is positive otherwise eviction won't work
         // TODO: Probably should use NonZeroUsize instead
-        let weight = if weight == 0 { 1 } else { weight };
+        let weight = weight.max(1);
 
         let old_weight = unit.admit(key, data, weight);
         if old_weight != weight {
@@ -307,6 +307,7 @@ impl<T> LruUnit<T> {
         self.lookup_table.remove(&key).map(|node| {
             let list_key = self.order.remove(node.list_index);
             assert_eq!(key, list_key);
+            self.used_weight -= node.weight;
             (node.data, node.weight)
         })
     }
@@ -322,6 +323,7 @@ impl<T> LruUnit<T> {
             weight,
         });
         self.lookup_table.insert(key, node);
+        self.used_weight += weight;
         true
     }
 
@@ -658,4 +660,64 @@ mod test_lru_unit {
         assert_eq!(lru.used_weight(), 0);
         assert_lru(&lru, &[]);
     }
+
+    #[test]
+    fn test_remove() {
+        let mut lru = LruUnit::with_capacity(10);
+
+        lru.admit(2, 2, 2);
+        lru.admit(3, 3, 3);
+        lru.admit(4, 4, 4);
+        lru.admit(5, 5, 5);
+        assert_lru(&lru, &[5, 4, 3, 2]);
+
+        assert!(lru.access(4));
+        assert!(lru.access(3));
+        assert!(lru.access(3));
+        assert!(lru.access(2));
+        assert_lru(&lru, &[2, 3, 4, 5]);
+
+        assert_eq!(lru.used_weight(), 2 + 3 + 4 + 5);
+        assert_eq!(lru.remove(2), Some((2, 2)));
+        assert_eq!(lru.used_weight(), 3 + 4 + 5);
+        assert_lru(&lru, &[3, 4, 5]);
+
+        assert_eq!(lru.remove(4), Some((4, 4)));
+        assert_eq!(lru.used_weight(), 3 + 5);
+        assert_lru(&lru, &[3, 5]);
+
+        assert_eq!(lru.remove(5), Some((5, 5)));
+        assert_eq!(lru.used_weight(), 3);
+        assert_lru(&lru, &[3]);
+
+        assert_eq!(lru.remove(1), None);
+        assert_eq!(lru.used_weight(), 3);
+        assert_lru(&lru, &[3]);
+
+        assert_eq!(lru.remove(3), Some((3, 3)));
+        assert_eq!(lru.used_weight(), 0);
+        assert_lru(&lru, &[]);
+    }
+
+    #[test]
+    fn test_insert_tail() {
+        let mut lru = LruUnit::with_capacity(10);
+        assert_eq!(lru.len(), 0);
+        assert!(lru.peek(0).is_none());
+
+        assert!(lru.insert_tail(2, 2, 1));
+        assert_eq!(lru.len(), 1);
+        assert_eq!(lru.peek(2).unwrap(), &2);
+        assert_eq!(lru.used_weight(), 1);
+
+        assert!(!lru.insert_tail(2, 2, 2));
+        assert!(lru.insert_tail(3, 3, 3));
+        assert_eq!(lru.used_weight(), 1 + 3);
+        assert_lru(&lru, &[2, 3]);
+
+        assert!(lru.insert_tail(4, 4, 4));
+        assert!(lru.insert_tail(5, 5, 5));
+        assert_eq!(lru.used_weight(), 1 + 3 + 4 + 5);
+        assert_lru(&lru, &[2, 3, 4, 5]);
+    }
 }
diff --git a/tinyufo/src/estimation.rs b/tinyufo/src/estimation.rs
index 18c2d4f67..6f165ac0e 100644
--- a/tinyufo/src/estimation.rs
+++ b/tinyufo/src/estimation.rs
@@ -183,20 +183,25 @@ mod tests {
         assert_eq!(tiny.incr(1), 2);
         assert_eq!(tiny.get(1), 2);
 
-        assert_eq!(tiny.get(2), 0);
-        assert_eq!(tiny.incr(2), 1);
-        assert_eq!(tiny.incr(2), 2);
-        assert_eq!(tiny.get(2), 2);
-
-        assert_eq!(tiny.incr(3), 1);
-        assert_eq!(tiny.incr(3), 2);
-        assert_eq!(tiny.incr(3), 3);
-        assert_eq!(tiny.incr(3), 4);
-
-        // 8 incr(), now reset
-
-        assert_eq!(tiny.incr(3), 3);
-        assert_eq!(tiny.incr(1), 2);
-        assert_eq!(tiny.incr(2), 2);
+        // Might have hash collisions for the others, need to
+        // get() before can assert on the incr() value.
+        let two = tiny.get(2);
+        assert_eq!(tiny.incr(2), two + 1);
+        assert_eq!(tiny.incr(2), two + 2);
+        assert_eq!(tiny.get(2), two + 2);
+
+        let three = tiny.get(3);
+        assert_eq!(tiny.incr(3), three + 1);
+        assert_eq!(tiny.incr(3), three + 2);
+        assert_eq!(tiny.incr(3), three + 3);
+        assert_eq!(tiny.incr(3), three + 4);
+
+        // 8 incr(), now resets on next incr
+        // can only assert they are greater than or equal
+        // to the incr() we do per key.
+
+        assert!(tiny.incr(3) >= 3); // had 4, reset to 2, added another.
+        assert!(tiny.incr(1) >= 2); // had 2, reset to 1, added another.
+        assert!(tiny.incr(2) >= 2); // had 2, reset to 1, added another.
     }
 }
diff --git a/tinyufo/src/lib.rs b/tinyufo/src/lib.rs
index 015afe0ef..248f7293c 100644
--- a/tinyufo/src/lib.rs
+++ b/tinyufo/src/lib.rs
@@ -121,7 +121,7 @@ pub struct KV<T> {
 }
 
 // the data and its metadata
-struct Bucket<T> {
+pub struct Bucket<T> {
     uses: Uses,
     queue: Location,
     weight: Weight,


Test patch: