From 7b2f616919a7018d6067f72cc44421249c5b891f Mon Sep 17 00:00:00 2001
From: "Michael P. Jung" <michael.jung@terreon.de>
Date: Wed, 18 Dec 2024 22:55:23 +0000
Subject: [PATCH] Add RetainResult as return value of Pool::retain method

This is related to #370
---
 CHANGELOG.md       |  1 +
 src/managed/mod.rs | 46 ++++++++++++++++++++++++++++++++++++++--------
 tests/managed.rs   | 13 ++++++++++---
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a7e3ec4..0bc82cd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Update `itertools` dependency to version `0.13.0`
 - Change predicate parameter of `Pool::retain` method to `FnMut`
+- Add `RetainResult` as return value of `Pool::retain` method
 
 ## [0.12.1] - 2024-05-07
 
diff --git a/src/managed/mod.rs b/src/managed/mod.rs
index c7dd35f..dd3373a 100644
--- a/src/managed/mod.rs
+++ b/src/managed/mod.rs
@@ -519,18 +519,30 @@ impl<M: Manager, W: From<Object<M>>> Pool<M, W> {
     ///     }
     /// });
     /// ```
-    pub fn retain(&self, mut f: impl FnMut(&M::Type, Metrics) -> bool) {
+    pub fn retain(
+        &self,
+        mut predicate: impl FnMut(&M::Type, Metrics) -> bool,
+    ) -> RetainResult<M::Type> {
+        let mut removed = Vec::with_capacity(self.status().size);
         let mut guard = self.inner.slots.lock().unwrap();
-        let len_before = guard.vec.len();
-        guard.vec.retain_mut(|obj| {
-            if f(&obj.obj, obj.metrics) {
-                true
+        let mut i = 0;
+        // This code can be simplified once `Vec::extract_if` lands in stable Rust.
+        // https://doc.rust-lang.org/std/vec/struct.Vec.html#method.extract_if
+        while i < guard.vec.len() {
+            let obj = &mut guard.vec[i];
+            if predicate(&mut obj.obj, obj.metrics) {
+                i += 1;
             } else {
+                let mut obj = guard.vec.remove(i).unwrap();
                 self.manager().detach(&mut obj.obj);
-                false
+                removed.push(obj.obj);
             }
-        });
-        guard.size -= len_before - guard.vec.len();
+        }
+        guard.size -= removed.len();
+        RetainResult {
+            retained: i,
+            removed,
+        }
     }
 
     /// Get current timeout configuration
@@ -661,3 +673,21 @@ async fn apply_timeout<O, E>(
         (None, Some(_)) => Err(PoolError::NoRuntimeSpecified),
     }
 }
+
+#[derive(Debug)]
+/// This is the result returned by `Pool::retain`
+pub struct RetainResult<T> {
+    /// Number of retained objects
+    pub retained: usize,
+    /// Objects that were removed from the pool
+    pub removed: Vec<T>,
+}
+
+impl<T> Default for RetainResult<T> {
+    fn default() -> Self {
+        Self {
+            retained: Default::default(),
+            removed: Default::default(),
+        }
+    }
+}
diff --git a/tests/managed.rs b/tests/managed.rs
index 3aeb750..0a538c4 100644
--- a/tests/managed.rs
+++ b/tests/managed.rs
@@ -291,10 +291,14 @@ async fn retain() {
         tokio::time::sleep(Duration::from_millis(5)).await;
     }
     assert_eq!(pool.status().size, 3);
-    pool.retain(|_, metrics| metrics.age() <= Duration::from_millis(10));
+    let retain_result = pool.retain(|_, metrics| metrics.age() <= Duration::from_millis(10));
+    assert_eq!(retain_result.retained, 1);
+    assert_eq!(retain_result.removed.len(), 2);
     assert_eq!(pool.status().size, 1);
     tokio::time::sleep(Duration::from_millis(5)).await;
-    pool.retain(|_, metrics| metrics.age() <= Duration::from_millis(10));
+    let retain_result = pool.retain(|_, metrics| metrics.age() <= Duration::from_millis(10));
+    assert_eq!(retain_result.retained, 0);
+    assert_eq!(retain_result.removed.len(), 1);
     assert_eq!(pool.status().size, 0);
 }
 
@@ -310,9 +314,12 @@ async fn retain_fnmut() {
     }
     let mut removed = 0;
     {
-        pool.retain(|_, _| {
+        let retain_result = pool.retain(|_, _| {
             removed += 1;
             false
         });
+        assert_eq!(retain_result.retained, 0);
+        assert_eq!(retain_result.removed.len(), 4);
     }
+    assert_eq!(pool.status().size, 0);
 }