@platformatic/watt-extra 1.6.2 → 1.6.3-alpha.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@platformatic/watt-extra",
-  "version": "1.6.2",
+  "version": "1.6.3-alpha.0",
   "description": "The Platformatic runtime manager",
   "type": "module",
   "scripts": {
@@ -19,10 +19,10 @@
   },
   "devDependencies": {
     "@fastify/websocket": "^11.1.0",
-    "@platformatic/composer": "^3.22.0",
-    "@platformatic/next": "^3.22.0",
-    "@platformatic/node": "^3.22.0",
-    "@platformatic/service": "^3.22.0",
+    "@platformatic/composer": "^3.23.0",
+    "@platformatic/next": "^3.23.0",
+    "@platformatic/node": "^3.23.0",
+    "@platformatic/service": "^3.23.0",
     "atomic-sleep": "^1.0.0",
     "borp": "^0.21.0",
     "eslint": "9",
@@ -30,16 +30,16 @@
     "fastify-plugin": "^5.0.1",
     "neostandard": "^0.12.0",
     "next": "^16.0.0",
-    "platformatic": "^3.22.0",
+    "platformatic": "^3.23.0",
     "pprof-format": "^2.1.0",
     "why-is-node-running": "^2.3.0"
   },
   "dependencies": {
     "@datadog/pprof": "^5.9.0",
     "@fastify/error": "^4.2.0",
-    "@platformatic/foundation": "^3.22.0",
-    "@platformatic/runtime": "^3.22.0",
-    "@platformatic/wattpm-pprof-capture": "^3.22.0",
+    "@platformatic/foundation": "^3.23.0",
+    "@platformatic/runtime": "^3.23.0",
+    "@platformatic/wattpm-pprof-capture": "^3.23.0",
     "avvio": "^9.1.0",
     "chalk": "^4.1.2",
     "commist": "^3.2.0",

package/plugins/alerts.js

@@ -40,6 +40,7 @@ async function alerts (app, _opts) {
       }
 
       const timestamp = Date.now()
+      const workerId = healthInfo.id
       const serviceId = healthInfo.application
       const healthWithTimestamp = { ...healthInfo, timestamp, service: serviceId }
       delete healthWithTimestamp.healthConfig // we don't need to store this
@@ -111,14 +112,12 @@ async function alerts (app, _opts) {
 
         const alert = await body.json()
 
-        try {
-          await app.sendFlamegraphs({
-            serviceIds: [serviceId],
-            alertId: alert.id
-          })
-        } catch (err) {
+        app.sendFlamegraphs({
+          workerIds: [workerId],
+          alertId: alert.id
+        }).catch(err => {
           app.log.error({ err }, 'Failed to send a flamegraph')
-        }
+        })
       }
     }
 

package/plugins/env.js

