Skip to content

Commit 9f3d243

Browse files
committed
refactor(execution): tighten removal stops after final audit
- disableUserResources reads owned workspaces once and stops earlier-archived ones up front, dropping the post-archive read that could fail the ban - Trigger.dev cancelByWorkspace reuses scanAndCancelTriggerRuns; workspace tag shares one builder with enqueue - database queue shares abortInlineJob with cancelJob and returns completeness from control flow - removal stop handler parses its payload with zod and signals while it sweeps - drop the unused event-id return and redundant bucketing state
1 parent 1feebd6 commit 9f3d243

5 files changed

Lines changed: 84 additions & 124 deletions

File tree

‎apps/sim/lib/core/async-jobs/backends/database.ts‎

Lines changed: 15 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,15 @@ function rowToJob(row: AsyncJobRow): Job {
4949
const inlineAbortControllers = new Map<string, AbortController>()
5050
const inlineExecutionControllers = new Map<string, Map<AbortController, JobType>>()
5151

52+
/** Aborts the runner this process holds for the job, if any. */
53+
function abortInlineJob(jobId: string): boolean {
54+
const controller = inlineAbortControllers.get(jobId)
55+
if (!controller) return false
56+
controller.abort(new DOMException('user', 'AbortError'))
57+
inlineAbortControllers.delete(jobId)
58+
return true
59+
}
60+
5261
/**
5362
* Per-cancel-key abort controllers for the `batchEnqueueAndWait` direct-call
5463
* path. Distinct from `inlineAbortControllers` (which keys by jobId) — this
@@ -418,13 +427,7 @@ export class DatabaseJobQueue implements JobQueueBackend {
418427
// Abort any in-process inline execution first so the running workflow
419428
// observes the signal and stops mid-flight. Then mark the row cancelled so
420429
// any future poller skips it without treating a user action as a failure.
421-
const controller = inlineAbortControllers.get(jobId)
422-
let aborted = false
423-
if (controller) {
424-
controller.abort(new DOMException('user', 'AbortError'))
425-
inlineAbortControllers.delete(jobId)
426-
aborted = true
427-
}
430+
const aborted = abortInlineJob(jobId)
428431

429432
const now = new Date()
430433
await db
@@ -532,7 +535,6 @@ export class DatabaseJobQueue implements JobQueueBackend {
532535
sql`${asyncJobs.metadata}->>'workspaceId' = ${workspaceId}`
533536
)
534537
let cancelledJobs = 0
535-
let complete = false
536538
while (cancelledJobs < options.maxRuns && !options.signal?.aborted) {
537539
const batchSize = Math.min(WORKSPACE_CANCELLATION_BATCH_SIZE, options.maxRuns - cancelledJobs)
538540
const now = new Date()
@@ -550,21 +552,16 @@ export class DatabaseJobQueue implements JobQueueBackend {
550552
)
551553
.returning({ id: asyncJobs.id })
552554

553-
for (const { id } of cancelled) {
554-
const controller = inlineAbortControllers.get(id)
555-
if (!controller) continue
556-
controller.abort(new DOMException('user', 'AbortError'))
557-
inlineAbortControllers.delete(id)
558-
}
555+
for (const { id } of cancelled) abortInlineJob(id)
559556
cancelledJobs += cancelled.length
560557
if (cancelled.length < batchSize) {
561-
complete = true
562-
break
558+
logger.info('Cancelled database jobs for workspace', { workspaceId, cancelledJobs })
559+
return { cancelledJobs, complete: true }
563560
}
564561
}
565562

566-
logger.info('Cancelled database jobs for workspace', { workspaceId, cancelledJobs, complete })
567-
return { cancelledJobs, complete }
563+
logger.info('Paused database job sweep for workspace', { workspaceId, cancelledJobs })
564+
return { cancelledJobs, complete: false }
568565
}
569566

570567
cancelByKey(cancelKey: string): boolean {

‎apps/sim/lib/core/async-jobs/backends/trigger-dev.ts‎

Lines changed: 34 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,10 @@ function buildExecutionTag(executionId: string): string {
127127
return tag.length <= 128 ? tag : `executionIdHash:${sha256Hex(executionId)}`
128128
}
129129

130+
function buildWorkspaceTag(workspaceId: string): string {
131+
return `workspaceId:${workspaceId}`
132+
}
133+
130134
function buildWorkflowTag(workflowId: string): string {
131135
const tag = `workflowId:${workflowId}`
132136
return tag.length <= 128 ? tag : `workflowIdHash:${sha256Hex(workflowId)}`
@@ -171,10 +175,13 @@ interface CancellationListRun {
171175
}
172176

173177
interface CancellationScanOptions {
174-
binding: ExecutionJobBinding
178+
/** Required by any candidate that sets `verifyPayload`. */
179+
binding?: ExecutionJobBinding
175180
cancelJob: (jobId: string) => Promise<void>
176181
listRuns: () => AsyncIterable<CancellationListRun>
177182
selectCandidate: (run: CancellationListRun) => CancellationCandidate | undefined
183+
/** Checked before each listed run; true ends the scan there. */
184+
shouldStop?: () => boolean
178185
state: CancellationScanState
179186
}
180187

