1
// Copyright (c) 2017-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 <workerd/api/actor-state.h>
8
#include <workerd/api/hibernatable-web-socket.h>
9
#include <workerd/api/web-socket.h>
10
#include <workerd/jsg/jsg.h>
11
 
12
#include <kj/exception.h>
13
 
14
#include <list>
15
 
16
namespace workerd {
17
 
18
// Implements the HibernationManager class.
19
class HibernationManagerImpl final: public Worker::Actor::HibernationManager {
20
 public:
21
  HibernationManagerImpl(kj::Own<Worker::Actor::Loopback> loopback, uint16_t hibernationEventType);
22
  ~HibernationManagerImpl() noexcept(false);
23
 
24
  // Tells the HibernationManager to create a new HibernatableWebSocket with the associated tags
25
  // and to initiate the `readLoop()` for this websocket. The `tags` array *must* contain only
26
  // unique elements.
27
  void acceptWebSocket(jsg::Ref<api::WebSocket> ws, kj::ArrayPtr<kj::String> tags) override;
28
 
29
  // Gets a collection of websockets associated with the given tag. Any hibernating websockets will
30
  // be woken up. If no tag is provided, we return all accepted websockets.
31
  kj::Vector<jsg::Ref<api::WebSocket>> getWebSockets(
32
      jsg::Lock& js, kj::Maybe<kj::StringPtr> tag) override;
33
 
34
  // Hibernates all the websockets held by the HibernationManager.
35
  // This converts our activeOrPackage from an api::WebSocket to a HibernationPackage.
36
  void hibernateWebSockets(Worker::Lock& lock) override;
37
 
38
  void setWebSocketAutoResponse(
39
      kj::Maybe<kj::StringPtr> request, kj::Maybe<kj::StringPtr> response) override;
40
  kj::Maybe<jsg::Ref<api::WebSocketRequestResponsePair>> getWebSocketAutoResponse(
41
      jsg::Lock& js) override;
42
  void setTimerChannel(TimerChannel& timerChannel) override;
43
 
44
  kj::Own<HibernationManager> addRef() override;
45
 
46
  friend class api::HibernatableWebSocketEvent;
47
 
48
  // Sets/Unset the maximum time in milliseconds that an hibernatable websocket event can run for.
49
  // If the timeout is reached, event is canceled.
50
  void setEventTimeout(kj::Maybe<uint32_t> timeoutMs) override;
51
 
52
  // Gets the event timeout if set.
53
  kj::Maybe<uint32_t> getEventTimeout() override;
54
 
55
 private:
56
  class HibernatableWebSocket;
57
 
58
  kj::Promise<void> handleReadLoop(HibernatableWebSocket& refToHibernatable);
59
 
60
  // Each HibernatableWebSocket can have multiple tags, so we want to store a reference
61
  // in our kj::List.
62
  struct TagListItem {
63
    kj::Maybe<HibernatableWebSocket&> hibWS;
64
    kj::ListLink<TagListItem> link;
65
    kj::StringPtr tag;
66
    // The List that refers to this TagListItem.
67
    // If `list` is null, we've already removed this item from the list.
68
    kj::Maybe<kj::List<TagListItem, &TagListItem::link>&> list;
69
  };
70
 
71
  // api::WebSockets cannot survive hibernation, but kj::WebSockets do. This class helps us
72
  // manage the transition of an api::WebSocket from its active state to a hibernated state
73
  // and vice versa.
74
  //
75
  // Some properties of the JS websocket object need to be retained throughout hibernation,
76
  // such as `attachment`, `url`, `extensions`, etc. These properties are only read/modified
77
  // when initiating, or waking from hibernation.
78
  class HibernatableWebSocket {
79
   public:
80
    HibernatableWebSocket(jsg::Ref<api::WebSocket> websocket,
81
        kj::ArrayPtr<kj::String> tags,
82
        HibernationManagerImpl& manager);
83
    ~HibernatableWebSocket() noexcept(false);
84
    KJ_DISALLOW_COPY_AND_MOVE(HibernatableWebSocket);
85
 
86
    // Returns the tags associated with this HibernatableWebSocket.
87
    kj::Array<kj::StringPtr> getTags();
88
 
89
    // Returns the tags associated with this HibernatableWebSocket.
90
    // Note that this returns an array of Strings, unlike `getTags()`.
91
    // Copying the strings each time tags are requested would be expensive,
92
    // so we only do it when we're delivering a close/error event because
93
    // we will be destroying the HibernatableWebSocket object,
94
    // which the tags need to outlive.
95
    kj::Array<kj::String> cloneTags();
96
 
97
    // Returns a reference to the active websocket. If the websocket is currently hibernating,
98
    // we have to unhibernate it first. The process moves values from the HibernatableWebSocket
99
    // to the api::WebSocket.
100
    jsg::Ref<api::WebSocket> getActiveOrUnhibernate(jsg::Lock& js);
101
 
102
    kj::ListLink<HibernatableWebSocket> link;
103
 
104
    // An array of all the items/nodes that refer to this HibernatableWebSocket.
105
    // Keeping track of these items allows us to quickly remove every reference from `tagToWs`
106
    // once the websocket disconnects -- rather than iterating through each relevant tag in the
107
    // hashmap and removing it from each kj::List.
108
    kj::Array<TagListItem> tagItems;
109
 
110
    // If active, we have an api::WebSocket reference, otherwise, we're hibernating, so we retain
111
    // the websocket's properties in a HibernationPackage until it's time to wake up.
112
    kj::OneOf<jsg::Ref<api::WebSocket>, api::WebSocket::HibernationPackage> activeOrPackage;
113
 
114
    // This is an owned websocket that we extract from the api::WebSocket after accepting as
115
    // hibernatable. It becomes null once we dispatch a close or error event because we want its
116
    // lifetime to be managed by IoContext's DeleteQueue. This helps prevent a situation where the
117
    // HibernationManager drops the websocket before all queued messages have sent.
118
    kj::Maybe<kj::Own<kj::WebSocket>> ws;
119
 
120
    HibernationManagerImpl& manager;
121
    // TODO(someday): We (currently) only use the HibernationManagerImpl reference to refer to
122
    // `tagToWs` when running the dtor for `HibernatableWebSocket`. This feels a bit excessive,
123
    // I would rather have the HibernationManager deal with its collections than have the
124
    // HibernatableWebSocket do so. Maybe come back to this at some point?
125
 
126
    // Reference to the Node in `allWs` that allows us to do fast deletion on disconnect.
127
    kj::Maybe<std::list<kj::Own<HibernatableWebSocket>>::iterator> node;
128
 
129
    // True once we have dispatched the close event.
130
    // This prevents us from dispatching it if we have already done so.
131
    bool hasDispatchedClose = false;
132
 
133
    // Stores the last received autoResponseRequest timestamp.
134
    kj::Maybe<kj::Date> autoResponseTimestamp;
135
 
136
    // Keeps track of the currently ongoing websocket auto-response send promise. This promise may
137
    // be moved to api::websocket if an hibernating websocket unhibernates.
138
    kj::Promise<void> autoResponsePromise = kj::READY_NOW;
139
 
140
    friend HibernationManagerImpl;
141
  };
142
 
143
  // Removes a HibernatableWebSocket from the HibernationManager's various collections.
144
  void dropHibernatableWebSocket(HibernatableWebSocket& hib);
145
 
146
  // Removes the HibernatableWebSocket from `allWs`.
147
  inline void removeFromAllWs(HibernatableWebSocket& hib);
148
 
149
  // Handles the termination of the websocket. If termination was not clean, we might try to
150
  // dispatch a close event (if we haven't already), or an error event.
151
  // We will also remove the HibernatableWebSocket from the HibernationManager's collections.
152
  kj::Promise<void> handleSocketTermination(
153
      HibernatableWebSocket& hib, kj::Maybe<kj::Exception>& maybeError) KJ_WARN_UNUSED_RESULT;
154
 
155
  // Like the api::WebSocket readLoop(), but we dispatch different types of events.
156
  kj::Promise<void> readLoop(HibernatableWebSocket& hib);
157
 
158
  // This struct is held by the `tagToWs` hashmap. The key is a StringPtr to tag, and the value
159
  // is this struct itself.
160
  struct TagCollection {
161
    kj::String tag;
162
    kj::Own<kj::List<TagListItem, &TagListItem::link>> list;
163
 
164
    TagCollection(kj::String tag, decltype(list) list): tag(kj::mv(tag)), list(kj::mv(list)) {}
165
    TagCollection(TagCollection&& other) = default;
166
  };
167
 
168
  // This structure will hold the request and corresponding response for hibernatable websockets
169
  // auto-response feature. Although we store 2 kj::Maybe strings, if we don't have a request set
170
  // we can't have a response, and vice versa.
171
  // TODO(cleanup): Remove kj::Maybe from request and response strings.
172
  struct AutoRequestResponsePair {
173
    kj::Maybe<kj::String> request = kj::none;
174
    kj::Maybe<kj::String> response = kj::none;
175
  };
176
 
177
  // A hashmap of tags to HibernatableWebSockets associated with the tag.
178
  // We use a kj::List so we can quickly remove websockets that have disconnected.
179
  // Also note that we box the keys and values such that in the event of a hashmap resizing we don't
180
  // move the underlying data (thereby keeping any references intact).
181
  kj::HashMap<kj::StringPtr, kj::Own<TagCollection>> tagToWs;
182
 
183
  // We store all of our HibernatableWebSockets in a doubly linked-list.
184
  std::list<kj::Own<HibernatableWebSocket>> allWs;
185
 
186
  // Used to obtain the worker so we can dispatch Hibernatable websocket events.
187
  kj::Own<Worker::Actor::Loopback> loopback;
188
 
189
  // Passed to HibernatableWebSocket custom event as the typeId.
190
  uint16_t hibernationEventType;
191
 
192
  // A map of { ID -> HibernatableWebSocket } that allows the event handler that is currently
193
  // running to access the HibernatableWebSocket that it needs to execute.
194
  //
195
  // Dispatching events tends to result in races when events are received on different websockets
196
  // around the same time. Suppose there are two websockets that disconnect at the same time.
197
  // It is possible that both of them will be added to the map (i.e. their `receive()`
198
  // will throw) before the first event is dispatched and manages to obtain its associated websocket.
199
  kj::HashMap<kj::String, HibernatableWebSocket*> webSocketsForEventHandler;
200
 
201
  // The maximum number of Hibernatable WebSocket connections a single HibernationManagerImpl
202
  // instance can manage.
203
  const size_t ACTIVE_CONNECTION_LIMIT = 1024 * 32;
204
 
205
  class DisconnectHandler: public kj::TaskSet::ErrorHandler {
206
   public:
207
    // We don't need to do anything here; we already handle disconnects in the callee of readLoop().
208
    void taskFailed(kj::Exception&& exception) override {};
209
  };
210
  DisconnectHandler onDisconnect;
211
  kj::TaskSet readLoopTasks;
212
  kj::Own<AutoRequestResponsePair> autoResponsePair = kj::heap<AutoRequestResponsePair>();
213
  kj::Maybe<TimerChannel&> timer;
214
  kj::Maybe<uint32_t> eventTimeoutMs;
215
};
216
};  // namespace workerd