Skip to content

Commit 59bc945

Browse files
committed
fix(oauth): take over a lapsed refresh lease instead of serving the stored token
- A follower whose wait ends with nothing now tries the lock once more: a leader whose process died leaves its lease to lapse, and the follower leads the refresh instead of failing. Losing that try still returns nothing. - Drops the stored-token fallback and its lease-held outcome, which served tokens below the five-minute floor a served token keeps. - The leader rereads the dead flag under the lock, so a follower taking over after a leader that failed terminally does not call the provider again. - The in-process refresh settles after the follower wait plus a full leader run and headroom.
1 parent 9a23d4b commit 59bc945

4 files changed

Lines changed: 96 additions & 104 deletions

File tree

‎apps/sim/lib/concurrency/leader-lock.test.ts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,32 @@ describe('withLeaderLock', () => {
119119
}
120120
})
121121

122+
describe('when the follower wait ends with nothing', () => {
123+
const waitedOut = {
124+
key: 'k',
125+
pollIntervalMs: 5,
126+
maxWaitMs: 20,
127+
onFollower: async () => null,
128+
}
129+
130+
it('leads once a leader that vanished leaves its lease to lapse', async () => {
131+
redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false)
132+
const onLeader = vi.fn(async () => 'took-over')
133+
await expect(withLeaderLock<string>({ ...waitedOut, onLeader })).resolves.toBe('took-over')
134+
expect(onLeader).toHaveBeenCalledOnce()
135+
expect(redisConfigMockFns.mockAcquireLock).toHaveBeenCalledTimes(2)
136+
expect(redisConfigMockFns.mockReleaseLock).toHaveBeenCalledOnce()
137+
})
138+
139+
it('returns the follower result when a live holder still has the lease', async () => {
140+
redisConfigMockFns.mockAcquireLock.mockResolvedValue(false)
141+
const onLeader = vi.fn(async () => 'should-not-run')
142+
await expect(withLeaderLock<string>({ ...waitedOut, onLeader })).resolves.toBeNull()
143+
expect(onLeader).not.toHaveBeenCalled()
144+
expect(redisConfigMockFns.mockAcquireLock).toHaveBeenCalledTimes(2)
145+
})
146+
})
147+
122148
it('only one of N concurrent callers acquires the lock', async () => {
123149
// Track which calls won the lock: first one returns true, rest return false.
124150
let acquired = false

‎apps/sim/lib/concurrency/leader-lock.ts‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,13 @@ export interface LeaderLockOptions<T> {
1919
onFollower: () => Promise<T | null>
2020
}
2121

22+
/**
23+
* Runs `onLeader` in the one caller that takes the lock for `key`; every other caller polls
24+
* `onFollower` until it reports the leader's result or `maxWaitMs` passes. A follower whose wait
25+
* ends with nothing tries the lock once more: a leader whose process died without releasing it
26+
* leaves the lease to lapse, and that follower then leads instead of failing. Losing that try
27+
* means a live holder has the lease, and the follower returns `null`.
28+
*/
2229
export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T | null> {
2330
const {
2431
key,
@@ -42,7 +49,7 @@ export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T |
4249
return onLeader()
4350
}
4451

45-
if (acquired) {
52+
const leadHoldingLock = async () => {
4653
try {
4754
return await onLeader()
4855
} finally {
@@ -57,6 +64,8 @@ export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T |
5764
}
5865
}
5966

67+
if (acquired) return leadHoldingLock()
68+
6069
const deadline = Date.now() + maxWaitMs
6170
while (Date.now() < deadline) {
6271
await sleep(pollIntervalMs)
@@ -68,6 +77,20 @@ export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T |
6877
const lastChance = await onFollower()
6978
if (lastChance !== null) return lastChance
7079

80+
let tookOver = false
81+
try {
82+
tookOver = await acquireLock(key, ownerToken, ttlSec)
83+
} catch (error) {
84+
logger.warn('Lock takeover failed after the follower wait', {
85+
key,
86+
error: toError(error).message,
87+
})
88+
}
89+
if (tookOver) {
90+
logger.warn('Follower took over a lease its leader left to lapse', { key, maxWaitMs })
91+
return leadHoldingLock()
92+
}
93+
7194
logger.warn('Follower timed out waiting for leader', { key, maxWaitMs })
7295
return null
7396
}

‎apps/sim/lib/oauth/credential-service.test.ts‎

Lines changed: 25 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -601,7 +601,7 @@ describe('OAuth access-token refresh headroom', () => {
601601
expect(guard).toContain('original-refresh-token')
602602
})
603603

604-
/** A Google token with a minute left: the refresh is awaited, and enough validity remains to serve it. */
604+
/** A Google token with a minute left, so its caller waits for the refresh instead of using it. */
605605
const WAITED_GOOGLE_REMAINING_MS = 60_000
606606

607607
it('returns no token when the rotation write finds the account gone', async () => {
@@ -810,62 +810,46 @@ describe('OAuth access-token refresh headroom', () => {
810810
redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false)
811811
})
812812

813-
/** Every follower poll, once the lock is contested, reads the unchanged row the caller read. */
814-
async function keepChainStoredDuringWait(stored: OAuthAccount) {
815-
await vi.waitFor(() => expect(redisConfigMockFns.mockAcquireLock).toHaveBeenCalled())
816-
dbChainMockFns.limit.mockResolvedValue([stored])
817-
}
818-
819813
it.each([
820-
{
821-
label: 'serves a Google token still valid after the wait',
822-
row: { providerId: 'google-drive', remainingMs: 100_000 },
823-
token: 'original-access-token',
824-
},
825-
{
826-
label: 'fails a Google token with under 30 seconds left after the wait',
827-
row: { providerId: 'google-drive', remainingMs: 50_000 },
828-
token: null,
829-
},
830-
{
831-
label: 'fails a provider whose refresh can invalidate the stored chain',
832-
row: { providerId: 'confluence', remainingMs: 100_000 },
833-
token: null,
834-
},
835-
])('$label', async ({ row, token }) => {
836-
const stored = { ...createOAuthAccount(row.remainingMs), providerId: row.providerId }
837-
queueCredentialAccount(stored)
814+
{ label: 'a Google token under the refresh floor', providerId: 'google-drive' },
815+
{ label: 'a provider that cannot refresh ahead', providerId: 'confluence' },
816+
])('refreshes $label once the follower takes over the lapsed lease', async ({ providerId }) => {
817+
queueCredentialAccount({ ...createOAuthAccount(240_000), providerId })
838818
const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
839-
await keepChainStoredDuringWait(stored)
840819
await vi.advanceTimersByTimeAsync(31_000)
841-
await expect(pending).resolves.toEqual(token === null ? null : { accessToken: token })
820+
await expect(pending).resolves.toEqual({ accessToken: 'refreshed-access-token' })
821+
expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()
822+
})
823+
824+
it('never serves the stored token under the refresh floor when a live holder keeps the lease', async () => {
825+
redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false)
826+
queueCredentialAccount(createOAuthAccount(240_000))
827+
const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
828+
await vi.advanceTimersByTimeAsync(31_000)
829+
await expect(pending).resolves.toBeNull()
842830
expect(mocks.refreshOAuthToken).not.toHaveBeenCalled()
843831
})
844832

845-
it.each([
846-
{ label: 'the account is deleted', deleted: true, deadCode: null },
847-
{ label: 'another leader flags it dead', deleted: false, deadCode: 'invalid_client' },
848-
])('fails a valid Google token when $label during the wait', async ({ deleted, deadCode }) => {
849-
const stored = createOAuthAccount(100_000)
850-
queueCredentialAccount(stored)
833+
it('does not refresh again after taking over from a leader that flagged the credential dead', async () => {
834+
mocks.getRecentTerminalError
835+
.mockResolvedValueOnce(null)
836+
.mockResolvedValueOnce('invalid_client')
837+
queueCredentialAccount(createOAuthAccount(240_000))
851838
const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
852-
if (!deleted) await keepChainStoredDuringWait(stored)
853-
if (deadCode) mocks.getRecentTerminalError.mockResolvedValue(deadCode)
854839
await vi.advanceTimersByTimeAsync(31_000)
855840
await expect(pending).resolves.toBeNull()
841+
expect(mocks.refreshOAuthToken).not.toHaveBeenCalled()
856842
})
857843

858-
it('lets the follower finish its wait before the in-process refresh times out', async () => {
844+
it('lets the follower take over before the in-process refresh times out', async () => {
859845
const { coalesceLocally } = await vi.importActual<
860846
typeof import('@/lib/concurrency/singleflight')
861847
>('@/lib/concurrency/singleflight')
862848
mocks.coalesceLocally.mockImplementation(coalesceLocally)
863-
const stored = createOAuthAccount(100_000)
864-
queueCredentialAccount(stored)
849+
queueCredentialAccount(createOAuthAccount(240_000))
865850
const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
866-
await keepChainStoredDuringWait(stored)
867851
await vi.advanceTimersByTimeAsync(31_000)
868-
await expect(pending).resolves.toEqual({ accessToken: 'original-access-token' })
852+
await expect(pending).resolves.toEqual({ accessToken: 'refreshed-access-token' })
869853
})
870854
})
871855

@@ -882,7 +866,7 @@ describe('OAuth access-token refresh headroom', () => {
882866
settled = true
883867
}
884868
)
885-
await vi.advanceTimersByTimeAsync(44_000)
869+
await vi.advanceTimersByTimeAsync(74_000)
886870
expect(settled).toBe(false)
887871
await vi.advanceTimersByTimeAsync(2_000)
888872
await expect(pending).resolves.toBeNull()

‎apps/sim/lib/oauth/credential-service.ts‎

Lines changed: 21 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -948,23 +948,16 @@ function isRevocationReprobeDue(chain: RevocableChain | null | undefined): boole
948948
return !revokedAt || Date.now() - revokedAt.getTime() >= REVOKED_REFRESH_REPROBE_INTERVAL_MS
949949
}
950950

951-
/** What a settled refresh produced: one this process led, or one a follower saw its leader store. */
952-
type SettledRefreshOutcome =
951+
/** What one coalesced refresh produced; `null` is a failure worth retrying later. */
952+
type RefreshOutcome =
953953
| { kind: 'token'; accessToken: string }
954954
| { kind: 'revoked'; errorCode: string }
955955

956-
/**
957-
* What one coalesced refresh produced; `null` is a failure worth retrying later. `lease-held`
958-
* means this caller followed another lease holder for its whole wait and nothing was stored, while
959-
* the account still holds the chain the caller read and no leader flagged it dead.
960-
*/
961-
type RefreshOutcome = SettledRefreshOutcome | { kind: 'lease-held' }
962-
963-
function tokenOutcome(accessToken: string | null): SettledRefreshOutcome | null {
956+
function tokenOutcome(accessToken: string | null): RefreshOutcome | null {
964957
return accessToken ? { kind: 'token', accessToken } : null
965958
}
966959

967-
function revokedOutcome(errorCode: string | null): SettledRefreshOutcome | null {
960+
function revokedOutcome(errorCode: string | null): RefreshOutcome | null {
968961
return errorCode ? { kind: 'revoked', errorCode } : null
969962
}
970963

@@ -1022,10 +1015,6 @@ function isOAuthAccessTokenExpiring(
10221015
return expiresAt != null && expiresAt.getTime() <= now.getTime() + refreshWindowMs
10231016
}
10241017

1025-
function remainingValidityMs(expiresAt: Date, now = new Date()): number {
1026-
return expiresAt.getTime() - now.getTime()
1027-
}
1028-
10291018
/**
10301019
* Where a provider can refresh ahead of expiry, a caller is served the stored token with this much
10311020
* validity left or less while a background refresh replaces it. It starts above
@@ -1039,12 +1028,6 @@ const OAUTH_REFRESH_AHEAD_WINDOW_MS = 10 * 60 * 1000
10391028
* the two windows so every account still refreshes ahead before callers have to wait.
10401029
*/
10411030
const OAUTH_REFRESH_AHEAD_JITTER_MS = 60 * 1000
1042-
/**
1043-
* The validity a stored token must still have, after a blocking refresh that never settled, to be
1044-
* served instead of failing the request. It sits below {@link OAUTH_ACCESS_TOKEN_REFRESH_WINDOW_MS}
1045-
* on purpose: the request can still use such a token, and its next caller refreshes again.
1046-
*/
1047-
const OAUTH_STORED_TOKEN_FALLBACK_MIN_MS = 30 * 1000
10481031

10491032
/**
10501033
* A Google refresh does not rotate the refresh token, and Microsoft does not revoke the old one
@@ -1106,7 +1089,7 @@ function planAccessTokenRefresh(
11061089
if (
11071090
refreshesAhead &&
11081091
accessTokenExpiresAt &&
1109-
remainingValidityMs(accessTokenExpiresAt, now) <=
1092+
accessTokenExpiresAt.getTime() - now.getTime() <=
11101093
OAUTH_REFRESH_AHEAD_WINDOW_MS - refreshAheadJitterMs(accountId)
11111094
) {
11121095
return { kind: 'refresh-ahead', refreshToken }
@@ -1146,11 +1129,13 @@ const REFRESH_LOCK_HEADROOM_MS = 15_000
11461129
const REFRESH_LOCK_TTL_SEC = Math.ceil((TOKEN_REFRESH_TIMEOUT_MS + REFRESH_LOCK_HEADROOM_MS) / 1000)
11471130
const REFRESH_FOLLOWER_MAX_WAIT_MS = REFRESH_LOCK_TTL_SEC * 1000
11481131
/**
1149-
* How long callers sharing one in-process refresh wait for it: past a follower's full wait (which
1150-
* starts once the lock is contested, and ends with one more read), so a follower concludes on its
1151-
* own; and still bounded, so a leader hung in this process fails its callers instead of wedging them.
1132+
* How long callers sharing one in-process refresh wait for it: a follower's full wait, then the
1133+
* leader run it may take over when that wait ends with a lapsed lease (budgeted like any lease),
1134+
* plus headroom for the lock round trips around them; so a follower concludes on its own, and a
1135+
* leader hung in this process still fails its callers instead of wedging them.
11521136
*/
1153-
const REFRESH_SETTLE_TIMEOUT_MS = REFRESH_FOLLOWER_MAX_WAIT_MS + REFRESH_LOCK_HEADROOM_MS
1137+
const REFRESH_SETTLE_TIMEOUT_MS =
1138+
REFRESH_FOLLOWER_MAX_WAIT_MS + REFRESH_LOCK_TTL_SEC * 1000 + REFRESH_LOCK_HEADROOM_MS
11541139

11551140
/**
11561141
* The raw scope one refresh coordinates on: the account row, or the installation for a Slack
@@ -1338,27 +1323,25 @@ async function performCoalescedRefresh({
13381323

13391324
const refreshPromise = coalesceLocally(
13401325
lockKey,
1341-
async (): Promise<RefreshOutcome | null> => {
1342-
let led = false
1343-
/** The refresh token the follower's latest read found, `undefined` when that read found nothing. */
1344-
let followerSawRefreshToken: string | null | undefined
1345-
const outcome = await withLeaderLock<SettledRefreshOutcome>({
1326+
() =>
1327+
withLeaderLock<RefreshOutcome>({
13461328
key: lockKey,
13471329
ttlSec: REFRESH_LOCK_TTL_SEC,
13481330
maxWaitMs: REFRESH_FOLLOWER_MAX_WAIT_MS,
13491331
onLeader: async () => {
1350-
led = true
13511332
try {
13521333
/**
1353-
* Reread under the lock: another process may have recorded a revocation after the
1354-
* caller read the row, and the version guards the record against a reconnect that
1355-
* lands while the provider call is in flight.
1334+
* Reread under the lock: another process may have recorded a revocation or flagged the
1335+
* credential dead after the caller read the row (a follower taking over the lock after
1336+
* a failed leader released it sees only these), and the version guards the record
1337+
* against a reconnect that lands while the provider call is in flight.
13561338
*/
13571339
const current = await readStoredChain(accountId)
13581340
const recordedMeanwhile = recordedRevocationCode(current)
13591341
if (recordedMeanwhile && !isRevocationReprobeDue(current)) {
13601342
return revokedOutcome(recordedMeanwhile)
13611343
}
1344+
if (await getRecentTerminalError(scopeKey)) return null
13621345
/**
13631346
* A writer that stored a new token after the caller's read (a refresh that finished just
13641347
* before this one took the lock, or a reconnect) already answers it, so this refresh
@@ -1578,7 +1561,6 @@ async function performCoalescedRefresh({
15781561
.from(account)
15791562
.where(eq(account.id, accountId))
15801563
.limit(1)
1581-
followerSawRefreshToken = row?.refreshToken
15821564
if (
15831565
row?.accessToken &&
15841566
row.accessTokenExpiresAt &&
@@ -1592,24 +1574,14 @@ async function performCoalescedRefresh({
15921574
? revokedOutcome(revokedByLeader)
15931575
: null
15941576
} catch (error) {
1595-
followerSawRefreshToken = undefined
15961577
logger.warn('Follower DB read failed during refresh poll', {
15971578
...logContext,
15981579
...(privacyMode === 'selector' ? {} : { error: toError(error).message }),
15991580
})
16001581
return null
16011582
}
16021583
},
1603-
})
1604-
if (outcome || led) return outcome
1605-
/**
1606-
* A follower's wait ran out with nothing stored. That counts against the chain unless the
1607-
* last read still found the chain this caller read and no leader flagged it dead meanwhile.
1608-
*/
1609-
const chainStanding =
1610-
followerSawRefreshToken === refreshToken && !(await getRecentTerminalError(scopeKey))
1611-
return chainStanding ? { kind: 'lease-held' } : null
1612-
},
1584+
}),
16131585
REFRESH_SETTLE_TIMEOUT_MS
16141586
)
16151587

@@ -1672,21 +1644,8 @@ async function resolveStoredAccessToken(
16721644
if (fresh?.kind === 'token') {
16731645
return { kind: 'token', accessToken: fresh.accessToken, refreshed: true }
16741646
}
1675-
/**
1676-
* A refresh ahead abandoned with its process (Trigger.dev runs have no keep-alive) holds its
1677-
* lease until the TTL, so a follower can wait it out with nothing; a chain that survives
1678-
* refreshes still serves its stored token if enough validity is left after that wait. Any other
1679-
* failure (a dead flag, a terminal error, a vanished account) says the chain itself is unusable.
1680-
*/
1681-
const orphaned =
1682-
fresh?.kind === 'lease-held' &&
1683-
refreshesAheadOfExpiry(refresh.providerId) &&
1684-
stored.accessTokenExpiresAt != null &&
1685-
remainingValidityMs(stored.accessTokenExpiresAt) > OAUTH_STORED_TOKEN_FALLBACK_MIN_MS
1686-
if (stored.accessToken && (plan.kind === 'refresh-or-reuse' || orphaned)) {
1687-
logger.info(
1688-
`${prefix}${orphaned ? 'Refresh lease still held elsewhere' : 'Refresh unavailable'}; reusing still-valid access token`
1689-
)
1647+
if (plan.kind === 'refresh-or-reuse' && stored.accessToken) {
1648+
logger.info(`${prefix}Refresh unavailable; reusing still-valid access token`)
16901649
return { kind: 'token', accessToken: stored.accessToken, refreshed: false }
16911650
}
16921651
return { kind: 'failed', outcome: fresh }

0 commit comments

Comments
 (0)