grpc 1.60.2 → 1.61.0.pre2

data/src/core/lib/promise/detail/status.h

@@ -45,7 +45,12 @@ inline absl::Status IntoStatus(absl::Status* status) {
 // can participate in TrySeq as result types that affect control flow.
 inline bool IsStatusOk(const absl::Status& status) { return status.ok(); }
 
-template <typename To, typename From>
+template <typename T>
+inline bool IsStatusOk(const absl::StatusOr<T>& status) {
+  return status.ok();
+}
+
+template <typename To, typename From, typename SfinaeVoid = void>
 struct StatusCastImpl;
 
 template <typename To>
@@ -59,20 +64,52 @@ struct StatusCastImpl<To, const To&> {
 };
 
 template <typename T>
-struct StatusCastImpl<absl::StatusOr<T>, absl::Status> {
-  static absl::StatusOr<T> Cast(absl::Status&& t) { return std::move(t); }
+struct StatusCastImpl<absl::Status, absl::StatusOr<T>> {
+  static absl::Status Cast(absl::StatusOr<T>&& t) {
+    return std::move(t.status());
+  }
 };
 
 template <typename T>
-struct StatusCastImpl<absl::StatusOr<T>, const absl::Status&> {
-  static absl::StatusOr<T> Cast(const absl::Status& t) { return t; }
+struct StatusCastImpl<absl::Status, absl::StatusOr<T>&> {
+  static absl::Status Cast(const absl::StatusOr<T>& t) { return t.status(); }
 };
 
+template <typename T>
+struct StatusCastImpl<absl::Status, const absl::StatusOr<T>&> {
+  static absl::Status Cast(const absl::StatusOr<T>& t) { return t.status(); }
+};
+
+// StatusCast<> allows casting from one status-bearing type to another,
+// regardless of whether the status indicates success or failure.
+// This means that we can go from StatusOr to Status safely, but not in the
+// opposite direction.
+// For cases where the status is guaranteed to be a failure (and hence not
+// needing to preserve values) see FailureStatusCast<> below.
 template <typename To, typename From>
 To StatusCast(From&& from) {
   return StatusCastImpl<To, From>::Cast(std::forward<From>(from));
 }
 
+template <typename To, typename From, typename SfinaeVoid = void>
+struct FailureStatusCastImpl : public StatusCastImpl<To, From> {};
+
+template <typename T>
+struct FailureStatusCastImpl<absl::StatusOr<T>, absl::Status> {
+  static absl::StatusOr<T> Cast(absl::Status&& t) { return std::move(t); }
+};
+
+template <typename T>
+struct FailureStatusCastImpl<absl::StatusOr<T>, const absl::Status&> {
+  static absl::StatusOr<T> Cast(const absl::Status& t) { return t; }
+};
+
+template <typename To, typename From>
+To FailureStatusCast(From&& from) {
+  GPR_DEBUG_ASSERT(!IsStatusOk(from));
+  return FailureStatusCastImpl<To, From>::Cast(std::forward<From>(from));
+}
+
 }  // namespace grpc_core
 
 #endif  // GRPC_SRC_CORE_LIB_PROMISE_DETAIL_STATUS_H

data/src/core/lib/promise/for_each.h

@@ -30,7 +30,9 @@
 #include "src/core/lib/gprpp/construct_destruct.h"
 #include "src/core/lib/promise/activity.h"
 #include "src/core/lib/promise/detail/promise_factory.h"
+#include "src/core/lib/promise/detail/status.h"
 #include "src/core/lib/promise/poll.h"
+#include "src/core/lib/promise/status_flag.h"
 #include "src/core/lib/promise/trace.h"
 
 namespace grpc_core {
@@ -48,6 +50,16 @@ struct Done<absl::Status> {
   static absl::Status Make() { return absl::OkStatus(); }
 };
 
+template <>
+struct Done<StatusFlag> {
+  static StatusFlag Make() { return StatusFlag(true); }
+};
+
+template <>
+struct Done<Success> {
+  static Success Make() { return Success{}; }
+};
+
 template <typename Reader, typename Action>
 class ForEach {
  private:
@@ -139,7 +151,7 @@ class ForEach {
     }
     auto r = in_action_.promise();
     if (auto* p = r.value_if_ready()) {
-      if (p->ok()) {
+      if (IsStatusOk(*p)) {
         Destruct(&in_action_);
         Construct(&reader_next_, reader_.Next());
         reading_next_ = true;

data/src/core/lib/promise/if.h

@@ -192,6 +192,10 @@ class If<bool, T, F> {
 // If it returns failure, returns failure for the entire combinator.
 // If it returns true, evaluates the second promise.
 // If it returns false, evaluates the third promise.
+// If C is a constant, it's guaranteed that one of the promise factories
+// if_true or if_false will be evaluated before returning from this function.
+// This makes it safe to capture lambda arguments in the promise factory by
+// reference.
 template <typename C, typename T, typename F>
 promise_detail::If<C, T, F> If(C condition, T if_true, F if_false) {
   return promise_detail::If<C, T, F>(std::move(condition), std::move(if_true),

data/src/core/lib/promise/latch.h

@@ -37,6 +37,7 @@ namespace grpc_core {
 // Initially the Latch is unset.
 // It can be waited upon by the Wait method, which produces a Promise that
 // resolves when the Latch is Set to a value of type T.
+// Latches only work correctly within a single activity.
 template <typename T>
 class Latch {
  public:
@@ -184,6 +185,8 @@ class Latch<void> {
     waiter_.Wake();
   }
 
+  bool is_set() const { return is_set_; }
+
  private:
   std::string DebugTag() {
     return absl::StrCat(Activity::current()->DebugTag(), " LATCH(void)[0x",
@@ -204,6 +207,9 @@ class Latch<void> {
   IntraActivityWaiter waiter_;
 };
 
+template <typename T>
+using LatchWaitPromise = decltype(std::declval<Latch<T>>().Wait());
+
 // A Latch that can have its value observed by outside threads, but only waited
 // upon from inside a single activity.
 template <typename T>
@@ -268,9 +274,6 @@ class ExternallyObservableLatch<void> {
   IntraActivityWaiter waiter_;
 };
 
-template <typename T>
-using LatchWaitPromise = decltype(std::declval<Latch<T>>().Wait());
-
 }  // namespace grpc_core
 
 #endif  // GRPC_SRC_CORE_LIB_PROMISE_LATCH_H

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

@@ -227,40 +227,42 @@ void Party::RunLocked() {
 bool Party::RunParty() {
   ScopedActivity activity(this);
   promise_detail::Context<Arena> arena_ctx(arena_);
-  return sync_.RunParty([this](int i) {
-    // 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_promise_primitives.enabled()) {
-        gpr_log(GPR_DEBUG, "%s[party] wakeup %d already complete",
-                DebugTag().c_str(), i);
-      }
-      return false;
-    }
-    absl::string_view name;
+  return sync_.RunParty([this](int i) { return RunOneParticipant(i); });
+}
+
+bool Party::RunOneParticipant(int i) {
+  // 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_promise_primitives.enabled()) {
-      name = participant->name();
-      gpr_log(GPR_DEBUG, "%s[%s] begin job %d", DebugTag().c_str(),
-              std::string(name).c_str(), i);
+      gpr_log(GPR_DEBUG, "%s[party] wakeup %d already complete",
+              DebugTag().c_str(), i);
     }
-    // Poll the participant.
-    currently_polling_ = i;
-    bool done = participant->Poll();
-    currently_polling_ = kNotPolling;
-    if (done) {
-      if (!name.empty()) {
-        gpr_log(GPR_DEBUG, "%s[%s] end poll and finish job %d",
-                DebugTag().c_str(), std::string(name).c_str(), i);
-      }
-      participants_[i].store(nullptr, std::memory_order_relaxed);
-    } else if (!name.empty()) {
-      gpr_log(GPR_DEBUG, "%s[%s] end poll", DebugTag().c_str(),
-              std::string(name).c_str());
+    return false;
+  }
+  absl::string_view name;
+  if (grpc_trace_promise_primitives.enabled()) {
+    name = participant->name();
+    gpr_log(GPR_DEBUG, "%s[%s] begin job %d", DebugTag().c_str(),
+            std::string(name).c_str(), i);
+  }
+  // Poll the participant.
+  currently_polling_ = i;
+  bool done = participant->PollParticipantPromise();
+  currently_polling_ = kNotPolling;
+  if (done) {
+    if (!name.empty()) {
+      gpr_log(GPR_DEBUG, "%s[%s] end poll and finish job %d",
+              DebugTag().c_str(), std::string(name).c_str(), i);
     }
-    return done;
-  });
+    participants_[i].store(nullptr, std::memory_order_relaxed);
+  } else if (!name.empty()) {
+    gpr_log(GPR_DEBUG, "%s[%s] end poll", DebugTag().c_str(),
+            std::string(name).c_str());
+  }
+  return done;
 }
 
 void Party::AddParticipants(Participant** participants, size_t count) {

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

@@ -24,6 +24,7 @@
 #include <string>
 #include <utility>
 
+#include "absl/base/attributes.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/strings/string_view.h"
 
@@ -38,6 +39,7 @@
 #include "src/core/lib/promise/activity.h"
 #include "src/core/lib/promise/context.h"
 #include "src/core/lib/promise/detail/promise_factory.h"
+#include "src/core/lib/promise/poll.h"
 #include "src/core/lib/promise/trace.h"
 #include "src/core/lib/resource_quota/arena.h"
 
@@ -100,7 +102,8 @@ class PartySyncUsingAtomics {
   template <typename F>
   GRPC_MUST_USE_RESULT bool RunParty(F poll_one_participant) {
     uint64_t prev_state;
-    do {
+    iteration_.fetch_add(1, std::memory_order_relaxed);
+    for (;;) {
       // Grab the current state, and clear the wakeup bits & add flag.
       prev_state = state_.fetch_and(kRefMask | kLocked | kAllocatedMask,
                                     std::memory_order_acquire);
@@ -131,9 +134,23 @@ class PartySyncUsingAtomics {
       // 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.
-    } while (!state_.compare_exchange_weak(
-        prev_state, (prev_state & (kRefMask | kAllocatedMask)),
-        std::memory_order_acq_rel, std::memory_order_acquire));
+      if (wake_after_poll_ == 0) {
+        if (state_.compare_exchange_weak(
+                prev_state, (prev_state & (kRefMask | kAllocatedMask)),
+                std::memory_order_acq_rel, std::memory_order_acquire)) {
+          return false;
+        }
+      } else {
+        if (state_.compare_exchange_weak(
+                prev_state,
+                (prev_state & (kRefMask | kAllocatedMask | kLocked)) |
+                    wake_after_poll_,
+                std::memory_order_acq_rel, std::memory_order_acquire)) {
+          iteration_.fetch_add(1, std::memory_order_relaxed);
+          wake_after_poll_ = 0;
+        }
+      }
+    }
     return false;
   }
 
@@ -184,6 +201,11 @@ class PartySyncUsingAtomics {
   // Returns true if the caller should run the party.
   GRPC_MUST_USE_RESULT bool ScheduleWakeup(WakeupMask mask);
 
+  void WakeAfterPoll(WakeupMask mask) { wake_after_poll_ |= mask; }
+  uint32_t iteration() const {
+    return iteration_.load(std::memory_order_relaxed);
+  }
+
  private:
   bool UnreffedLast();
 
@@ -223,6 +245,8 @@ class PartySyncUsingAtomics {
   static constexpr uint64_t kOneRef = 1ull << kRefShift;
 
   std::atomic<uint64_t> state_;
+  std::atomic<uint32_t> iteration_{0};
+  WakeupMask wake_after_poll_ = 0;
 };
 
 class PartySyncUsingMutex {
@@ -298,7 +322,7 @@ class Party : public Activity, private Wakeable {
     explicit Participant(absl::string_view name) : name_(name) {}
     // Poll the participant. Return true if complete.
     // Participant should take care of its own deallocation in this case.
-    virtual bool Poll() = 0;
+    virtual bool PollParticipantPromise() = 0;
 
     // Destroy the participant before finishing.
     virtual void Destroy() = 0;
@@ -330,6 +354,9 @@ class Party : public Activity, private Wakeable {
   void Spawn(absl::string_view name, Factory promise_factory,
              OnComplete on_complete);
 
+  template <typename Factory>
+  auto SpawnWaitable(absl::string_view name, Factory factory);
+
   void Orphan() final { Crash("unused"); }
 
   // Activity implementation: not allowed to be overridden by derived types.
@@ -353,6 +380,20 @@ class Party : public Activity, private Wakeable {
 
   Arena* arena() const { return arena_; }
 
+  // Return a promise that resolves to Empty{} when the current party poll is
+  // complete.
+  // This is useful for implementing batching and the like: we can hold some
+  // action until the rest of the party resolves itself.
+  auto AfterCurrentPoll() {
+    GPR_DEBUG_ASSERT(Activity::current() == this);
+    sync_.WakeAfterPoll(CurrentParticipant());
+    return [this, iteration = sync_.iteration()]() -> Poll<Empty> {
+      GPR_DEBUG_ASSERT(Activity::current() == this);
+      if (iteration == sync_.iteration()) return Pending{};
+      return Empty{};
+    };
+  }
+
   class BulkSpawner {
    public:
     explicit BulkSpawner(Party* party) : party_(party) {}
@@ -414,7 +455,7 @@ class Party : public Activity, private Wakeable {
       }
     }
 
-    bool Poll() override {
+    bool PollParticipantPromise() override {
       if (!started_) {
         auto p = factory_.Make();
         Destruct(&factory_);
@@ -441,6 +482,89 @@ class Party : public Activity, private Wakeable {
     bool started_ = false;
   };
 
+  template <typename SuppliedFactory>
+  class PromiseParticipantImpl final
+      : public RefCounted<PromiseParticipantImpl<SuppliedFactory>,
+                          NonPolymorphicRefCount>,
+        public Participant {
+    using Factory = promise_detail::OncePromiseFactory<void, SuppliedFactory>;
+    using Promise = typename Factory::Promise;
+    using Result = typename Promise::Result;
+
+   public:
+    PromiseParticipantImpl(absl::string_view name,
+                           SuppliedFactory promise_factory)
+        : Participant(name) {
+      Construct(&factory_, std::move(promise_factory));
+    }
+
+    ~PromiseParticipantImpl() {
+      switch (state_.load(std::memory_order_acquire)) {
+        case State::kFactory:
+          Destruct(&factory_);
+          break;
+        case State::kPromise:
+          Destruct(&promise_);
+          break;
+        case State::kResult:
+          Destruct(&result_);
+          break;
+      }
+    }
+
+    // Inside party poll: drive from factory -> promise -> result
+    bool PollParticipantPromise() override {
+      switch (state_.load(std::memory_order_relaxed)) {
+        case State::kFactory: {
+          auto p = factory_.Make();
+          Destruct(&factory_);
+          Construct(&promise_, std::move(p));
+          state_.store(State::kPromise, std::memory_order_relaxed);
+        }
+          ABSL_FALLTHROUGH_INTENDED;
+        case State::kPromise: {
+          auto p = promise_();
+          if (auto* r = p.value_if_ready()) {
+            Destruct(&promise_);
+            Construct(&result_, std::move(*r));
+            state_.store(State::kResult, std::memory_order_release);
+            waiter_.Wakeup();
+            this->Unref();
+            return true;
+          }
+          return false;
+        }
+        case State::kResult:
+          Crash(
+              "unreachable: promises should not be repolled after completion");
+      }
+    }
+
+    // Outside party poll: check whether the spawning party has completed this
+    // promise.
+    Poll<Result> PollCompletion() {
+      switch (state_.load(std::memory_order_acquire)) {
+        case State::kFactory:
+        case State::kPromise:
+          return Pending{};
+        case State::kResult:
+          return std::move(result_);
+      }
+    }
+
+    void Destroy() override { this->Unref(); }
+
+   private:
+    enum class State : uint8_t { kFactory, kPromise, kResult };
+    union {
+      GPR_NO_UNIQUE_ADDRESS Factory factory_;
+      GPR_NO_UNIQUE_ADDRESS Promise promise_;
+      GPR_NO_UNIQUE_ADDRESS Result result_;
+    };
+    Waker waiter_{Activity::current()->MakeOwningWaker()};
+    std::atomic<State> state_{State::kFactory};
+  };
+
   // Notification that the party has finished and this instance can be deleted.
   // Derived types should arrange to call CancelRemainingParticipants during
   // this sequence.
@@ -460,6 +584,7 @@ class Party : public Activity, private Wakeable {
 
   // Add a participant (backs Spawn, after type erasure to ParticipantFactory).
   void AddParticipants(Participant** participant, size_t count);
+  bool RunOneParticipant(int i);
 
   virtual grpc_event_engine::experimental::EventEngine* event_engine()
       const = 0;
@@ -502,6 +627,17 @@ void Party::Spawn(absl::string_view name, Factory promise_factory,
                           std::move(on_complete));
 }
 
+template <typename Factory>
+auto Party::SpawnWaitable(absl::string_view name, Factory promise_factory) {
+  auto participant = MakeRefCounted<PromiseParticipantImpl<Factory>>(
+      name, std::move(promise_factory));
+  Participant* p = participant->Ref().release();
+  AddParticipants(&p, 1);
+  return [participant = std::move(participant)]() mutable {
+    return participant->PollCompletion();
+  };
+}
+
 }  // namespace grpc_core
 
 #endif  // GRPC_SRC_CORE_LIB_PROMISE_PARTY_H

data/src/core/lib/promise/poll.h

@@ -186,25 +186,13 @@ class Poll<Pending>;
 template <class T>
 class Poll<Poll<T>>;
 
-template <typename T>
-bool operator==(const Poll<T>& a, const Poll<T>& b) {
-  if (a.pending() && b.pending()) return true;
-  if (a.ready() && b.ready()) return a.value() == b.value();
-  return false;
-}
-
-template <typename T, typename U>
-Poll<T> poll_cast(Poll<U> poll) {
-  if (poll.pending()) return Pending{};
-  return static_cast<T>(std::move(poll.value()));
-}
-
 // PollTraits tells us whether a type is Poll<> or some other type, and is
 // leveraged in the PromiseLike/PromiseFactory machinery to select the
 // appropriate implementation of those concepts based upon the return type of a
 // lambda, for example (via enable_if).
 template <typename T>
 struct PollTraits {
+  using Type = T;
   static constexpr bool is_poll() { return false; }
 };
 
@@ -214,6 +202,44 @@ struct PollTraits<Poll<T>> {
   static constexpr bool is_poll() { return true; }
 };
 
+template <typename T>
+bool operator==(const Poll<T>& a, const Poll<T>& b) {
+  if (a.pending() && b.pending()) return true;
+  if (a.ready() && b.ready()) return a.value() == b.value();
+  return false;
+}
+
+template <typename T, typename U, typename SfinaeVoid = void>
+struct PollCastImpl;
+
+template <typename T, typename U>
+struct PollCastImpl<T, Poll<U>> {
+  static Poll<T> Cast(Poll<U>&& poll) {
+    if (poll.pending()) return Pending{};
+    return static_cast<T>(std::move(poll.value()));
+  }
+};
+
+template <typename T, typename U>
+struct PollCastImpl<T, U, std::enable_if<!PollTraits<U>::is_poll()>> {
+  static Poll<T> Cast(U&& poll) { return Poll<T>(T(std::move(poll))); }
+};
+
+template <typename T>
+struct PollCastImpl<T, T> {
+  static Poll<T> Cast(T&& poll) { return Poll<T>(std::move(poll)); }
+};
+
+template <typename T>
+struct PollCastImpl<T, Poll<T>> {
+  static Poll<T> Cast(Poll<T>&& poll) { return std::move(poll); }
+};
+
+template <typename T, typename U>
+Poll<T> poll_cast(U poll) {
+  return PollCastImpl<T, U>::Cast(std::move(poll));
+}
+
 // Convert a poll to a string
 template <typename T, typename F>
 std::string PollToString(

data/src/core/lib/promise/promise.h

@@ -89,6 +89,10 @@ auto WithResult(F f) ->
   return f;
 }
 
+template <typename Promise>
+using PromiseResult = typename PollTraits<
+    typename promise_detail::PromiseLike<Promise>::Result>::Type;
+
 }  // namespace grpc_core
 
 #endif  // GRPC_SRC_CORE_LIB_PROMISE_PROMISE_H