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
#pragma once
6
 
7
#include <kj/async.h>
8
#include <kj/list.h>
9
#include <kj/map.h>
10
#include <kj/mutex.h>
11
 
12
namespace workerd {
13
 
14
using kj::uint;
15
 
16
// A class that allows multiple threads to wait for an event, and for any thread to later trigger
17
// that event. This is like using kj::newPromiseAndCrossThreadFulfiller<void>() and forking the
18
// promise, except:
19
// * Normally, a ForkedPromise's addBranch() can only be called in the thread that created the
20
//   fork. `CrossThreadWaitList` can be awaited from any thread.
21
// * CrossThreadWaitList is one object, not a promise/fulfiller pair. In many use cases, this
22
//   turns out to be most convenient. But if you want a separate fulfiller, you can call the
23
//   `makeSeparateFulfiller()` method.
24
class CrossThreadWaitList {
25
 public:
26
  struct Options {
27
    // Enable this if it is common for there to be multiple waiters in the same thread. This avoids
28
    // sending multiple cross-thread signals in this case, instead sending one signal that all
29
    // waiters in the thread wait on.
30
    bool useThreadLocalOptimization = false;
31
  };
32
 
33
  CrossThreadWaitList(): CrossThreadWaitList(Options()) {}
34
  CrossThreadWaitList(Options options);
35
  CrossThreadWaitList(CrossThreadWaitList&& other) = default;
36
  ~CrossThreadWaitList() noexcept(false) {
37
    // Check if moved away.
38
    if (state.get() != nullptr) destroyed();
39
  }
40
 
41
  kj::Promise<void> addWaiter() const;
42
 
43
  // Wake all current *and future* waiters.
44
  void fulfill() const {
45
    KJ_IREQUIRE(!createdFulfiller);
46
    state->fulfill();
47
  }
48
 
49
  // Causes all past and future `addWaiter()` calls to reject with the given exception.
50
  void reject(kj::Exception&& e) const {
51
    KJ_IREQUIRE(!createdFulfiller);
52
    state->reject(kj::mv(e));
53
  }
54
 
55
  // Has `fulfill()` or `reject()` been called? Of course, the caller should consider if
56
  // `fulfill()` might be called in another thread concurrently.
57
  bool isDone() const {
58
    return __atomic_load_n(&state->done, __ATOMIC_ACQUIRE);
59
  }
60
 
61
  // Creates a PromiseFulfiller that will fulfill this wait list. Once this is called, it is no
62
  // longer the CrossThreadWaitList's responsibility to fulfill the waiters.
63
  //
64
  // Arguably, we should always make people create a PromiseFulfiller-CrossThreadWaitList pair,
65
  // like kj::newPromiseAndFulfiller, instead of having methods directly on CrossThreadWaitList
66
  // to fulfill/reject. However, in practice, in many use cases the fulfiller would be stored
67
  // right next to the wait list, so it's convenient to let people opt into having two parts
68
  // explicitly.
69
  kj::Own<kj::CrossThreadPromiseFulfiller<void>> makeSeparateFulfiller();
70
 
71
 private:
72
  // Forward declare our private structs so we can name the Map for public use in the source file.
73
  struct State;
74
  struct Waiter;
75
 
76
 public:
77
  using WaiterMap = kj::HashMap<const CrossThreadWaitList::State*, Waiter*>;
78
 
79
 private:
80
  struct Waiter: public kj::Refcounted {
81
    Waiter(const State& state, kj::Own<kj::CrossThreadPromiseFulfiller<void>> fulfiller);
82
    ~Waiter() noexcept(false);
83
 
84
    kj::Own<const State> state;
85
    kj::Own<kj::CrossThreadPromiseFulfiller<void>> fulfiller;
86
 
87
    // Protected by list mutex.
88
    kj::ListLink<Waiter> link;
89
 
90
    // Optimization: This is atomically set true when the waiter is removed from the list so that
91
    // we don't have to redundantly take the lock.
92
    bool unlinked = false;
93
 
94
    // Only initialized if useThreadLocalOptimization is enabled.
95
    kj::ForkedPromise<void> forkedPromise = nullptr;
96
  };
97
 
98
  struct State: public kj::AtomicRefcounted {
99
    kj::MutexGuarded<kj::List<Waiter, &Waiter::link>> waiters;
100
 
101
    const bool useThreadLocalOptimization = false;
102
 
103
    // Atomically set true at the start of fulfill() or reject(). This can be checked before taking
104
    // the lock, but if false, it must be checked again after taking the lock, to avoid a race.
105
    mutable bool done = false;
106
 
107
    // If `done` is true due to `reject()` being called, this is the exception. This field
108
    // does not change after `done` is set true.
109
    mutable kj::Maybe<kj::Exception> exception;
110
 
111
    bool wakeNext() const;
112
    void fulfill() const;
113
    void reject(kj::Exception&& e) const;
114
    void lostFulfiller() const;
115
 
116
    explicit State(const Options& options)
117
        : useThreadLocalOptimization(options.useThreadLocalOptimization) {}
118
  };
119
 
120
  kj::Own<const State> state;
121
  bool createdFulfiller = false;
122
 
123
  void destroyed();
124
};
125
 
126
}  // namespace workerd