1 // Copyright 2022 gRPC authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #ifndef GRPC_SRC_CORE_LIB_PROMISE_MPSC_H
16 #define GRPC_SRC_CORE_LIB_PROMISE_MPSC_H
17 
18 #include <grpc/support/port_platform.h>
19 
20 #include <stddef.h>
21 
22 #include <algorithm>
23 #include <utility>
24 #include <vector>
25 
26 #include "absl/base/thread_annotations.h"
27 
28 #include <grpc/support/log.h>
29 
30 #include "src/core/lib/gprpp/ref_counted.h"
31 #include "src/core/lib/gprpp/ref_counted_ptr.h"
32 #include "src/core/lib/gprpp/sync.h"
33 #include "src/core/lib/promise/activity.h"
34 #include "src/core/lib/promise/poll.h"
35 #include "src/core/lib/promise/wait_set.h"
36 
37 // Multi producer single consumer inter-activity comms.
38 
39 namespace grpc_core {
40 
41 namespace mpscpipe_detail {
42 
43 // "Center" of the communication pipe.
44 // Contains sent but not received messages, and open/close state.
45 template <typename T>
46 class Center : public RefCounted<Center<T>> {
47  public:
48   // Construct the center with a maximum queue size.
Center(size_t max_queued)49   explicit Center(size_t max_queued) : max_queued_(max_queued) {}
50 
51   // Poll for new items.
52   // - Returns true if new items were obtained, in which case they are contained
53   //   in dest in the order they were added. Wakes up all pending senders since
54   //   there will now be space to send.
55   // - If no new items are available, returns
56   //   false and sets up a waker to be awoken when more items are available.
57   // TODO(ctiller): consider the problem of thundering herds here. There may be
58   // more senders than there are queue spots, and so the strategy of waking up
59   // all senders is ill-advised.
60   // That said, some senders may have been cancelled by the time we wake them,
61   // and so waking a subset could cause starvation.
PollReceiveBatch(std::vector<T> & dest)62   bool PollReceiveBatch(std::vector<T>& dest) {
63     ReleasableMutexLock lock(&mu_);
64     if (queue_.empty()) {
65       receive_waker_ = Activity::current()->MakeNonOwningWaker();
66       return false;
67     }
68     dest.swap(queue_);
69     queue_.clear();
70     auto wakeups = send_wakers_.TakeWakeupSet();
71     lock.Release();
72     wakeups.Wakeup();
73     return true;
74   }
75 
76   // Poll to send one item.
77   // Returns pending if no send slot was available.
78   // Returns true if the item was sent.
79   // Returns false if the receiver has been closed.
PollSend(T & t)80   Poll<bool> PollSend(T& t) {
81     ReleasableMutexLock lock(&mu_);
82     if (receiver_closed_) return Poll<bool>(false);
83     if (queue_.size() < max_queued_) {
84       queue_.push_back(std::move(t));
85       auto receive_waker = std::move(receive_waker_);
86       lock.Release();
87       receive_waker.Wakeup();
88       return Poll<bool>(true);
89     }
90     send_wakers_.AddPending(Activity::current()->MakeNonOwningWaker());
91     return Pending{};
92   }
93 
94   // Mark that the receiver is closed.
ReceiverClosed()95   void ReceiverClosed() {
96     MutexLock lock(&mu_);
97     receiver_closed_ = true;
98   }
99 
100  private:
101   Mutex mu_;
102   const size_t max_queued_;
103   std::vector<T> queue_ ABSL_GUARDED_BY(mu_);
104   bool receiver_closed_ ABSL_GUARDED_BY(mu_) = false;
105   Waker receive_waker_ ABSL_GUARDED_BY(mu_);
106   WaitSet send_wakers_ ABSL_GUARDED_BY(mu_);
107 };
108 
109 }  // namespace mpscpipe_detail
110 
111 template <typename T>
112 class MpscReceiver;
113 
114 // Send half of an mpsc pipe.
115 template <typename T>
116 class MpscSender {
117  public:
118   MpscSender(const MpscSender&) = delete;
119   MpscSender& operator=(const MpscSender&) = delete;
120   MpscSender(MpscSender&&) noexcept = default;
121   MpscSender& operator=(MpscSender&&) noexcept = default;
122 
123   // Return a promise that will send one item.
124   // Resolves to true if sent, false if the receiver was closed (and the value
125   // will never be successfully sent).
Send(T t)126   auto Send(T t) {
127     return [this, t = std::move(t)]() mutable { return center_->PollSend(t); };
128   }
129 
130  private:
131   friend class MpscReceiver<T>;
MpscSender(RefCountedPtr<mpscpipe_detail::Center<T>> center)132   explicit MpscSender(RefCountedPtr<mpscpipe_detail::Center<T>> center)
133       : center_(std::move(center)) {}
134   RefCountedPtr<mpscpipe_detail::Center<T>> center_;
135 };
136 
137 // Receive half of an mpsc pipe.
138 template <typename T>
139 class MpscReceiver {
140  public:
141   // max_buffer_hint is the maximum number of elements we'd like to buffer.
142   // We half this before passing to Center so that the number there is the
143   // maximum number of elements that can be queued in the center of the pipe.
144   // The receiver also holds some of the buffered elements (up to half of them!)
145   // so the total outstanding is equal to max_buffer_hint (unless it's 1 in
146   // which case instantaneosly we may have two elements buffered).
MpscReceiver(size_t max_buffer_hint)147   explicit MpscReceiver(size_t max_buffer_hint)
148       : center_(MakeRefCounted<mpscpipe_detail::Center<T>>(
149             std::max(static_cast<size_t>(1), max_buffer_hint / 2))) {}
~MpscReceiver()150   ~MpscReceiver() {
151     if (center_ != nullptr) center_->ReceiverClosed();
152   }
153   MpscReceiver(const MpscReceiver&) = delete;
154   MpscReceiver& operator=(const MpscReceiver&) = delete;
155   // Only movable until it's first polled, and so we don't need to contend with
156   // a non-empty buffer during a legal move!
MpscReceiver(MpscReceiver && other)157   MpscReceiver(MpscReceiver&& other) noexcept
158       : center_(std::move(other.center_)) {
159     GPR_DEBUG_ASSERT(other.buffer_.empty());
160   }
161   MpscReceiver& operator=(MpscReceiver&& other) noexcept {
162     GPR_DEBUG_ASSERT(other.buffer_.empty());
163     center_ = std::move(other.center_);
164     return *this;
165   }
166 
167   // Construct a new sender for this receiver.
MakeSender()168   MpscSender<T> MakeSender() { return MpscSender<T>(center_); }
169 
170   // Return a promise that will resolve to the next item (and remove said item).
Next()171   auto Next() {
172     return [this]() -> Poll<T> {
173       if (buffer_it_ != buffer_.end()) {
174         return Poll<T>(std::move(*buffer_it_++));
175       }
176       if (center_->PollReceiveBatch(buffer_)) {
177         buffer_it_ = buffer_.begin();
178         return Poll<T>(std::move(*buffer_it_++));
179       }
180       return Pending{};
181     };
182   }
183 
184  private:
185   // Received items. We move out of here one by one, but don't resize the
186   // vector. Instead, when we run out of items, we poll the center for more -
187   // which swaps this buffer in for the new receive queue and clears it.
188   // In this way, upon hitting a steady state the queue ought to be allocation
189   // free.
190   std::vector<T> buffer_;
191   typename std::vector<T>::iterator buffer_it_ = buffer_.end();
192   RefCountedPtr<mpscpipe_detail::Center<T>> center_;
193 };
194 
195 }  // namespace grpc_core
196 
197 #endif  // GRPC_SRC_CORE_LIB_PROMISE_MPSC_H
198