Skip to content
rachel/workerd
File

Blob: src/workerd/io/io-gate.c++

11.5 KB
1// Copyright (c) 2017-2022 Cloudflare, Inc.
2// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3//     https://opensource.org/licenses/Apache-2.0
4 
5#include <workerd/io/io-gate.h>
6 
7#include <kj/debug.h>
8 
9namespace workerd {
10 
11const InputGate::Hooks InputGate::Hooks::DEFAULT;
12 
13InputGate::InputGate(Hooks& hooks): InputGate(hooks, kj::newPromiseAndFulfiller<void>()) {}
14InputGate::InputGate(Hooks& hooks, kj::PromiseFulfillerPair<void> paf)
15    : hooks(hooks),
16      brokenPromise(paf.promise.fork()),
17      brokenState(kj::mv(paf.fulfiller)) {}
18InputGate::~InputGate() noexcept {
19  // Intentionally `noexcept` because if this throws then there are dangling references.
20  KJ_ASSERT(lockCount == 0,
21      "destroying InputGate when locks are still present; they would become dangling references");
22 
23  // If the lock count is zero, then the waiters must be empty.
24  KJ_ASSERT(waiters.empty());
25}
26 
27InputGate::Waiter::Waiter(kj::PromiseFulfiller<Lock>& fulfiller,
28    InputGate& gate,
29    bool isChildWaiter,
30    SpanParent parentSpan)
31    : fulfiller(fulfiller),
32      gate(&gate),
33      isChildWaiter(isChildWaiter),
34      waitSpan(parentSpan.newChild("input_gate_lock_wait"_kjc)),
35      lockSpanParent(kj::mv(parentSpan)) {
36  gate.hooks.inputGateWaiterAdded();
37  if (isChildWaiter) {
38    gate.waitingChildren.add(*this);
39  } else {
40    gate.waiters.add(*this);
41  }
42}
43InputGate::Waiter::~Waiter() noexcept(false) {
44  gate->hooks.inputGateWaiterRemoved();
45  if (link.isLinked()) {
46    if (isChildWaiter) {
47      gate->waitingChildren.remove(*this);
48    } else {
49      gate->waiters.remove(*this);
50    }
51  }
52}
53 
54kj::Promise<InputGate::Lock> InputGate::wait(SpanParent parentSpan) {
55  auto methodSpan = parentSpan.newChild("input_gate_wait_attempt"_kjc);
56  KJ_IF_SOME(e, brokenState.tryGet<kj::Exception>()) {
57    return e.clone();
58  } else if (lockCount == 0) {
59    return Lock(*this, methodSpan);
60  } else {
61    return kj::newAdaptedPromise<Lock, Waiter>(*this, false, methodSpan);
62  }
63}
64 
65kj::Promise<void> InputGate::onBroken() {
66  KJ_IF_SOME(e, brokenState.tryGet<kj::Exception>()) {
67    return e.clone();
68  } else {
69    return brokenPromise.addBranch();
70  }
71}
72 
73InputGate::Lock::Lock(InputGate& gate, SpanParent parentSpan)
74    : gate(&gate),
75      cs(gate.isCriticalSection ? kj::Maybe(kj::addRef(static_cast<CriticalSection&>(gate)))
76                                : kj::none),
77      lockSpan(parentSpan.newChild("input_gate_lock_hold"_kjc)) {
78  InputGate* gateToLock = &gate;
79 
80  KJ_IF_SOME(c, cs) {
81    if (c.get()->state == CriticalSection::REPARENTED) {
82      gateToLock = &c.get()->parentAsInputGate();
83    }
84  }
85 
86  if (++gateToLock->lockCount == 1) {
87    gateToLock->hooks.inputGateLocked();
88  }
89}
90 
91void InputGate::releaseLock() {
92  if (isCriticalSection) {
93    auto& self = static_cast<CriticalSection&>(*this);
94    if (self.state == CriticalSection::REPARENTED) {
95      // This lock was for a critical section that has already completed, therefore the lock
96      // should be considered "reparented", and we should forward the release to the parent.
97 
98      // Ensure any waiters on us have already been reparented.
99      KJ_DASSERT(self.waitingChildren.size() == 0);
100      KJ_DASSERT(self.waiters.size() == 0);
101      KJ_DASSERT(lockCount == 0);
102 
103      self.parentAsInputGate().releaseLock();
104      return;
105    }
106  }
107 
108  KJ_ASSERT(lockCount-- > 0);
109 
110  // Check if any waiters can be released.
111  if (lockCount == 0) {
112    hooks.inputGateReleased();
113    if (!waitingChildren.empty()) {
114      auto& waiter = waitingChildren.front();
115      waitingChildren.remove(waiter);
116      waiter.fulfiller.fulfill(Lock(*this, kj::mv(waiter.lockSpanParent)));
117    } else if (!waiters.empty()) {
118      auto& waiter = waiters.front();
119      waiters.remove(waiter);
120      waiter.fulfiller.fulfill(Lock(*this, kj::mv(waiter.lockSpanParent)));
121    }
122  }
123}
124 
125kj::Own<InputGate::CriticalSection> InputGate::Lock::startCriticalSection() {
126  return kj::refcounted<CriticalSection>(*gate);
127}
128 
129kj::Maybe<InputGate::CriticalSection&> InputGate::Lock::getCriticalSection() {
130  if (gate->isCriticalSection) {
131    return static_cast<CriticalSection&>(*gate);
132  } else {
133    return kj::none;
134  }
135}
136 
137bool InputGate::Lock::isFor(const InputGate& otherGate) const {
138  KJ_ASSERT(!otherGate.isCriticalSection);
139 
140  InputGate* ptr = gate;
141  while (ptr->isCriticalSection) {
142    ptr = &static_cast<CriticalSection&>(*ptr).parentAsInputGate();
143  }
144  return ptr == &otherGate;
145}
146 
147InputGate::CriticalSection::CriticalSection(InputGate& parent) {
148  isCriticalSection = true;
149  if (parent.isCriticalSection) {
150    this->parent = kj::addRef(static_cast<CriticalSection&>(parent));
151  } else {
152    this->parent = &parent;
153  }
154}
155InputGate::CriticalSection::~CriticalSection() noexcept(false) {
156  switch (state) {
157    case NOT_STARTED:
158      // Oh well.
159      break;
160    case INITIAL_WAIT:
161      // The initial wait() had better have been canceled... but we have no way to tell here.
162      break;
163    case RUNNING:
164      failed(KJ_EXCEPTION(FAILED,
165          "jsg.Error: A critical section within this Durable Object awaited a Promise that "
166          "apparently will never complete. This could happen in particular if a critical section "
167          "awaits a task that was initiated outside of the critical section. Since a critical "
168          "section blocks all other tasks from completing, this leads to deadlock."));
169      break;
170    case REPARENTED:
171      // Common case.
172      break;
173  }
174}
175 
176kj::Promise<InputGate::Lock> InputGate::CriticalSection::wait(SpanParent parentSpan) {
177  auto methodSpan = parentSpan.newChild("input_gate_critical_section_wait_attempt"_kjc);
178  for (;;) {
179    switch (state) {
180      case NOT_STARTED: {
181        state = INITIAL_WAIT;
182 
183        auto& target = parentAsInputGate();
184        KJ_IF_SOME(e, target.brokenState.tryGet<kj::Exception>()) {
185          // Oops, we're broken.
186          setBroken(e);
187          kj::throwFatalException(e.clone());
188        }
189 
190        // Add ourselves to this parent's child waiter list.
191        if (target.lockCount == 0) {
192          state = RUNNING;
193          parentLock = Lock(target, methodSpan);
194          continue;
195        } else {
196          try {
197            auto lock = co_await kj::newAdaptedPromise<Lock, Waiter>(target, true, methodSpan);
198            state = RUNNING;
199            parentLock = kj::mv(lock);
200            continue;
201          } catch (...) {
202            auto exception = kj::getCaughtExceptionAsKj();
203            state = RUNNING;
204            setBroken(exception);
205            kj::throwFatalException(kj::mv(exception));
206          }
207        }
208      }
209      case INITIAL_WAIT:
210        // To avoid the need for a ForkedPromise, we assume wait() is called once initially to
211        // get things started. This is the case in practice because any further tasks would be
212        // started only after some code runs under the initial lock.
213        KJ_FAIL_REQUIRE("CriticalSection::wait() should be called once initially");
214      case RUNNING:
215        // CriticalSection is active, so defer to InputGate implementation.
216        co_return co_await InputGate::wait(methodSpan);
217      case REPARENTED:
218        // Once the CriticalSection has declared itself done, then any straggler tasks it initiated
219        // are adopted by the parent.
220        // WARNING: Don't use parentAsInputGate() here as that'll bypass the override of wait() if
221        //   the parent is a CriticalSection itself.
222        KJ_SWITCH_ONEOF(parent) {
223          KJ_CASE_ONEOF(p, InputGate*) {
224            co_return co_await p->wait(methodSpan);
225          }
226          KJ_CASE_ONEOF(c, kj::Own<CriticalSection>) {
227            co_return co_await c->wait(methodSpan);
228          }
229        }
230        KJ_UNREACHABLE;
231    }
232    KJ_UNREACHABLE;
233  }
234}
235 
236InputGate::Lock InputGate::CriticalSection::succeeded() {
237  KJ_REQUIRE(state == RUNNING);
238 
239  // Once the CriticalSection has declared itself done, then any straggler tasks it initiated are
240  // adopted by the parent.
241  auto& parentGate = parentAsInputGate();
242  for (auto& waiter: waitingChildren) {
243    waitingChildren.remove(waiter);
244    parentGate.waitingChildren.add(waiter);
245    waiter.gate = &parentGate;
246  }
247  for (auto& waiter: waiters) {
248    waiters.remove(waiter);
249    parentGate.waiters.add(waiter);
250    waiter.gate = &parentGate;
251  }
252  parentGate.lockCount += lockCount;
253  lockCount = 0;
254 
255  state = REPARENTED;
256  auto result = KJ_ASSERT_NONNULL(kj::mv(parentLock));
257  parentLock = kj::none;
258  return result;
259}
260 
261void InputGate::CriticalSection::failed(const kj::Exception& e) {
262  if (brokenState.is<kj::Exception>()) {
263    // Already failed I guess.
264    return;
265  }
266 
267  setBroken(e);
268  KJ_SWITCH_ONEOF(parent) {
269    KJ_CASE_ONEOF(p, InputGate*) {
270      p->setBroken(e);
271    }
272    KJ_CASE_ONEOF(c, kj::Own<CriticalSection>) {
273      c->failed(e);
274    }
275  }
276}
277 
278void InputGate::setBroken(const kj::Exception& e) {
279  for (auto& waiter: waitingChildren) {
280    waiter.fulfiller.reject(e.clone());
281    waitingChildren.remove(waiter);
282  }
283  for (auto& waiter: waiters) {
284    waiter.fulfiller.reject(e.clone());
285    waiters.remove(waiter);
286  }
287  KJ_IF_SOME(f, brokenState.tryGet<kj::Own<kj::PromiseFulfiller<void>>>()) {
288    f.get()->reject(e.clone());
289  }
290  brokenState = e.clone();
291}
292 
293InputGate& InputGate::CriticalSection::parentAsInputGate() {
294  CriticalSection* ptr = this;
295  for (;;) {
296    KJ_SWITCH_ONEOF(ptr->parent) {
297      KJ_CASE_ONEOF(p, InputGate*) {
298        return *p;
299      }
300      KJ_CASE_ONEOF(c, kj::Own<CriticalSection>) {
301        if (c.get()->state == REPARENTED) {
302          // Keep looping...
303          ptr = c;
304        } else {
305          return *c.get();
306        }
307      }
308    }
309  }
310}
311 
312// =======================================================================================
313 
314OutputGate::OutputGate(Hooks& hooks)
315    : hooks(hooks),
316      pastLocksPromise(kj::Promise<void>(kj::READY_NOW).fork()) {}
317OutputGate::~OutputGate() noexcept(false) {}
318 
319const OutputGate::Hooks OutputGate::Hooks::DEFAULT;
320 
321kj::Own<kj::PromiseFulfiller<void>> OutputGate::lock() {
322  auto paf = kj::newPromiseAndFulfiller<void>();
323  auto joined = kj::joinPromises(kj::arr(pastLocksPromise.addBranch(), kj::mv(paf.promise)));
324  pastLocksPromise = joined.fork();
325  return kj::mv(paf.fulfiller);
326}
327 
328kj::Promise<void> OutputGate::wait(SpanParent parentSpan) {
329  hooks.outputGateWaiterAdded();
330  SpanBuilder waitSpan = parentSpan.newChild("output_gate_lock_wait"_kjc);
331  return pastLocksPromise.addBranch().attach(
332      kj::defer([this]() { hooks.outputGateWaiterRemoved(); }), kj::mv(waitSpan));
333}
334 
335kj::Promise<void> OutputGate::onBroken() {
336  KJ_REQUIRE(
337      !brokenState.is<kj::Own<kj::PromiseFulfiller<void>>>(), "onBroken() can only be called once");
338 
339  KJ_IF_SOME(e, brokenState.tryGet<kj::Exception>()) {
340    return e.clone();
341  } else {
342    auto paf = kj::newPromiseAndFulfiller<void>();
343    brokenState = kj::mv(paf.fulfiller);
344    return kj::mv(paf.promise);
345  }
346}
347 
348bool OutputGate::isBroken() {
349  return brokenState.is<kj::Exception>();
350}
351 
352namespace {
353 
354void END_OUTPUT_LOCK_CANCELATION_STACK_START_WAITER_STACK() {}
355 
356}  // namespace
357 
358kj::Exception OutputGate::makeUnfulfilledException() {
359  return kj::getDestructionReason(
360      reinterpret_cast<void*>(&END_OUTPUT_LOCK_CANCELATION_STACK_START_WAITER_STACK),
361      kj::Exception::Type::FAILED, __FILE__, __LINE__,
362      "output lock was canceled before completion"_kj);
363}
364 
365void OutputGate::setBroken(const kj::Exception& e) {
366  // We assume the exception is already propagated into `pastLocksPromise`, so all we need to do
367  // is handle onBroken().
368  KJ_IF_SOME(f, brokenState.tryGet<kj::Own<kj::PromiseFulfiller<void>>>()) {
369    f.get()->reject(e.clone());
370  }
371  brokenState = e.clone();
372}
373 
374}  // namespace workerd