Defer TCP socket connect() until DO output gate clears

Adds TCP_SOCKET_CONNECT_OUTPUT_GATE autogate that defers httpClient->connect()
until IoContext::waitForOutputLocks() resolves. Socket is returned immediately
backed by kj::newPromisedStream so reads/writes queue until the real connection
arrives.

Author: vaibhavSharma-cf

src/workerd/api/sockets.c++

@@ -11,6 +11,7 @@
 #include <workerd/io/worker-interface.h>
 #include <workerd/jsg/exception.h>
 #include <workerd/jsg/url.h>
+#include <workerd/util/autogate.h>
 
 namespace workerd::api {
 
@@ -90,7 +91,8 @@ jsg::Ref<Socket> setupSocket(jsg::Lock& js,
     SecureTransportKind secureTransport,
     kj::Maybe<kj::String> domain,
     bool isDefaultFetchPort,
-    kj::Maybe<jsg::PromiseResolverPair<SocketInfo>> maybeOpenedPrPair) {
+    kj::Maybe<jsg::PromiseResolverPair<SocketInfo>> maybeOpenedPrPair,
+    kj::Maybe<kj::Promise<void>> connectTask) {
   auto& ioContext = IoContext::current();
 
   // Disconnection handling is annoyingly complicated:
@@ -173,7 +175,7 @@ jsg::Ref<Socket> setupSocket(jsg::Lock& js,
   auto result = js.alloc<Socket>(js, ioContext, kj::mv(refcountedConnection), kj::mv(remoteAddress),
       kj::mv(readable), kj::mv(writable), kj::mv(closedPrPair), kj::mv(watchForDisconnectTask),
       kj::mv(options), kj::mv(tlsStarter), secureTransport, kj::mv(domain), isDefaultFetchPort,
-      kj::mv(openedPrPair));
+      kj::mv(openedPrPair), kj::mv(connectTask));
 
   KJ_IF_SOME(p, eofPromise) {
     result->handleReadableEof(js, kj::mv(p));
@@ -256,25 +258,99 @@ jsg::Ref<Socket> connectImplNoOutputLock(jsg::Lock& js,
   }
   kj::Own<kj::TlsStarterCallback> tlsStarter = kj::heap<kj::TlsStarterCallback>();
   httpConnectSettings.tlsStarter = tlsStarter;
-  auto request = httpClient->connect(addressStr, *headers, httpConnectSettings);
-  request.connection = request.connection.attach(kj::mv(httpClient));
-
-  auto result = setupSocket(js, kj::mv(request.connection), kj::mv(addressStr), kj::mv(options),
-      kj::mv(tlsStarter), secureTransport, kj::mv(domain), isDefaultFetchPort,
-      kj::none /* maybeOpenedPrPair */);
-  // `handleProxyStatus` needs an initialized refcount to use `JSG_THIS`, hence it cannot be
-  // called in Socket's constructor. Also it's only necessary when creating a Socket as a result of
-  // a `connect`.
-  result->handleProxyStatus(js, kj::mv(request.status));
-  return result;
+
+  if (util::Autogate::isEnabled(util::AutogateKey::TCP_SOCKET_CONNECT_OUTPUT_GATE)) {
+    // Deferred-connect path: return a Socket backed by a promised stream immediately, deferring
+    // the actual httpClient->connect() until after the DO output gate clears. This prevents
+    // premature network output while storage writes are still pending.
+    //
+    // Two promise-fulfiller pairs bridge the deferred connect to the Socket:
+    //   connStreamPaf — fulfilled with the raw connection stream once connect() returns
+    //   statusPaf     — fulfilled with the proxy status once connect()'s status promise resolves
+    //
+    // The connect task is stored inside ConnectionData alongside tlsStarter. This ensures
+    // that if the Socket is GC'd (destroying ConnectionData), the connect task is cancelled
+    // before tlsStarter is destroyed, preventing a use-after-free on the TlsStarterCallback
+    // pointer embedded in httpConnectSettings.
+
+    // Promise-fulfiller for the raw connection stream.
+    auto connStreamPaf = kj::newPromiseAndFulfiller<kj::Own<kj::AsyncIoStream>>();
+    auto& connStreamFulfiller = *connStreamPaf.fulfiller;
+    auto deferredCancelStream = kj::defer([fulfiller = kj::mv(connStreamPaf.fulfiller)]() mutable {
+      fulfiller->reject(KJ_EXCEPTION(DISCONNECTED, "socket connect cancelled"));
+    });
+
+    // Promise-fulfiller for the proxy status.
+    auto statusPaf = kj::newPromiseAndFulfiller<kj::HttpClient::ConnectRequest::Status>();
+    auto& statusFulfiller = *statusPaf.fulfiller;
+    auto deferredCancelStatus = kj::defer([fulfiller = kj::mv(statusPaf.fulfiller)]() mutable {
+      fulfiller->reject(KJ_EXCEPTION(DISCONNECTED, "socket connect cancelled"));
+    });
+
+    // The connect task: waits for the output gate, performs the connect, then forwards results.
+    static auto constexpr doConnect =
+        [](IoContext& ioContext, kj::StringPtr addressStr, kj::HttpHeaders& headers,
+            kj::HttpConnectSettings httpConnectSettings, kj::Own<kj::HttpClient> httpClient,
+            kj::PromiseFulfiller<kj::Own<kj::AsyncIoStream>>& connFulfiller,
+            kj::PromiseFulfiller<kj::HttpClient::ConnectRequest::Status>& statusFulfiller)
+        -> kj::Promise<void> {
+      try {
+        co_await ioContext.waitForOutputLocks();
+        auto request = httpClient->connect(addressStr, headers, httpConnectSettings);
+        request.connection = request.connection.attach(kj::mv(httpClient));
+        auto status = kj::mv(request.status);
+        connFulfiller.fulfill(kj::mv(request.connection));
+        auto resolvedStatus = co_await kj::mv(status);
+        statusFulfiller.fulfill(kj::mv(resolvedStatus));
+      } catch (...) {
+        auto e = kj::getCaughtExceptionAsKj();
+        if (!connFulfiller.isWaiting()) {
+          // Connection was already fulfilled; only reject status.
+          statusFulfiller.reject(kj::mv(e));
+        } else {
+          connFulfiller.reject(kj::cp(e));
+          statusFulfiller.reject(kj::mv(e));
+        }
+      }
+    };
+
+    auto connectTaskPromise =
+        doConnect(ioContext, addressStr, *headers, httpConnectSettings, kj::mv(httpClient),
+            connStreamFulfiller, statusFulfiller)
+            .attach(kj::mv(deferredCancelStream), kj::mv(deferredCancelStatus), kj::mv(headers));
+
+    auto result = setupSocket(js, kj::newPromisedStream(kj::mv(connStreamPaf.promise)),
+        kj::mv(addressStr), kj::mv(options), kj::mv(tlsStarter), secureTransport, kj::mv(domain),
+        isDefaultFetchPort, kj::none /* maybeOpenedPrPair */, kj::mv(connectTaskPromise));
+
+    // `handleProxyStatus` needs an initialized refcount to use `JSG_THIS`, hence it cannot be
+    // called in Socket's constructor. Also it's only necessary when creating a Socket as a result
+    // of a `connect`.
+    result->handleProxyStatus(js, kj::mv(statusPaf.promise));
+    return result;
+  } else {
+    // Original synchronous-connect path (no output gate wait).
+    auto request = httpClient->connect(addressStr, *headers, httpConnectSettings);
+    request.connection = request.connection.attach(kj::mv(httpClient));
+
+    auto result = setupSocket(js, kj::mv(request.connection), kj::mv(addressStr), kj::mv(options),
+        kj::mv(tlsStarter), secureTransport, kj::mv(domain), isDefaultFetchPort,
+        kj::none /* maybeOpenedPrPair */);
+    // `handleProxyStatus` needs an initialized refcount to use `JSG_THIS`, hence it cannot be
+    // called in Socket's constructor. Also it's only necessary when creating a Socket as a result
+    // of a `connect`.
+    result->handleProxyStatus(js, kj::mv(request.status));
+    return result;
+  }
 }
 
 jsg::Ref<Socket> connectImpl(jsg::Lock& js,
     kj::Maybe<jsg::Ref<Fetcher>> fetcher,
     AnySocketAddress address,
     jsg::Optional<SocketOptions> options) {
-  // TODO(soon): Doesn't this need to check for the presence of an output lock, and if it finds one
-  // then wait on it, before calling into connectImplNoOutputLock?
+  // When the TCP_SOCKET_CONNECT_OUTPUT_GATE autogate is enabled, the output gate wait is
+  // handled inside connectImplNoOutputLock via a deferred connect task, so no separate wait
+  // is needed here.
   return connectImplNoOutputLock(js, kj::mv(fetcher), kj::mv(address), kj::mv(options));
 }
 

src/workerd/api/sockets.h

@@ -70,9 +70,12 @@ class Socket: public jsg::Object {
       SecureTransportKind secureTransport,
       kj::Maybe<kj::String> domain,
       bool isDefaultFetchPort,
-      jsg::PromiseResolverPair<SocketInfo> openedPrPair)
-      : connectionData(context.addObject(kj::heap<ConnectionData>(
-            kj::mv(tlsStarter), kj::mv(connectionStream), kj::mv(watchForDisconnectTask)))),
+      jsg::PromiseResolverPair<SocketInfo> openedPrPair,
+      kj::Maybe<kj::Promise<void>> connectTask = kj::none)
+      : connectionData(context.addObject(kj::heap<ConnectionData>(kj::mv(tlsStarter),
+            kj::mv(connectionStream),
+            kj::mv(watchForDisconnectTask),
+            kj::mv(connectTask)))),
         readable(kj::mv(readableParam)),
         writable(kj::mv(writable)),
         closedResolver(kj::mv(closedPrPair.resolver)),
@@ -179,14 +182,21 @@ class Socket: public jsg::Object {
   struct ConnectionData {
     kj::Own<kj::RefcountedWrapper<kj::Own<kj::AsyncIoStream>>> connectionStream;
     kj::Maybe<kj::Promise<void>> watchForDisconnectTask;
+    // When the deferred-connect autogate is enabled, holds the task that waits for the
+    // output gate then calls httpClient->connect(). Lives here so that cancelling the
+    // Socket (destroying ConnectionData) also cancels the connect task, preventing a
+    // TlsStarterCallback use-after-free.
+    kj::Maybe<kj::Promise<void>> connectTask;
     // tlsStarter must be declared after connectionStream so that it is destroyed first,
     // since it holds a reference that keeps the connection alive.
     kj::Own<kj::TlsStarterCallback> tlsStarter;
     ConnectionData(kj::Own<kj::TlsStarterCallback> tlsStarter,
         kj::Own<kj::RefcountedWrapper<kj::Own<kj::AsyncIoStream>>> connStream,
-        kj::Promise<void> disconnectTask)
+        kj::Promise<void> disconnectTask,
+        kj::Maybe<kj::Promise<void>> connectTask = kj::none)
         : connectionStream(kj::mv(connStream)),
           watchForDisconnectTask(kj::mv(disconnectTask)),
+          connectTask(kj::mv(connectTask)),
           tlsStarter(kj::mv(tlsStarter)) {}
   };
   kj::Maybe<IoOwn<ConnectionData>> connectionData;
@@ -251,7 +261,8 @@ jsg::Ref<Socket> setupSocket(jsg::Lock& js,
     SecureTransportKind secureTransport,
     kj::Maybe<kj::String> domain,
     bool isDefaultFetchPort,
-    kj::Maybe<jsg::PromiseResolverPair<SocketInfo>> maybeOpenedPrPair);
+    kj::Maybe<jsg::PromiseResolverPair<SocketInfo>> maybeOpenedPrPair,
+    kj::Maybe<kj::Promise<void>> connectTask = kj::none);
 
 jsg::Ref<Socket> connectImplNoOutputLock(jsg::Lock& js,
     kj::Maybe<jsg::Ref<Fetcher>> fetcher,

src/workerd/api/tests/BUILD.bazel

@@ -630,6 +630,19 @@ wd_test(
     ],
 )
 
