move maxMessageSize and reading feature flag into startReadLoop

src/workerd/api/web-socket.c++

@@ -259,10 +259,6 @@ jsg::Ref<WebSocket> WebSocket::constructor(
     headers.unset(kj::HttpHeaderId::SEC_WEBSOCKET_EXTENSIONS);
   }
 
-  if (FeatureFlags::get(js).getIncreaseWebsocketMessageSize()) {
-    ws->maxMessageSize = 128u << 20;
-  }
-
   auto prom = ([](auto& context, auto connUrl, auto headers, auto client)
       -> kj::Promise<PackedWebSocket> {
     auto response = co_await client->openWebSocket(connUrl, headers);
@@ -454,14 +450,19 @@ WebSocket::Accepted::~Accepted() noexcept(false) {
 }
 
 void WebSocket::startReadLoop(jsg::Lock& js, kj::Maybe<kj::Own<InputGate::CriticalSection>> cs) {
+  size_t maxMessageSize = 1u << 20;
+  if (FeatureFlags::get(js).getIncreaseWebsocketMessageSize()) {
+    maxMessageSize = 128u << 20;
+  }
+
   // If the kj::WebSocket happens to be an AbortableWebSocket (see util/abortable.h), then
   // calling readLoop here could throw synchronously if the canceler has already been tripped.
   // Using kj::evalNow() here let's us capture that and handle correctly.
   //
   // We catch exceptions and return Maybe<Exception> instead since we want to handle the exceptions
   // in awaitIo() below, but we don't want the KJ exception converted to JavaScript before we can
   // examine it.
-  kj::Promise<kj::Maybe<kj::Exception>> promise = readLoop(kj::mv(cs));
+  kj::Promise<kj::Maybe<kj::Exception>> promise = readLoop(kj::mv(cs), maxMessageSize);
 
   auto& context = IoContext::current();
 
@@ -908,14 +909,15 @@ kj::Array<kj::StringPtr> WebSocket::getHibernatableTags() {
 }
 
 kj::Promise<kj::Maybe<kj::Exception>> WebSocket::readLoop(
-    kj::Maybe<kj::Own<InputGate::CriticalSection>> cs) {
+    kj::Maybe<kj::Own<InputGate::CriticalSection>> cs,
+    size_t maxMessageSize) {
   try {
     // Note that we'll throw if the websocket has enabled hibernation.
     auto& ws = *KJ_REQUIRE_NONNULL(
         KJ_ASSERT_NONNULL(farNative->state.tryGet<Accepted>()).ws.getIfNotHibernatable());
     auto& context = IoContext::current();
     while (true) {
-      auto message = co_await ws.receive(ws.maxMessageSize);
+      auto message = co_await ws.receive(maxMessageSize);
 
       context.getLimitEnforcer().topUpActor();
       KJ_IF_SOME(a, context.getActor()) {

src/workerd/api/web-socket.h

@@ -425,9 +425,6 @@ class WebSocket: public EventTarget {
   // `close()`, thereby preventing calls to `send()` even after we wake from hibernation.
   bool closedOutgoingForHib = false;
 
-  // Maximum allowed size for WebSocket messages
-  size_t maxMessageSize = 1u << 20;
-
   // Maximum size of a WebSocket attachment.
   inline static const size_t MAX_ATTACHMENT_SIZE = 1024 * 2;