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 "wait-list.h"
6
 
7
#include <kj/debug.h>
8
 
9
namespace workerd {
10
 
11
namespace {
12
// Optimization: If the same wait list is waited multiple times in the same thread, we want to
13
// share the signal rather than send two cross-thread signals.
14
static const kj::EventLoopLocal<CrossThreadWaitList::WaiterMap> threadLocalWaiters;
15
 
16
void END_WAIT_LIST_CANCELER_STACK_START_CANCELEE_STACK() {}
17
}  // namespace
18
 
19
CrossThreadWaitList::CrossThreadWaitList(Options options)
20
    : state(kj::atomicRefcounted<State>(options)) {}
21
 
22
void CrossThreadWaitList::destroyed() {
23
  if (!createdFulfiller) state->lostFulfiller();
24
}
25
 
26
CrossThreadWaitList::Waiter::Waiter(
27
    const State& state, kj::Own<kj::CrossThreadPromiseFulfiller<void>> fulfillerArg)
28
    : state(kj::atomicAddRef(state)),
29
      fulfiller(kj::mv(fulfillerArg)) {
30
  auto lock = state.waiters.lockExclusive();
31
  if (__atomic_load_n(&state.done, __ATOMIC_ACQUIRE)) {
32
    KJ_IF_SOME(e, state.exception) {
33
      fulfiller->reject(e.clone());
34
    } else {
35
      fulfiller->fulfill();
36
    }
37
  } else {
38
    lock->add(*this);
39
  }
40
}
41
CrossThreadWaitList::Waiter::~Waiter() noexcept(false) {
42
  if (__atomic_load_n(&unlinked, __ATOMIC_ACQUIRE)) {
43
    // No need to take a lock, already unlinked.
44
    KJ_ASSERT(!link.isLinked());
45
  } else {
46
    auto lock = state->waiters.lockExclusive();
47
    if (link.isLinked()) {
48
      lock->remove(*this);
49
    }
50
  }
51
 
52
  if (state->useThreadLocalOptimization) {
53
    auto& entry = KJ_ASSERT_NONNULL(threadLocalWaiters->findEntry(state.get()));
54
    KJ_ASSERT(entry.value == this);
55
    threadLocalWaiters->erase(entry);
56
  }
57
}
58
 
59
kj::Promise<void> CrossThreadWaitList::addWaiter() const {
60
  if (__atomic_load_n(&state->done, __ATOMIC_ACQUIRE)) {
61
    KJ_IF_SOME(e, state->exception) {
62
      return e.clone();
63
    } else {
64
      return kj::READY_NOW;
65
    }
66
  }
67
 
68
  if (state->useThreadLocalOptimization) {
69
    kj::Own<Waiter> ownWaiter;
70
 
71
    auto& waiter = threadLocalWaiters->findOrCreate(
72
        state.get(), [&]() -> CrossThreadWaitList::WaiterMap::Entry {
73
      auto paf = kj::newPromiseAndCrossThreadFulfiller<void>();
74
      ownWaiter = kj::refcounted<Waiter>(*state, kj::mv(paf.fulfiller));
75
      ownWaiter->forkedPromise = paf.promise.fork();
76
      return {state.get(), ownWaiter.get()};
77
    });
78
 
79
    if (ownWaiter.get() == nullptr) {
80
      ownWaiter = kj::addRef(*waiter);
81
    }
82
 
83
    return waiter->forkedPromise.addBranch().attach(kj::mv(ownWaiter));
84
  } else {
85
    // No refcounting, no forked promise.
86
    auto paf = kj::newPromiseAndCrossThreadFulfiller<void>();
87
    auto waiter = kj::heap<Waiter>(*state, kj::mv(paf.fulfiller));
88
    return paf.promise.attach(kj::mv(waiter));
89
  }
90
}
91
 
92
kj::Own<kj::CrossThreadPromiseFulfiller<void>> CrossThreadWaitList::makeSeparateFulfiller() {
93
  class FulfillerImpl final: public kj::CrossThreadPromiseFulfiller<void> {
94
   public:
95
    FulfillerImpl(kj::Own<const State> state): state(kj::mv(state)) {}
96
    ~FulfillerImpl() noexcept(false) {
97
      state->lostFulfiller();
98
    }
99
    void fulfill(kj::_::Void&&) const override {
100
      state->fulfill();
101
    }
102
    void reject(kj::Exception&& exception) const override {
103
      state->reject(kj::mv(exception));
104
    }
105
    bool isWaiting() const override {
106
      // Note that it would be incorrect for isWaiting() to return false when `done` is false
107
      // even if the waiter list is empty, because the waiter list could become non-empty later.
108
      // In theory if we could determine that there will never be a waiter, then isWaiting()
109
      // could return false.
110
      return !__atomic_load_n(&state->done, __ATOMIC_ACQUIRE);
111
    }
112
 
113
   private:
114
    kj::Own<const State> state;
115
  };
116
 
117
  KJ_REQUIRE(!createdFulfiller, "makeSeparateFulfiller() can only be called once");
118
  createdFulfiller = true;
119
  return kj::heap<FulfillerImpl>(kj::atomicAddRef(*state));
120
}
121
 
122
void CrossThreadWaitList::State::fulfill() const {
123
  if (__atomic_load_n(&done, __ATOMIC_ACQUIRE)) return;
124
  auto lock = waiters.lockExclusive();
125
  if (done) return;
126
  __atomic_store_n(&done, true, __ATOMIC_RELEASE);
127
 
128
  for (auto& waiter: *lock) {
129
    lock->remove(waiter);
130
    waiter.fulfiller->fulfill();
131
    __atomic_store_n(&waiter.unlinked, true, __ATOMIC_RELEASE);
132
  }
133
}
134
 
135
void CrossThreadWaitList::State::reject(kj::Exception&& e) const {
136
  if (__atomic_load_n(&done, __ATOMIC_ACQUIRE)) return;
137
  auto lock = waiters.lockExclusive();
138
  if (done) return;
139
  auto& exceptionRef = exception.emplace(kj::mv(e));
140
  __atomic_store_n(&done, true, __ATOMIC_RELEASE);
141
 
142
  for (auto& waiter: *lock) {
143
    lock->remove(waiter);
144
    waiter.fulfiller->reject(exceptionRef.clone());
145
    __atomic_store_n(&waiter.unlinked, true, __ATOMIC_RELEASE);
146
  }
147
}
148
 
149
void CrossThreadWaitList::State::lostFulfiller() const {
150
  if (__atomic_load_n(&done, __ATOMIC_ACQUIRE)) return;
151
  auto lock = waiters.lockExclusive();
152
  if (done) return;
153
  auto& exceptionRef = exception.emplace(kj::getDestructionReason(
154
      reinterpret_cast<void*>(&END_WAIT_LIST_CANCELER_STACK_START_CANCELEE_STACK),
155
      kj::Exception::Type::FAILED, __FILE__, __LINE__, "wait list was never fulfilled"_kj));
156
  __atomic_store_n(&done, true, __ATOMIC_RELEASE);
157
 
158
  if (!lock->empty()) {
159
    for (auto& waiter: *lock) {
160
      lock->remove(waiter);
161
      waiter.fulfiller->reject(exceptionRef.clone());
162
      __atomic_store_n(&waiter.unlinked, true, __ATOMIC_RELEASE);
163
    }
164
  }
165
}
166
 
167
}  // namespace workerd