1
// Copyright (c) 2023 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/refcount.h>
8
 
9
namespace workerd {
10
 
11
// This class is used to track a number of associated requests so that some desired behavior
12
// is carried out once all requests have completed. `activeRequests` is incremented each time a
13
// new request is created, and then decremented once it completes.
14
class RequestTracker final: public kj::Refcounted {
15
 public:
16
  class Hooks {
17
   public:
18
    virtual void active() = 0;
19
    virtual void inactive() = 0;
20
  };
21
 
22
  // An object that should be associated with (attached to) a request.
23
  class ActiveRequest final {
24
   public:
25
    // On creation, if the parent RequestTracker has 0 active requests, we call the `active()` hook.
26
    // On destruction, if the RequestTracker has 0 active requests, we call the `inactive()` hook.
27
    // Otherwise, we just increment/decrement the count on creation/destruction respectively.
28
    ActiveRequest(kj::Badge<RequestTracker>, RequestTracker& parent);
29
    ActiveRequest(ActiveRequest&& other) = default;
30
    KJ_DISALLOW_COPY(ActiveRequest);
31
    ~ActiveRequest() noexcept(false);
32
 
33
   private:
34
    kj::Maybe<kj::Own<RequestTracker>> maybeParent;
35
  };
36
 
37
  RequestTracker(Hooks& hooks);
38
  ~RequestTracker() noexcept(false);
39
  KJ_DISALLOW_COPY(RequestTracker);
40
 
41
  // Returns a new ActiveRequest, thereby bumping the count of active requests associated with the
42
  // RequestTracker. The ActiveRequest must be attached to the lifetime of the request such that we
43
  // destroy the ActiveRequest when the request is finished. On destruction, we decrement the count
44
  // of active requests associated with the RequestTracker, and if there are no more active requests
45
  // we call the `inactive()` hook.
46
  ActiveRequest startRequest();
47
 
48
  void shutdown() {
49
    // We want to prevent any hooks from running after this point.
50
    hooks = kj::none;
51
  }
52
 
53
  kj::Own<RequestTracker> addRef() {
54
    return kj::addRef(*this);
55
  }
56
 
57
 private:
58
  void requestActive();
59
  void requestInactive();
60
 
61
  int activeRequests = 0;
62
  kj::Maybe<Hooks&> hooks;
63
};
64
 
65
}  // namespace workerd