@@ -21,6 +21,7 @@ const schema = {
     PLT_FLAMEGRAPHS_INTERVAL_SEC: { type: 'number', default: 60 },
     PLT_FLAMEGRAPHS_ELU_THRESHOLD: { type: 'number', default: 0.4 },
     PLT_FLAMEGRAPHS_GRACE_PERIOD: { type: 'number', default: 3000 },
+    PLT_FLAMEGRAPHS_ATTEMPT_TIMEOUT: { type: 'number', default: 10000 },
     PLT_JWT_EXPIRATION_OFFSET_SEC: { type: 'number', default: 60 },
     PLT_UPDATES_RECONNECT_INTERVAL_SEC: { type: 'number', default: 1 },
     PLT_ELU_HEALTH_SIGNAL_THRESHOLD: { type: 'number', default: 0.8 },

package/plugins/flamegraphs.js

@@ -8,10 +8,13 @@ async function flamegraphs (app, _opts) {
   const flamegraphsIntervalSec = app.env.PLT_FLAMEGRAPHS_INTERVAL_SEC
   const flamegraphsELUThreshold = app.env.PLT_FLAMEGRAPHS_ELU_THRESHOLD
   const flamegraphsGracePeriod = app.env.PLT_FLAMEGRAPHS_GRACE_PERIOD
+  const flamegraphsAttemptTimeout = app.env.PLT_FLAMEGRAPHS_ATTEMPT_TIMEOUT
 
   const durationMillis = parseInt(flamegraphsIntervalSec) * 1000
   const eluThreshold = parseFloat(flamegraphsELUThreshold)
   const gracePeriod = parseInt(flamegraphsGracePeriod)
+  const attemptTimeout = Math.min(parseInt(flamegraphsAttemptTimeout), durationMillis)
+  const maxAttempts = Math.ceil(durationMillis / attemptTimeout) + 1
 
   let workerStartedListener = null
 
@@ -125,13 +128,75 @@ async function flamegraphs (app, _opts) {
     }
   }
 
+  async function getServiceFlamegraph (workerId, profileType, attempt = 1) {
+    const runtime = app.watt.runtime
+
+    try {
+      const profile = await runtime.sendCommandToApplication(workerId, 'getLastProfile', { type: profileType })
+      return profile
+    } catch (err) {
+      if (err.code === 'PLT_PPROF_NO_PROFILE_AVAILABLE') {
+        app.log.info(
+          { workerId, attempt, maxAttempts, attemptTimeout },
+          'No profile available for the service. Waiting for profiling to complete.'
+        )
+        if (attempt <= maxAttempts) {
+          await sleep(attemptTimeout)
+          return getServiceFlamegraph(workerId, profileType, attempt + 1)
+        }
+      } else if (err.code === 'PLT_PPROF_NOT_ENOUGH_ELU') {
+        app.log.info({ workerId }, 'ELU low, CPU profiling not active')
+      } else {
+        app.log.warn({ err, workerId }, 'Failed to get profile from a worker')
+
+        const [serviceId, workerIndex] = workerId.split(':')
+        if (workerIndex) {
+          app.log.warn('Worker not available, trying to get profile from another worker')
+          return getServiceFlamegraph(serviceId, profileType)
+        }
+      }
+    }
+  }
+
+  async function sendServiceFlamegraph (scalerUrl, serviceId, profile, profileType, alertId) {
+    const podId = app.instanceId
+    const url = `${scalerUrl}/pods/${podId}/services/${serviceId}/flamegraph`
+    app.log.info({ serviceId, podId, profileType }, 'Sending flamegraph')
+
+    const query = { profileType }
+    if (alertId) {
+      query.alertId = alertId
+    }
+
+    try {
+      const authHeaders = await app.getAuthorizationHeader()
+      const { statusCode, body } = await request(url, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/octet-stream',
+          ...authHeaders
+        },
+        query,
+        body: profile
+      })
+
+      if (statusCode !== 200) {
+        const error = await body.text()
+        app.log.error({ error }, 'Failed to send flamegraph')
+        throw new Error(`Failed to send flamegraph: ${error}`)
+      }
+    } catch (err) {
+      app.log.warn({ err, serviceId, podId }, 'Failed to send flamegraph from service')
+    }
+  }
+
   app.sendFlamegraphs = async (options = {}) => {
     if (isFlamegraphsDisabled) {
       app.log.info('PLT_DISABLE_FLAMEGRAPHS is set, flamegraphs are disabled')
       return
     }
 
-    let { serviceIds, alertId, profileType = 'cpu' } = options
+    let { workerIds, alertId, profileType = 'cpu' } = options
 
     const scalerUrl = app.instanceConfig?.iccServices?.scaler?.url
     if (!scalerUrl) {
@@ -139,57 +204,22 @@ async function flamegraphs (app, _opts) {
       throw new Error('No scaler URL found in ICC services, cannot send flamegraph')
     }
 
-    const podId = app.instanceId
     const runtime = app.watt.runtime
 
-    if (!serviceIds) {
+    if (!workerIds) {
       const { applications } = await runtime.getApplications()
-      serviceIds = applications.map(app => app.id)
+      workerIds = applications.map(app => `${app.id}:0`)
     }
 
-    const authHeaders = await app.getAuthorizationHeader()
-
-    const uploadPromises = serviceIds.map(async (serviceId) => {
-      try {
-        const profile = await runtime.sendCommandToApplication(serviceId, 'getLastProfile', { type: profileType })
-        if (!profile || !(profile instanceof Uint8Array)) {
-          app.log.error({ serviceId }, 'Failed to get profile from service')
-          return
-        }
-
-        const url = `${scalerUrl}/pods/${podId}/services/${serviceId}/flamegraph`
-
-        app.log.info({ serviceId, podId, profileType }, 'Sending flamegraph')
-
-        const query = { profileType }
-        if (alertId) {
-          query.alertId = alertId
-        }
-
-        const { statusCode, body } = await request(url, {
-          method: 'POST',
-          headers: {
-            'Content-Type': 'application/octet-stream',
-            ...authHeaders
-          },
-          query,
-          body: profile
-        })
-
-        if (statusCode !== 200) {
-          const error = await body.text()
-          app.log.error({ error }, 'Failed to send flamegraph')
-          throw new Error(`Failed to send flamegraph: ${error}`)
-        }
-      } catch (err) {
-        if (err.code === 'PLT_PPROF_NO_PROFILE_AVAILABLE') {
-          app.log.info({ serviceId, podId }, 'No profile available for the service')
-        } else if (err.code === 'PLT_PPROF_NOT_ENOUGH_ELU') {
-          app.log.info({ serviceId, podId }, 'ELU low, CPU profiling not active')
-        } else {
-          app.log.warn({ err, serviceId, podId }, 'Failed to send flamegraph from service')
-        }
+    const uploadPromises = workerIds.map(async (workerId) => {
+      const profile = await getServiceFlamegraph(workerId, profileType)
+      if (!profile || !(profile instanceof Uint8Array)) {
+        app.log.error({ workerId }, 'Failed to get profile from service')
+        return
       }
+
+      const serviceId = workerId.split(':')[0]
+      await sendServiceFlamegraph(scalerUrl, serviceId, profile, profileType, alertId)
     })
 
     await Promise.all(uploadPromises)

package/plugins/health-signals.js

@@ -80,6 +80,7 @@ async function healthSignals (app, _opts) {
       }
 
       const {
+        id: workerId,
         application: serviceId,
         currentHealth,
         healthSignals
@@ -125,13 +126,13 @@ async function healthSignals (app, _opts) {
       }
 
       if (healthSignals.length > 0) {
-        await sendHealthSignalsWithTimeout(serviceId, healthSignals)
+        await sendHealthSignalsWithTimeout(serviceId, workerId, healthSignals)
       }
     })
   }
   app.setupHealthSignals = setupHealthSignals
 
-  async function sendHealthSignalsWithTimeout (serviceId, signals) {
+  async function sendHealthSignalsWithTimeout (serviceId, workerId, signals) {
     signalsCaches[serviceId] ??= new HealthSignalsCache()
     servicesSendingStatuses[serviceId] ??= false
 
@@ -148,7 +149,7 @@ async function healthSignals (app, _opts) {
 
         try {
           const signals = signalsCache.getAll()
-          await sendHealthSignals(serviceId, signals, metrics)
+          await sendHealthSignals(serviceId, workerId, signals, metrics)
         } catch (err) {
           app.log.error({ err }, 'Failed to send health signals to scaler')
         }
@@ -156,7 +157,7 @@ async function healthSignals (app, _opts) {
     }
   }
 
-  async function sendHealthSignals (serviceId, signals, metrics) {
+  async function sendHealthSignals (serviceId, workerId, signals, metrics) {
     const scalerUrl = app.instanceConfig?.iccServices?.scaler?.url
     const applicationId = app.instanceConfig?.applicationId
     const authHeaders = await app.getAuthorizationHeader()
@@ -184,11 +185,13 @@ async function healthSignals (app, _opts) {
 
     const alert = await body.json()
 
-    try {
-      await app.sendFlamegraphs({ serviceIds: [serviceId], alertId: alert.id })
-    } catch (err) {
+    app.sendFlamegraphs({
+      serviceIds: [serviceId],
+      workerIds: [workerId],
+      alertId: alert.id
+    }).catch(err => {
       app.log.error({ err }, 'Failed to send a flamegraph')
-    }
+    })
   }
 }
 

package/plugins/update.js

@@ -23,14 +23,14 @@ async function updatePlugin (app) {
       // Handle trigger-flamegraph command from ICC
       if (command === 'trigger-flamegraph') {
         app.log.info({ command }, 'Received trigger-flamegraph command from ICC')
-        await app.sendFlamegraphs({ profileType: 'cpu' })
+        app.sendFlamegraphs({ profileType: 'cpu' })
         return
       }
 
       // Handle trigger-heapprofile command from ICC
       if (command === 'trigger-heapprofile') {
         app.log.info({ command }, 'Received trigger-heapprofile command from ICC')
-        await app.sendFlamegraphs({ profileType: 'heap' })
+        app.sendFlamegraphs({ profileType: 'heap' })
         return
       }
 

package/test/trigger-flamegraphs.test.js

@@ -35,7 +35,7 @@ function setupMockIccServer (wss, receivedMessages, validateAuth = false) {
   return { waitForClientSubscription, getWs: () => ws }
 }
 
-function createMockApp (port, includeScalerUrl = true) {
+function createMockApp (port, includeScalerUrl = true, env = {}) {
   const eventListeners = new Map()
 
   const mockWatt = {
@@ -100,7 +100,9 @@ function createMockApp (port, includeScalerUrl = true) {
       PLT_DISABLE_FLAMEGRAPHS: false,
       PLT_FLAMEGRAPHS_INTERVAL_SEC: 1,
       PLT_FLAMEGRAPHS_ELU_THRESHOLD: 0,
-      PLT_FLAMEGRAPHS_GRACE_PERIOD: 0
+      PLT_FLAMEGRAPHS_GRACE_PERIOD: 0,
+      PLT_FLAMEGRAPHS_ATTEMPT_TIMEOUT: 1000,
+      ...env
     },
     watt: mockWatt
   }
@@ -333,15 +335,15 @@ test('sendFlamegraphs should handle missing profile data', async (t) => {
   equal(errors.length, 2, 'Should log errors for both services with missing profiles')
 })
 
-test('sendFlamegraphs should filter by serviceIds when provided', async (t) => {
+test('sendFlamegraphs should filter by workerIds when provided', async (t) => {
   setUpEnvironment()
 
   const app = createMockApp(port + 12)
   const getProfileCalls = []
 
-  app.watt.runtime.sendCommandToApplication = async (serviceId, command) => {
+  app.watt.runtime.sendCommandToApplication = async (workerId, command) => {
     if (command === 'getLastProfile') {
-      getProfileCalls.push(serviceId)
+      getProfileCalls.push(workerId)
       return new Uint8Array([1, 2, 3])
     }
     return { success: false }
@@ -362,10 +364,49 @@ test('sendFlamegraphs should filter by serviceIds when provided', async (t) => {
   t.after(() => server.close())
 
   await flamegraphsPlugin(app)
-  await app.sendFlamegraphs({ serviceIds: ['service-1'] })
+  await app.sendFlamegraphs({ workerIds: ['service-1:0'] })
 
   equal(getProfileCalls.length, 1, 'Should only request profile for specified service')
-  equal(getProfileCalls[0], 'service-1', 'Should request profile for service-1')
+  equal(getProfileCalls[0], 'service-1:0', 'Should request profile for service-1')
+})
+
+test('sendFlamegraphs should try to get the profile from a service if worker is not available', async (t) => {
+  setUpEnvironment()
+
+  const app = createMockApp(port + 12)
+  const getProfileCalls = []
+
+  app.watt.runtime.sendCommandToApplication = async (workerId, command) => {
+    if (command === 'getLastProfile') {
+      getProfileCalls.push(workerId)
+      if (workerId === 'service-1:2') {
+        throw new Error('Worker not available')
+      }
+      return new Uint8Array([1, 2, 3])
+    }
+    return { success: false }
+  }
+
+  // Mock HTTP server
+  const { createServer } = await import('node:http')
+  const server = createServer((req, res) => {
+    const body = []
+    req.on('data', chunk => body.push(chunk))
+    req.on('end', () => {
+      res.writeHead(200)
+      res.end()
+    })
+  })
+
+  await new Promise(resolve => server.listen(port + 12, resolve))
+  t.after(() => server.close())
+
+  await flamegraphsPlugin(app)
+  await app.sendFlamegraphs({ workerIds: ['service-1:2'] })
+
+  equal(getProfileCalls.length, 2)
+  equal(getProfileCalls[0], 'service-1:2')
+  equal(getProfileCalls[1], 'service-1')
 })
 
 test('sendFlamegraphs should skip when PLT_DISABLE_FLAMEGRAPHS is set', async (t) => {
@@ -376,9 +417,9 @@ test('sendFlamegraphs should skip when PLT_DISABLE_FLAMEGRAPHS is set', async (t
 
   const getProfileCalls = []
 
-  app.watt.runtime.sendCommandToApplication = async (serviceId, command) => {
+  app.watt.runtime.sendCommandToApplication = async (workerId, command) => {
     if (command === 'getLastProfile') {
-      getProfileCalls.push(serviceId)
+      getProfileCalls.push(workerId)
       return new Uint8Array([1, 2, 3])
     }
     return { success: false }
@@ -465,14 +506,14 @@ test('should handle trigger-flamegraph command and upload flamegraphs from servi
   equal(getFlamegraphReqs.length, 2)
 
   const service1Req = getFlamegraphReqs.find(
-    (f) => f.serviceId === 'service-1'
+    (f) => f.serviceId === 'service-1:0'
   )
   const service2Req = getFlamegraphReqs.find(
-    (f) => f.serviceId === 'service-2'
+    (f) => f.serviceId === 'service-2:0'
   )
 
-  equal(service1Req.serviceId, 'service-1')
-  equal(service2Req.serviceId, 'service-2')
+  equal(service1Req.serviceId, 'service-1:0')
+  equal(service2Req.serviceId, 'service-2:0')
 
   if (app.cleanupFlamegraphs) app.cleanupFlamegraphs()
   await app.closeUpdates()
@@ -622,14 +663,14 @@ test('should handle trigger-heapprofile command and upload heap profiles from se
   equal(getHeapProfileReqs.length, 2)
 
   const service1Req = getHeapProfileReqs.find(
-    (f) => f.serviceId === 'service-1'
+    (f) => f.serviceId === 'service-1:0'
   )
   const service2Req = getHeapProfileReqs.find(
-    (f) => f.serviceId === 'service-2'
+    (f) => f.serviceId === 'service-2:0'
   )
 
-  equal(service1Req.serviceId, 'service-1')
-  equal(service2Req.serviceId, 'service-2')
+  equal(service1Req.serviceId, 'service-1:0')
+  equal(service2Req.serviceId, 'service-2:0')
 
   if (app.cleanupFlamegraphs) app.cleanupFlamegraphs()
   await app.closeUpdates()
@@ -640,11 +681,6 @@ test('should handle PLT_PPROF_NO_PROFILE_AVAILABLE error with info log', async (
 
   const receivedMessages = []
   const infoLogs = []
-  let errorCount = 0
-  let uploadResolve
-  const allUploadsComplete = new Promise((resolve) => {
-    uploadResolve = resolve
-  })
 
   const wss = new WebSocketServer({ port: port + 4 })
   t.after(async () => wss.close())
@@ -655,19 +691,21 @@ test('should handle PLT_PPROF_NO_PROFILE_AVAILABLE error with info log', async (
     true
   )
 
-  const app = createMockApp(port + 4)
+  const app = createMockApp(port + 4, true, {
+    PLT_FLAMEGRAPHS_INTERVAL_SEC: 10,
+    PLT_FLAMEGRAPHS_ATTEMPT_TIMEOUT: 1000
+  })
+
   const originalInfo = app.log.info
   app.log.info = (...args) => {
     originalInfo(...args)
-    if (args[1] && args[1].includes('No profile available for the service')) {
-      infoLogs.push(args)
-      errorCount++
-      if (errorCount === 2) {
-        uploadResolve()
-      }
-    }
+    infoLogs.push(args)
   }
 
+  // Profile will be generated in 10s
+  const profileGenerationDate = Date.now() + 10000
+  const mockProfile = new Uint8Array([1, 2, 3, 4, 5])
+
   app.watt.runtime.sendCommandToApplication = async (
     serviceId,
     command
@@ -676,9 +714,13 @@ test('should handle PLT_PPROF_NO_PROFILE_AVAILABLE error with info log', async (
       return { success: true }
     }
     if (command === 'getLastProfile') {
-      const error = new Error('No profile available - wait for profiling to complete or trigger manual capture')
-      error.code = 'PLT_PPROF_NO_PROFILE_AVAILABLE'
-      throw error
+      const now = Date.now()
+      if (now < profileGenerationDate) {
+        const error = new Error('No profile available - wait for profiling to complete or trigger manual capture')
+        error.code = 'PLT_PPROF_NO_PROFILE_AVAILABLE'
+        throw error
+      }
+      return mockProfile
     }
     return { success: false }
   }
@@ -689,6 +731,13 @@ test('should handle PLT_PPROF_NO_PROFILE_AVAILABLE error with info log', async (
   await app.connectToUpdates()
   await app.setupFlamegraphs()
 
+  t.after(async () => {
+    if (app.cleanupFlamegraphs) {
+      app.cleanupFlamegraphs()
+    }
+    await app.closeUpdates()
+  })
+
   await waitForClientSubscription
 
   const triggerFlamegraphMessage = {
@@ -697,15 +746,47 @@ test('should handle PLT_PPROF_NO_PROFILE_AVAILABLE error with info log', async (
 
   getWs().send(JSON.stringify(triggerFlamegraphMessage))
 
-  await allUploadsComplete
+  await sleep(15000)
 
-  equal(infoLogs.length, 2)
-  equal(infoLogs[0][0].serviceId, 'service-1')
-  equal(infoLogs[0][0].podId, 'test-pod-123')
-  equal(infoLogs[0][1], 'No profile available for the service')
+  const service1AttemptLogs = []
+  const service2AttemptLogs = []
+  const service1SuccessLogs = []
+  const service2SuccessLogs = []
 
-  if (app.cleanupFlamegraphs) app.cleanupFlamegraphs()
-  await app.closeUpdates()
+  for (const infoLog of infoLogs) {
+    if (infoLog.length !== 2) continue
+    const [options, message] = infoLog
+
+    if (message.includes('No profile available for the service')) {
+      const { workerId, attempt, maxAttempts, attemptTimeout } = options
+
+      equal(maxAttempts, 11)
+      equal(attemptTimeout, 1000)
+
+      if (workerId === 'service-1:0') {
+        service1AttemptLogs.push(infoLog)
+        equal(attempt, service1AttemptLogs.length)
+      }
+      if (workerId === 'service-2:0') {
+        service2AttemptLogs.push(infoLog)
+        equal(attempt, service2AttemptLogs.length)
+      }
+      continue
+    }
+
+    if (message.includes('Sending flamegraph')) {
+      if (options.serviceId === 'service-1') {
+        service1SuccessLogs.push(infoLog)
+      } else if (options.serviceId === 'service-2') {
+        service2SuccessLogs.push(infoLog)
+      }
+    }
+  }
+
+  equal(service1AttemptLogs.length, 10)
+  equal(service2AttemptLogs.length, 10)
+  equal(service1SuccessLogs.length, 1)
+  equal(service2SuccessLogs.length, 1)
 })
 
 test('should handle PLT_PPROF_NOT_ENOUGH_ELU error with info log', async (t) => {
@@ -762,6 +843,13 @@ test('should handle PLT_PPROF_NOT_ENOUGH_ELU error with info log', async (t) =>
   await app.connectToUpdates()
   await app.setupFlamegraphs()
 
+  t.after(async () => {
+    if (app.cleanupFlamegraphs) {
+      app.cleanupFlamegraphs()
+    }
+    await app.closeUpdates()
+  })
+
   await waitForClientSubscription
 
   const triggerFlamegraphMessage = {
@@ -773,12 +861,8 @@ test('should handle PLT_PPROF_NOT_ENOUGH_ELU error with info log', async (t) =>
   await allUploadsComplete
 
   equal(infoLogs.length, 2)
-  equal(infoLogs[0][0].serviceId, 'service-1')
-  equal(infoLogs[0][0].podId, 'test-pod-123')
+  equal(infoLogs[0][0].workerId, 'service-1:0')
   equal(infoLogs[0][1], 'ELU low, CPU profiling not active')
-
-  if (app.cleanupFlamegraphs) app.cleanupFlamegraphs()
-  await app.closeUpdates()
 })
 
 test('should start profiling on new workers that start after initial setup', async (t) => {
@@ -815,6 +899,13 @@ test('should start profiling on new workers that start after initial setup', asy
   await app.connectToUpdates()
   await app.setupFlamegraphs()
 
+  t.after(async () => {
+    if (app.cleanupFlamegraphs) {
+      app.cleanupFlamegraphs()
+    }
+    await app.closeUpdates()
+  })
+
   await waitForClientSubscription
 
   equal(startProfilingCalls.length, 4)
@@ -844,9 +935,6 @@ test('should start profiling on new workers that start after initial setup', asy
   equal(startProfilingCalls[5].options.durationMillis, 1000)
   equal(startProfilingCalls[5].options.eluThreshold, 0)
   equal(startProfilingCalls[5].options.type, 'heap')
-
-  if (app.cleanupFlamegraphs) app.cleanupFlamegraphs()
-  await app.closeUpdates()
 })
 
 test('should not start profiling on new workers when flamegraphs are disabled', async (t) => {
@@ -884,6 +972,13 @@ test('should not start profiling on new workers when flamegraphs are disabled',
   await app.connectToUpdates()
   await app.setupFlamegraphs()
 
+  t.after(async () => {
+    if (app.cleanupFlamegraphs) {
+      app.cleanupFlamegraphs()
+    }
+    await app.closeUpdates()
+  })
+
   await waitForClientSubscription
 
   equal(startProfilingCalls.length, 0)
@@ -897,9 +992,6 @@ test('should not start profiling on new workers when flamegraphs are disabled',
   await sleep(10)
 
   equal(startProfilingCalls.length, 0)
-
-  if (app.cleanupFlamegraphs) app.cleanupFlamegraphs()
-  await app.closeUpdates()
 })
 
 test('sendFlamegraphs should include alertId in query params when provided', async (t) => {