grpc 1.66.0 → 1.67.0.pre1

data/src/core/lib/promise/party.cc

@@ -15,6 +15,7 @@
 #include "src/core/lib/promise/party.h"
 
 #include <atomic>
+#include <cstdint>
 
 #include "absl/base/thread_annotations.h"
 #include "absl/log/check.h"
@@ -39,7 +40,7 @@ namespace grpc_core {
 ///////////////////////////////////////////////////////////////////////////////
 // PartySyncUsingAtomics
 
-GRPC_MUST_USE_RESULT bool PartySyncUsingAtomics::RefIfNonZero() {
+GRPC_MUST_USE_RESULT bool Party::RefIfNonZero() {
   auto count = state_.load(std::memory_order_relaxed);
   do {
     // If zero, we are done (without an increment). If not, we must do a CAS
@@ -55,33 +56,6 @@ GRPC_MUST_USE_RESULT bool PartySyncUsingAtomics::RefIfNonZero() {
   return true;
 }
 
-bool PartySyncUsingAtomics::UnreffedLast() {
-  uint64_t prev_state =
-      state_.fetch_or(kDestroying | kLocked, std::memory_order_acq_rel);
-  LogStateChange("UnreffedLast", prev_state,
-                 prev_state | kDestroying | kLocked);
-  return (prev_state & kLocked) == 0;
-}
-
-bool PartySyncUsingAtomics::ScheduleWakeup(WakeupMask mask) {
-  // Or in the wakeup bit for the participant, AND the locked bit.
-  uint64_t prev_state = state_.fetch_or((mask & kWakeupMask) | kLocked,
-                                        std::memory_order_acq_rel);
-  LogStateChange("ScheduleWakeup", prev_state,
-                 prev_state | (mask & kWakeupMask) | kLocked);
-  // If the lock was not held now we hold it, so we need to run.
-  return ((prev_state & kLocked) == 0);
-}
-
-///////////////////////////////////////////////////////////////////////////////
-// PartySyncUsingMutex
-
-bool PartySyncUsingMutex::ScheduleWakeup(WakeupMask mask) {
-  MutexLock lock(&mu_);
-  wakeups_ |= mask;
-  return !std::exchange(locked_, true);
-}
-
 ///////////////////////////////////////////////////////////////////////////////
 // Party::Handle
 
@@ -175,7 +149,7 @@ Party::Participant::~Participant() {
 Party::~Party() {}
 
 void Party::CancelRemainingParticipants() {
-  if (!sync_.has_participants()) return;
+  if ((state_.load(std::memory_order_relaxed) & kAllocatedMask) == 0) return;
   ScopedActivity activity(this);
   promise_detail::Context<Arena> arena_ctx(arena_.get());
   for (size_t i = 0; i < party_detail::kMaxParticipants; i++) {
@@ -206,37 +180,42 @@ Waker Party::MakeNonOwningWaker() {
 
 void Party::ForceImmediateRepoll(WakeupMask mask) {
   DCHECK(is_current());
-  sync_.ForceImmediateRepoll(mask);
+  wakeup_mask_ |= mask;
 }
 
-void Party::RunLocked(Party* party) {
+void Party::RunLockedAndUnref(Party* party, uint64_t prev_state) {
   GRPC_LATENT_SEE_PARENT_SCOPE("Party::RunLocked");
 #ifdef GRPC_MAXIMIZE_THREADYNESS
   Thread thd(
       "RunParty",
-      [party]() {
+      [party, prev_state]() {
         ApplicationCallbackExecCtx app_exec_ctx;
         ExecCtx exec_ctx;
-        if (party->RunParty()) party->PartyIsOver();
+        party->RunPartyAndUnref(prev_state);
       },
       nullptr, Thread::Options().set_joinable(false));
   thd.Start();
 #else
   struct RunState;
   static thread_local RunState* g_run_state = nullptr;
+  struct PartyWakeup {
+    PartyWakeup() : party{nullptr} {}
+    PartyWakeup(Party* party, uint64_t prev_state)
+        : party{party}, prev_state{prev_state} {}
+    Party* party;
+    uint64_t prev_state;
+  };
   struct RunState {
-    explicit RunState(Party* party) : running(party), next(nullptr) {}
-    Party* running;
-    Party* next;
-    void Run() {
+    explicit RunState(PartyWakeup first) : first{first}, next{} {}
+    PartyWakeup first;
+    PartyWakeup next;
+    GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION void Run() {
       g_run_state = this;
       do {
         GRPC_LATENT_SEE_INNER_SCOPE("run_one_party");
-        if (running->RunParty()) {
-          running->PartyIsOver();
-        }
-        running = std::exchange(next, nullptr);
-      } while (running != nullptr);
+        first.party->RunPartyAndUnref(first.prev_state);
+        first = std::exchange(next, PartyWakeup{});
+      } while (first.party != nullptr);
       DCHECK(g_run_state == this);
       g_run_state = nullptr;
     }
@@ -245,119 +224,264 @@ void Party::RunLocked(Party* party) {
   // but instead add it to the end of the list of parties to run.
   // This enables a fairly straightforward batching of work from a
   // call to a transport (or back again).
-  if (g_run_state != nullptr) {
-    if (g_run_state->running == party || g_run_state->next == party) {
-      // Already running or already queued.
+  if (GPR_UNLIKELY(g_run_state != nullptr)) {
+    if (g_run_state->first.party == party) {
+      g_run_state->first.prev_state = prev_state;
+      party->Unref();
       return;
     }
-    if (g_run_state->next != nullptr) {
+    if (g_run_state->next.party == party) {
+      g_run_state->next.prev_state = prev_state;
+      party->Unref();
+      return;
+    }
+    if (g_run_state->next.party != nullptr) {
       // If there's already a different party queued, we're better off asking
       // event engine to run it so we can spread load.
       // We swap the oldest party to run on the event engine so that we don't
       // accidentally end up with a tail latency problem whereby one party
       // gets held for a really long time.
-      std::swap(g_run_state->next, party);
+      auto wakeup =
+          std::exchange(g_run_state->next, PartyWakeup{party, prev_state});
       party->arena_->GetContext<grpc_event_engine::experimental::EventEngine>()
-          ->Run([party]() {
+          ->Run([wakeup]() {
             GRPC_LATENT_SEE_PARENT_SCOPE("Party::RunLocked offload");
             ApplicationCallbackExecCtx app_exec_ctx;
             ExecCtx exec_ctx;
-            RunState{party}.Run();
+            RunState{wakeup}.Run();
           });
       return;
     }
-    g_run_state->next = party;
+    g_run_state->next = PartyWakeup{party, prev_state};
     return;
   }
-  RunState{party}.Run();
+  RunState{{party, prev_state}}.Run();
 #endif
 }
 
-bool Party::RunParty() {
+void Party::RunPartyAndUnref(uint64_t prev_state) {
   ScopedActivity activity(this);
   promise_detail::Context<Arena> arena_ctx(arena_.get());
-  return sync_.RunParty([this](int i) { return RunOneParticipant(i); });
-}
-
-bool Party::RunOneParticipant(int i) {
-  GRPC_LATENT_SEE_INNER_SCOPE("Party::RunOneParticipant");
-  // If the participant is null, skip.
-  // This allows participants to complete whilst wakers still exist
-  // somewhere.
-  auto* participant = participants_[i].load(std::memory_order_acquire);
-  if (participant == nullptr) {
-    if (GRPC_TRACE_FLAG_ENABLED(promise_primitives)) {
-      LOG(INFO) << DebugTag() << "[party] wakeup " << i << " already complete";
-    }
-    return false;
-  }
-  absl::string_view name;
-  if (GRPC_TRACE_FLAG_ENABLED(promise_primitives)) {
-    name = participant->name();
-    LOG(INFO) << DebugTag() << "[" << name << "] begin job " << i;
+  DCHECK_EQ(prev_state & kLocked, 0u)
+      << "Party should be unlocked prior to first wakeup";
+  DCHECK_GE(prev_state & kRefMask, kOneRef);
+  // Now update prev_state to be what we want the CAS to see below.
+  DCHECK_EQ(prev_state & ~(kRefMask | kAllocatedMask), 0u)
+      << "Party should have contained no wakeups on lock";
+  prev_state |= kLocked;
+  absl::optional<ScopedTimeCache> time_cache;
+#if !TARGET_OS_IPHONE
+  if (IsTimeCachingInPartyEnabled()) {
+    time_cache.emplace();
   }
-  // Poll the participant.
-  currently_polling_ = i;
-  bool done = participant->PollParticipantPromise();
-  currently_polling_ = kNotPolling;
-  if (done) {
-    if (!name.empty()) {
-      GRPC_TRACE_LOG(promise_primitives, INFO)
-          << DebugTag() << "[" << name << "] end poll and finish job " << i;
+#endif
+  for (;;) {
+    uint64_t keep_allocated_mask = kAllocatedMask;
+    // For each wakeup bit...
+    while (wakeup_mask_ != 0) {
+      auto wakeup_mask = std::exchange(wakeup_mask_, 0);
+      while (wakeup_mask != 0) {
+        const uint64_t t = LowestOneBit(wakeup_mask);
+        const int i = absl::countr_zero(t);
+        wakeup_mask ^= t;
+        // If the participant is null, skip.
+        // This allows participants to complete whilst wakers still exist
+        // somewhere.
+        auto* participant = participants_[i].load(std::memory_order_acquire);
+        if (GPR_UNLIKELY(participant == nullptr)) {
+          GRPC_TRACE_LOG(promise_primitives, INFO)
+              << "Party " << this << "                 Run:Wakeup " << i
+              << " already complete";
+          continue;
+        }
+        GRPC_TRACE_LOG(promise_primitives, INFO)
+            << "Party " << this << "                 Run:Wakeup " << i;
+        // Poll the participant.
+        currently_polling_ = i;
+        if (participant->PollParticipantPromise()) {
+          participants_[i].store(nullptr, std::memory_order_relaxed);
+          const uint64_t allocated_bit = (1u << i << kAllocatedShift);
+          keep_allocated_mask &= ~allocated_bit;
+        }
+      }
+    }
+    currently_polling_ = kNotPolling;
+    // Try to CAS the state we expected to have (with no wakeups or adds)
+    // back to unlocked (by masking in only the ref mask - sans locked bit).
+    // If this succeeds then no wakeups were added, no adds were added, and we
+    // have successfully unlocked.
+    // Otherwise, we need to loop again.
+    // Note that if an owning waker is created or the weak cas spuriously
+    // fails we will also loop again, but in that case see no wakeups or adds
+    // and so will get back here fairly quickly.
+    // TODO(ctiller): consider mitigations for the accidental wakeup on owning
+    // waker creation case -- I currently expect this will be more expensive
+    // than this quick loop.
+    if (state_.compare_exchange_weak(
+            prev_state,
+            (prev_state & (kRefMask | keep_allocated_mask)) - kOneRef,
+            std::memory_order_acq_rel, std::memory_order_acquire)) {
+      LogStateChange("Run:End", prev_state,
+                     prev_state & (kRefMask | kAllocatedMask) - kOneRef);
+      if ((prev_state & kRefMask) == kOneRef) {
+        // We're done with the party.
+        PartyIsOver();
+      }
+      return;
+    }
+    // CAS out (but retrieve) any allocations and wakeups that occurred during
+    // the run.
+    while (!state_.compare_exchange_weak(
+        prev_state, prev_state & (kRefMask | kLocked | keep_allocated_mask))) {
+      // Nothing to do here.
     }
-    participants_[i].store(nullptr, std::memory_order_relaxed);
-  } else if (!name.empty()) {
-    GRPC_TRACE_LOG(promise_primitives, INFO)
-        << DebugTag() << "[" << name << "] end poll";
+    LogStateChange("Run:Continue", prev_state,
+                   prev_state & (kRefMask | kLocked | keep_allocated_mask));
+    DCHECK(prev_state & kLocked)
+        << "Party should be locked; prev_state=" << prev_state;
+    DCHECK_GE(prev_state & kRefMask, kOneRef);
+    // From the previous state, extract which participants we're to wakeup.
+    wakeup_mask_ |= prev_state & kWakeupMask;
+    // Now update prev_state to be what we want the CAS to see once wakeups
+    // complete next iteration.
+    prev_state &= kRefMask | kLocked | keep_allocated_mask;
   }
-  return done;
 }
 
 void Party::AddParticipants(Participant** participants, size_t count) {
-  bool run_party = sync_.AddParticipantsAndRef(count, [this, participants,
-                                                       count](size_t* slots) {
+  uint64_t state = state_.load(std::memory_order_acquire);
+  uint64_t allocated;
+
+  size_t slots[party_detail::kMaxParticipants];
+
+  // Find slots for each new participant, ordering them from lowest available
+  // slot upwards to ensure the same poll ordering as presentation ordering to
+  // this function.
+  WakeupMask wakeup_mask;
+  uint64_t new_state;
+  do {
+    wakeup_mask = 0;
+    allocated = (state & kAllocatedMask) >> kAllocatedShift;
     for (size_t i = 0; i < count; i++) {
-      if (GRPC_TRACE_FLAG_ENABLED(party_state)) {
-        LOG(INFO) << "Party " << &sync_ << "                 AddParticipant: "
-                  << participants[i]->name() << " @ " << slots[i]
-                  << " [participant=" << participants[i] << "]";
+      auto new_mask = LowestOneBit(~allocated);
+      if (GPR_UNLIKELY((new_mask & kWakeupMask) == 0)) {
+        DelayAddParticipants(participants, count);
+        return;
       }
-      participants_[slots[i]].store(participants[i], std::memory_order_release);
+      wakeup_mask |= new_mask;
+      allocated |= new_mask;
+      slots[i] = absl::countr_zero(new_mask);
+    }
+    // Try to allocate this slot and take a ref (atomically).
+    // Ref needs to be taken because once we store the participant it could be
+    // spuriously woken up and unref the party.
+    new_state = (state | (allocated << kAllocatedShift)) + kOneRef;
+  } while (!state_.compare_exchange_weak(
+      state, new_state, std::memory_order_acq_rel, std::memory_order_acquire));
+  LogStateChange("AddParticipantsAndRef", state, new_state);
+
+  for (size_t i = 0; i < count; i++) {
+    GRPC_TRACE_LOG(party_state, INFO)
+        << "Party " << this << "                 AddParticipant: " << slots[i]
+        << " " << participants[i];
+    participants_[slots[i]].store(participants[i], std::memory_order_release);
+  }
+
+  // Now we need to wake up the party.
+  WakeupFromState(new_state, wakeup_mask);
+}
+
+void Party::AddParticipant(Participant* participant) {
+  uint64_t state = state_.load(std::memory_order_acquire);
+  uint64_t allocated;
+  size_t slot;
+
+  // Find slots for each new participant, ordering them from lowest available
+  // slot upwards to ensure the same poll ordering as presentation ordering to
+  // this function.
+  uint64_t wakeup_mask;
+  uint64_t new_state;
+  do {
+    allocated = (state & kAllocatedMask) >> kAllocatedShift;
+    wakeup_mask = LowestOneBit(~allocated);
+    if (GPR_UNLIKELY((wakeup_mask & kWakeupMask) == 0)) {
+      DelayAddParticipants(&participant, 1);
+      return;
     }
-  });
-  if (run_party) RunLocked(this);
-  Unref();
+    DCHECK_NE(wakeup_mask & kWakeupMask, 0u)
+        << "No available slots for new participant; allocated=" << allocated
+        << " state=" << state << " wakeup_mask=" << wakeup_mask;
+    allocated |= wakeup_mask;
+    slot = absl::countr_zero(wakeup_mask);
+    // Try to allocate this slot and take a ref (atomically).
+    // Ref needs to be taken because once we store the participant it could be
+    // spuriously woken up and unref the party.
+    new_state = (state | (allocated << kAllocatedShift)) + kOneRef;
+  } while (!state_.compare_exchange_weak(
+      state, new_state, std::memory_order_acq_rel, std::memory_order_acquire));
+  LogStateChange("AddParticipantsAndRef", state, new_state);
+  GRPC_TRACE_LOG(party_state, INFO)
+      << "Party " << this << "                 AddParticipant: " << slot
+      << " [participant=" << participant << "]";
+  participants_[slot].store(participant, std::memory_order_release);
+  // Now we need to wake up the party.
+  WakeupFromState(new_state, wakeup_mask);
 }
 
-void Party::Wakeup(WakeupMask wakeup_mask) {
-  if (sync_.ScheduleWakeup(wakeup_mask)) RunLocked(this);
-  Unref();
+void Party::DelayAddParticipants(Participant** participants, size_t count) {
+  // We need to delay the addition of participants.
+  IncrementRefCount();
+  VLOG_EVERY_N_SEC(2, 10) << "Delaying addition of " << count
+                          << " participants to party " << this
+                          << " because it is full.";
+  std::vector<Participant*> delayed_participants{participants,
+                                                 participants + count};
+  arena_->GetContext<grpc_event_engine::experimental::EventEngine>()->Run(
+      [this, delayed_participants = std::move(delayed_participants)]() mutable {
+        ApplicationCallbackExecCtx app_exec_ctx;
+        ExecCtx exec_ctx;
+        AddParticipants(delayed_participants.data(),
+                        delayed_participants.size());
+        Unref();
+      });
 }
 
 void Party::WakeupAsync(WakeupMask wakeup_mask) {
-  if (sync_.ScheduleWakeup(wakeup_mask)) {
-    arena_->GetContext<grpc_event_engine::experimental::EventEngine>()->Run(
-        [this]() {
-          ApplicationCallbackExecCtx app_exec_ctx;
-          ExecCtx exec_ctx;
-          RunLocked(this);
-          Unref();
-        });
-  } else {
-    Unref();
+  // Or in the wakeup bit for the participant, AND the locked bit.
+  uint64_t prev_state = state_.load(std::memory_order_relaxed);
+  LogStateChange("ScheduleWakeup", prev_state,
+                 prev_state | (wakeup_mask & kWakeupMask) | kLocked);
+  while (true) {
+    if ((prev_state & kLocked) == 0) {
+      if (state_.compare_exchange_weak(prev_state, prev_state | kLocked,
+                                       std::memory_order_acq_rel,
+                                       std::memory_order_acquire)) {
+        LogStateChange("WakeupAsync", prev_state, prev_state | kLocked);
+        wakeup_mask_ |= wakeup_mask;
+        arena_->GetContext<grpc_event_engine::experimental::EventEngine>()->Run(
+            [this, prev_state]() {
+              ApplicationCallbackExecCtx app_exec_ctx;
+              ExecCtx exec_ctx;
+              RunLockedAndUnref(this, prev_state);
+            });
+        return;
+      }
+    } else {
+      if (state_.compare_exchange_weak(
+              prev_state, (prev_state | wakeup_mask) - kOneRef,
+              std::memory_order_acq_rel, std::memory_order_acquire)) {
+        LogStateChange("WakeupAsync", prev_state, prev_state | wakeup_mask);
+        return;
+      }
+    }
   }
 }
 
 void Party::Drop(WakeupMask) { Unref(); }
 
 void Party::PartyIsOver() {
-  auto arena = arena_;
-  {
-    ScopedActivity activity(this);
-    promise_detail::Context<Arena> arena_ctx(arena_.get());
-    CancelRemainingParticipants();
-    arena->DestroyManagedNewObjects();
-  }
+  CancelRemainingParticipants();
+  auto arena = std::move(arena_);
   this->~Party();
 }