Extract TSlotMap and use it in TFuture

commit_hash:624d9e33790566421f01090a99e63155da5027f8

8 files changed, 449 insertions, 72 deletions

diff --git a/library/cpp/yt/containers/slot_map-inl.h b/library/cpp/yt/containers/slot_map-inl.h
new file mode 100644
index 00000000000..931678e3c40
--- /dev/null
+++ b/library/cpp/yt/containers/slot_map-inl.h
@@ -0,0 +1,96 @@
+#ifndef SLOT_MAP_INL_H_
+#error "Direct inclusion of this file is not allowed, include slot_map.h"
+// For the sake of sane code completion.
+#include "slot_map.h"
+#endif
+
+#include <library/cpp/yt/assert/assert.h>
+
+#include <iterator>
+#include <utility>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+bool TDefaultSlotMapTraits::IsEmpty(const T& value) const
+{
+    return !static_cast<bool>(value);
+}
+
+template <class T>
+T TDefaultSlotMapTraits::MakeEmpty() const
+{
+    return T();
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T, template <class> class TVector, class TTraits>
+constexpr TSlotMap<T, TVector, TTraits>::TSlotMap(TTraits traits)
+    : Traits_(std::move(traits))
+{ }
+
+template <class T, template <class> class TVector, class TTraits>
+TSlotMapIndex TSlotMap<T, TVector, TTraits>::Insert(T value)
+{
+    YT_VERIFY(!Traits_.IsEmpty(value));
+
+    if (!FreeList_.empty()) {
+        auto index = FreeList_.back();
+        FreeList_.pop_back();
+        auto& slot = Slots_[index.Underlying()];
+        YT_ASSERT(Traits_.IsEmpty(slot));
+        slot = std::move(value);
+        return index;
+    }
+
+    auto index = TSlotMapIndex(std::ssize(Slots_));
+    Slots_.push_back(std::move(value));
+    return index;
+}
+
+template <class T, template <class> class TVector, class TTraits>
+T TSlotMap<T, TVector, TTraits>::Extract(TSlotMapIndex index)
+{
+    auto rawIndex = index.Underlying();
+    YT_ASSERT(rawIndex < std::ssize(Slots_));
+
+    auto& slot = Slots_[rawIndex];
+    YT_ASSERT(!Traits_.IsEmpty(slot));
+
+    auto value = std::move(slot);
+    slot = Traits_.template MakeEmpty<T>();
+    FreeList_.push_back(index);
+    return value;
+}
+
+template <class T, template <class> class TVector, class TTraits>
+template <class F>
+void TSlotMap<T, TVector, TTraits>::ExtractAll(F&& func)
+{
+    for (auto& slot : Slots_) {
+        if (!Traits_.IsEmpty(slot)) {
+            func(std::move(slot));
+        }
+    }
+    Slots_.clear();
+    FreeList_.clear();
+}
+
+template <class T, template <class> class TVector, class TTraits>
+bool TSlotMap<T, TVector, TTraits>::IsEmpty() const
+{
+    return Slots_.size() == FreeList_.size();
+}
+
+template <class T, template <class> class TVector, class TTraits>
+int TSlotMap<T, TVector, TTraits>::GetSize() const
+{
+    return std::ssize(Slots_) - std::ssize(FreeList_);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/library/cpp/yt/containers/slot_map.h b/library/cpp/yt/containers/slot_map.h
new file mode 100644
index 00000000000..f506c0cb42d
--- /dev/null
+++ b/library/cpp/yt/containers/slot_map.h
@@ -0,0 +1,77 @@
+#pragma once
+
+#include <library/cpp/yt/misc/strong_typedef.h>
+
+#include <vector>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! A stable handle into a #TSlotMap; remains valid until the value is extracted.
+YT_DEFINE_STRONG_TYPEDEF(TSlotMapIndex, ui32);
+
+//! A sentinel that never refers to a live slot.
+constexpr auto InvalidSlotMapIndex = TSlotMapIndex(-1);
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! Defines the tombstone marking free slots: a value-initialized |T| convertible
+//! to |false| (empty callback, null pointer). Provide a custom traits object for
+//! types lacking a null state.
+struct TDefaultSlotMapTraits
+{
+    template <class T>
+    bool IsEmpty(const T& value) const;
+
+    template <class T>
+    T MakeEmpty() const;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! Maps stable indices to values, recycling freed slots via a free list.
+/*!
+ *  Values live in a dense vector; each #Insert returns an index valid until the
+ *  value is extracted. Freed slots hold a tombstone (see #TDefaultSlotMapTraits).
+ *  Inserted values must be non-empty. Not thread-safe.
+ */
+template <
+    class T,
+    template <class> class TVector = std::vector,
+    class TTraits = TDefaultSlotMapTraits>
+class TSlotMap
+{
+public:
+    explicit constexpr TSlotMap(TTraits traits = {});
+
+    //! Inserts a non-empty #value; returns its stable index.
+    TSlotMapIndex Insert(T value);
+
+    //! Extracts and returns the value at #index, freeing the slot.
+    /*!
+     *  #index must refer to a live slot (as returned by a prior #Insert and not
+     *  yet extracted).
+     */
+    T Extract(TSlotMapIndex index);
+
+    //! Moves each live value into #func, then clears the map.
+    template <class F>
+    void ExtractAll(F&& func);
+
+    bool IsEmpty() const;
+    int GetSize() const;
+
+private:
+    Y_NO_UNIQUE_ADDRESS TTraits Traits_;
+    TVector<T> Slots_;
+    TVector<TSlotMapIndex> FreeList_;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+#define SLOT_MAP_INL_H_
+#include "slot_map-inl.h"
+#undef SLOT_MAP_INL_H_
diff --git a/library/cpp/yt/containers/unittests/slot_map_ut.cpp b/library/cpp/yt/containers/unittests/slot_map_ut.cpp
new file mode 100644
index 00000000000..68488b6c6eb
--- /dev/null
+++ b/library/cpp/yt/containers/unittests/slot_map_ut.cpp
@@ -0,0 +1,206 @@
+#include <library/cpp/yt/containers/slot_map.h>
+
+#include <library/cpp/yt/compact_containers/compact_vector.h>
+
+#include <library/cpp/testing/gtest/gtest.h>
+
+#include <memory>
+#include <vector>
+
+namespace NYT {
+namespace {
+
+////////////////////////////////////////////////////////////////////////////////
+
+using TIntPtr = std::shared_ptr<int>;
+
+TIntPtr MakeInt(int value)
+{
+    return std::make_shared<int>(value);
+}
+
+std::vector<int> DrainSorted(TSlotMap<TIntPtr>& map)
+{
+    std::vector<int> result;
+    map.ExtractAll([&] (TIntPtr value) {
+        result.push_back(*value);
+    });
+    std::sort(result.begin(), result.end());
+    return result;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+TEST(TSlotMapTest, Empty)
+{
+    TSlotMap<TIntPtr> map;
+    EXPECT_TRUE(map.IsEmpty());
+    EXPECT_EQ(0, map.GetSize());
+}
+
+TEST(TSlotMapTest, InsertAndExtract)
+{
+    TSlotMap<TIntPtr> map;
+
+    auto index = map.Insert(MakeInt(42));
+    EXPECT_FALSE(map.IsEmpty());
+    EXPECT_EQ(1, map.GetSize());
+
+    auto value = map.Extract(index);
+    EXPECT_EQ(42, *value);
+    EXPECT_TRUE(map.IsEmpty());
+    EXPECT_EQ(0, map.GetSize());
+}
+
+TEST(TSlotMapTest, DistinctIndices)
+{
+    TSlotMap<TIntPtr> map;
+
+    auto a = map.Insert(MakeInt(1));
+    auto b = map.Insert(MakeInt(2));
+    auto c = map.Insert(MakeInt(3));
+
+    EXPECT_NE(a, b);
+    EXPECT_NE(b, c);
+    EXPECT_NE(a, c);
+    EXPECT_EQ(3, map.GetSize());
+    EXPECT_THAT(DrainSorted(map), ::testing::ElementsAre(1, 2, 3));
+}
+
+TEST(TSlotMapTest, SlotReuse)
+{
+    TSlotMap<TIntPtr> map;
+
+    auto a = map.Insert(MakeInt(1));
+    map.Insert(MakeInt(2));
+
+    EXPECT_EQ(1, *map.Extract(a));
+
+    // The freed slot should be recycled by the next insert.
+    auto c = map.Insert(MakeInt(3));
+    EXPECT_EQ(a, c);
+    EXPECT_EQ(2, map.GetSize());
+    EXPECT_THAT(DrainSorted(map), ::testing::ElementsAre(2, 3));
+}
+
+TEST(TSlotMapTest, ExtractAllSkipsHoles)
+{
+    TSlotMap<TIntPtr> map;
+
+    auto a = map.Insert(MakeInt(1));
+    map.Insert(MakeInt(2));
+    auto c = map.Insert(MakeInt(3));
+    map.Insert(MakeInt(4));
+
+    map.Extract(a);
+    map.Extract(c);
+
+    EXPECT_THAT(DrainSorted(map), ::testing::ElementsAre(2, 4));
+    EXPECT_TRUE(map.IsEmpty());
+
+    // The map is fully reusable after draining.
+    auto index = map.Insert(MakeInt(5));
+    EXPECT_EQ(TSlotMapIndex(0), index);
+    EXPECT_EQ(5, *map.Extract(index));
+}
+
+TEST(TSlotMapTest, DestroysValueOnExtract)
+{
+    TSlotMap<TIntPtr> map;
+
+    auto value = MakeInt(1);
+    std::weak_ptr<int> weak = value;
+    auto index = map.Insert(std::move(value));
+    EXPECT_FALSE(weak.expired());
+
+    // The map no longer holds a reference; only the extracted value does.
+    auto extracted = map.Extract(index);
+    extracted.reset();
+    EXPECT_TRUE(weak.expired());
+}
+
+TEST(TSlotMapTest, DestroysValuesOnExtractAll)
+{
+    TSlotMap<TIntPtr> map;
+
+    auto value = MakeInt(1);
+    std::weak_ptr<int> weak = value;
+    map.Insert(std::move(value));
+
+    map.ExtractAll([] (TIntPtr) { });
+    EXPECT_TRUE(weak.expired());
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// A stateful traits object using a caller-supplied sentinel as the tombstone.
+struct TSentinelTraits
+{
+    int Sentinel = 0;
+
+    template <class U>
+    bool IsEmpty(const U& value) const
+    {
+        return value == Sentinel;
+    }
+
+    template <class U>
+    U MakeEmpty() const
+    {
+        return Sentinel;
+    }
+};
+
+TEST(TSlotMapTest, CustomTraits)
+{
+    TSlotMap<int, std::vector, TSentinelTraits> map(TSentinelTraits{.Sentinel = -1});
+
+    // Zero is an ordinary value here since the sentinel is -1.
+    auto a = map.Insert(0);
+    map.Insert(5);
+    EXPECT_EQ(2, map.GetSize());
+
+    EXPECT_EQ(0, map.Extract(a));
+    auto c = map.Insert(7);
+    EXPECT_EQ(a, c);
+
+    std::vector<int> values;
+    map.ExtractAll([&] (int value) {
+        values.push_back(value);
+    });
+    std::sort(values.begin(), values.end());
+    EXPECT_THAT(values, ::testing::ElementsAre(5, 7));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class U>
+using TCompactVector2 = TCompactVector<U, 2>;
+
+TEST(TSlotMapTest, CustomVector)
+{
+    TSlotMap<TIntPtr, TCompactVector2> map;
+
+    auto a = map.Insert(MakeInt(1));
+    map.Insert(MakeInt(2));
+    // Force reallocation past the inline capacity.
+    map.Insert(MakeInt(3));
+
+    EXPECT_EQ(3, map.GetSize());
+    EXPECT_EQ(1, *map.Extract(a));
+
+    auto d = map.Insert(MakeInt(4));
+    EXPECT_EQ(a, d);
+
+    std::vector<int> values;
+    map.ExtractAll([&] (TIntPtr value) {
+        values.push_back(*value);
+    });
+    std::sort(values.begin(), values.end());
+    EXPECT_THAT(values, ::testing::ElementsAre(2, 3, 4));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+} // namespace NYT
diff --git a/library/cpp/yt/containers/unittests/ya.make b/library/cpp/yt/containers/unittests/ya.make
index aa5aac875c8..4ec4f5ef158 100644
--- a/library/cpp/yt/containers/unittests/ya.make
+++ b/library/cpp/yt/containers/unittests/ya.make
@@ -11,10 +11,12 @@ SRCS(
     ordered_hash_map_ut.cpp
     sentinel_optional_ut.cpp
     sharded_set_ut.cpp
+    slot_map_ut.cpp
 )
 
 PEERDIR(
     library/cpp/yt/containers
+    library/cpp/yt/compact_containers
 
     library/cpp/testing/gtest
 )
diff --git a/library/cpp/yt/containers/ya.make b/library/cpp/yt/containers/ya.make
index fbc70dbd7c8..8dcea5d170b 100644
--- a/library/cpp/yt/containers/ya.make
+++ b/library/cpp/yt/containers/ya.make
@@ -7,6 +7,7 @@ SRCS(
 
 PEERDIR(
     library/cpp/yt/assert
+    library/cpp/yt/misc
 )
 
 END()
diff --git a/yt/yt/core/actions/future-inl.h b/yt/yt/core/actions/future-inl.h
index c95ef391a54..d4544dba01d 100644
--- a/yt/yt/core/actions/future-inl.h
+++ b/yt/yt/core/actions/future-inl.h
@@ -17,6 +17,8 @@
 
 #include <library/cpp/yt/compact_containers/compact_vector.h>
 
+#include <library/cpp/yt/containers/slot_map.h>
+
 #include <algorithm>
 #include <atomic>
 #include <type_traits>
@@ -80,70 +82,19 @@ inline TError TryExtractCancelationError()
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <class T, TFutureCallbackCookie MinCookie, TFutureCallbackCookie MaxCookie>
-class TFutureCallbackList
-{
-public:
-    static bool IsValidCookie(TFutureCallbackCookie cookie)
-    {
-        return cookie >= MinCookie && cookie <= MaxCookie;
-    }
-
-    TFutureCallbackCookie Add(T callback)
-    {
-        YT_ASSERT(callback);
-        TFutureCallbackCookie cookie;
-        if (SpareCookies_.empty()) {
-            cookie = static_cast<TFutureCallbackCookie>(Callbacks_.size());
-            Callbacks_.push_back(std::move(callback));
-        } else {
-            cookie = SpareCookies_.back();
-            SpareCookies_.pop_back();
-            YT_ASSERT(!Callbacks_[cookie]);
-            Callbacks_[cookie] = std::move(callback);
-        }
-        cookie += MinCookie;
-        YT_ASSERT(cookie <= MaxCookie);
-        return cookie;
-    }
-
-    bool TryRemove(TFutureCallbackCookie cookie, TGuard<NThreading::TSpinLock>* guard)
-    {
-        if (!IsValidCookie(cookie)) {
-            return false;
-        }
-        cookie -= MinCookie;
-        YT_ASSERT(cookie >= 0 && cookie < std::ssize(Callbacks_));
-        YT_ASSERT(Callbacks_[cookie]);
-        SpareCookies_.push_back(cookie);
-        auto callback = std::move(Callbacks_[cookie]);
-        // Make sure callback is not being destroyed under spinlock.
-        guard->Release();
-        return true;
-    }
-
-    template <class... As>
-    void RunAndClear(const As&... args)
-    {
-        for (const auto& callback : Callbacks_) {
-            if (callback) {
-                RunFutureHandler(callback, args...);
-            }
-        }
-        Callbacks_.clear();
-        SpareCookies_.clear();
-    }
+// Future callback cookies are partitioned into two ranges so that a single
+// cookie identifies which handler list owns it: void-result handlers occupy
+// [VoidResultHandlerCookieBase, ...Base + Span), typed-result handlers occupy
+// [ResultHandlerCookieBase, ...Base + Span).
+constexpr ui32 VoidResultHandlerCookieBase = 0;
+constexpr ui32 ResultHandlerCookieBase = 1u << 30;
+constexpr ui32 FutureCallbackCookieSpan = 1u << 30;
 
-    bool IsEmpty() const
-    {
-        return Callbacks_.size() == SpareCookies_.size();
-    }
+template <class T>
+using TFutureCallbackVector = TCompactVector<T, 2>;
 
-private:
-    static constexpr int TypicalCount = 2;
-    TCompactVector<T, TypicalCount> Callbacks_;
-    TCompactVector<TFutureCallbackCookie, TypicalCount> SpareCookies_;
-};
+template <class T>
+using TFutureCallbackMap = TSlotMap<T, TFutureCallbackVector>;
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -223,7 +174,7 @@ class TFutureState<void>
 {
 public:
     using TVoidResultHandler = TCallback<void(const TError&)>;
-    using TVoidResultHandlers = TFutureCallbackList<TVoidResultHandler, 0, (1ULL << 30) - 1>;
+    using TVoidResultHandlers = TFutureCallbackMap<TVoidResultHandler>;
 
     using TUniqueVoidResultHandler = TCallback<void(TError&&)>;
 
@@ -432,7 +383,9 @@ protected:
             CancelHandlers_.clear();
         }
 
-        VoidResultHandlers_.RunAndClear(ResultError_);
+        VoidResultHandlers_.ExtractAll([&] (const TVoidResultHandler& handler) {
+            RunFutureHandler(handler, ResultError_);
+        });
 
         return true;
     }
@@ -455,6 +408,42 @@ protected:
     void WaitUntilSet() const;
     bool CheckIfSet() const;
 
+    static TFutureCallbackCookie EncodeFutureCallbackCookie(TSlotMapIndex index, ui32 base)
+    {
+        auto offset = static_cast<ui32>(index.Underlying());
+        YT_ASSERT(offset < FutureCallbackCookieSpan);
+        return TFutureCallbackCookie(base + offset);
+    }
+
+    static TSlotMapIndex TryDecodeFutureCallbackCookie(TFutureCallbackCookie cookie, ui32 base)
+    {
+        // NB: Unsigned wraparound also rejects cookies below #base.
+        auto offset = cookie.Underlying() - base;
+        if (offset >= FutureCallbackCookieSpan) {
+            return InvalidSlotMapIndex;
+        }
+        return TSlotMapIndex(offset);
+    }
+
+    // Extracts a handler from #map by #cookie, releasing #guard before the handler
+    // is destroyed. Returns |false| if #cookie does not belong to #map's range.
+    template <class T>
+    static bool TryUnsubscribe(
+        TFutureCallbackMap<T>* map,
+        TFutureCallbackCookie cookie,
+        ui32 base,
+        TGuard<NThreading::TSpinLock>* guard)
+    {
+        auto index = TryDecodeFutureCallbackCookie(cookie, base);
+        if (index == InvalidSlotMapIndex) {
+            return false;
+        }
+        auto handler = map->Extract(index);
+        // Make sure handler is not being destroyed under spinlock.
+        guard->Release();
+        return true;
+    }
+
 private:
     void OnLastFutureRefLost();
     void OnLastPromiseRefLost();
@@ -468,7 +457,7 @@ class TFutureState
 {
 public:
     using TResultHandler = TCallback<void(const TErrorOr<T>&)>;
-    using TResultHandlers = TFutureCallbackList<TResultHandler, (1ULL << 30), (1ULL << 31) - 1>;
+    using TResultHandlers = TFutureCallbackMap<TResultHandler>;
 
     using TUniqueResultHandler = TCallback<void(TErrorOr<T>&&)>;
 
@@ -522,7 +511,10 @@ private:
         // It is possible that the result has already been moved out by, e.g., GetUnique.
         // Hence GetResult must only be called when we actually have handlers to invoke.
         if (!ResultHandlers_.IsEmpty()) {
-            ResultHandlers_.RunAndClear(GetResult());
+            const auto& result = GetResult();
+            ResultHandlers_.ExtractAll([&] (const TResultHandler& handler) {
+                RunFutureHandler(handler, result);
+            });
         }
 
         if (UniqueResultHandler_) {
@@ -580,7 +572,7 @@ private:
     {
         YT_ASSERT_SPINLOCK_AFFINITY(SpinLock_);
         return
-            ResultHandlers_.TryRemove(cookie, guard) ||
+            TryUnsubscribe(&ResultHandlers_, cookie, ResultHandlerCookieBase, guard) ||
             TFutureState<void>::DoUnsubscribe(cookie, guard);
     }
 
@@ -680,7 +672,7 @@ public:
                 return NullFutureCallbackCookie;
             } else {
                 HasHandlers_ = true;
-                return ResultHandlers_.Add(std::move(handler));
+                return EncodeFutureCallbackCookie(ResultHandlers_.Insert(std::move(handler)), ResultHandlerCookieBase);
             }
         }
     }
diff --git a/yt/yt/core/actions/future.cpp b/yt/yt/core/actions/future.cpp
index d0c3b634827..6702134b98f 100644
--- a/yt/yt/core/actions/future.cpp
+++ b/yt/yt/core/actions/future.cpp
@@ -51,7 +51,7 @@ TFutureCallbackCookie TFutureState<void>::Subscribe(TVoidResultHandler handler)
             return NullFutureCallbackCookie;
         } else {
             HasHandlers_ = true;
-            return VoidResultHandlers_.Add(std::move(handler));
+            return EncodeFutureCallbackCookie(VoidResultHandlers_.Insert(std::move(handler)), VoidResultHandlerCookieBase);
         }
     }
 }
@@ -205,7 +205,7 @@ void TFutureState<void>::SetErrorGuarded(const TError& error, TGuard<NThreading:
 bool TFutureState<void>::DoUnsubscribe(TFutureCallbackCookie cookie, TGuard<NThreading::TSpinLock>* guard)
 {
     YT_ASSERT_SPINLOCK_AFFINITY(SpinLock_);
-    return VoidResultHandlers_.TryRemove(cookie, guard);
+    return TryUnsubscribe(&VoidResultHandlers_, cookie, VoidResultHandlerCookieBase, guard);
 }
 
 void TFutureState<void>::WaitUntilSet() const
diff --git a/yt/yt/core/actions/future.h b/yt/yt/core/actions/future.h
index 837266a17da..653d811762b 100644
--- a/yt/yt/core/actions/future.h
+++ b/yt/yt/core/actions/future.h
@@ -6,6 +6,9 @@
 
 #include <yt/yt/core/misc/error.h>
 
+#include <library/cpp/yt/misc/strong_typedef.h>
+
+#include <limits>
 #include <optional>
 #include <type_traits>
 
@@ -159,8 +162,8 @@ private:
 ////////////////////////////////////////////////////////////////////////////////
 
 //! An opaque future callback id.
-using TFutureCallbackCookie = int;
-constexpr TFutureCallbackCookie NullFutureCallbackCookie = -1;
+YT_DEFINE_STRONG_TYPEDEF(TFutureCallbackCookie, ui32);
+constexpr auto NullFutureCallbackCookie = TFutureCallbackCookie(-1);
 
 ////////////////////////////////////////////////////////////////////////////////