Skip to content

Commit 656271b

Browse files
authored
[use-cache] track cache read earlier when encrypting bound args (#81427)
bound args need to be encrypted for use cache functions and currently we track the read slightly later than will be necessary when we sync the latest React. This change moves the tracking to be as early as possible (as soon as the input signal is aborted or when the bound args are finished being serialized, whichever is first).
1 parent bdb87a1 commit 656271b

File tree

3 files changed

+47
-10
lines changed
  • packages/next/src/server/app-render
  • test/e2e/app-dir/use-cache-hanging-inputs/app

3 files changed

+47
-10
lines changed

packages/next/src/server/app-render/encryption.ts

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -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.
8995
export 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

test/e2e/app-dir/use-cache-hanging-inputs/app/uncached-promise-nested/page.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import React from 'react'
22
import { setTimeout } from 'timers/promises'
33

44
async function getUncachedData() {
5-
await setTimeout(100)
5+
await setTimeout(0)
66

77
return Math.random()
88
}

test/e2e/app-dir/use-cache-hanging-inputs/app/uncached-promise/page.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import React from 'react'
22
import { setTimeout } from 'timers/promises'
33

44
async function fetchUncachedData() {
5-
await setTimeout(100)
5+
await setTimeout(0)
66

77
return Math.random()
88
}

0 commit comments

Comments
 (0)