Skip to content

Commit 197e901

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 86faeab commit 197e901

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
@@ -55,6 +55,15 @@ function rowToJob(row: AsyncJobRow): Job {
5555
const inlineAbortControllers = new Map<string, AbortController>()
5656
const inlineExecutionControllers = new Map<string, Map<AbortController, JobType>>()
5757

58+
/** Aborts the runner this process holds for the job, if any. */
59+
function abortInlineJob(jobId: string): boolean {
60+
const controller = inlineAbortControllers.get(jobId)
61+
if (!controller) return false
62+
controller.abort(new DOMException('user', 'AbortError'))
63+
inlineAbortControllers.delete(jobId)
64+
return true
65+
}
66+
5867
/**
5968
* Per-cancel-key abort controllers for the `batchEnqueueAndWait` direct-call
6069
* path. Distinct from `inlineAbortControllers` (which keys by jobId) — this
@@ -424,13 +433,7 @@ export class DatabaseJobQueue implements JobQueueBackend {
424433
// Abort any in-process inline execution first so the running workflow
425434
// observes the signal and stops mid-flight. Then mark the row cancelled so
426435
// any future poller skips it without treating a user action as a failure.
427-
const controller = inlineAbortControllers.get(jobId)
428-
let aborted = false
429-
if (controller) {
430-
controller.abort(new DOMException('user', 'AbortError'))
431-
inlineAbortControllers.delete(jobId)
432-
aborted = true
433-
}
436+
const aborted = abortInlineJob(jobId)
434437

435438
const now = new Date()
436439
await db
@@ -538,7 +541,6 @@ export class DatabaseJobQueue implements JobQueueBackend {
538541
sql`${asyncJobs.metadata}->>'workspaceId' = ${workspaceId}`
539542
)
540543
let cancelledJobs = 0
541-
let complete = false
542544
while (cancelledJobs < options.maxRuns && !options.signal?.aborted) {
543545
const batchSize = Math.min(WORKSPACE_CANCELLATION_BATCH_SIZE, options.maxRuns - cancelledJobs)
544546
const now = new Date()
@@ -556,21 +558,16 @@ export class DatabaseJobQueue implements JobQueueBackend {
556558
)
557559
.returning({ id: asyncJobs.id })
558560

559-
for (const { id } of cancelled) {
560-
const controller = inlineAbortControllers.get(id)
561-
if (!controller) continue
562-
controller.abort(new DOMException('user', 'AbortError'))
563-
inlineAbortControllers.delete(id)
564-
}
561+
for (const { id } of cancelled) abortInlineJob(id)
565562
cancelledJobs += cancelled.length
566563
if (cancelled.length < batchSize) {
567-
complete = true
568-
break
564+
logger.info('Cancelled database jobs for workspace', { workspaceId, cancelledJobs })
565+
return { cancelledJobs, complete: true }
569566
}
570567
}
571568

572-
logger.info('Cancelled database jobs for workspace', { workspaceId, cancelledJobs, complete })
573-
return { cancelledJobs, complete }
569+
logger.info('Paused database job sweep for workspace', { workspaceId, cancelledJobs })
570+
return { cancelledJobs, complete: false }
574571
}
575572

576573
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
@@ -128,6 +128,10 @@ function buildExecutionTag(executionId: string): string {
128128
return tag.length <= 128 ? tag : `executionIdHash:${sha256Hex(executionId)}`
129129
}
130130

131+
function buildWorkspaceTag(workspaceId: string): string {
132+
return `workspaceId:${workspaceId}`
133+
}
134+
131135
function buildWorkflowTag(workflowId: string): string {
132136
const tag = `workflowId:${workflowId}`
133137
return tag.length <= 128 ? tag : `workflowIdHash:${sha256Hex(workflowId)}`
@@ -172,10 +176,13 @@ interface CancellationListRun {
172176
}
173177

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

@@ -186,7 +193,7 @@ function recordCancellationCandidateFailure(state: CancellationScanState, error:
186193

187194
async function processCancellationCandidateBatch(
188195
candidates: readonly CancellationCandidate[],
189-
binding: ExecutionJobBinding,
196+
binding: ExecutionJobBinding | undefined,
190197
state: CancellationScanState,
191198
cancelJob: (jobId: string) => Promise<void>
192199
): Promise<void> {
@@ -200,7 +207,7 @@ async function processCancellationCandidateBatch(
200207
uniqueCandidates.map(async (candidate) => {
201208
if (candidate.verifyPayload) {
202209
const legacyRun = await runs.retrieve(candidate.id)
203-
if (!payloadMatchesExecution(legacyRun.payload, binding)) return false
210+
if (!binding || !payloadMatchesExecution(legacyRun.payload, binding)) return false
204211
}
205212
await cancelJob(candidate.id)
206213
return true
@@ -232,6 +239,7 @@ async function scanAndCancelTriggerRuns(options: CancellationScanOptions): Promi
232239

233240
try {
234241
for await (const run of options.listRuns()) {
242+
if (options.shouldStop?.()) break
235243
const candidate = options.selectCandidate(run)
236244
if (!candidate) continue
237245
candidates.push(candidate)
@@ -665,44 +673,33 @@ export class TriggerDevJobQueue implements JobQueueBackend {
665673
workspaceId: string,
666674
options: WorkspaceJobCancellationOptions
667675
): Promise<WorkspaceJobCancellationResult> {
668-
const workspaceTag = `workspaceId:${workspaceId}`
676+
const workspaceTag = buildWorkspaceTag(workspaceId)
669677
const state: CancellationScanState = { cancelledJobs: 0, failureCount: 0 }
670-
let pending: string[] = []
671678
let inspectedRuns = 0
672679
let complete = true
673680

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

707704
if (state.firstError) {
708705
logger.error('Failed to cancel trigger.dev runs for workspace', {
@@ -748,7 +745,7 @@ function buildTags(options?: EnqueueOptions): string[] {
748745
const workflowId =
749746
meta?.correlation?.workflowId ??
750747
(typeof meta?.workflowId === 'string' ? meta.workflowId : undefined)
751-
if (meta?.workspaceId) tags.push(`workspaceId:${meta.workspaceId}`)
748+
if (meta?.workspaceId) tags.push(buildWorkspaceTag(meta.workspaceId))
752749
if (workflowId) tags.push(buildWorkflowTag(workflowId))
753750
if (meta?.userId) tags.push(`userId:${meta.userId}`)
754751

‎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)