fix: stop accumulating WhenAny continuations on the Conditions aggregate

Author: beekld

libs/server-sdk/src/data_systems/fdv2/conditions.cpp

@@ -2,6 +2,7 @@
 
 #include <launchdarkly/async/timer.hpp>
 
+#include <algorithm>
 #include <utility>
 #include <variant>
 
@@ -133,15 +134,73 @@ async::Future<IFDv2Condition::Type> MakeAggregateFuture(
 }  // namespace
 
 Conditions::Conditions(std::vector<std::unique_ptr<IFDv2Condition>> conditions)
-    : conditions_(std::move(conditions)),
-      future_(MakeAggregateFuture(conditions_)) {}
+    : conditions_(std::move(conditions)), state_(std::make_shared<State>()) {
+    MakeAggregateFuture(conditions_)
+        .Then(
+            [state =
+                 state_](IFDv2Condition::Type const& type) -> std::monostate {
+                std::vector<PendingEntry> drained;
+                {
+                    std::lock_guard lock(state->mutex);
+                    state->aggregate_result = type;
+                    drained = std::move(state->pending);
+                }
+                for (auto& entry : drained) {
+                    entry.promise.Resolve(type);
+                }
+                return {};
+            },
+            async::kInlineExecutor);
+}
 
 Conditions::~Conditions() {
     Close();
 }
 
