@shipfox/api-runners 12.3.0 → 12.4.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (153) hide show
  1. package/.turbo/turbo-build.log +4 -4
  2. package/CHANGELOG.md +36 -0
  3. package/dist/config.d.ts +1 -0
  4. package/dist/config.d.ts.map +1 -1
  5. package/dist/config.js +9 -2
  6. package/dist/config.js.map +1 -1
  7. package/dist/core/demand.d.ts +4 -2
  8. package/dist/core/demand.d.ts.map +1 -1
  9. package/dist/core/demand.js +22 -6
  10. package/dist/core/demand.js.map +1 -1
  11. package/dist/core/entities/reservation.d.ts +1 -0
  12. package/dist/core/entities/reservation.d.ts.map +1 -1
  13. package/dist/core/entities/reservation.js.map +1 -1
  14. package/dist/core/entities/runner-instance.d.ts +3 -0
  15. package/dist/core/entities/runner-instance.d.ts.map +1 -1
  16. package/dist/core/entities/runner-instance.js.map +1 -1
  17. package/dist/core/errors.d.ts +5 -1
  18. package/dist/core/errors.d.ts.map +1 -1
  19. package/dist/core/errors.js +9 -2
  20. package/dist/core/errors.js.map +1 -1
  21. package/dist/core/runner-activation.d.ts +4 -1
  22. package/dist/core/runner-activation.d.ts.map +1 -1
  23. package/dist/core/runner-activation.js +32 -3
  24. package/dist/core/runner-activation.js.map +1 -1
  25. package/dist/core/runner-control-sessions.d.ts +8 -4
  26. package/dist/core/runner-control-sessions.d.ts.map +1 -1
  27. package/dist/core/runner-control-sessions.js +116 -20
  28. package/dist/core/runner-control-sessions.js.map +1 -1
  29. package/dist/core/runner-instances.d.ts +1 -0
  30. package/dist/core/runner-instances.d.ts.map +1 -1
  31. package/dist/core/runner-instances.js +1 -0
  32. package/dist/core/runner-instances.js.map +1 -1
  33. package/dist/db/db.d.ts +68 -0
  34. package/dist/db/db.d.ts.map +1 -1
  35. package/dist/db/ephemeral-registration-tokens.d.ts.map +1 -1
  36. package/dist/db/ephemeral-registration-tokens.js +1 -1
  37. package/dist/db/ephemeral-registration-tokens.js.map +1 -1
  38. package/dist/db/job-executions.d.ts.map +1 -1
  39. package/dist/db/job-executions.js +35 -4
  40. package/dist/db/job-executions.js.map +1 -1
  41. package/dist/db/reservation-locks.d.ts +6 -0
  42. package/dist/db/reservation-locks.d.ts.map +1 -0
  43. package/dist/db/reservation-locks.js +15 -0
  44. package/dist/db/reservation-locks.js.map +1 -0
  45. package/dist/db/reservations.d.ts +17 -3
  46. package/dist/db/reservations.d.ts.map +1 -1
  47. package/dist/db/reservations.js +246 -51
  48. package/dist/db/reservations.js.map +1 -1
  49. package/dist/db/runner-assignments.d.ts +25 -1
  50. package/dist/db/runner-assignments.d.ts.map +1 -1
  51. package/dist/db/runner-assignments.js +160 -18
  52. package/dist/db/runner-assignments.js.map +1 -1
  53. package/dist/db/runner-instances.d.ts +11 -2
  54. package/dist/db/runner-instances.d.ts.map +1 -1
  55. package/dist/db/runner-instances.js +167 -24
  56. package/dist/db/runner-instances.js.map +1 -1
  57. package/dist/db/runner-sessions.d.ts.map +1 -1
  58. package/dist/db/runner-sessions.js +14 -1
  59. package/dist/db/runner-sessions.js.map +1 -1
  60. package/dist/db/runner-states.d.ts +2 -0
  61. package/dist/db/runner-states.d.ts.map +1 -0
  62. package/dist/db/runner-states.js +7 -0
  63. package/dist/db/runner-states.js.map +1 -0
  64. package/dist/db/schema/reservations.d.ts +18 -0
  65. package/dist/db/schema/reservations.d.ts.map +1 -1
  66. package/dist/db/schema/reservations.js +7 -1
  67. package/dist/db/schema/reservations.js.map +1 -1
  68. package/dist/db/schema/runner-instances.d.ts +18 -0
  69. package/dist/db/schema/runner-instances.d.ts.map +1 -1
  70. package/dist/db/schema/runner-instances.js +8 -0
  71. package/dist/db/schema/runner-instances.js.map +1 -1
  72. package/dist/metrics/index.d.ts +2 -2
  73. package/dist/metrics/index.d.ts.map +1 -1
  74. package/dist/metrics/index.js +1 -1
  75. package/dist/metrics/index.js.map +1 -1
  76. package/dist/metrics/instance.d.ts +72 -2
  77. package/dist/metrics/instance.d.ts.map +1 -1
  78. package/dist/metrics/instance.js +161 -1
  79. package/dist/metrics/instance.js.map +1 -1
  80. package/dist/metrics/service.d.ts.map +1 -1
  81. package/dist/metrics/service.js +25 -4
  82. package/dist/metrics/service.js.map +1 -1
  83. package/dist/presentation/dto/poll-demand.d.ts.map +1 -1
  84. package/dist/presentation/dto/poll-demand.js +3 -0
  85. package/dist/presentation/dto/poll-demand.js.map +1 -1
  86. package/dist/presentation/dto/runner-instances.d.ts +4 -1
  87. package/dist/presentation/dto/runner-instances.d.ts.map +1 -1
  88. package/dist/presentation/dto/runner-instances.js +4 -1
  89. package/dist/presentation/dto/runner-instances.js.map +1 -1
  90. package/dist/presentation/routes/assign-runner-instances.d.ts.map +1 -1
  91. package/dist/presentation/routes/assign-runner-instances.js +10 -4
  92. package/dist/presentation/routes/assign-runner-instances.js.map +1 -1
  93. package/dist/presentation/routes/poll-demand.d.ts.map +1 -1
  94. package/dist/presentation/routes/poll-demand.js +1 -0
  95. package/dist/presentation/routes/poll-demand.js.map +1 -1
  96. package/dist/presentation/routes/reconcile-runner-instances.d.ts.map +1 -1
  97. package/dist/presentation/routes/reconcile-runner-instances.js +4 -2
  98. package/dist/presentation/routes/reconcile-runner-instances.js.map +1 -1
  99. package/dist/presentation/routes/runner-enrollment.d.ts +4 -1
  100. package/dist/presentation/routes/runner-enrollment.d.ts.map +1 -1
  101. package/dist/presentation/routes/runner-enrollment.js +11 -6
  102. package/dist/presentation/routes/runner-enrollment.js.map +1 -1
  103. package/dist/tsconfig.test.tsbuildinfo +1 -1
  104. package/drizzle/0005_previous_lenny_balinger.sql +5 -0
  105. package/drizzle/meta/0005_snapshot.json +2220 -0
  106. package/drizzle/meta/_journal.json +7 -0
  107. package/package.json +6 -6
  108. package/src/config.test.ts +22 -0
  109. package/src/config.ts +15 -2
  110. package/src/core/demand.test.ts +192 -3
  111. package/src/core/demand.ts +44 -4
  112. package/src/core/entities/reservation.ts +1 -0
  113. package/src/core/entities/runner-instance.ts +4 -0
  114. package/src/core/errors.ts +30 -1
  115. package/src/core/runner-activation.ts +48 -2
  116. package/src/core/runner-control-sessions.test.ts +135 -3
  117. package/src/core/runner-control-sessions.ts +141 -19
  118. package/src/core/runner-instances.test.ts +19 -0
  119. package/src/core/runner-instances.ts +2 -0
  120. package/src/core/runner-sessions.test.ts +121 -1
  121. package/src/db/ephemeral-registration-tokens.test.ts +20 -1
  122. package/src/db/ephemeral-registration-tokens.ts +1 -0
  123. package/src/db/job-executions.test.ts +98 -12
  124. package/src/db/job-executions.ts +52 -4
  125. package/src/db/reservation-locks.ts +22 -0
  126. package/src/db/reservations.test.ts +733 -22
  127. package/src/db/reservations.ts +408 -144
  128. package/src/db/runner-assignments.test.ts +105 -1
  129. package/src/db/runner-assignments.ts +245 -25
  130. package/src/db/runner-instances.test.ts +357 -2
  131. package/src/db/runner-instances.ts +289 -55
  132. package/src/db/runner-sessions.ts +19 -1
  133. package/src/db/runner-states.ts +7 -0
  134. package/src/db/schema/reservations.ts +5 -1
  135. package/src/db/schema/runner-instances.ts +9 -0
  136. package/src/metrics/index.ts +24 -1
  137. package/src/metrics/instance.test.ts +118 -0
  138. package/src/metrics/instance.ts +247 -3
  139. package/src/metrics/service.test.ts +46 -1
  140. package/src/metrics/service.ts +55 -8
  141. package/src/presentation/dto/poll-demand.ts +3 -0
  142. package/src/presentation/dto/runner-instances.ts +8 -1
  143. package/src/presentation/routes/assign-runner-instances.ts +14 -2
  144. package/src/presentation/routes/late-runner-enrollment.test.ts +219 -2
  145. package/src/presentation/routes/poll-demand.test.ts +56 -0
  146. package/src/presentation/routes/poll-demand.ts +1 -0
  147. package/src/presentation/routes/reconcile-runner-instances.test.ts +34 -2
  148. package/src/presentation/routes/reconcile-runner-instances.ts +7 -1
  149. package/src/presentation/routes/runner-enrollment.test.ts +316 -3
  150. package/src/presentation/routes/runner-enrollment.ts +8 -5
  151. package/test/factories/reservation.ts +2 -0
  152. package/test/factories/runner-instance.ts +6 -0
  153. package/tsconfig.build.tsbuildinfo +1 -1
