diff --git a/packages/services/service-common/src/iam-aws.ts b/packages/services/service-common/src/iam-aws.ts index c03b0d1aac..df69274d88 100644 --- a/packages/services/service-common/src/iam-aws.ts +++ b/packages/services/service-common/src/iam-aws.ts @@ -1,3 +1,4 @@ +import { setTimeout as sleep } from 'node:timers/promises'; import { Sha256 } from '@aws-crypto/sha256-js'; import { fromNodeProviderChain } from '@aws-sdk/credential-providers'; import { formatUrl } from '@aws-sdk/util-format-url'; @@ -114,45 +115,51 @@ export function startTokenRefreshTimer( const retryBackoffMs = options?.retryBackoffMs ?? 5_000; // Track shutdown state to prevent in-flight callbacks and retries during graceful shutdown - let isShuttingDown = false; + const abortController = new AbortController(); let currentTimer: ReturnType | undefined; function scheduleNext() { - if (isShuttingDown) return; - + if (abortController.signal.aborted) return; const jitterMs = Math.floor(Math.random() * maxJitterMs); + currentTimer = setTimeout(() => { - if (isShuttingDown) return; + if (abortController.signal.aborted) return; void (async () => { for (let attempt = 1; attempt <= maxRetries; attempt++) { - if (isShuttingDown) return; + if (abortController.signal.aborted) return; try { await refreshFn(attempt, maxRetries); break; } catch { - if (attempt < maxRetries) { - await new Promise(r => setTimeout(r, attempt * retryBackoffMs)); + if (attempt >= maxRetries) { + return; + } + } + + try { + await sleep(attempt * retryBackoffMs, { + signal: abortController.signal, + }); + } catch (err) { + if (abortController.signal.aborted) { + return; } + throw err; } } // Schedule next cycle only after current refresh (including retries) completes scheduleNext(); })(); }, intervalMs + jitterMs); - - // Allow clean process shutdown even if the refresh timer is still active. - if (typeof currentTimer.unref === 'function') { - currentTimer.unref(); - } } scheduleNext(); // Return a cleanup function that sets shutdown flag and clears the pending timer return () => { - isShuttingDown = true; + abortController.abort(); if (currentTimer !== undefined) { clearTimeout(currentTimer); } diff --git a/packages/services/service-common/src/redis-client.ts b/packages/services/service-common/src/redis-client.ts index 46bce572c1..c500b23c4b 100644 --- a/packages/services/service-common/src/redis-client.ts +++ b/packages/services/service-common/src/redis-client.ts @@ -94,7 +94,8 @@ export async function createRedisClient( // Start IAM token refresh if enabled // Each client gets its own independent token lifecycle if (redisIamConfig) { - startIamTokenRefresh(redis, redisIamConfig, options.logger); + const stopTokenRefresh = startIamTokenRefresh(redis, redisIamConfig, options.logger); + redis.on('end', stopTokenRefresh); } return redis; diff --git a/packages/services/workflows/src/index.ts b/packages/services/workflows/src/index.ts index 66140317e4..5fa3992295 100644 --- a/packages/services/workflows/src/index.ts +++ b/packages/services/workflows/src/index.ts @@ -215,7 +215,7 @@ registerShutdown({ await pg.end(); logger.info('Shutdown postgres connection successful.'); logger.info('Shutdown redis connection.'); - redis.disconnect(false); + await Promise.all([redis.quit(), redisSubscriber.quit()]); if (shutdownMetrics) { logger.info('Stopping prometheus endpoint'); await shutdownMetrics();