-async::Future<IFDv2Condition::Type> Conditions::GetFuture() const {
-    return future_;
+async::Future<IFDv2Condition::Type> Conditions::GetFuture(
+    async::CancellationToken token) {
+    async::Promise<IFDv2Condition::Type> promise;
+    auto future = promise.GetFuture();
+    std::int64_t id;
+
+    {
+        std::lock_guard lock(state_->mutex);
+        if (state_->aggregate_result) {
+            promise.Resolve(*state_->aggregate_result);
+            return future;
+        }
+        id = state_->next_id++;
+        state_->pending.push_back({id, std::move(promise), nullptr});
+    }
+
+    // Construct cb outside the lock: if the token is already cancelled, the
+    // callback fires synchronously inside the ctor and needs to acquire
+    // state_->mutex.
+    auto cb = std::make_unique<async::CancellationCallback>(
+        token, [state = state_, id]() {
+            std::lock_guard lock(state->mutex);
+            state->pending.erase(
+                std::remove_if(
+                    state->pending.begin(), state->pending.end(),
+                    [id](PendingEntry const& e) { return e.id == id; }),
+                state->pending.end());
+        });
+
+    // Re-find the entry under the lock and attach cb. Between the push above
+    // and here, the aggregate may have fired and drained the vector, or cb's
+    // callback may have fired synchronously during construction and erased
+    // the entry. If gone, drop cb.
+    {
+        std::lock_guard lock(state_->mutex);
+        auto it =
+            std::find_if(state_->pending.begin(), state_->pending.end(),
+                         [id](PendingEntry const& e) { return e.id == id; });
+        if (it != state_->pending.end()) {
+            it->cancel_cb = std::move(cb);
+        }
+    }
+
+    return future;
 }
 
 void Conditions::Inform(FDv2SourceResult const& result) {

libs/server-sdk/src/data_systems/fdv2/conditions.hpp

@@ -8,9 +8,11 @@
 #include <boost/asio/any_io_executor.hpp>
 
 #include <chrono>
+#include <cstdint>
 #include <memory>
 #include <mutex>
 #include <optional>
+#include <vector>
 
 namespace launchdarkly::server_side::data_systems {
 
@@ -136,7 +138,7 @@ class RecoveryConditionFactory final
  * Aggregates a set of conditions into a single Future that resolves with the
  * type of the first condition to fire. Inform() and Close() forward to every
  * underlying condition. If constructed with no conditions, GetFuture()
- * returns a Future that never resolves.
+ * returns a Future that never resolves (until Close).
  *
  * Thread-safe: GetFuture, Inform, and Close may be called from any thread.
  */
@@ -153,16 +155,35 @@ class Conditions final {
     Conditions& operator=(Conditions const&) = delete;
     Conditions& operator=(Conditions&&) = delete;
 
+    /**
+     * Returns a fresh Future that resolves with the type of the first
+     * condition to fire. The caller must cancel the source corresponding to
+     * `token` once the result is no longer needed, so that the per-call
+     * Promise (and its registered continuations) can be released.
+     */
     [[nodiscard]] async::Future<data_interfaces::IFDv2Condition::Type>
-    GetFuture() const;
+    GetFuture(async::CancellationToken token);
 
     void Inform(data_interfaces::FDv2SourceResult const& result);
 
     void Close();
 
    private:
+    struct PendingEntry {
+        std::int64_t id;
+        async::Promise<data_interfaces::IFDv2Condition::Type> promise;
+        std::unique_ptr<async::CancellationCallback> cancel_cb;
+    };
+
+    struct State {
+        std::mutex mutex;
+        std::int64_t next_id = 0;
+        std::vector<PendingEntry> pending;
+        std::optional<data_interfaces::IFDv2Condition::Type> aggregate_result;
+    };
+
     std::vector<std::unique_ptr<data_interfaces::IFDv2Condition>> conditions_;
-    async::Future<data_interfaces::IFDv2Condition::Type> future_;
+    std::shared_ptr<State> const state_;
 };
 
 }  // namespace launchdarkly::server_side::data_systems

libs/server-sdk/src/data_systems/fdv2/fdv2_data_system.cpp

@@ -274,11 +274,13 @@ void FDv2DataSystem::RunSynchronizerNext() {
         return;
     }
     auto next_future = active_synchronizer_->Next(selector_);
-    auto cond_future = active_conditions_->GetFuture();
+    auto cond_cancel = std::make_shared<async::CancellationSource>();
+    auto cond_future = active_conditions_->GetFuture(cond_cancel->GetToken());
     async::WhenAny(cond_future, next_future)
         .Then(
-            [this, cond_future,
-             next_future](std::size_t const& idx) -> std::monostate {
+            [this, cond_future, next_future,
+             cond_cancel](std::size_t const& idx) -> std::monostate {
+                cond_cancel->Cancel();
                 if (idx == 0) {
                     OnConditionFired(*cond_future.GetResult());
                 } else {

libs/server-sdk/tests/conditions_test.cpp

@@ -12,6 +12,8 @@ using namespace launchdarkly::server_side::data_interfaces;
 using namespace launchdarkly::server_side::data_systems;
 using namespace std::chrono_literals;
 
+using launchdarkly::async::CancellationToken;
+
 namespace {
 
 // Holds an io_context running on a worker thread; the executor produced by
@@ -162,7 +164,7 @@ TEST(RecoveryConditionTest, CloseCancelsActiveTimerAndResolvesWithCancelled) {
 TEST(ConditionsTest, EmptyAggregateNeverResolves) {
     Conditions conditions({});
 
-    auto result = conditions.GetFuture().WaitForResult(50ms);
+    auto result = conditions.GetFuture(CancellationToken{}).WaitForResult(50ms);
 
     EXPECT_FALSE(result.has_value());
 }
@@ -178,7 +180,7 @@ TEST(ConditionsTest, AggregateResolvesWithTypeOfFirstFiringCondition) {
                                                         /*timeout=*/100ms));
     Conditions conditions(std::move(conds));
 
-    auto result = conditions.GetFuture().WaitForResult(1s);
+    auto result = conditions.GetFuture(CancellationToken{}).WaitForResult(1s);
 
     ASSERT_TRUE(result.has_value());
     EXPECT_EQ(IFDv2Condition::Type::kRecovery, *result);
@@ -202,7 +204,7 @@ TEST(ConditionsTest, InformForwardsToAllUnderlyingConditions) {
             /*status_code=*/0, "boom", std::chrono::system_clock::now()},
     }});
 
-    auto result = conditions.GetFuture().WaitForResult(1s);
+    auto result = conditions.GetFuture(CancellationToken{}).WaitForResult(1s);
 
     ASSERT_TRUE(result.has_value());
     EXPECT_EQ(IFDv2Condition::Type::kFallback, *result);
@@ -220,7 +222,8 @@ TEST(ConditionsTest, CloseForwardsToAllUnderlyingConditions) {
 
     conditions.Close();
 
-    auto result = conditions.GetFuture().WaitForResult(200ms);
+    auto result =
+        conditions.GetFuture(CancellationToken{}).WaitForResult(200ms);
     ASSERT_TRUE(result.has_value());
     EXPECT_EQ(IFDv2Condition::Type::kCancelled, *result);
 }