1
#pragma once
2
 
3
#include <kj/mutex.h>
4
#include <kj/refcount.h>
5
 
6
namespace workerd {
7
 
8
// Represents a weak reference back to an object that code can use as an indirect pointer when
9
// they want to be able to race destruction safely. A caller wishing to use a weak reference to
10
// the object should acquire a strong reference. It's always safe to invoke `tryAddStrongRef` to
11
// try to obtain a strong reference of the underlying object. This is because the Object's
12
// destructor will explicitly clear the underlying pointer that would be dereferenced by
13
// `tryAddStrongRef`. This means that after the refcount reaches 0, `tryAddStrongRef` is always
14
// still safe to invoke even if the underlying object memory has been deallocated (provided
15
// ownership of the weak object reference is retained).
16
// T must itself extend from kj::AtomicRefcounted
17
template <typename T>
18
class AtomicWeakRef final: public kj::AtomicRefcounted {
19
 public:
20
  inline static kj::Own<const AtomicWeakRef<T>> wrap(T* this_) {
21
    return kj::atomicRefcounted<AtomicWeakRef<T>>(this_);
22
  }
23
 
24
  inline explicit AtomicWeakRef(T* thisArg): this_(thisArg) {}
25
 
26
  // This tries to materialize a strong reference to the owner. It will fail if the owner's
27
  // refcount has already dropped to 0. As discussed in the class, the lifetime of this weak
28
  // reference can exceed the lifetime of the object it's tracking.
29
  inline kj::Maybe<kj::Own<const T>> tryAddStrongRef() const {
30
    auto lock = this_.lockShared();
31
    if (*lock == nullptr) return kj::none;
32
    return kj::atomicAddRefWeak(**lock);
33
  }
34
 
35
  inline kj::Own<const AtomicWeakRef<T>> addRef() const {
36
    return kj::atomicAddRef(*this);
37
  }
38
 
39
 private:
40
  kj::MutexGuarded<const T*> this_;
41
 
42
  // This is invoked by the owner destructor to clear the pointer. That means that any racing
43
  // code will never try to invoke `atomicAddRefWeak` on the instance any more. Any code racing
44
  // in between the refcount dropping to 0 and the invalidation getting invoked will still fail
45
  // to acquire a strong reference. Any code acquiring a strong reference prior to the refcount
46
  // dropping to 0 will prevent invalidation until that extra reference is dropped.
47
  inline void invalidate() const {
48
    *this_.lockExclusive() = nullptr;
49
  }
50
 
51
  friend T;
52
};
53
 
54
// A WeakRef is a weak reference to a thing. Note that because T may not itself be ref-counted,
55
// we cannot follow the usual pattern of a weak reference that potentially converts to a strong
56
// reference. Instead, intended usage looks like so:
57
// ```
58
//   kj::Own<WeakRef<Foo>> weakFoo = getWeakRefSomehow();
59
//
60
//   auto wasValid = weak->runIfAlive([](Foo& thing){
61
//     // Use thing
62
//   });
63
// ```
64
//
65
// TODO(cleanup): It would eventually be nice to replace kj::Own<WeakRef<T>> with a
66
// kj::WeakOwn<T> type with the same basic characteristics.
67
template <typename T>
68
class WeakRef final: public kj::Refcounted {
69
 public:
70
  inline WeakRef(kj::Badge<T>, T& thing): maybeThing(thing) {}
71
 
72
  // The use of the kj::Badge<T> in the constructor ensures that the initial instances
73
  // of WeakRef<T> can only be created within an instance of T. The instance T is responsible
74
  // for creating the initial refcounted kj::Own<WeakRef<T>>, and is responsible for calling
75
  // invalidate() in the destructor.
76
 
77
  KJ_DISALLOW_COPY_AND_MOVE(WeakRef);
78
 
79
  // Run the functor and return true if the context is alive, otherwise return false. Note that
80
  // since the `IoContext` might not be alive for any async continuation, we do not provide
81
  // a `kj::Maybe<IoContext&> tryGet()` function. You are expected to invoke this function
82
  // again in the next continuation to re-check if the `IoContext` is still around.
83
  template <typename F>
84
  inline bool runIfAlive(F&& f) const {
85
    KJ_IF_SOME(thing, maybeThing) {
86
      kj::fwd<F>(f)(thing);
87
      return true;
88
    }
89
 
90
    return false;
91
  }
92
 
93
  // Note: This is safe to call on a const WeakRef because the WeakRef doesn't own
94
  // the target - it just observes it. The returned reference allows mutation of
95
  // the target, which is intentional (the target's lifetime is managed elsewhere).
96
  inline kj::Maybe<T&> tryGet() const {
97
    KJ_IF_SOME(thing, maybeThing) {
98
      return thing;
99
    }
100
    return kj::none;
101
  }
102
  inline kj::Own<WeakRef> addRef() {
103
    return kj::addRef(*this);
104
  }
105
  inline bool isValid() const {
106
    return maybeThing != kj::none;
107
  }
108
 
109
 private:
110
  friend T;
111
 
112
  inline void invalidate() {
113
    maybeThing = kj::none;
114
  }
115
 
116
  kj::Maybe<T&> maybeThing;
117
};
118
 
119
}  // namespace workerd