@@ -11,18 +11,22 @@ import {
11
11
  isNull,
12
12
  lt,
13
13
  notExists,
14
+ notInArray,
14
15
  or,
15
16
  sql,
16
17
  } from 'drizzle-orm';
18
+ import {recordProviderRunnerActivationOutcome} from '#metrics/instance.js';
17
19
  import type {Tx} from './db.js';
18
20
  import {db} from './db.js';
21
+ import {lockRunnerReservationAdvisoryKeysTx} from './reservation-locks.js';
22
+ import {terminalStates} from './runner-states.js';
19
23
  import {pendingJobExecutions} from './schema/pending-job-executions.js';
20
24
  import {provisionerCapabilitySnapshots} from './schema/provisioner-capability-snapshots.js';
21
25
  import {provisionerTokens} from './schema/provisioner-tokens.js';
22
26
  import {reservations} from './schema/reservations.js';
23
27
  import {runnerActivationTokens} from './schema/runner-activation-tokens.js';
24
28
  import {runnerControlSessions} from './schema/runner-control-sessions.js';
25
- import {providerRunners} from './schema/runner-instances.js';
29
+ import {type providerRunnerLaunchKindEnum, providerRunners} from './schema/runner-instances.js';
26
30
 
27
31
  export interface ReservationTemplate {
28
32
  templateKey: string;
@@ -53,6 +57,8 @@ export interface PollDemandAndReserveParams {
53
57
  provisionerId: string;
54
58
  maxReservations: number;
55
59
  ttlSeconds: number;
60
+ /** Lifetime of the short reservation that gives rebound runners time to activate. */
61
+ activationGraceSeconds?: number;
56
62
  templates: ReservationTemplate[];
57
63
  capabilityWindowSeconds?: number;
58
64
  }
@@ -61,6 +67,8 @@ export interface InstallationPollDemandAndReserveParams {
61
67
  provisionerId: string;
62
68
  maxReservations: number;
63
69
  ttlSeconds: number;
70
+ /** Lifetime of the short reservation that gives rebound runners time to activate. */
71
+ activationGraceSeconds?: number;
64
72
  templates: ReservationTemplate[];
65
73
  capabilityWindowSeconds: number;
66
74
  eligibleWorkspaceIds: ReadonlySet<string>;
@@ -70,6 +78,22 @@ export interface InstallationPollDemandAndReserveParams {
70
78
 
71
79
  type DemandScope = 'installation' | 'workspace';
72
80
 
81
+ interface IdleRunnerCandidate {
82
+ id: string;
83
+ launchKind: (typeof providerRunnerLaunchKindEnum.enumValues)[number];
84
+ }
85
+
86
+ interface BindableRunnerParams {
87
+ provisionerId: string;
88
+ workspaceId: string;
89
+ requiredLabels: string[];
90
+ scope: DemandScope;
91
+ }
92
+
93
+ interface IdleRunnerSelectionParams extends BindableRunnerParams {
94
+ count: number;
95
+ }
96
+
73
97
  type PollDemandAndReserveLockedParams = PollDemandAndReserveParams & {
74
98
  scope: DemandScope;
75
99
  };
@@ -86,29 +110,47 @@ interface DemandRow {
86
110
  oldestQueuedAt: Date;
87
111
  }
88
112
 
113
+ interface NewReservationUnits {
114
+ labels: string[];
115
+ count: number;
116
+ }
117
+
118
+ interface PollDemandAndReserveResult {
119
+ stats: DemandStat[];
120
+ /** Only launch reservations are exposed to the provisioner. */
121
+ reservations: ReservationGrant[];
122
+ /** Internal allocation accounting for bound and launch rows together. */
123
+ newlyReservedUnits: NewReservationUnits[];
124
+ }
125
+
89
126
  export async function pollDemandAndReserve(
90
127
  params: PollDemandAndReserveParams,
91
128
  ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
92
- return await db().transaction(async (tx) => {
129
+ const result = await db().transaction(async (tx) => {
93
130
  return await pollDemandAndReserveTx(tx, params);
94
131
  });
132
+ return {stats: result.stats, reservations: result.reservations};
95
133
  }
96
134
 
97
135
  export async function pollDemandAndReserveTx(
98
136
  tx: Tx,
99
137
  params: PollDemandAndReserveParams,
100
- ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
138
+ ): Promise<PollDemandAndReserveResult> {
101
139
  await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${params.workspaceId}))`);
102
140
  return await pollDemandAndReserveLockedTx(tx, {...params, scope: 'workspace'});
103
141
  }
104
142
 
105
143
  export async function pollInstallationDemandAndReserve(
106
144
  params: InstallationPollDemandAndReserveParams,
107
- ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
145
+ ): Promise<{
146
+ stats: DemandStat[];
147
+ reservations: ReservationGrant[];
148
+ newlyReservedCount?: number;
149
+ }> {
108
150
  const candidateWorkspaceIds = await listInstallationDemandWorkspaceIds(
109
151
  params.eligibleWorkspaceIds,
110
152
  );
111
- const results: Array<{stats: DemandStat[]; reservations: ReservationGrant[]}> = [];
153
+ const results: PollDemandAndReserveResult[] = [];
112
154
  let remainingMaxReservations = params.maxReservations;
113
155
  const remainingTemplates = params.templates.map((template) => ({
114
156
  ...template,
@@ -122,12 +164,15 @@ export async function pollInstallationDemandAndReserve(
122
164
  sql`select pg_try_advisory_xact_lock(hashtext(${workspaceId})) as locked`,
123
165
  );
124
166
  const locked = lockResult.rows[0];
125
- if (!locked?.locked) return {stats: [], reservations: []};
167
+ if (!locked?.locked) return {stats: [], reservations: [], newlyReservedUnits: []};
126
168
  return await pollDemandAndReserveLockedTx(tx, {
127
169
  workspaceId,
128
170
  provisionerId: params.provisionerId,
129
171
  maxReservations: remainingMaxReservations,
130
172
  ttlSeconds: params.ttlSeconds,
173
+ ...(params.activationGraceSeconds !== undefined
174
+ ? {activationGraceSeconds: params.activationGraceSeconds}
175
+ : {}),
131
176
  templates: remainingTemplates.map((template) => ({
132
177
  ...template,
133
178
  availableSlots: template.remainingSlots,
@@ -139,22 +184,31 @@ export async function pollInstallationDemandAndReserve(
139
184
  results.push(result);
140
185
  consumeInstallationTemplateSlots(remainingTemplates, result.reservations);
141
186
  params.onReservations?.(result.reservations);
142
- remainingMaxReservations -= result.reservations.reduce(
187
+ remainingMaxReservations -= result.newlyReservedUnits.reduce(
143
188
  (total, reservation) => total + reservation.count,
144
189
  0,
145
190
  );
146
191
  }
192
+ const newlyReservedCount = results.reduce(
193
+ (total, result) =>
194
+ total +
195
+ result.newlyReservedUnits.reduce((subtotal, reservation) => subtotal + reservation.count, 0),
196
+ 0,
197
+ );
147
198
  return {
148
199
  stats: results.flatMap((result) => result.stats),
149
200
  reservations: results.flatMap((result) => result.reservations),
201
+ ...(newlyReservedCount > 0 ? {newlyReservedCount} : {}),
150
202
  };
151
203
  }
152
204
 
153
205
  function consumeInstallationTemplateSlots(
154
206
  templates: NormalizedTemplate[],
155
- reservations: ReservationGrant[],
207
+ launchGrants: ReservationGrant[],
156
208
  ): void {
157
- for (const reservation of reservations) {
209
+ // Adopted runners are already included in the running count behind availableSlots.
210
+ // Only units that still need a launch consume advertised template capacity.
211
+ for (const reservation of launchGrants) {
158
212
  const satisfyingTemplates = templates
159
213
  .filter((template) => isSubset(reservation.labels, template.labels))
160
214
  .sort(
@@ -167,7 +221,7 @@ function consumeInstallationTemplateSlots(
167
221
  async function pollDemandAndReserveLockedTx(
168
222
  tx: Tx,
169
223
  params: PollDemandAndReserveLockedParams,
170
- ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
224
+ ): Promise<PollDemandAndReserveResult> {
171
225
  let demandRows = (
172
226
  await tx
173
227
  .select({
@@ -206,20 +260,10 @@ async function pollDemandAndReserveLockedTx(
206
260
  const reservedByLabels = new Map(
207
261
  activeReservationRows.map((row) => [labelKey(row.requiredLabels), row.reserved]),
208
262
  );
209
- const activeProvisionerReservationRows = await tx
210
- .select({
211
- requiredLabels: reservations.requiredLabels,
212
- reserved: sql<number>`coalesce(sum(${reservations.count}), 0)::int`,
213
- })
214
- .from(reservations)
215
- .where(
216
- and(
217
- eq(reservations.workspaceId, params.workspaceId),
218
- eq(reservations.provisionerId, params.provisionerId),
219
- gt(reservations.expiresAt, sql`now()`),
220
- ),
221
- )
222
- .groupBy(reservations.requiredLabels);
263
+ const activeProvisionerReservationRows = await listActiveProvisionerReservationRowsTx(tx, {
264
+ workspaceId: params.workspaceId,
265
+ provisionerId: params.provisionerId,
266
+ });
223
267
 
224
268
  const templates = params.templates.map((template) => ({
225
269
  templateKey: template.templateKey,
@@ -229,6 +273,7 @@ async function pollDemandAndReserveLockedTx(
229
273
  deductProvisionerReservations(templates, activeProvisionerReservationRows);
230
274
  const stats: DemandStat[] = [];
231
275
  const grants: ReservationGrant[] = [];
276
+ const newlyReservedUnits: NewReservationUnits[] = [];
232
277
  let remainingMaxReservations = params.maxReservations;
233
278
  const workspaceIdField =
234
279
  params.capabilityWindowSeconds === undefined ? {} : {workspaceId: params.workspaceId};
@@ -252,37 +297,73 @@ async function pollDemandAndReserveLockedTx(
252
297
  let reservedAfterGrant = reserved;
253
298
 
254
299
  if (grant > 0 && params.maxReservations > 0) {
255
- const [inserted] = await tx
256
- .insert(reservations)
257
- .values({
258
- workspaceId: params.workspaceId,
259
- provisionerId: params.provisionerId,
260
- requiredLabels: demand.requiredLabels,
261
- count: grant,
262
- expiresAt: sql`now() + (${params.ttlSeconds} || ' seconds')::interval`,
263
- })
264
- .returning({id: reservations.id, expiresAt: reservations.expiresAt});
265
-
266
- if (!inserted) throw new Error('Insert returned no rows');
267
-
268
- remainingMaxReservations -= grant;
269
- drawSlots(satisfyingTemplates, grant);
270
- await bindIdleRunnerInstancesTx(tx, {
300
+ const idleRunners = await listIdleRunnerInstancesTx(tx, {
271
301
  provisionerId: params.provisionerId,
272
- reservationId: inserted.id,
273
302
  requiredLabels: demand.requiredLabels,
274
303
  count: grant,
275
304
  workspaceId: params.workspaceId,
276
305
  scope: params.scope,
277
306
  });
307
+
308
+ remainingMaxReservations -= grant;
309
+
310
+ const boundCount = idleRunners.length;
311
+ if (boundCount > 0) {
312
+ const [boundReservation] = await tx
313
+ .insert(reservations)
314
+ .values({
315
+ workspaceId: params.workspaceId,
316
+ provisionerId: params.provisionerId,
317
+ requiredLabels: demand.requiredLabels,
318
+ count: boundCount,
319
+ kind: 'bound',
320
+ expiresAt: sql`now() + (${params.activationGraceSeconds ?? params.ttlSeconds} || ' seconds')::interval`,
321
+ })
322
+ .returning({id: reservations.id});
323
+
324
+ if (!boundReservation) throw new Error('Insert returned no rows');
325
+
326
+ await bindIdleRunnerInstancesTx(tx, {
327
+ provisionerId: params.provisionerId,
328
+ reservationId: boundReservation.id,
329
+ workspaceId: params.workspaceId,
330
+ requiredLabels: demand.requiredLabels,
331
+ scope: params.scope,
332
+ idleRunners,
333
+ });
334
+ }
335
+
336
+ const launchCount = grant - boundCount;
337
+ drawSlots(satisfyingTemplates, launchCount);
338
+ let launchReservation: {id: string; expiresAt: Date} | undefined;
339
+ if (launchCount > 0) {
340
+ const [inserted] = await tx
341
+ .insert(reservations)
342
+ .values({
343
+ workspaceId: params.workspaceId,
344
+ provisionerId: params.provisionerId,
345
+ requiredLabels: demand.requiredLabels,
346
+ count: launchCount,
347
+ kind: 'launch',
348
+ expiresAt: sql`now() + (${params.ttlSeconds} || ' seconds')::interval`,
349
+ })
350
+ .returning({id: reservations.id, expiresAt: reservations.expiresAt});
351
+
352
+ if (!inserted) throw new Error('Insert returned no rows');
353
+ launchReservation = inserted;
354
+ }
355
+
356
+ newlyReservedUnits.push({labels: demand.requiredLabels, count: grant});
278
357
  reservedAfterGrant += grant;
279
- grants.push({
280
- reservationId: inserted.id,
281
- ...workspaceIdField,
282
- labels: demand.requiredLabels,
283
- count: grant,
284
- expiresAt: inserted.expiresAt,
285
- });
358
+ if (launchReservation) {
359
+ grants.push({
360
+ reservationId: launchReservation.id,
361
+ ...workspaceIdField,
362
+ labels: demand.requiredLabels,
363
+ count: launchCount,
364
+ expiresAt: launchReservation.expiresAt,
365
+ });
366
+ }
286
367
  }
287
368
 
288
369
  stats.push({
@@ -294,111 +375,276 @@ async function pollDemandAndReserveLockedTx(
294
375
  });
295
376
  }
296
377
 
297
- return {stats, reservations: grants};
378
+ return {stats, reservations: grants, newlyReservedUnits};
298
379
  }
299
380
 
300
- async function bindIdleRunnerInstancesTx(
381
+ async function listActiveProvisionerReservationRowsTx(
301
382
  tx: Tx,
302
- params: {
303
- provisionerId: string;
304
- reservationId: string;
305
- workspaceId: string;
306
- requiredLabels: string[];
307
- count: number;
308
- scope: DemandScope;
309
- },
310
- ): Promise<void> {
311
- // Reservation expiry is the activation grace deadline. A live reservation
312
- // protects a runner that is still booting; an expired or missing row is stale.
313
- const canBindRunner = () =>
314
- and(
315
- or(
316
- and(isNull(providerRunners.workspaceId), isNull(providerRunners.reservationId)),
317
- and(
318
- isNotNull(providerRunners.reservationId),
319
- params.scope === 'installation'
320
- ? undefined
321
- : or(
322
- isNull(providerRunners.workspaceId),
323
- eq(providerRunners.workspaceId, params.workspaceId),
324
- ),
325
- notExists(
326
- tx
327
- .select({id: reservations.id})
328
- .from(reservations)
329
- .where(
330
- and(
331
- eq(reservations.id, providerRunners.reservationId),
332
- gt(reservations.expiresAt, sql`now()`),
333
- ),
334
- ),
383
+ params: {workspaceId: string; provisionerId: string},
384
+ ): Promise<Array<{requiredLabels: string[]; reserved: number}>> {
385
+ const candidateRows = await tx
386
+ .select({id: reservations.id})
387
+ .from(reservations)
388
+ .where(
389
+ and(
390
+ eq(reservations.workspaceId, params.workspaceId),
391
+ eq(reservations.provisionerId, params.provisionerId),
392
+ gt(reservations.expiresAt, sql`now()`),
393
+ ),
394
+ );
395
+ const candidateIds = candidateRows.map((row) => row.id);
396
+ if (candidateIds.length === 0) return [];
397
+
398
+ // Assignment and provider-report transactions use this key before changing runner
399
+ // links. Lock it before counting so an in-flight assignment cannot expose its unit twice.
400
+ await lockRunnerReservationAdvisoryKeysTx(tx, {
401
+ provisionerId: params.provisionerId,
402
+ reservationIds: candidateIds,
403
+ });
404
+
405
+ const activeRows = await tx
406
+ .select({
407
+ id: reservations.id,
408
+ requiredLabels: reservations.requiredLabels,
409
+ count: reservations.count,
410
+ })
411
+ .from(reservations)
412
+ .where(
413
+ and(
414
+ eq(reservations.workspaceId, params.workspaceId),
415
+ eq(reservations.provisionerId, params.provisionerId),
416
+ gt(reservations.expiresAt, sql`now()`),
417
+ ),
418
+ )
419
+ .orderBy(asc(reservations.id))
420
+ .for('update');
421
+ if (activeRows.length === 0) return [];
422
+
423
+ const activeIds = activeRows.map((row) => row.id);
424
+ const activeIdSet = new Set(activeIds);
425
+ const usedRunnerRows = await tx
426
+ .select({
427
+ id: providerRunners.id,
428
+ reservationId: providerRunners.reservationId,
429
+ intendedReservationId: providerRunners.intendedReservationId,
430
+ state: providerRunners.state,
431
+ })
432
+ .from(providerRunners)
433
+ .where(
434
+ and(
435
+ eq(providerRunners.provisionerId, params.provisionerId),
436
+ isNull(providerRunners.reservationReleasedAt),
437
+ or(
438
+ inArray(providerRunners.reservationId, activeIds),
439
+ and(
440
+ inArray(providerRunners.intendedReservationId, activeIds),
441
+ notInArray(providerRunners.state, [...terminalStates]),
335
442
  ),
336
443
  ),
337
444
  ),
338
- or(
339
- isNull(providerRunners.intendedReservationId),
445
+ );
446
+ const usedByReservation = new Map<string, Set<string>>();
447
+ for (const runner of usedRunnerRows) {
448
+ if (runner.reservationId && activeIdSet.has(runner.reservationId)) {
449
+ addUsedRunner(usedByReservation, runner.reservationId, runner.id);
450
+ }
451
+ if (
452
+ runner.intendedReservationId &&
453
+ activeIdSet.has(runner.intendedReservationId) &&
454
+ !terminalStates.includes(runner.state as (typeof terminalStates)[number])
455
+ ) {
456
+ addUsedRunner(usedByReservation, runner.intendedReservationId, runner.id);
457
+ }
458
+ }
459
+
460
+ return activeRows.flatMap((reservation) => {
461
+ const used = usedByReservation.get(reservation.id)?.size ?? 0;
462
+ const pending = Math.max(0, reservation.count - used);
463
+ return pending > 0 ? [{requiredLabels: reservation.requiredLabels, reserved: pending}] : [];
464
+ });
465
+ }
466
+
467
+ function addUsedRunner(
468
+ usedByReservation: Map<string, Set<string>>,
469
+ reservationId: string,
470
+ runnerId: string,
471
+ ): void {
472
+ const runnerIds = usedByReservation.get(reservationId);
473
+ if (runnerIds) runnerIds.add(runnerId);
474
+ else usedByReservation.set(reservationId, new Set([runnerId]));
475
+ }
476
+
477
+ // An expired or missing reservation is stale. A live reservation protects a runner that is
478
+ // still booting or waiting for its activation grace period.
479
+ function canBindRunner(tx: Tx, params: BindableRunnerParams) {
480
+ return and(
481
+ or(
482
+ and(isNull(providerRunners.workspaceId), isNull(providerRunners.reservationId)),
483
+ and(
484
+ isNotNull(providerRunners.reservationId),
485
+ params.scope === 'installation'
486
+ ? undefined
487
+ : or(
488
+ isNull(providerRunners.workspaceId),
489
+ eq(providerRunners.workspaceId, params.workspaceId),
490
+ ),
340
491
  notExists(
341
492
  tx
342
493
  .select({id: reservations.id})
343
494
  .from(reservations)
344
495
  .where(
345
496
  and(
346
- eq(reservations.id, providerRunners.intendedReservationId),
497
+ eq(reservations.id, providerRunners.reservationId),
347
498
  gt(reservations.expiresAt, sql`now()`),
348
499
  ),
349
500
  ),
350
501
  ),
351
502
  ),
352
- isNull(providerRunners.reservationReleasedAt),
353
- isNull(providerRunners.runnerSessionId),
354
- );
355
-
356
- const idleRunners = await tx
357
- .select({id: providerRunners.id})
358
- .from(providerRunners)
359
- .where(
360
- and(
361
- eq(providerRunners.provisionerId, params.provisionerId),
362
- canBindRunner(),
363
- isNotNull(providerRunners.providerRunnerId),
364
- eq(providerRunners.state, 'running'),
365
- arrayContains(providerRunners.labels, params.requiredLabels),
366
- exists(
367
- tx
368
- .select({id: runnerControlSessions.id})
369
- .from(runnerControlSessions)
370
- .where(
371
- and(
372
- eq(runnerControlSessions.runnerInstanceId, providerRunners.id),
373
- eq(runnerControlSessions.provisionerId, params.provisionerId),
374
- isNull(runnerControlSessions.closedAt),
375
- gt(runnerControlSessions.expiresAt, sql`now()`),
376
- ),
503
+ ),
504
+ or(
505
+ isNull(providerRunners.intendedReservationId),
506
+ notExists(
507
+ tx
508
+ .select({id: reservations.id})
509
+ .from(reservations)
510
+ .where(
511
+ and(
512
+ eq(reservations.id, providerRunners.intendedReservationId),
513
+ gt(reservations.expiresAt, sql`now()`),
377
514
  ),
378
- ),
515
+ ),
379
516
  ),
380
- )
381
- .orderBy(asc(providerRunners.createdAt), asc(providerRunners.id))
382
- .limit(params.count)
383
- .for('update');
517
+ ),
518
+ isNull(providerRunners.reservationReleasedAt),
519
+ isNull(providerRunners.runnerSessionId),
520
+ );
521
+ }
384
522
 
385
- if (idleRunners.length === 0) return;
523
+ function isBindableRunner(tx: Tx, params: BindableRunnerParams) {
524
+ return and(
525
+ eq(providerRunners.provisionerId, params.provisionerId),
526
+ canBindRunner(tx, params),
527
+ isNotNull(providerRunners.providerRunnerId),
528
+ eq(providerRunners.state, 'running'),
529
+ arrayContains(providerRunners.labels, params.requiredLabels),
530
+ exists(
531
+ tx
532
+ .select({id: runnerControlSessions.id})
533
+ .from(runnerControlSessions)
534
+ .where(
535
+ and(
536
+ eq(runnerControlSessions.runnerInstanceId, providerRunners.id),
537
+ eq(runnerControlSessions.provisionerId, params.provisionerId),
538
+ isNull(runnerControlSessions.closedAt),
539
+ gt(runnerControlSessions.expiresAt, sql`now()`),
540
+ ),
541
+ ),
542
+ ),
543
+ );
544
+ }
545
+
546
+ async function listIdleRunnerInstancesTx(
547
+ tx: Tx,
548
+ params: IdleRunnerSelectionParams & {count: number},
549
+ ): Promise<IdleRunnerCandidate[]> {
550
+ const lockedRunners = await selectIdleRunnerInstancesTx(tx, params);
551
+ if (lockedRunners.length === 0) return lockedRunners;
552
+
553
+ // A SELECT ... FOR UPDATE can wake with a fresh target row but the original
554
+ // statement snapshot for reservation subqueries. Re-check the full predicate
555
+ // in a new statement while retaining the row locks from the first query.
556
+ return await selectIdleRunnerInstancesTx(tx, {
557
+ ...params,
558
+ runnerIds: lockedRunners.map((runner) => runner.id),
559
+ count: lockedRunners.length,
560
+ });
561
+ }
562
+
563
+ async function selectIdleRunnerInstancesTx(
564
+ tx: Tx,
565
+ params: IdleRunnerSelectionParams & {count: number; runnerIds?: string[]},
566
+ ): Promise<IdleRunnerCandidate[]> {
567
+ if (params.runnerIds) {
568
+ // The first pass retains the row locks. Recheck the full predicate in a fresh statement so
569
+ // reservation subqueries see the state that was current when the lock was acquired.
570
+ return await tx
571
+ .select({id: providerRunners.id, launchKind: providerRunners.launchKind})
572
+ .from(providerRunners)
573
+ .where(and(inArray(providerRunners.id, params.runnerIds), isBindableRunner(tx, params)))
574
+ .orderBy(asc(providerRunners.createdAt), asc(providerRunners.id))
575
+ .limit(params.count);
576
+ }
577
+
578
+ // Select the oldest candidates first, but leave row locking to a second pass. The cleanup
579
+ // path locks runner rows by id, so the locking pass must use that same order. A skipped row
580
+ // rolls back only this nested savepoint and retries the bounded candidate scan, allowing the
581
+ // next-oldest eligible runner to refill the grant without retaining a partial lock set.
582
+ const retrySelection = Symbol('retry bindable runner selection');
583
+ const idleRunners = await (async () => {
584
+ while (true) {
585
+ try {
586
+ return await tx.transaction(async (lockTx) => {
587
+ const candidateRunners = await lockTx
588
+ .select({id: providerRunners.id, launchKind: providerRunners.launchKind})
589
+ .from(providerRunners)
590
+ .where(isBindableRunner(lockTx, params))
591
+ .orderBy(asc(providerRunners.createdAt), asc(providerRunners.id))
592
+ .limit(params.count);
593
+
594
+ if (candidateRunners.length === 0) return candidateRunners;
595
+
596
+ // PostgreSQL can acquire row locks before sorting a multi-row FOR UPDATE result. Lock
597
+ // each selected candidate individually so both polling and cleanup use id order.
598
+ const lockedRunners: typeof candidateRunners = [];
599
+ for (const candidate of [...candidateRunners].sort(compareRunnerIds)) {
600
+ const [runner] = await lockTx
601
+ .select({id: providerRunners.id, launchKind: providerRunners.launchKind})
602
+ .from(providerRunners)
603
+ .where(and(eq(providerRunners.id, candidate.id), isBindableRunner(lockTx, params)))
604
+ .for('update');
605
+ if (runner) lockedRunners.push(runner);
606
+ }
607
+
608
+ if (lockedRunners.length !== candidateRunners.length) throw retrySelection;
609
+ return lockedRunners;
610
+ });
611
+ } catch (error) {
612
+ if (error !== retrySelection) throw error;
613
+ }
614
+ }
615
+ })();
616
+
617
+ return idleRunners;
618
+ }
619
+
620
+ async function bindIdleRunnerInstancesTx(
621
+ tx: Tx,
622
+ params: {
623
+ provisionerId: string;
624
+ reservationId: string;
625
+ workspaceId: string;
626
+ requiredLabels: string[];
627
+ scope: DemandScope;
628
+ idleRunners: IdleRunnerCandidate[];
629
+ },
630
+ ): Promise<void> {
631
+ const idleRunnerIds = params.idleRunners.map((runner) => runner.id);
386
632
 
387
633
  await tx
388
634
  .update(runnerActivationTokens)
389
635
  .set({revokedAt: sql`now()`})
390
636
  .where(
391
637
  and(
392
- inArray(
393
- runnerActivationTokens.runnerInstanceId,
394
- idleRunners.map((runner) => runner.id),
395
- ),
638
+ inArray(runnerActivationTokens.runnerInstanceId, idleRunnerIds),
396
639
  isNull(runnerActivationTokens.consumedAt),
397
640
  isNull(runnerActivationTokens.revokedAt),
398
641
  ),
399
642
  );
400
643
 
401
- await tx
644
+ // listIdleRunnerInstancesTx re-checks the full eligibility predicate in a fresh statement
645
+ // after acquiring row locks. Those locks prevent the validated set from changing before this
646
+ // update runs.
647
+ const boundRunners = await tx
402
648
  .update(providerRunners)
403
649
  .set({
404
650
  workspaceId: params.workspaceId,
@@ -407,15 +653,15 @@ async function bindIdleRunnerInstancesTx(
407
653
  assignedAt: sql`now()`,
408
654
  updatedAt: sql`now()`,
409
655
  })
410
- .where(
411
- and(
412
- inArray(
413
- providerRunners.id,
414
- idleRunners.map((runner) => runner.id),
415
- ),
416
- canBindRunner(),
417
- ),
418
- );
656
+ .where(and(inArray(providerRunners.id, idleRunnerIds), isBindableRunner(tx, params)))
657
+ .returning({id: providerRunners.id, launchKind: providerRunners.launchKind});
658
+
659
+ if (boundRunners.length !== params.idleRunners.length)
660
+ throw new Error('Locked idle runner set changed before binding');
661
+
662
+ const reboundCount = boundRunners.filter((runner) => runner.launchKind === 'demand').length;
663
+ if (reboundCount > 0)
664
+ recordProviderRunnerActivationOutcome({outcome: 'rebound', count: reboundCount});
419
665
  }
420
666
 
421
667
  async function listInstallationDemandWorkspaceIds(eligibleWorkspaceIds: ReadonlySet<string>) {
@@ -498,7 +744,16 @@ async function deleteReservationsWithCleanupTx(
498
744
  ): Promise<number> {
499
745
  if (ids.length === 0) return 0;
500
746
 
501
- const affectedRunners = await tx
747
+ const isAffectedRunner = () =>
748
+ or(
749
+ and(
750
+ inArray(providerRunners.reservationId, ids),
751
+ isNull(providerRunners.runnerSessionId),
752
+ isNull(providerRunners.reservationReleasedAt),
753
+ ),
754
+ inArray(providerRunners.intendedReservationId, ids),
755
+ );
756
+ const candidateRunners = await tx
502
757
  .select({
503
758
  id: providerRunners.id,
504
759
  reservationId: providerRunners.reservationId,
@@ -507,18 +762,23 @@ async function deleteReservationsWithCleanupTx(
507
762
  reservationReleasedAt: providerRunners.reservationReleasedAt,
508
763
  })
509
764
  .from(providerRunners)
510
- .where(
511
- or(
512
- and(
513
- inArray(providerRunners.reservationId, ids),
514
- isNull(providerRunners.runnerSessionId),
515
- isNull(providerRunners.reservationReleasedAt),
516
- ),
517
- inArray(providerRunners.intendedReservationId, ids),
518
- ),
519
- )
520
- .orderBy(asc(providerRunners.id))
521
- .for('update');
765
+ .where(isAffectedRunner())
766
+ .orderBy(asc(providerRunners.id));
767
+ const affectedRunners: typeof candidateRunners = [];
768
+ for (const candidate of [...candidateRunners].sort(compareRunnerIds)) {
769
+ const [runner] = await tx
770
+ .select({
771
+ id: providerRunners.id,
772
+ reservationId: providerRunners.reservationId,
773
+ intendedReservationId: providerRunners.intendedReservationId,
774
+ runnerSessionId: providerRunners.runnerSessionId,
775
+ reservationReleasedAt: providerRunners.reservationReleasedAt,
776
+ })
777
+ .from(providerRunners)
778
+ .where(and(eq(providerRunners.id, candidate.id), isAffectedRunner()))
779
+ .for('update');
780
+ if (runner) affectedRunners.push(runner);
781
+ }
522
782
  const reservationRows = await tx
523
783
  .select({id: reservations.id})
524
784
  .from(reservations)
@@ -660,6 +920,10 @@ function deductProvisionerReservations(
660
920
  }
661
921
  }
662
922
 
923
+ function compareRunnerIds(left: {id: string}, right: {id: string}): number {
924
+ return left.id < right.id ? -1 : left.id > right.id ? 1 : 0;
925
+ }
926
+
663
927
  function sortDemandRows(rows: DemandRow[]): DemandRow[] {
664
928
  return [...rows].sort((a, b) => {
665
929
  const specificity = b.requiredLabels.length - a.requiredLabels.length;