Implement some bulk methods (Retain and ApproximateSize) for TConcurrentHash

commit_hash:d1b4740ab0854d5cd2baaff61c17709d97f398f0

2 files changed, 62 insertions, 0 deletions

diff --git a/library/cpp/containers/concurrent_hash/concurrent_hash.h b/library/cpp/containers/concurrent_hash/concurrent_hash.h
index 204434a9783..298d348d539 100644
--- a/library/cpp/containers/concurrent_hash/concurrent_hash.h
+++ b/library/cpp/containers/concurrent_hash/concurrent_hash.h
@@ -83,6 +83,22 @@ public:
             typename TActualMap::iterator it = Map.find(key);
             return it == Map.end() ? nullptr : &it->second;
         }
+
+        template <typename Predicate>
+        void RetainUnsafe(Predicate predicate) {
+            typename TActualMap::iterator it = Map.begin();
+            while (it != Map.end()) {
+                if (predicate(*it)) {
+                    it++;
+                } else {
+                    Map.erase(it++);
+                }
+            }
+        }
+
+        size_t SizeUnsafe() const {
+            return Map.size();
+        }
     };
 
     std::array<TBucket, BucketCount> Buckets;
@@ -182,4 +198,21 @@ public:
         TBucketGuard guard(bucket.Mutex);
         return bucket.HasUnsafe(key);
     }
+
+    template <typename Predicate>
+    void Retain(Predicate predicate) {
+        for (auto& bucket: Buckets) {
+            TBucketGuard guard(bucket.Mutex);
+            bucket.RetainUnsafe(predicate);
+        }
+    }
+
+    size_t ApproximateSize() const {
+        size_t total = 0;
+        for (const auto& bucket: Buckets) {
+            TBucketGuard guard(bucket.Mutex);
+            total += bucket.SizeUnsafe();
+        }
+        return total;
+    }
 };
diff --git a/library/cpp/containers/concurrent_hash/concurrent_hash_ut.cpp b/library/cpp/containers/concurrent_hash/concurrent_hash_ut.cpp
index 6427adf92c6..618cdf5f8f0 100644
--- a/library/cpp/containers/concurrent_hash/concurrent_hash_ut.cpp
+++ b/library/cpp/containers/concurrent_hash/concurrent_hash_ut.cpp
@@ -196,3 +196,32 @@ TEST(TConcurrentHashTest, TGetBucketTest) {
         EXPECT_EQ(&bucket1, &bucket2);
     }
 }
+
+TEST(TConcurrentHashTest, TRetainTest) {
+    TConcurrentHashMap<ui32, ui32> h;
+
+    for (ui32 i = 0; i < 100; ++i) {
+        h.InsertIfAbsent(i, i);
+    }
+
+    h.Retain([](const auto& item) {
+        return item.second == 77;
+    });
+
+    for (ui32 i = 0; i < 100; ++i) {
+        EXPECT_EQ(h.Has(i), i == 77);
+    }
+}
+
+TEST(TConcurrentHashTest, TApproximateSize) {
+    TConcurrentHashMap<ui32, ui32> h;
+    EXPECT_EQ(h.ApproximateSize(), 0);
+    h.InsertIfAbsent(1, 1);
+    EXPECT_EQ(h.ApproximateSize(), 1);
+    h.InsertIfAbsent(2, 2);
+    EXPECT_EQ(h.ApproximateSize(), 2);
+    h.Remove(1);
+    EXPECT_EQ(h.ApproximateSize(), 1);
+    h.Remove(2);
+    EXPECT_EQ(h.ApproximateSize(), 0);
+}