refactor: optimize RSC payload handling and improve hook management for better performance and clarity

Author: AbanoubGhadban

lib/react_on_rails/utils.rb

@@ -216,8 +216,10 @@ def self.react_on_rails_pro_version
     def self.rsc_support_enabled?
       return false unless react_on_rails_pro?
 
+      return @rsc_support_enabled if defined?(@rsc_support_enabled)
+
       rorp_config = ReactOnRailsPro.configuration
-      rorp_config.respond_to?(:enable_rsc_support) && rorp_config.enable_rsc_support
+      @rsc_support_enabled = rorp_config.respond_to?(:enable_rsc_support) && rorp_config.enable_rsc_support
     end
 
     def self.full_text_errors_enabled?

node_package/src/RSCPayloadGenerator.ts

@@ -13,7 +13,7 @@ declare global {
   ): Promise<NodeJS.ReadableStream>;
 }
 
-const mapRailsContextToRSCPayloadStreams = new Map<string, RSCPayloadStreamInfo[]>();
+const rscPayloadStreams = new Map<string, RSCPayloadStreamInfo[]>();
 
 const rscPayloadCallbacks = new Map<string, RSCPayloadCallback[]>();
 
@@ -44,7 +44,7 @@ export const onRSCPayloadGenerated = (
   }
 
   // Call callback for any existing streams for this context
-  const existingStreams = mapRailsContextToRSCPayloadStreams.get(renderRequestId);
+  const existingStreams = rscPayloadStreams.get(renderRequestId);
   if (existingStreams) {
     existingStreams.forEach((streamInfo) => callback(streamInfo));
   }
@@ -95,11 +95,11 @@ export const getRSCPayloadStream = async (
     props,
     stream: stream2,
   };
-  const streams = mapRailsContextToRSCPayloadStreams.get(renderRequestId);
+  const streams = rscPayloadStreams.get(renderRequestId);
   if (streams) {
     streams.push(streamInfo);
   } else {
-    mapRailsContextToRSCPayloadStreams.set(renderRequestId, [streamInfo]);
+    rscPayloadStreams.set(renderRequestId, [streamInfo]);
   }
 
   // Notify callbacks about the new stream in a sync manner to maintain proper hydration timing
@@ -120,11 +120,11 @@ export const getRSCPayloadStreams = (
   stream: NodeJS.ReadableStream;
 }[] => {
   const { renderRequestId } = railsContext.componentSpecificMetadata;
-  return mapRailsContextToRSCPayloadStreams.get(renderRequestId) ?? [];
+  return rscPayloadStreams.get(renderRequestId) ?? [];
 };
 
 export const clearRSCPayloadStreams = (railsContext: RailsContextWithServerComponentCapabilities) => {
   const { renderRequestId } = railsContext.componentSpecificMetadata;
-  mapRailsContextToRSCPayloadStreams.delete(renderRequestId);
+  rscPayloadStreams.delete(renderRequestId);
   rscPayloadCallbacks.delete(renderRequestId);
 };

node_package/src/injectRSCPayload.ts

@@ -78,15 +78,12 @@ export default function injectRSCPayload(
         // The client-side RSCProvider actively monitors the array for new chunks, processing them as they arrive and forwarding them to the RSC payload stream, regardless of whether the array is initially empty.
         initializeCacheKeyJSArray(cacheKey, resultStream);
         rscPromises.push(
-          new Promise((resolve, reject) => {
-            (async () => {
-              for await (const chunk of stream ?? []) {
-                const decodedChunk = typeof chunk === 'string' ? chunk : decoder.decode(chunk);
-                writeChunk(JSON.stringify(decodedChunk), resultStream, cacheKey);
-              }
-              resolve();
-            })().catch(reject);
-          }),
+          (async () => {
+            for await (const chunk of stream ?? []) {
+              const decodedChunk = typeof chunk === 'string' ? chunk : decoder.decode(chunk);
+              writeChunk(JSON.stringify(decodedChunk), resultStream, cacheKey);
+            }
+          })(),
         );
       });
 

node_package/src/postSSRHooks.ts

@@ -7,14 +7,18 @@ export const addPostSSRHook = (
   railsContext: RailsContextWithServerComponentCapabilities,
   hook: PostSSRHook,
 ) => {
-  const hooks = postSSRHooks.get(railsContext.componentSpecificMetadata.renderRequestId) || [];
-  hooks.push(hook);
-  postSSRHooks.set(railsContext.componentSpecificMetadata.renderRequestId, hooks);
+  const hooks = postSSRHooks.get(railsContext.componentSpecificMetadata.renderRequestId);
+  if (hooks) {
+    hooks.push(hook);
+  } else {
+    postSSRHooks.set(railsContext.componentSpecificMetadata.renderRequestId, [hook]);
+  }
 };
 
 export const notifySSREnd = (railsContext: RailsContextWithServerComponentCapabilities) => {
   const hooks = postSSRHooks.get(railsContext.componentSpecificMetadata.renderRequestId);
   if (hooks) {
     hooks.forEach((hook) => hook());
   }
+  postSSRHooks.delete(railsContext.componentSpecificMetadata.renderRequestId);
 };