@@ -82,12 +82,24 @@ async function encodeActionBoundArg(actionId: string, arg: string) {
8282  return  btoa ( ivValue  +  arrayBufferToString ( encrypted ) ) 
8383} 
8484
85+ enum  ReadStatus  { 
86+   Ready , 
87+   Pending , 
88+   Complete , 
89+ } 
90+ 
8591// Encrypts the action's bound args into a string. For the same combination of 
8692// actionId and args the same cached promise is returned. This ensures reference 
8793// equality for returned objects from "use cache" functions when they're invoked 
8894// multiple times within one render pass using the same bound args. 
8995export  const  encryptActionBoundArgs  =  React . cache ( 
9096  async  function  encryptActionBoundArgs ( actionId : string ,  ...args : any [ ] )  { 
97+     const  workUnitStore  =  workUnitAsyncStorage . getStore ( ) 
98+     const  cacheSignal  = 
99+       workUnitStore ?. type  ===  'prerender' 
100+         ? workUnitStore . cacheSignal 
101+         : undefined 
102+ 
91103    const  {  clientModules }  =  getClientReferenceManifestForRsc ( ) 
92104
93105    // Create an error before any asynchronous calls, to capture the original 
@@ -97,13 +109,38 @@ export const encryptActionBoundArgs = React.cache(
97109
98110    let  didCatchError  =  false 
99111
100-     const  workUnitStore  =  workUnitAsyncStorage . getStore ( ) 
101- 
102112    const  hangingInputAbortSignal  = 
103113      workUnitStore ?. type  ===  'prerender' 
104114        ? createHangingInputAbortSignal ( workUnitStore ) 
105115        : undefined 
106116
117+     let  readStatus  =  ReadStatus . Ready 
118+     function  startReadOnce ( )  { 
119+       if  ( readStatus  ===  ReadStatus . Ready )  { 
120+         readStatus  =  ReadStatus . Pending 
121+         cacheSignal ?. beginRead ( ) 
122+       } 
123+     } 
124+ 
125+     function  endReadIfStarted ( )  { 
126+       if  ( readStatus  ===  ReadStatus . Pending )  { 
127+         cacheSignal ?. endRead ( ) 
128+       } 
129+       readStatus  =  ReadStatus . Complete 
130+     } 
131+ 
132+     // streamToString might take longer than a microtask to resolve and then other things 
133+     // waiting on the cache signal might not realize there is another cache to fill so if 
134+     // we are no longer waiting on the bound args serialization via the hangingInputAbortSignal 
135+     // we should eagerly start the cache read to prevent other readers of the cache signal from 
136+     // missing this cache fill. We use a idempotent function to only start reading once because 
137+     // it's also possible that streamToString finishes before the hangingInputAbortSignal aborts. 
138+     if  ( hangingInputAbortSignal  &&  cacheSignal )  { 
139+       hangingInputAbortSignal . addEventListener ( 'abort' ,  startReadOnce ,  { 
140+         once : true , 
141+       } ) 
142+     } 
143+ 
107144    // Using Flight to serialize the args into a string. 
108145    const  serialized  =  await  streamToString ( 
109146      renderToReadableStream ( args ,  clientModules ,  { 
@@ -139,13 +176,18 @@ export const encryptActionBoundArgs = React.cache(
139176        console . error ( error ) 
140177      } 
141178
179+       endReadIfStarted ( ) 
142180      throw  error 
143181    } 
144182
145183    if  ( ! workUnitStore )  { 
184+       // We don't need to call cacheSignal.endRead here because we can't have a cacheSignal 
185+       // if we do not have a workUnitStore. 
146186      return  encodeActionBoundArg ( actionId ,  serialized ) 
147187    } 
148188
189+     startReadOnce ( ) 
190+ 
149191    const  prerenderResumeDataCache  =  getPrerenderResumeDataCache ( workUnitStore ) 
150192    const  renderResumeDataCache  =  getRenderResumeDataCache ( workUnitStore ) 
151193    const  cacheKey  =  actionId  +  serialized 
@@ -158,14 +200,9 @@ export const encryptActionBoundArgs = React.cache(
158200      return  cachedEncrypted 
159201    } 
160202
161-     const  cacheSignal  = 
162-       workUnitStore . type  ===  'prerender'  ? workUnitStore . cacheSignal  : undefined 
163- 
164-     cacheSignal ?. beginRead ( ) 
165- 
166203    const  encrypted  =  await  encodeActionBoundArg ( actionId ,  serialized ) 
167204
168-     cacheSignal ?. endRead ( ) 
205+     endReadIfStarted ( ) 
169206    prerenderResumeDataCache ?. encryptedBoundArgs . set ( cacheKey ,  encrypted ) 
170207
171208    return  encrypted 
0 commit comments