@convex-dev/workpool 0.1.3-alpha.0 → 0.2.0-beta.0

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 (113) hide show
  1. package/README.md +144 -4
  2. package/dist/commonjs/client/index.d.ts +123 -35
  3. package/dist/commonjs/client/index.d.ts.map +1 -1
  4. package/dist/commonjs/client/index.js +122 -15
  5. package/dist/commonjs/client/index.js.map +1 -1
  6. package/dist/commonjs/client/utils.d.ts +16 -0
  7. package/dist/commonjs/client/utils.d.ts.map +1 -0
  8. package/dist/commonjs/client/utils.js +2 -0
  9. package/dist/commonjs/client/utils.js.map +1 -0
  10. package/dist/commonjs/component/convex.config.d.ts.map +1 -1
  11. package/dist/commonjs/component/convex.config.js +0 -2
  12. package/dist/commonjs/component/convex.config.js.map +1 -1
  13. package/dist/commonjs/component/kick.d.ts +9 -0
  14. package/dist/commonjs/component/kick.d.ts.map +1 -0
  15. package/dist/commonjs/component/kick.js +97 -0
  16. package/dist/commonjs/component/kick.js.map +1 -0
  17. package/dist/commonjs/component/lib.d.ts +23 -32
  18. package/dist/commonjs/component/lib.d.ts.map +1 -1
  19. package/dist/commonjs/component/lib.js +70 -564
  20. package/dist/commonjs/component/lib.js.map +1 -1
  21. package/dist/commonjs/component/logging.d.ts +6 -4
  22. package/dist/commonjs/component/logging.d.ts.map +1 -1
  23. package/dist/commonjs/component/logging.js +13 -2
  24. package/dist/commonjs/component/logging.js.map +1 -1
  25. package/dist/commonjs/component/loop.d.ts +26 -0
  26. package/dist/commonjs/component/loop.d.ts.map +1 -0
  27. package/dist/commonjs/component/loop.js +453 -0
  28. package/dist/commonjs/component/loop.js.map +1 -0
  29. package/dist/commonjs/component/recovery.d.ts +8 -0
  30. package/dist/commonjs/component/recovery.d.ts.map +1 -0
  31. package/dist/commonjs/component/recovery.js +74 -0
  32. package/dist/commonjs/component/recovery.js.map +1 -0
  33. package/dist/commonjs/component/schema.d.ts +163 -93
  34. package/dist/commonjs/component/schema.d.ts.map +1 -1
  35. package/dist/commonjs/component/schema.js +54 -65
  36. package/dist/commonjs/component/schema.js.map +1 -1
  37. package/dist/commonjs/component/shared.d.ts +130 -0
  38. package/dist/commonjs/component/shared.d.ts.map +1 -0
  39. package/dist/commonjs/component/shared.js +65 -0
  40. package/dist/commonjs/component/shared.js.map +1 -0
  41. package/dist/commonjs/component/stats.d.ts +3 -2
  42. package/dist/commonjs/component/stats.d.ts.map +1 -1
  43. package/dist/commonjs/component/stats.js +17 -3
  44. package/dist/commonjs/component/stats.js.map +1 -1
  45. package/dist/commonjs/component/worker.d.ts +25 -0
  46. package/dist/commonjs/component/worker.d.ts.map +1 -0
  47. package/dist/commonjs/component/worker.js +86 -0
  48. package/dist/commonjs/component/worker.js.map +1 -0
  49. package/dist/esm/client/index.d.ts +123 -35
  50. package/dist/esm/client/index.d.ts.map +1 -1
  51. package/dist/esm/client/index.js +122 -15
  52. package/dist/esm/client/index.js.map +1 -1
  53. package/dist/esm/client/utils.d.ts +16 -0
  54. package/dist/esm/client/utils.d.ts.map +1 -0
  55. package/dist/esm/client/utils.js +2 -0
  56. package/dist/esm/client/utils.js.map +1 -0
  57. package/dist/esm/component/convex.config.d.ts.map +1 -1
  58. package/dist/esm/component/convex.config.js +0 -2
  59. package/dist/esm/component/convex.config.js.map +1 -1
  60. package/dist/esm/component/kick.d.ts +9 -0
  61. package/dist/esm/component/kick.d.ts.map +1 -0
  62. package/dist/esm/component/kick.js +97 -0
  63. package/dist/esm/component/kick.js.map +1 -0
  64. package/dist/esm/component/lib.d.ts +23 -32
  65. package/dist/esm/component/lib.d.ts.map +1 -1
  66. package/dist/esm/component/lib.js +70 -564
  67. package/dist/esm/component/lib.js.map +1 -1
  68. package/dist/esm/component/logging.d.ts +6 -4
  69. package/dist/esm/component/logging.d.ts.map +1 -1
  70. package/dist/esm/component/logging.js +13 -2
  71. package/dist/esm/component/logging.js.map +1 -1
  72. package/dist/esm/component/loop.d.ts +26 -0
  73. package/dist/esm/component/loop.d.ts.map +1 -0
  74. package/dist/esm/component/loop.js +453 -0
  75. package/dist/esm/component/loop.js.map +1 -0
  76. package/dist/esm/component/recovery.d.ts +8 -0
  77. package/dist/esm/component/recovery.d.ts.map +1 -0
  78. package/dist/esm/component/recovery.js +74 -0
  79. package/dist/esm/component/recovery.js.map +1 -0
  80. package/dist/esm/component/schema.d.ts +163 -93
  81. package/dist/esm/component/schema.d.ts.map +1 -1
  82. package/dist/esm/component/schema.js +54 -65
  83. package/dist/esm/component/schema.js.map +1 -1
  84. package/dist/esm/component/shared.d.ts +130 -0
  85. package/dist/esm/component/shared.d.ts.map +1 -0
  86. package/dist/esm/component/shared.js +65 -0
  87. package/dist/esm/component/shared.js.map +1 -0
  88. package/dist/esm/component/stats.d.ts +3 -2
  89. package/dist/esm/component/stats.d.ts.map +1 -1
  90. package/dist/esm/component/stats.js +17 -3
  91. package/dist/esm/component/stats.js.map +1 -1
  92. package/dist/esm/component/worker.d.ts +25 -0
  93. package/dist/esm/component/worker.d.ts.map +1 -0
  94. package/dist/esm/component/worker.js +86 -0
  95. package/dist/esm/component/worker.js.map +1 -0
  96. package/package.json +6 -5
  97. package/src/client/index.ts +231 -70
  98. package/src/client/utils.ts +45 -0
  99. package/src/component/README.md +73 -0
  100. package/src/component/_generated/api.d.ts +36 -66
  101. package/src/component/convex.config.ts +0 -3
  102. package/src/component/kick.test.ts +286 -0
  103. package/src/component/kick.ts +118 -0
  104. package/src/component/lib.test.ts +203 -0
  105. package/src/component/lib.ts +80 -671
  106. package/src/component/logging.ts +24 -10
  107. package/src/component/loop.ts +579 -0
  108. package/src/component/recovery.ts +79 -0
  109. package/src/component/schema.ts +59 -77
  110. package/src/component/setup.test.ts +5 -0
  111. package/src/component/shared.ts +127 -0
  112. package/src/component/stats.ts +20 -6
  113. package/src/component/worker.ts +94 -0
