Extract TExpiringSet to library/cpp/containers

commit_hash:dd1c08771b1d4865d03a492927afa0f9895a5f44

4 files changed, 239 insertions, 0 deletions

diff --git a/library/cpp/yt/containers/expiring_set-inl.h b/library/cpp/yt/containers/expiring_set-inl.h
new file mode 100644
index 00000000000..799a0680a02
--- /dev/null
+++ b/library/cpp/yt/containers/expiring_set-inl.h
@@ -0,0 +1,82 @@
+#ifndef EXPIRING_SET_INL_H_
+#error "Direct inclusion of this file is not allowed, include expiring_set.h"
+// For the sake of sane code completion.
+#include "expiring_set.h"
+#endif
+
+#include <array>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class TItem, class THash, class TEqual>
+void TExpiringSet<TItem, THash, TEqual>::SetTTl(TDuration ttl)
+{
+    TTl_ = ttl;
+}
+
+template <class TItem, class THash, class TEqual>
+void TExpiringSet<TItem, THash, TEqual>::Insert(TInstant now, const TItem& item)
+{
+    std::array<std::reference_wrapper<const TItem>, 1> items{std::cref(item)};
+    InsertMany(now, items);
+}
+
+template <class TItem, class THash, class TEqual>
+template <class TItems>
+void TExpiringSet<TItem, THash, TEqual>::InsertMany(TInstant now, const TItems& items)
+{
+    Expire(now);
+    auto deadline = now + TTl_;
+    for (const auto& item : items) {
+        ItemToDeadline_[item] = deadline;
+    }
+    ExpirationQueue_.push(TItemPack{.Items = {items.begin(), items.end()}, .Deadline = deadline});
+}
+
+template <class TItem, class THash, class TEqual>
+void TExpiringSet<TItem, THash, TEqual>::Expire(TInstant now)
+{
+    while (!ExpirationQueue_.empty() && ExpirationQueue_.top().Deadline <= now) {
+        for (const auto& item : ExpirationQueue_.top().Items) {
+            if (auto it = ItemToDeadline_.find(item); it != ItemToDeadline_.end()) {
+                if (it->second <= now) {
+                    ItemToDeadline_.erase(it);
+                }
+            }
+        }
+        ExpirationQueue_.pop();
+    }
+}
+
+template <class TItem, class THash, class TEqual>
+void TExpiringSet<TItem, THash, TEqual>::Clear()
+{
+    ItemToDeadline_ = {};
+    ExpirationQueue_ = {};
+}
+
+template <class TItem, class THash, class TEqual>
+template <class TItemLike>
+bool TExpiringSet<TItem, THash, TEqual>::Contains(const TItemLike& item) const
+{
+    return ItemToDeadline_.contains(item);
+}
+
+template <class TItem, class THash, class TEqual>
+int TExpiringSet<TItem, THash, TEqual>::GetSize() const
+{
+    return std::ssize(ItemToDeadline_);
+}
+
+template <class TItem, class THash, class TEqual>
+bool TExpiringSet<TItem, THash, TEqual>::TItemPack::operator<(const TItemPack& other) const
+{
+    // Reversed ordering for the priority queue.
+    return Deadline > other.Deadline;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/library/cpp/yt/containers/expiring_set.h b/library/cpp/yt/containers/expiring_set.h
new file mode 100644
index 00000000000..f8a1e336309
--- /dev/null
+++ b/library/cpp/yt/containers/expiring_set.h
@@ -0,0 +1,54 @@
+#pragma once
+
+#include <util/generic/hash.h>
+
+#include <util/datetime/base.h>
+
+#include <queue>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! Maintains a set of items that expire after a certain time.
+template <class TItem, class THash = THash<TItem>, class TEqual = TEqualTo<TItem>>
+class TExpiringSet
+{
+public:
+    void SetTTl(TDuration ttl);
+
+    void Insert(TInstant now, const TItem& item);
+    template <class TItems>
+    void InsertMany(TInstant now, const TItems& items);
+
+    void Expire(TInstant now);
+
+    void Clear();
+
+    template <class TItemLike>
+    bool Contains(const TItemLike& item) const;
+
+    int GetSize() const;
+
+private:
+    TDuration TTl_;
+
+    struct TItemPack
+    {
+        std::vector<TItem> Items;
+        TInstant Deadline;
+
+        bool operator<(const TItemPack& other) const;
+    };
+
+    THashMap<TItem, TInstant, THash, TEqual> ItemToDeadline_;
+    std::priority_queue<TItemPack> ExpirationQueue_;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+#define EXPIRING_SET_INL_H_
+#include "expiring_set-inl.h"
+#undef EXPIRING_SET_INL_H_
diff --git a/library/cpp/yt/containers/unittests/expiring_set_ut.cpp b/library/cpp/yt/containers/unittests/expiring_set_ut.cpp
new file mode 100644
index 00000000000..e9f2743a807
--- /dev/null
+++ b/library/cpp/yt/containers/unittests/expiring_set_ut.cpp
@@ -0,0 +1,102 @@
+#include <library/cpp/yt/containers/expiring_set.h>
+
+#include <library/cpp/testing/gtest/gtest.h>
+
+namespace NYT {
+namespace {
+
+////////////////////////////////////////////////////////////////////////////////
+
+TInstant operator""_ts(unsigned long long seconds)
+{
+    return TInstant::Zero() + TDuration::Seconds(seconds);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+TEST(TExpiringSetTest, Empty)
+{
+    TExpiringSet<int> set;
+    EXPECT_EQ(set.GetSize(), 0);
+}
+
+TEST(TExpiringSetTest, ExpireSingle)
+{
+    TExpiringSet<int> set;
+    set.SetTTl(TDuration::Seconds(2));
+
+    set.Insert(0_ts, 1);
+    EXPECT_EQ(set.GetSize(), 1);
+
+    set.Expire(1_ts);
+    EXPECT_EQ(set.GetSize(), 1);
+
+    set.Expire(2_ts);
+    EXPECT_EQ(set.GetSize(), 0);
+}
+
+TEST(TExpiringSetTest, ExpireBatch)
+{
+    TExpiringSet<int> set;
+    set.SetTTl(TDuration::Seconds(2));
+
+    set.InsertMany(0_ts, std::vector<int>{1, 2, 3});
+    EXPECT_EQ(set.GetSize(), 3);
+
+    set.Expire(1_ts);
+    EXPECT_EQ(set.GetSize(), 3);
+
+    set.Expire(2_ts);
+    EXPECT_EQ(set.GetSize(), 0);
+}
+
+TEST(TExpiringSetTest, Reinsert)
+{
+    TExpiringSet<int> set;
+    set.SetTTl(TDuration::Seconds(2));
+
+    set.Insert(0_ts, 1);
+    EXPECT_EQ(set.GetSize(), 1);
+
+    set.Insert(1_ts, 1);
+    EXPECT_EQ(set.GetSize(), 1);
+
+    set.Expire(2_ts);
+    EXPECT_EQ(set.GetSize(), 1);
+
+    set.Expire(3_ts);
+    EXPECT_EQ(set.GetSize(), 0);
+}
+
+TEST(TExpiringSetTest, Contains)
+{
+    TExpiringSet<int> set;
+    set.SetTTl(TDuration::Seconds(1));
+
+    EXPECT_FALSE(set.Contains(1));
+
+    set.Insert(0_ts, 1);
+    EXPECT_TRUE(set.Contains(1));
+
+    set.Expire(1_ts);
+    EXPECT_FALSE(set.Contains(1));
+}
+
+TEST(TExpiringSetTest, Clear)
+{
+    TExpiringSet<int> set;
+    set.SetTTl(TDuration::Seconds(1));
+
+    set.Insert(0_ts, 1);
+    EXPECT_EQ(set.GetSize(), 1);
+    EXPECT_TRUE(set.Contains(1));
+
+    set.Clear();
+    EXPECT_EQ(set.GetSize(), 0);
+    EXPECT_FALSE(set.Contains(1));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+} // namespace NYT
diff --git a/library/cpp/yt/containers/unittests/ya.make b/library/cpp/yt/containers/unittests/ya.make
index 3ffc4206589..3732116a513 100644
--- a/library/cpp/yt/containers/unittests/ya.make
+++ b/library/cpp/yt/containers/unittests/ya.make
@@ -6,6 +6,7 @@ SIZE(MEDIUM)
 
 SRCS(
     enum_indexed_array_ut.cpp
+    expiring_set_ut.cpp
     sharded_set_ut.cpp
 )