@@ -185,7 +192,7 @@ function recordCancellationCandidateFailure(state: CancellationScanState, error:
185192

186193
async function processCancellationCandidateBatch(
187194
candidates: readonly CancellationCandidate[],
188-
binding: ExecutionJobBinding,
195+
binding: ExecutionJobBinding | undefined,
189196
state: CancellationScanState,
190197
cancelJob: (jobId: string) => Promise<void>
191198
): Promise<void> {
@@ -199,7 +206,7 @@ async function processCancellationCandidateBatch(
199206
uniqueCandidates.map(async (candidate) => {
200207
if (candidate.verifyPayload) {
201208
const legacyRun = await runs.retrieve(candidate.id)
202-
if (!payloadMatchesExecution(legacyRun.payload, binding)) return false
209+
if (!binding || !payloadMatchesExecution(legacyRun.payload, binding)) return false
203210
}
204211
await cancelJob(candidate.id)
205212
return true
@@ -231,6 +238,7 @@ async function scanAndCancelTriggerRuns(options: CancellationScanOptions): Promi
231238

232239
try {
233240
for await (const run of options.listRuns()) {
241+
if (options.shouldStop?.()) break
234242
const candidate = options.selectCandidate(run)
235243
if (!candidate) continue
236244
candidates.push(candidate)
@@ -659,44 +667,33 @@ export class TriggerDevJobQueue implements JobQueueBackend {
659667
workspaceId: string,
660668
options: WorkspaceJobCancellationOptions
661669
): Promise<WorkspaceJobCancellationResult> {
662-
const workspaceTag = `workspaceId:${workspaceId}`
670+
const workspaceTag = buildWorkspaceTag(workspaceId)
663671
const state: CancellationScanState = { cancelledJobs: 0, failureCount: 0 }
664-
let pending: string[] = []
665672
let inspectedRuns = 0
666673
let complete = true
667674

668-
const flush = async (): Promise<void> => {
669-
const runIds = pending
670-
pending = []
671-
const results = await Promise.allSettled(runIds.map((runId) => this.cancelJob(runId)))
672-
for (const result of results) {
673-
if (result.status === 'rejected') recordCancellationCandidateFailure(state, result.reason)
674-
else state.cancelledJobs += 1
675-
}
676-
}
677-
678-
try {
679-
for await (const run of runs.list({
680-
tag: workspaceTag,
681-
taskIdentifier: [...WORKSPACE_EXECUTION_JOB_TYPES],
682-
status: [...ACTIVE_TRIGGER_RUN_STATUSES],
683-
from: new Date(Date.now() - JOB_PENDING_RETENTION_HOURS * 60 * 60 * 1000),
684-
limit: WORKSPACE_CANCELLATION_LIST_PAGE_SIZE,
685-
})) {
686-
if (inspectedRuns >= options.maxRuns || options.signal?.aborted) {
687-
complete = false
688-
break
689-
}
690-
if (!run.tags.includes(workspaceTag)) continue
675+
await scanAndCancelTriggerRuns({
676+
cancelJob: (jobId) => this.cancelJob(jobId),
677+
listRuns: () =>
678+
runs.list({
679+
tag: workspaceTag,
680+
taskIdentifier: [...WORKSPACE_EXECUTION_JOB_TYPES],
681+
status: [...ACTIVE_TRIGGER_RUN_STATUSES],
682+
from: new Date(Date.now() - JOB_PENDING_RETENTION_HOURS * 60 * 60 * 1000),
683+
limit: WORKSPACE_CANCELLATION_LIST_PAGE_SIZE,
684+
}),
685+
shouldStop: () => {
686+
if (inspectedRuns < options.maxRuns && !options.signal?.aborted) return false
687+
complete = false
688+
return true
689+
},
690+
selectCandidate: (run) => {
691+
if (!run.tags.includes(workspaceTag)) return undefined
691692
inspectedRuns += 1
692-
pending.push(run.id)
693-
if (pending.length >= CANCELLATION_OPERATION_CHUNK_SIZE) await flush()
694-
}
695-
} catch (error) {
696-
recordCancellationCandidateFailure(state, error)
697-
} finally {
698-
await flush()
699-
}
693+
return { id: run.id, verifyPayload: false }
694+
},
695+
state,
696+
})
700697

701698
if (state.firstError) {
702699
logger.error('Failed to cancel trigger.dev runs for workspace', {
@@ -742,7 +739,7 @@ function buildTags(options?: EnqueueOptions): string[] {
742739
const workflowId =
743740
meta?.correlation?.workflowId ??
744741
(typeof meta?.workflowId === 'string' ? meta.workflowId : undefined)
745-
if (meta?.workspaceId) tags.push(`workspaceId:${meta.workspaceId}`)
742+
if (meta?.workspaceId) tags.push(buildWorkspaceTag(meta.workspaceId))
746743
if (workflowId) tags.push(buildWorkflowTag(workflowId))
747744
if (meta?.userId) tags.push(`userId:${meta.userId}`)
748745

‎apps/sim/lib/execution/removal-stop-records.ts‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ export interface ExecutionStopPlan {
4343
executions: PlannedExecutionStop[]
4444
}
4545

46-
export interface ExecutionStopPayload {
46+
interface ExecutionStopPayload {
4747
/** Workspace whose queued and running jobs this event sweeps. */
4848
workspaceId?: string
4949
executions: { executionId: string; executionDeadlineAt: string | null }[]
@@ -89,7 +89,6 @@ export async function collectExecutionStops(
8989
})
9090
.from(workflowExecutionLogs)
9191
.where(and(eq(workflowExecutionLogs.status, 'running'), inScope))
92-
/** Driven from the scope's log rows, so `resume_queue` is read through its parent index. */
9392
const resuming = await executor
9493
.select({
9594
executionId: resumeQueue.newExecutionId,
@@ -142,17 +141,14 @@ function toPayloadExecutions(
142141

143142
/** One event per removed workspace carries its sweep; executions beyond one event's share spill over. */
144143
function buildPayloads(plan: ExecutionStopPlan): ExecutionStopPayload[] {
145-
const removedWorkspaces = new Set(plan.workspaceIds)
146144
const byWorkspace = new Map<string, PlannedExecutionStop[]>(
147145
plan.workspaceIds.map((workspaceId) => [workspaceId, []])
148146
)
149147
const elsewhere: PlannedExecutionStop[] = []
150148
for (const execution of plan.executions) {
151-
if (execution.workspaceId && removedWorkspaces.has(execution.workspaceId)) {
152-
byWorkspace.get(execution.workspaceId)?.push(execution)
153-
} else {
154-
elsewhere.push(execution)
155-
}
149+
const removed = execution.workspaceId ? byWorkspace.get(execution.workspaceId) : undefined
150+
if (removed) removed.push(execution)
151+
else elsewhere.push(execution)
156152
}
157153

158154
const payloads: ExecutionStopPayload[] = []
@@ -174,12 +170,10 @@ function buildPayloads(plan: ExecutionStopPlan): ExecutionStopPayload[] {
174170
export async function enqueueExecutionStops(
175171
executor: Pick<typeof db, 'insert'>,
176172
plan: ExecutionStopPlan
177-
): Promise<string[]> {
178-
const eventIds: string[] = []
173+
): Promise<void> {
179174
for (const chunk of chunkArray(buildPayloads(plan), MAX_BULK_ENQUEUE_EVENTS)) {
180-
eventIds.push(...(await enqueueOutboxEvents(executor, EXECUTION_REMOVAL_STOP_EVENT, chunk)))
175+
await enqueueOutboxEvents(executor, EXECUTION_REMOVAL_STOP_EVENT, chunk)
181176
}
182-
return eventIds
183177
}
184178

185179
/**

‎apps/sim/lib/execution/removal-stops.ts‎

Lines changed: 22 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { db } from '@sim/db'
2-
import { isRecordLike } from '@sim/utils/object'
2+
import { z } from 'zod'
33
import { getInlineJobQueue, getJobQueue } from '@/lib/core/async-jobs'
44
import {
55
continueOutboxHandler,
@@ -8,40 +8,24 @@ import {
88
withOutboxHandlerTimeout,
99
} from '@/lib/core/outbox/service'
1010
import { EXECUTION_REMOVAL_STOP_EVENT } from '@/lib/execution/removal-stop-event'
11-
import {
12-
collectExecutionStops,
13-
type ExecutionStopPayload,
14-
signalExecutionStops,
15-
} from '@/lib/execution/removal-stop-records'
11+
import { collectExecutionStops, signalExecutionStops } from '@/lib/execution/removal-stop-records'
1612

1713
/** Queue runs one handler attempt cancels before handing the rest to its next run. */
1814
const SWEEP_MAX_RUNS_PER_ATTEMPT = 1_000
1915
const HANDLER_TIMEOUT_MS = 5 * 60 * 1000
2016

21-
function parsePayload(payload: unknown): ExecutionStopPayload {
22-
if (!isRecordLike(payload) || !Array.isArray(payload.executions)) {
23-
throw new Error('Execution removal stop payload is malformed')
24-
}
25-
let workspaceId: string | undefined
26-
if (payload.workspaceId !== undefined) {
27-
if (typeof payload.workspaceId !== 'string' || payload.workspaceId.length === 0) {
28-
throw new Error('Execution removal stop payload has an invalid workspaceId')
29-
}
30-
workspaceId = payload.workspaceId
31-
}
32-
const executions = payload.executions.map((entry: unknown) => {
33-
if (
34-
!isRecordLike(entry) ||
35-
typeof entry.executionId !== 'string' ||
36-
entry.executionId.length === 0 ||
37-
(entry.executionDeadlineAt !== null && typeof entry.executionDeadlineAt !== 'string')
38-
) {
39-
throw new Error('Execution removal stop payload has an invalid execution')
40-
}
41-
return { executionId: entry.executionId, executionDeadlineAt: entry.executionDeadlineAt }
42-
})
43-
return { ...(workspaceId ? { workspaceId } : {}), executions }
44-
}
17+
const payloadSchema = z.object({
18+
workspaceId: z.string().min(1).optional(),
19+
executions: z.array(
20+
z.object({
21+
executionId: z.string().min(1),
22+
executionDeadlineAt: z
23+
.string()
24+
.nullable()
25+
.transform((deadline) => (deadline ? new Date(deadline) : null)),
26+
})
27+
),
28+
})
4529

4630
/**
4731
* Sweeps the primary queue and the database queue that inline webhook runs use even when
@@ -70,15 +54,9 @@ async function sweepWorkspaceJobs(workspaceId: string, signal: AbortSignal) {
7054
*/
7155
const stopRemovedExecutions: OutboxHandler<unknown> = withOutboxHandlerTimeout(
7256
async (rawPayload, context) => {
73-
const payload = parsePayload(rawPayload)
57+
const payload = payloadSchema.parse(rawPayload)
7458
const executions = new Map(
75-
payload.executions.map(({ executionId, executionDeadlineAt }) => [
76-
executionId,
77-
{
78-
executionId,
79-
executionDeadlineAt: executionDeadlineAt ? new Date(executionDeadlineAt) : null,
80-
},
81-
])
59+
payload.executions.map((execution) => [execution.executionId, execution])
8260
)
8361
/**
8462
* An archived workspace keeps its rows, so a run that passed admission before the archive
@@ -90,10 +68,12 @@ const stopRemovedExecutions: OutboxHandler<unknown> = withOutboxHandlerTimeout(
9068
})
9169
for (const execution of live) executions.set(execution.executionId, execution)
9270
}
93-
const unrecorded = await signalExecutionStops([...executions.values()], context.signal)
94-
const sweep = payload.workspaceId
95-
? await sweepWorkspaceJobs(payload.workspaceId, context.signal)
96-
: { cancelledJobs: 0, complete: true }
71+
const [unrecorded, sweep] = await Promise.all([
72+
signalExecutionStops([...executions.values()], context.signal),
73+
payload.workspaceId
74+
? sweepWorkspaceJobs(payload.workspaceId, context.signal)
75+
: { cancelledJobs: 0, complete: true },
76+
])
9777
if (unrecorded > 0) {
9878
throw new Error(`${unrecorded} execution cancellations were not recorded`)
9979
}

0 commit comments

Comments
 (0)