@@ -1,25 +1,20 @@
1
1
  import { v } from "convex/values";
2
+ import { mutation, query } from "./_generated/server.js";
2
3
  import {
3
- ActionCtx,
4
- internalAction,
5
- internalMutation,
6
- mutation,
7
- MutationCtx,
8
- query,
9
- QueryCtx,
10
- } from "./_generated/server.js";
11
- import { FunctionHandle, WithoutSystemFields } from "convex/server";
12
- import { Doc, Id } from "./_generated/dataModel.js";
13
- import { api, internal } from "./_generated/api.js";
14
- import { createLogger, logLevel } from "./logging.js";
15
- import { components } from "./_generated/api.js";
16
- import { Crons } from "@convex-dev/crons";
17
- import { recordCompleted, recordStarted } from "./stats.js";
18
- import { completionStatus } from "./schema.js";
19
-
20
- const crons = new Crons(components.crons);
21
-
22
- const ACTION_TIMEOUT_MS = 15 * 60 * 1000;
4
+ nextSegment,
5
+ onComplete,
6
+ retryBehavior,
7
+ config,
8
+ status as statusValidator,
9
+ toSegment,
10
+ boundScheduledTime,
11
+ } from "./shared.js";
12
+ import { logLevel } from "./logging.js";
13
+ import { kickMainLoop } from "./kick.js";
14
+ import { api } from "./_generated/api.js";
15
+ import { createLogger } from "./logging.js";
16
+
17
+ const MAX_POSSIBLE_PARALLELISM = 100;
23
18
 