+js_binary(
+    name = "socket-output-gate-server",
+    entry_point = "socket-output-gate-server.js",
+)
+
+wd_test(
+    src = "socket-output-gate-test.wd-test",
+    args = ["--experimental"],
+    data = ["socket-output-gate-test.js"],
+    sidecar = "socket-output-gate-server",
+    sidecar_port_bindings = ["ECHO_SERVER_PORT"],
+)
+
 wd_test(
     src = "js-rpc-socket-test.wd-test",
     args = [

src/workerd/api/tests/socket-output-gate-server.js

@@ -0,0 +1,11 @@
+// Simple TCP echo server for socket-output-gate test.
+const net = require('node:net');
+
+const server = net.createServer((socket) => {
+  socket.on('data', (chunk) => socket.write(chunk));
+  socket.on('error', () => {});
+});
+
+server.listen(process.env.ECHO_SERVER_PORT, () => {
+  console.info(`Echo server on port ${server.address().port}`);
+});

src/workerd/api/tests/socket-output-gate-test.js

@@ -0,0 +1,44 @@
+// Tests that TCP socket connect() respects the DO output gate when
+// TCP_SOCKET_CONNECT_OUTPUT_GATE is enabled.
+//
+// Invariant: by the time socket.opened resolves, any storage write started
+// in the same turn must already be committed.
+
+import { connect } from 'cloudflare:sockets';
+import { strict as assert } from 'node:assert';
+import { DurableObject } from 'cloudflare:workers';
+
+export class SocketTestDO extends DurableObject {
+  async fetch(request) {
+    const port = new URL(request.url).searchParams.get('port');
+
+    let putCompleted = false;
+
+    // Storage write — no await. This creates a pending output gate lock.
+    this.ctx.storage.put('key', 'value').then(() => {
+      putCompleted = true;
+    });
+
+    // Connect immediately while the write may still be in-flight.
+    const socket = connect(`localhost:${port}`);
+
+    // socket.opened only resolves after handleProxyStatus confirms a successful
+    // proxy status — which, with the autogate, cannot happen until
+    // waitForOutputLocks() resolves, i.e. after the put is committed.
+    await socket.opened;
+
+    assert.ok(putCompleted, 'storage put must complete before socket opens');
+
+    await socket.close();
+    return new Response('ok');
+  }
+}
+
+export const connectRespectsOutputGate = {
+  async test(ctrl, env) {
+    const id = env.SOCKET_TEST_DO.newUniqueId();
+    const stub = env.SOCKET_TEST_DO.get(id);
+    const resp = await stub.fetch(`http://do/?port=${env.ECHO_SERVER_PORT}`);
+    assert.equal(resp.status, 200);
+  },
+};

src/workerd/api/tests/socket-output-gate-test.wd-test

@@ -0,0 +1,25 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  autogates = ["workerd-autogate-tcp-socket-connect-output-gate"],
+  services = [
+    (
+      name = "main",
+      worker = (
+        modules = [
+          (name = "worker", esModule = embed "socket-output-gate-test.js"),
+        ],
+        compatibilityFlags = ["nodejs_compat", "experimental"],
+        bindings = [
+          (name = "ECHO_SERVER_PORT", fromEnvironment = "ECHO_SERVER_PORT"),
+          (name = "SOCKET_TEST_DO", durableObjectNamespace = "SocketTestDO"),
+        ],
+        durableObjectNamespaces = [
+          (className = "SocketTestDO", uniqueKey = "210bd0cbd803ef7883a1ee9d86cce06e"),
+        ],
+        durableObjectStorage = (inMemory = void),
+      )
+    ),
+    ( name = "internet", network = ( allow = ["private"] ) ),
+  ],
+);

src/workerd/util/autogate.c++

@@ -35,6 +35,8 @@ kj::StringPtr KJ_STRINGIFY(AutogateKey key) {
       return "enable-fast-textencoder"_kj;
     case AutogateKey::ENABLE_DRAINING_READ_ON_STANDARD_STREAMS:
       return "enable-draining-read-on-standard-streams"_kj;
+    case AutogateKey::TCP_SOCKET_CONNECT_OUTPUT_GATE:
+      return "tcp-socket-connect-output-gate"_kj;
     case AutogateKey::NumOfKeys:
       KJ_FAIL_ASSERT("NumOfKeys should not be used in getName");
   }

src/workerd/util/autogate.h

@@ -39,6 +39,9 @@ enum class AutogateKey {
   ENABLE_FAST_TEXTENCODER,
   // Enable draining read on standard streams
   ENABLE_DRAINING_READ_ON_STANDARD_STREAMS,
+  // Defers TCP socket connect() to wait for DO output gate, preventing
+  // network outputs while storage writes are pending.
+  TCP_SOCKET_CONNECT_OUTPUT_GATE,
   NumOfKeys  // Reserved for iteration.
 };