[pull] main from activepieces:main

packages/server/sandbox/src/lib/cache/pieces/piece-cache.ts

@@ -1,12 +1,14 @@
 import path from 'path'
-import { ActivepiecesError, ErrorCode } from '@activepieces/core-utils'
+import { ActivepiecesError, ErrorCode, isNil } from '@activepieces/core-utils'
 import { type ApLogger, wideEvent } from '@activepieces/server-utils'
 import { ApEnvironment, EXACT_VERSION_REGEX, PackageType, PiecePackage, PieceType, WorkerToApiContract } from '@activepieces/shared'
 import { SandboxSettings } from '../../types'
 import { cacheUtils } from '../cache-paths'
 import { cacheState, NO_SAVE_GUARD } from '../cache-state'
 import { isValidPackageName } from './piece-installer'
 
+const inFlightPieceReads = new Map<string, Promise<PiecePackage>>()
+
 export const pieceCache = (log: ApLogger, apiClient: WorkerToApiContract, basePath: string, getSettings: () => SandboxSettings) => ({
     async getPiece({ pieceName, pieceVersion, platformId }: PieceCacheKey): Promise<PiecePackage> {
         if (!isValidPackageName(pieceName)) {
@@ -22,40 +24,51 @@ export const pieceCache = (log: ApLogger, apiClient: WorkerToApiContract, basePa
         }
 
         const cacheKey = `${pieceName}-${pieceVersion}-${platformId}`
-        const cache = cacheState(path.join(cacheUtils(basePath).getGlobalCachePiecesPath(), cacheKey))
-
-        const { state, cacheHit } = await cache.getOrSetCache({
-            key: cacheKey,
-            cacheMiss: (_: string) => {
-                const environment = getSettings().ENVIRONMENT
-                if (environment === ApEnvironment.TESTING) {
-                    return true
-                }
-                const devPieces = getSettings().DEV_PIECES
-                if (devPieces.includes(pieceName)) {
-                    return true
-                }
-                return false
-            },
-            installFn: async () => {
-                return wideEvent.timed({
-                    name: 'pieceFetch',
-                    fn: async () => {
-                        const piecePackage = await getPiecePackage({ pieceName, pieceVersion, platformId }, apiClient)
-                        log.info({ piece: { name: pieceName, version: pieceVersion }, platform: { id: platformId } }, 'Cached piece')
-                        return JSON.stringify(piecePackage)
-                    },
-                })
-            },
-            skipSave: NO_SAVE_GUARD,
-        })
-
-        wideEvent.set({ pieceCacheHit: cacheHit })
-
-        return JSON.parse(state as string) as PiecePackage
+        const inFlight = inFlightPieceReads.get(cacheKey)
+        if (!isNil(inFlight)) {
+            return inFlight
+        }
+        const read = readPieceThroughCache({ cacheKey, pieceName, pieceVersion, platformId, log, apiClient, basePath, getSettings })
+            .finally(() => inFlightPieceReads.delete(cacheKey))
+        inFlightPieceReads.set(cacheKey, read)
+        return read
     },
 })
 
+async function readPieceThroughCache({ cacheKey, pieceName, pieceVersion, platformId, log, apiClient, basePath, getSettings }: ReadPieceThroughCacheParams): Promise<PiecePackage> {
+    const cache = cacheState(path.join(cacheUtils(basePath).getGlobalCachePiecesPath(), cacheKey))
+
+    const { state, cacheHit } = await cache.getOrSetCache({
+        key: cacheKey,
+        cacheMiss: (_: string) => {
+            const environment = getSettings().ENVIRONMENT
+            if (environment === ApEnvironment.TESTING) {
+                return true
+            }
+            const devPieces = getSettings().DEV_PIECES
+            if (devPieces.includes(pieceName)) {
+                return true
+            }
+            return false
+        },
+        installFn: async () => {
+            return wideEvent.timed({
+                name: 'pieceFetch',
+                fn: async () => {
+                    const piecePackage = await getPiecePackage({ pieceName, pieceVersion, platformId }, apiClient)
+                    log.info({ piece: { name: pieceName, version: pieceVersion }, platform: { id: platformId } }, 'Cached piece')
+                    return JSON.stringify(piecePackage)
+                },
+            })
+        },
+        skipSave: NO_SAVE_GUARD,
+    })
+
+    wideEvent.set({ pieceCacheHit: cacheHit })
+
+    return JSON.parse(state as string) as PiecePackage
+}
+
 async function getPiecePackage(query: PieceCacheKey, apiClient: WorkerToApiContract): Promise<PiecePackage> {
     const pieceMetadata = await apiClient.getPiece({
         name: query.pieceName,
@@ -104,3 +117,11 @@ type PieceCacheKey = {
     pieceVersion: string
     platformId: string
 }
+
+type ReadPieceThroughCacheParams = PieceCacheKey & {
+    cacheKey: string
+    log: ApLogger
+    apiClient: WorkerToApiContract
+    basePath: string
+    getSettings: () => SandboxSettings
+}

packages/server/sandbox/src/lib/sandbox.ts

@@ -1,19 +1,21 @@
-import { ActivepiecesError, ErrorCode, isNil, tryCatch } from '@activepieces/core-utils'
+import { ActivepiecesError, chunk, ErrorCode, isNil, tryCatch } from '@activepieces/core-utils'
 import { type ApLogger, wideEvent } from '@activepieces/server-utils'
-import { PiecePackage } from '@activepieces/shared'
 import { localExecutionCache } from './cache/local-execution-cache'
 import { createResolver } from './resolver'
 import { createSandboxManager, SandboxManager } from './sandbox-manager'
 import {
-    CodeArtifact,
     ExecuteParams,
     PreWarmSandboxParams,
+    ProvisionInput,
+    Resolver,
     Runtime,
     RuntimeExecutionResult,
     RuntimeExecutorInfo,
     SandboxSettings,
 } from './types'
 
+const PREWARM_RESOLVE_CONCURRENCY = 5
+
 // One box per worker at the destination (concurrency 1), or N independent boxes in the transitional
 // compatibility mode that honors AP_WORKER_CONCURRENCY. Each box is its own manager, holding one
 // in-flight operation at a time; the worker runs one poll loop per box and routes each execute to its
@@ -114,29 +116,19 @@ export function createSandboxRuntime({ concurrency = 1, basePath, getSettings }:
             if (isNil(apiClient) || isNil(publicApiUrl)) {
                 return
             }
+            const startedAt = Date.now()
             const { error } = await tryCatch(async () => {
                 const { flows, platformId, engineToken } = await apiClient.getPrewarmData({
                     workerGroupId: getSettings().WORKER_GROUP_ID,
                     projectWorker: getSettings().PROJECT_WORKER,
                     flow,
                 })
                 const resolver = createResolver({ apiClient, basePath, getSettings, log })
-                const pieces: PiecePackage[] = []
-                const codeSteps: CodeArtifact[] = []
-                for (const flow of flows) {
-                    const { data: resolved, error: flowError } = await tryCatch(() => resolver.resolve({ flow, platformId, publicApiUrl, engineToken }))
-                    if (flowError) {
-                        log.warn({ error: String(flowError), flow: { id: flow.id } }, 'Failed to resolve flow for prewarm')
-                        continue
-                    }
-                    if (resolved.kind !== 'ready') {
-                        continue
-                    }
-                    pieces.push(...resolved.provision.pieces)
-                    codeSteps.push(...resolved.provision.codes)
-                }
+                const provisions = await resolveFlowsForPrewarm({ resolver, flows, platformId, publicApiUrl, engineToken, log })
+                const pieces = provisions.flatMap((provision) => provision.pieces)
+                const codeSteps = provisions.flatMap((provision) => provision.codes)
                 await localExecutionCache(log, basePath, getSettings).provision({ pieces, codeSteps, publicApiUrl, engineToken })
-                log.info({ flowCount: flows.length, pieceCount: pieces.length }, 'Prewarmed sandbox cache')
+                log.info({ flowCount: flows.length, pieceCount: pieces.length, durationMs: Date.now() - startedAt }, 'Prewarmed sandbox cache')
             })
             if (error) {
                 log.warn({ error: String(error) }, 'Cache prewarm failed')
@@ -148,13 +140,39 @@ export function createSandboxRuntime({ concurrency = 1, basePath, getSettings }:
     }
 }
 
+async function resolveFlowsForPrewarm({ resolver, flows, platformId, publicApiUrl, engineToken, log }: ResolveFlowsForPrewarmParams): Promise<ProvisionInput[]> {
+    const provisions: ProvisionInput[] = []
+    for (const batch of chunk(flows, PREWARM_RESOLVE_CONCURRENCY)) {
+        const resolvedBatch = await Promise.all(batch.map(async (flow) => {
+            const { data: resolved, error: flowError } = await tryCatch(() => resolver.resolve({ flow, platformId, publicApiUrl, engineToken }))
+            if (flowError) {
+                log.warn({ error: String(flowError), flow: { id: flow.id } }, 'Failed to resolve flow for prewarm')
+                return null
+            }
+            return resolved.kind === 'ready' ? resolved.provision : null
+        }))
+        provisions.push(...resolvedBatch.filter((provision) => !isNil(provision)))
+    }
+    return provisions
+}
+
+
 function remainingTimeoutInSeconds({ timeoutInSeconds, expiresAt }: { timeoutInSeconds: number, expiresAt?: number }): number {
     if (isNil(expiresAt)) {
         return timeoutInSeconds
     }
     return Math.min(timeoutInSeconds, Math.floor((expiresAt - Date.now()) / 1000))
 }
 
+type ResolveFlowsForPrewarmParams = {
+    resolver: Resolver
+    flows: { id: string, versionId: string, projectId: string }[]
+    platformId: string
+    publicApiUrl: string
+    engineToken: string
+    log: ApLogger
+}
+
 type CreateSandboxRuntimeParams = {
     concurrency?: number
     basePath: string