24
19
  export const enqueue = mutation({
25
20
  args: {
@@ -27,31 +22,32 @@ export const enqueue = mutation({
27
22
  fnName: v.string(),
28
23
  fnArgs: v.any(),
29
24
  fnType: v.union(v.literal("action"), v.literal("mutation")),
30
- options: v.object({
31
- maxParallelism: v.number(),
32
- logLevel: v.optional(logLevel),
33
- statusTtl: v.optional(v.number()),
34
- }),
25
+ runAt: v.number(),
26
+ // TODO: annotation?
27
+ onComplete: v.optional(onComplete),
28
+ retryBehavior: v.optional(retryBehavior),
29
+ config,
35
30
  },
36
31
  returns: v.id("work"),
37
- handler: async (ctx, { fnHandle, fnName, options, fnArgs, fnType }) => {
38
- await ensurePoolAndLoopExist(
39
- ctx,
40
- {
41
- maxParallelism: options.maxParallelism,
42
- statusTtl: options.statusTtl ?? 24 * 60 * 60 * 1000,
43
- logLevel: options.logLevel ?? "WARN",
44
- },
45
- "enqueue"
46
- );
32
+ handler: async (ctx, { config, runAt, ...workArgs }) => {
33
+ const console = createLogger(config.logLevel);
34
+ if (config.maxParallelism > MAX_POSSIBLE_PARALLELISM) {
35
+ throw new Error(`maxParallelism must be <= ${MAX_POSSIBLE_PARALLELISM}`);
36
+ }
37
+ if (config.maxParallelism < 1) {
38
+ throw new Error("maxParallelism must be >= 1");
39
+ }
40
+ runAt = boundScheduledTime(runAt, console);
47
41
  const workId = await ctx.db.insert("work", {
48
- fnHandle,
49
- fnName,
50
- fnArgs,
51
- fnType,
42
+ ...workArgs,
43
+ attempts: 0,
44
+ });
45
+ await ctx.db.insert("pendingStart", {
46
+ workId,
47
+ segment: toSegment(runAt),
52
48
  });
53
- await ctx.db.insert("pendingStart", { workId });
54
- await kickMainLoop(ctx, "enqueue");
49
+ await kickMainLoop(ctx, "enqueue", config);
50
+ // TODO: stats event
55
51
  return workId;
56
52
  },
57
53
  });
@@ -59,658 +55,71 @@ export const enqueue = mutation({
59
55
  export const cancel = mutation({
60
56
  args: {
61
57
  id: v.id("work"),
58
+ logLevel,
62
59
  },
63
- handler: async (ctx, { id }) => {
64
- await ctx.db.insert("pendingCancelation", { workId: id });
60
+ handler: async (ctx, { id, logLevel }) => {
61
+ await ctx.db.insert("pendingCancelation", {
62
+ workId: id,
63
+ segment: nextSegment(),
64
+ });
65
+ await kickMainLoop(ctx, "cancel", { logLevel });
66
+ // TODO: stats event
65
67
  },
66
68
  });
67
69
 
68
- async function console(ctx: QueryCtx | ActionCtx) {
69
- if ("runAction" in ctx) {
70
- return globalThis.console;
71
- }
72
- const pool = await ctx.db.query("pool").unique();
73
- if (!pool) {
74
- return globalThis.console;
75
- }
76
- return createLogger(pool.logLevel);
77
- }
78
-
79
- const BATCH_SIZE = 10;
80
-
81
- // There should only ever be at most one of these scheduled or running.
82
- // The scheduled one is in the "mainLoop" table.
83
- export const mainLoop = internalMutation({
84
- args: {
85
- generation: v.number(),
86
- },
87
- handler: async (ctx, args) => {
88
- const console_ = await console(ctx);
89
-
90
- const options = (await ctx.db.query("pool").unique())!;
91
- if (!options) {
92
- throw new Error("no pool in mainLoop");
93
- }
94
- const { maxParallelism } = options;
95
- let didSomething = false;
96
-
97
- let inProgressCountChange = 0;
98
-
99
- // Move from pendingCompletion to completedWork, deleting from inProgressWork.
100
- // Generation is used to avoid OCCs with work completing.
101
- console_.time("[mainLoop] pendingCompletion");
102
- const generation = await ctx.db.query("completionGeneration").unique();
103
- const generationNumber = generation?.generation ?? 0;
104
- if (generationNumber !== args.generation) {
105
- throw new Error(
106
- `generation mismatch: ${generationNumber} !== ${args.generation}`
107
- );
108
- }
109
- // Collect all pending completions for the previous generation.
110
- // This won't be too many because the jobs all correspond to being scheduled
111
- // by a single mainLoop (the previous one), so they're limited by MAX_PARALLELISM.
112
- const completed = await ctx.db
113
- .query("pendingCompletion")
114
- .withIndex("generation", (q) => q.eq("generation", generationNumber - 1))
115
- .collect();
116
- console_.debug(`[mainLoop] completing ${completed.length}`);
117
- await Promise.all(
118
- completed.map(async (pendingCompletion) => {
119
- const inProgressWork = await ctx.db
120
- .query("inProgressWork")
121
- .withIndex("workId", (q) => q.eq("workId", pendingCompletion.workId))
122
- .unique();
123
- if (inProgressWork) {
124
- await ctx.db.delete(inProgressWork._id);
125
- inProgressCountChange--;
126
- await ctx.db.insert("completedWork", {
127
- completionStatus: pendingCompletion.completionStatus,
128
- workId: pendingCompletion.workId,
129
- });
130
- const work = (await ctx.db.get(pendingCompletion.workId))!;
131
- console_.info(
132
- recordCompleted(work, pendingCompletion.completionStatus)
133
- );
134
- await ctx.db.delete(work._id);
135
- }
136
- await ctx.db.delete(pendingCompletion._id);
137
- didSomething = true;
138
- })
139
- );
140
- console_.timeEnd("[mainLoop] pendingCompletion");
141
-
142
- console_.time("[mainLoop] inProgress count");
143
- // This is the only function reading and writing inProgressWork,
144
- // and it's bounded by MAX_POSSIBLE_PARALLELISM, so we can
145
- // read it all into memory. BUT we don't have to -- we can just read
146
- // the count from the inProgressCount table.
147
- const inProgressCount = await ctx.db.query("inProgressCount").unique();
148
- const inProgressBefore =
149
- (inProgressCount?.count ?? 0) + inProgressCountChange;
150
- console_.debug(`[mainLoop] ${inProgressBefore} in progress`);
151
- console_.timeEnd("[mainLoop] inProgress count");
152
-
153
- // Move from pendingStart to inProgressWork.
154
- console_.time("[mainLoop] pendingStart");
155
-
156
- // Start reading from the latest cursor _creationTime, which allows us to
157
- // skip over tombstones of pendingStart documents which haven't been cleaned up yet.
158
- // WARNING: this might skip over pendingStart documents if their _creationTime
159
- // was assigned out of order. We handle that below.
160
- const pendingStartCursorDoc = await ctx.db
161
- .query("pendingStartCursor")
162
- .unique();
163
- const pendingStartCursor = pendingStartCursorDoc?.cursor ?? 0;
164
-
165
- // Schedule as many as needed to reach maxParallelism.
166
- const toSchedule = maxParallelism - inProgressBefore;
167
-
168
- const pending = await ctx.db
169
- .query("pendingStart")
170
- .withIndex("by_creation_time", (q) =>
171
- q.gt("_creationTime", pendingStartCursor)
172
- )
173
- .take(toSchedule);
174
- console_.debug(`[mainLoop] scheduling ${pending.length} pending work`);
175
- await Promise.all(
176
- pending.map(async (pendingWork) => {
177
- const { scheduledId, timeoutMs } = await beginWork(ctx, pendingWork);
178
- await ctx.db.insert("inProgressWork", {
179
- running: scheduledId,
180
- timeoutMs,
181
- workId: pendingWork.workId,
182
- });
183
- inProgressCountChange++;
184
- await ctx.db.delete(pendingWork._id);
185
- didSomething = true;
186
- })
187
- );
188
- const newPendingStartCursor =
189
- pending.length > 0
190
- ? pending[pending.length - 1]._creationTime
191
- : pendingStartCursor;
192
- if (!pendingStartCursorDoc) {
193
- await ctx.db.insert("pendingStartCursor", {
194
- cursor: newPendingStartCursor,
195
- });
196
- } else {
197
- await ctx.db.patch(pendingStartCursorDoc._id, {
198
- cursor: newPendingStartCursor,
199
- });
200
- }
201
- console_.timeEnd("[mainLoop] pendingStart");
202
-
203
- console_.time("[mainLoop] pendingCancelation");
204
- const canceled = await ctx.db.query("pendingCancelation").take(BATCH_SIZE);
205
- console_.debug(`[mainLoop] canceling ${canceled.length}`);
70
+ const PAGE_SIZE = 64;
71
+ export const cancelAll = mutation({
72
+ args: { logLevel, before: v.optional(v.number()) },
73
+ handler: async (ctx, { logLevel, before }) => {
74
+ const beforeTime = before ?? Date.now();
75
+ const segment = nextSegment();
76
+ const pageOfWork = await ctx.db
77
+ .query("work")
78
+ .withIndex("by_creation_time", (q) => q.lte("_creationTime", beforeTime))
79
+ .order("desc")
80
+ .take(PAGE_SIZE);
206
81
  await Promise.all(
207
- canceled.map(async (pendingCancelation) => {
208
- const inProgressWork = await ctx.db
209
- .query("inProgressWork")
210
- .withIndex("workId", (q) => q.eq("workId", pendingCancelation.workId))
211
- .unique();
212
- if (inProgressWork) {
213
- await ctx.scheduler.cancel(inProgressWork.running);
214
- await ctx.db.delete(inProgressWork._id);
215
- inProgressCountChange--;
216
- await ctx.db.insert("completedWork", {
217
- workId: pendingCancelation.workId,
218
- completionStatus: "canceled",
219
- });
220
- const work = (await ctx.db.get(pendingCancelation.workId))!;
221
- console_.info(recordCompleted(work, "canceled"));
222
- await ctx.db.delete(work._id);
82
+ pageOfWork.map(async ({ _id }) => {
83
+ if (
84
+ await ctx.db
85
+ .query("pendingCancelation")
86
+ .withIndex("workId", (q) => q.eq("workId", _id))
87
+ .first()
88
+ ) {
89
+ return;
223
90
  }
224
- await ctx.db.delete(pendingCancelation._id);
225
- didSomething = true;
91
+ await ctx.db.insert("pendingCancelation", { workId: _id, segment });
226
92
  })
227
93
  );
228
- console_.timeEnd("[mainLoop] pendingCancelation");
229
-
230
- // In case there are more pending completions at higher generation numbers,
231
- // there's more to do.
232
- if (!didSomething) {
233
- const nextPendingCompletion = await ctx.db
234
- .query("pendingCompletion")
235
- .withIndex("generation", (q) => q.eq("generation", generationNumber))
236
- .first();
237
- didSomething = nextPendingCompletion !== null;
238
- }
239
-
240
- // In case there are more "pendingStart" items we missed due to out-of-order
241
- // _creationTime, we need to check for them. Do that here, when we're
242
- // otherwise idle so it doesn't matter if this function takes a while walking
243
- // tombstones.
244
- if (!didSomething && pendingStartCursorDoc) {
245
- const pendingStartDoc = await ctx.db
246
- .query("pendingStart")
247
- .order("desc")
248
- .first();
249
- if (pendingStartDoc) {
250
- console_.warn(`[mainLoop] missed pendingStart docs; discarding cursor`);
251
- await ctx.db.delete(pendingStartCursorDoc._id);
252
- didSomething = true;
253
- }
254
- }
255
-
256
- if (!didSomething) {
257
- console_.time("[mainLoop] inProgressWork check for unclean exits");
258
- // If all completions are handled, check everything in inProgressWork.
259
- // This will find everything that timed out, failed ungracefully, was
260
- // cancelled, or succeeded without a return value.
261
- const inProgress = await ctx.db.query("inProgressWork").collect();
262
- await Promise.all(
263
- inProgress.map(async (inProgressWork) => {
264
- const result = await checkInProgressWork(ctx, inProgressWork);
265
- if (result !== null) {
266
- console_.warn(
267
- "[mainLoop] inProgressWork finished uncleanly",
268
- inProgressWork.workId,
269
- result
270
- );
271
- inProgressCountChange--;
272
- await ctx.db.delete(inProgressWork._id);
273
- await ctx.db.insert("completedWork", {
274
- workId: inProgressWork.workId,
275
- completionStatus: result.completionStatus,
276
- });
277
- const work = (await ctx.db.get(inProgressWork.workId))!;
278
- console_.info(recordCompleted(work, result.completionStatus));
279
- didSomething = true;
280
- }
281
- })
282
- );
283
- console_.timeEnd("[mainLoop] inProgressWork check for unclean exits");
284
- }
285
-
286
- if (inProgressCountChange !== 0) {
287
- if (inProgressCount) {
288
- await ctx.db.patch(inProgressCount._id, {
289
- count: inProgressCount.count + inProgressCountChange,
290
- });
291
- } else {
292
- await ctx.db.insert("inProgressCount", {
293
- count: inProgressCountChange,
294
- });
295
- }
296
- }
297
-
298
- console_.time("[mainLoop] kickMainLoop");
299
- if (didSomething) {
300
- // There might be more to do.
301
- await loopFromMainLoop(ctx, 0);
302
- } else {
303
- // Decide when to wake up.
304
- const allInProgressWork = await ctx.db.query("inProgressWork").collect();
305
- const nextPending = await ctx.db.query("pendingStart").first();
306
- const nextPendingTime = nextPending
307
- ? nextPending._creationTime
308
- : Number.POSITIVE_INFINITY;
309
- const nextInProgress = allInProgressWork.length
310
- ? Math.min(
311
- ...allInProgressWork
312
- .filter((w) => w.timeoutMs !== null)
313
- .map((w) => w._creationTime + w.timeoutMs!)
314
- )
315
- : Number.POSITIVE_INFINITY;
316
- const nextTime = Math.min(nextPendingTime, nextInProgress);
317
- await loopFromMainLoop(ctx, nextTime - Date.now());
318
- }
319
- console_.timeEnd("[mainLoop] kickMainLoop");
320
- },
321
- });
322
-
323
- async function beginWork(
324
- ctx: MutationCtx,
325
- pendingStart: Doc<"pendingStart">
326
- ): Promise<{
327
- scheduledId: Id<"_scheduled_functions">;
328
- timeoutMs: number | null;
329
- }> {
330
- const console_ = await console(ctx);
331
- const work = await ctx.db.get(pendingStart.workId);
332
- if (!work) {
333
- throw new Error("work not found");
334
- }
335
- console_.info(recordStarted(work));
336
- if (work.fnType === "action") {
337
- return {
338
- scheduledId: await ctx.scheduler.runAfter(
339
- 0,
340
- internal.lib.runActionWrapper,
341
- {
342
- workId: work._id,
343
- fnHandle: work.fnHandle,
344
- fnArgs: work.fnArgs,
345
- }
346
- ),
347
- timeoutMs: ACTION_TIMEOUT_MS,
348
- };
349
- } else if (work.fnType === "mutation") {
350
- return {
351
- scheduledId: await ctx.scheduler.runAfter(
352
- 0,
353
- internal.lib.runMutationWrapper,
354
- {
355
- workId: work._id,
356
- fnHandle: work.fnHandle,
357
- fnArgs: work.fnArgs,
358
- }
359
- ),
360
- timeoutMs: null, // Mutations cannot timeout
361
- };
362
- } else {
363
- throw new Error(`Unexpected fnType ${work.fnType}`);
364
- }
365
- }
366
-
367
- async function checkInProgressWork(
368
- ctx: MutationCtx,
369
- doc: Doc<"inProgressWork">
370
- ): Promise<{
371
- completionStatus: "success" | "error" | "canceled" | "timeout";
372
- } | null> {
373
- const workStatus = await ctx.db.system.get(doc.running);
374
- if (workStatus === null) {
375
- return { completionStatus: "timeout" };
376
- } else if (
377
- workStatus.state.kind === "pending" ||
378
- workStatus.state.kind === "inProgress"
379
- ) {
380
- if (
381
- doc.timeoutMs !== null &&
382
- Date.now() - workStatus._creationTime > doc.timeoutMs
383
- ) {
384
- await ctx.scheduler.cancel(doc.running);
385
- return { completionStatus: "timeout" };
386
- }
387
- } else if (workStatus.state.kind === "success") {
388
- // Usually this would be handled by pendingCompletion, but for "unknown"
389
- // functions, this is how we know that they're done, and we can't get their
390
- // return values.
391
- return { completionStatus: "success" };
392
- } else if (workStatus.state.kind === "canceled") {
393
- return { completionStatus: "canceled" };
394
- } else if (workStatus.state.kind === "failed") {
395
- return { completionStatus: "error" };
396
- }
397
- return null;
398
- }
399
-
400
- export const runActionWrapper = internalAction({
401
- args: {
402
- workId: v.id("work"),
403
- fnHandle: v.string(),
404
- fnArgs: v.any(),
405
- },
406
- handler: async (ctx, { workId, fnHandle: handleStr, fnArgs }) => {
407
- const console_ = await console(ctx);
408
- const fnHandle = handleStr as FunctionHandle<"action">;
409
- try {
410
- await ctx.runAction(fnHandle, fnArgs);
411
- // NOTE: we could run `ctx.runMutation`, but we want to guarantee execution,
412
- // and `ctx.scheduler.runAfter` won't OCC.
413
- await ctx.scheduler.runAfter(0, internal.lib.saveResult, {
414
- workId,
415
- completionStatus: "success",
416
- });
417
- } catch (e: unknown) {
418
- console_.error(e);
419
- await ctx.scheduler.runAfter(0, internal.lib.saveResult, {
420
- workId,
421
- completionStatus: "error",
94
+ if (pageOfWork.length === PAGE_SIZE) {
95
+ await ctx.scheduler.runAfter(0, api.lib.cancelAll, {
96
+ logLevel,
97
+ before: pageOfWork[pageOfWork.length - 1]._creationTime,
422
98
  });
423
99
  }
100
+ await kickMainLoop(ctx, "cancel", { logLevel });
424
101
  },
425
102
  });
426
103
 
427
- export const saveResult = internalMutation({
428
- args: {
429
- workId: v.id("work"),
430
- completionStatus,
431
- },
432
- handler: async (ctx, args) => {
433
- const currentGeneration = await ctx.db
434
- .query("completionGeneration")
435
- .unique();
436
- const generation = currentGeneration?.generation ?? 0;
437
- await ctx.db.insert("pendingCompletion", {
438
- completionStatus: args.completionStatus,
439
- workId: args.workId,
440
- generation,
441
- });
442
- await kickMainLoop(ctx, "saveResult");
443
- },
444
- });
445
-
446
- export const bumpGeneration = internalMutation({
447
- args: {},
448
- handler: async (ctx) => {
449
- const currentGeneration = await ctx.db
450
- .query("completionGeneration")
451
- .unique();
452
- const generation = (currentGeneration?.generation ?? 0) + 1;
453
- if (!currentGeneration) {
454
- await ctx.db.insert("completionGeneration", { generation });
455
- } else {
456
- await ctx.db.patch(currentGeneration._id, { generation });
457
- }
458
- await ctx.scheduler.runAfter(0, internal.lib.mainLoop, { generation });
459
- },
460
- });
461
-
462
- export const runMutationWrapper = internalMutation({
463
- args: {
464
- workId: v.id("work"),
465
- fnHandle: v.string(),
466
- fnArgs: v.any(),
467
- },
468
- handler: async (ctx, { workId, fnHandle: handleStr, fnArgs }) => {
469
- const console_ = await console(ctx);
470
- const fnHandle = handleStr as FunctionHandle<"mutation">;
471
- try {
472
- await ctx.runMutation(fnHandle, fnArgs);
473
- // NOTE: we could run the `saveResult` handler here, or call `ctx.runMutation`,
474
- // but we want the mutation to be a separate transaction to reduce the window for OCCs.
475
- await ctx.scheduler.runAfter(0, internal.lib.saveResult, {
476
- workId,
477
- completionStatus: "success",
478
- });
479
- } catch (e: unknown) {
480
- console_.error(e);
481
- await ctx.scheduler.runAfter(0, internal.lib.saveResult, {
482
- workId,
483
- completionStatus: "error",
484
- });
485
- }
486
- },
487
- });
488
-
489
- async function getMainLoop(ctx: QueryCtx) {
490
- const mainLoop = await ctx.db.query("mainLoop").unique();
491
- if (!mainLoop) {
492
- throw new Error("mainLoop doesn't exist");
493
- }
494
- return mainLoop;
495
- }
496
-
497
- export const stopCleanup = mutation({
498
- args: {},
499
- handler: async (ctx) => {
500
- const cron = await crons.get(ctx, { name: CLEANUP_CRON_NAME });
501
- if (cron) {
502
- await crons.delete(ctx, { id: cron.id });
503
- }
504
- },
505
- });
506
-
507
- async function loopFromMainLoop(ctx: MutationCtx, delayMs: number) {
508
- const console_ = await console(ctx);
509
- const mainLoop = await getMainLoop(ctx);
510
- if (mainLoop.state.kind === "idle") {
511
- throw new Error("mainLoop is idle but `loopFromMainLoop` was called");
512
- }
513
- if (delayMs <= 0) {
514
- console_.debug(
515
- "[mainLoop] mainLoop is actively running and wants to keep running"
516
- );
517
- await ctx.scheduler.runAfter(0, internal.lib.bumpGeneration, {});
518
- if (mainLoop.state.kind !== "running") {
519
- await ctx.db.patch(mainLoop._id, { state: { kind: "running" } });
520
- }
521
- } else if (delayMs < Number.POSITIVE_INFINITY) {
522
- console_.debug(`[mainLoop] mainLoop wants to run after ${delayMs}ms`);
523
- const runAtTime = Date.now() + delayMs;
524
- const fn = await ctx.scheduler.runAt(
525
- runAtTime,
526
- internal.lib.bumpGeneration,
527
- {}
528
- );
529
- await ctx.db.patch(mainLoop._id, {
530
- state: {
531
- kind: "scheduled",
532
- fn,
533
- runAtTime,
534
- },
535
- });
536
- } else {
537
- console_.debug("[mainLoop] mainLoop wants to become idle");
538
- await ctx.db.patch(mainLoop._id, { state: { kind: "idle" } });
539
- }
540
- }
541
-
542
- async function kickMainLoop(
543
- ctx: MutationCtx,
544
- source: "saveResult" | "enqueue"
545
- ): Promise<void> {
546
- const console_ = await console(ctx);
547
-
548
- // Look for mainLoop documents that we want to reschedule.
549
- // Only kick to run now if we're scheduled or idle.
550
- const mainLoop = await getMainLoop(ctx);
551
- if (mainLoop.state.kind === "running") {
552
- console_.debug(
553
- `[${source}] mainLoop is actively running, so we don't need to do anything`
554
- );
555
- return;
556
- }
557
- // mainLoop is scheduled to run later, so we should cancel it and reschedule.
558
- if (mainLoop.state.kind === "scheduled") {
559
- await ctx.scheduler.cancel(mainLoop.state.fn);
560
- }
561
- const currentGeneration = await ctx.db.query("completionGeneration").unique();
562
- const generation = currentGeneration?.generation ?? 0;
563
- await ctx.scheduler.runAfter(0, internal.lib.mainLoop, { generation });
564
- console_.debug(
565
- `[${source}] mainLoop was scheduled later, so reschedule it to run now`
566
- );
567
- await ctx.db.patch(mainLoop._id, { state: { kind: "running" } });
568
- }
569
-
570
104
  export const status = query({
571
- args: {
572
- id: v.id("work"),
573
- },
574
- returns: v.union(
575
- v.object({
576
- kind: v.literal("pending"),
577
- }),
578
- v.object({
579
- kind: v.literal("inProgress"),
580
- }),
581
- v.object({
582
- kind: v.literal("completed"),
583
- completionStatus,
584
- })
585
- ),
105
+ args: { id: v.id("work") },
106
+ returns: statusValidator,
586
107
  handler: async (ctx, { id }) => {
587
- const completedWork = await ctx.db
588
- .query("completedWork")
589
- .withIndex("workId", (q) => q.eq("workId", id))
590
- .unique();
591
- if (completedWork) {
592
- return {
593
- kind: "completed",
594
- completionStatus: completedWork.completionStatus,
595
- } as const;
108
+ const work = await ctx.db.get(id);
109
+ if (!work) {
110
+ return { state: "finished" } as const;
596
111
  }
597
112
  const pendingStart = await ctx.db
598
113
  .query("pendingStart")
599
114
  .withIndex("workId", (q) => q.eq("workId", id))
600
115
  .unique();
601
116
  if (pendingStart) {
602
- return { kind: "pending" } as const;
603
- }
604
- // If it's not pending or completed, it must be in progress.
605
- // Note we do not check inProgressWork, because we don't want to intersect
606
- // mainLoop.
607
- return { kind: "inProgress" } as const;
608
- },
609
- });
610
-
611
- export const MAX_CLEANUP_DOCS = 1000;
612
-
613
- export const cleanup = mutation({
614
- args: {
615
- maxAgeMs: v.number(),
616
- },
617
- handler: async (ctx, { maxAgeMs }) => {
618
- const old = Date.now() - maxAgeMs;
619
- const docs = await ctx.db
620
- .query("completedWork")
621
- .withIndex("by_creation_time", (q) => q.lte("_creationTime", old))
622
- .order("desc")
623
- .take(MAX_CLEANUP_DOCS);
624
- await Promise.all(
625
- docs.map(async (doc) => {
626
- await ctx.db.delete(doc._id);
627
- const work = await ctx.db.get(doc.workId);
628
- if (work) {
629
- await ctx.db.delete(work._id);
630
- }
631
- })
632
- );
633
- if (docs.length === MAX_CLEANUP_DOCS) {
634
- // Schedule the next cleanup to run starting from the oldest document.
635
- await ctx.scheduler.runAfter(0, api.lib.cleanup, {
636
- maxAgeMs: docs[docs.length - 1]._creationTime,
637
- });
117
+ return { state: "pending", attempt: work.attempts } as const;
638
118
  }
119
+ // Assume it's in progress. It could be pending cancelation
120
+ return { state: "running", attempt: work.attempts } as const;
639
121
  },
640
122
  });
641
123
 
642
- const MAX_POSSIBLE_PARALLELISM = 300;
643
- const CLEANUP_CRON_NAME = "cleanup";
644
-
645
- async function ensurePoolAndLoopExist(
646
- ctx: MutationCtx,
647
- opts: WithoutSystemFields<Doc<"pool">>,
648
- source: "enqueue" | "saveResult" | "mainLoop"
649
- ) {
650
- if (opts.maxParallelism > MAX_POSSIBLE_PARALLELISM) {
651
- throw new Error(`maxParallelism must be <= ${MAX_POSSIBLE_PARALLELISM}`);
652
- }
653
- if (opts.maxParallelism < 1) {
654
- throw new Error("maxParallelism must be >= 1");
655
- }
656
- const pool = await ctx.db.query("pool").unique();
657
- if (pool) {
658
- let update = false;
659
- for (const key in opts) {
660
- if (pool[key as keyof typeof opts] !== opts[key as keyof typeof opts]) {
661
- update = true;
662
- }
663
- }
664
- if (update) {
665
- await ctx.db.patch(pool._id, opts);
666
- }
667
- } else {
668
- const console_ = await console(ctx);
669
- await ctx.db.insert("pool", opts);
670
- console_.debug(`[${source}] starting mainLoop`);
671
- const exists = await ctx.db.query("mainLoop").unique();
672
- if (exists) {
673
- throw new Error("mainLoop already exists");
674
- }
675
- await ctx.db.insert("mainLoop", { state: { kind: "running" } });
676
- const currentGeneration = await ctx.db
677
- .query("completionGeneration")
678
- .unique();
679
- if (currentGeneration) {
680
- throw new Error("completionGeneration already exists");
681
- }
682
- await ctx.db.insert("completionGeneration", { generation: 0 });
683
- await ctx.scheduler.runAfter(0, internal.lib.mainLoop, { generation: 0 });
684
- }
685
- await ensureCleanupCron(ctx, opts.statusTtl);
686
- }
687
-
688
- async function ensureCleanupCron(ctx: MutationCtx, ttl: number) {
689
- let cleanupCron = await crons.get(ctx, { name: CLEANUP_CRON_NAME });
690
- if (ttl === Number.POSITIVE_INFINITY) {
691
- if (cleanupCron) {
692
- await crons.delete(ctx, { id: cleanupCron.id });
693
- }
694
- return;
695
- }
696
- const cronFrequencyMs = Math.min(ttl, 24 * 60 * 60 * 1000);
697
- if (
698
- cleanupCron !== null &&
699
- !(
700
- cleanupCron.schedule.kind === "interval" &&
701
- cleanupCron.schedule.ms === cronFrequencyMs
702
- )
703
- ) {
704
- await crons.delete(ctx, { id: cleanupCron.id });
705
- cleanupCron = null;
706
- }
707
- if (cleanupCron === null) {
708
- await crons.register(
709
- ctx,
710
- { kind: "interval", ms: ttl },
711
- api.lib.cleanup,
712
- { maxAgeMs: ttl },
713
- CLEANUP_CRON_NAME
714
- );
715
- }
716
- }
124
+ // eslint-disable-next-line @typescript-eslint/no-unused-vars
125
+ const console = "THIS IS A REMINDER TO USE createLogger";