Additional checked queue cleanups/improvements

Author: jasnell

src/workerd/api/streams/internal.c++

@@ -2024,9 +2024,8 @@ void WritableStreamInternalController::drain(jsg::Lock& js, v8::Local<v8::Value>
 }
 
 void WritableStreamInternalController::visitForGc(jsg::GcVisitor& visitor) {
-  queue.forEach([&](const WriteEvent& event) {
-    WriteEvent& e = const_cast<WriteEvent&>(event);
-    KJ_SWITCH_ONEOF(e.event) {
+  queue.forEach([&](auto& event) {
+    KJ_SWITCH_ONEOF(event.event) {
       KJ_CASE_ONEOF(write, Write) {
         visitor.visit(write.promise);
       }

src/workerd/api/streams/queue.h

@@ -835,14 +835,14 @@ class ByteQueue final {
 
     JSG_MEMORY_INFO(ByteQueue::ByobRequest) {}
 
+    inline ReadRequest& getRequest() {
+      return KJ_ASSERT_NONNULL(request->tryGet(), "The ByobRequest has been invalidated.");
+    }
+
    private:
     kj::Own<WeakRef<ReadRequest>> request;
     ConsumerImpl& consumer;
     QueueImpl& queue;
-
-    inline ReadRequest& getRequest() {
-      return KJ_ASSERT_NONNULL(request->tryGet(), "The ByobRequest has been invalidated.");
-    }
   };
 
   struct State {

src/workerd/api/streams/standard.c++

@@ -1556,11 +1556,7 @@ void WritableImpl<Self>::visitForGc(jsg::GcVisitor& visitor) {
   KJ_IF_SOME(pendingAbort, maybePendingAbort) {
     visitor.visit(*pendingAbort);
   }
-  writeRequests.forEach([&](const WriteRequest& write) {
-    // const_cast is ugly here but GcVisitor only supports non-const references.
-    auto& w = const_cast<WriteRequest&>(write);
-    visitor.visit(w.resolver, w.value);
-  });
+  writeRequests.forEach([&](auto& write) { visitor.visit(write.resolver, write.value); });
 }
 
 template <typename Self>
@@ -3971,7 +3967,7 @@ void WritableImpl<Self>::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
   tracker.trackField("writeAlgorithm", algorithms.write);
   tracker.trackField("sizeAlgorithm", algorithms.size);
 
-  writeRequests.forEach([&](const auto& req) { tracker.trackField("writeRequest", req); });
+  writeRequests.forEach([&](auto& req) { tracker.trackField("writeRequest", req); });
 
   tracker.trackField("inFlightWrite", inFlightWrite);
   tracker.trackField("inFlightClose", inFlightClose);

src/workerd/jsg/BUILD.bazel

@@ -45,6 +45,7 @@ wd_cc_library(
         ":jsg-core",
         ":memory-tracker",
         "//src/workerd/util",
+        "//src/workerd/util:checked-queue",
         "//src/workerd/util:sentry",
         "//src/workerd/util:thread-scopes",
         "@capnp-cpp//src/kj",

src/workerd/jsg/iterator.c++

@@ -7,24 +7,19 @@
 namespace workerd::jsg {
 
 kj::Maybe<jsg::Promise<void>&> AsyncIteratorImpl::maybeCurrent() {
-  if (!pendingStack.empty()) {
-    return pendingStack.back();
-  }
-  return kj::none;
+  return pendingStack.peekBack();
 }
 
 void AsyncIteratorImpl::pushCurrent(Promise<void> promise) {
-  pendingStack.push_back(kj::mv(promise));
+  pendingStack.push(kj::mv(promise));
 }
 
 void AsyncIteratorImpl::popCurrent() {
-  if (!pendingStack.empty()) {
-    pendingStack.pop_front();
-  }
+  auto dropped KJ_UNUSED = KJ_ASSERT_NONNULL(pendingStack.pop());
 }
 
 void AsyncIteratorImpl::visitForGc(jsg::GcVisitor& visitor) {
-  visitor.visitAll(pendingStack);
+  pendingStack.forEach([&](auto& p) { visitor.visit(p); });
 }
 
 }  // namespace workerd::jsg

src/workerd/jsg/iterator.h

@@ -7,6 +7,7 @@
 #include <workerd/jsg/jsg.h>
 #include <workerd/jsg/memory.h>
 #include <workerd/jsg/struct.h>
+#include <workerd/util/checked-queue.h>
 
 #include <concepts>
 #include <list>
@@ -800,7 +801,7 @@ class AsyncIteratorImpl {
   }
 
  private:
-  std::list<Promise<void>> pendingStack;
+  workerd::util::Queue<Promise<void>> pendingStack;
 };
 
 // Provides the base implementation of JSG_ASYNC_ITERATOR types. See the documentation

src/workerd/util/checked-queue.h

@@ -141,6 +141,22 @@ class Queue final {
     return count;
   }
 
+  // Applies the callback to each element in the queue.
+  // Returns the number of elements processed.
+  // If the callback returns false, the iteration stops.
+  inline size_t forEach(auto callback) {
+    size_t count = 0;
+    for (auto& item: inner) {
+      count++;
+      if constexpr (std::is_void_v<decltype(callback(item))>) {
+        callback(item);
+      } else {
+        if (!callback(item)) break;
+      }
+    }
+    return count;
+  }
+
   // Checks if the queue is empty.
   inline bool empty() const {
     return inner.empty();