Merge pull request #2404 from cloudflare/jsnell/websocketpair-iterable

cloudflare · Jul 19, 2024 · d46550a · d46550a
2 parents f8c915c + 3b72ddb
commit d46550a
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 7 deletions.
diff --git a/src/workerd/api/tests/global-scope-test.js b/src/workerd/api/tests/global-scope-test.js
@@ -570,3 +570,11 @@ export const base64 = {
     generate_tests(testAtob, tests);
   }
 };
+
+export const webSocketPairIterable = {
+  test() {
+    const [a, b] = new WebSocketPair();
+    ok(a instanceof WebSocket);
+    ok(b instanceof WebSocket);
+  }
+};
diff --git a/src/workerd/api/web-socket.c++ b/src/workerd/api/web-socket.c++
@@ -1015,6 +1015,13 @@ jsg::Ref<WebSocketPair> WebSocketPair::constructor() {
   return kj::mv(pair);
 }
 
+jsg::Ref<WebSocketPair::PairIterator> WebSocketPair::entries(jsg::Lock&) {
+  return jsg::alloc<PairIterator>(IteratorState {
+    .pair = JSG_THIS,
+    .index = 0,
+   });
+}
+
 void WebSocket::reportError(jsg::Lock& js, kj::Exception&& e) {
   reportError(js, js.exceptionToJsValue(kj::cp(e)));
 }

diff --git a/src/workerd/api/web-socket.h b/src/workerd/api/web-socket.h
@@ -134,6 +134,19 @@ class CloseEvent: public Event {
 class WebSocket;
 
 class WebSocketPair: public jsg::Object {
+private:
+  struct IteratorState final {
+    jsg::Ref<WebSocketPair> pair;
+    size_t index = 0;
+
+    void visitForGc(jsg::GcVisitor& visitor) {
+      visitor.visit(pair);
+    }
+
+    JSG_MEMORY_INFO(IteratorState) {
+      tracker.trackField("pair", pair);
+    }
+  };
 public:
   WebSocketPair(jsg::Ref<WebSocket> first, jsg::Ref<WebSocket> second)
       : sockets { kj::mv(first), kj::mv(second) } {}
@@ -143,11 +156,14 @@ class WebSocketPair: public jsg::Object {
   jsg::Ref<WebSocket> getFirst() { return sockets[0].addRef(); }
   jsg::Ref<WebSocket> getSecond() { return sockets[1].addRef(); }
 
+  JSG_ITERATOR(PairIterator, entries, jsg::Ref<WebSocket>, IteratorState, iteratorNext);
+
   JSG_RESOURCE_TYPE(WebSocketPair) {
     // TODO(soon): These really should be using an indexed property handler rather
     // than named instance properties but jsg does not yet have support for that.
     JSG_READONLY_INSTANCE_PROPERTY(0, getFirst);
     JSG_READONLY_INSTANCE_PROPERTY(1, getSecond);
+    JSG_ITERABLE(entries);
 
     JSG_TS_OVERRIDE(const WebSocketPair: {
       new (): { 0: WebSocket; 1: WebSocket };
@@ -179,6 +195,13 @@ class WebSocketPair: public jsg::Object {
 private:
   jsg::Ref<WebSocket> sockets[2];
 
+  static kj::Maybe<jsg::Ref<WebSocket>> iteratorNext(jsg::Lock& js, IteratorState& state) {
+    if (state.index >= 2) {
+      return kj::none;
+    }
+    return state.pair->sockets[state.index++].addRef();
+  }
+
   void visitForGc(jsg::GcVisitor& visitor) {
     visitor.visit(sockets[0]);
     visitor.visit(sockets[1]);
@@ -633,13 +656,15 @@ class WebSocket: public EventTarget {
   void assertNoError(jsg::Lock& js);
 };
 
-#define EW_WEBSOCKET_ISOLATE_TYPES \
-  api::CloseEvent,                 \
-  api::CloseEvent::Initializer,    \
-  api::MessageEvent,               \
-  api::MessageEvent::Initializer,  \
-  api::WebSocket,                  \
-  api::WebSocketPair
+#define EW_WEBSOCKET_ISOLATE_TYPES       \
+  api::CloseEvent,                       \
+  api::CloseEvent::Initializer,          \
+  api::MessageEvent,                     \
+  api::MessageEvent::Initializer,        \
+  api::WebSocket,                        \
+  api::WebSocketPair,                    \
+  api::WebSocketPair::PairIterator,      \
+  api::WebSocketPair::PairIterator::Next \
 // The list of websocket.h types that are added to worker.c++'s JSG_DECLARE_ISOLATE_TYPE
 
 }  // namespace workerd::api