UNPKG

durable-execution-storage-convex

Version:

Convex storage implementation for durable-execution

724 lines (718 loc) 22.3 kB
// src/component/lib.ts import { v as v3 } from "convex/values"; // src/common.ts import { v as v2 } from "convex/values"; import { applyTaskExecutionStorageUpdate } from "durable-execution"; import { omitUndefinedValues } from "@gpahal/std/objects"; // src/component/schema.ts import { defineSchema, defineTable } from "convex/server"; import { v } from "convex/values"; var vDurableExecutionErrorType = v.union( v.literal("generic"), v.literal("not_found"), v.literal("timed_out"), v.literal("cancelled") ); var vDurableExecutionError = v.object({ errorType: vDurableExecutionErrorType, message: v.string(), isRetryable: v.boolean(), isInternal: v.boolean() }); var vTaskExecutionStatus = v.union( v.literal("ready"), v.literal("running"), v.literal("failed"), v.literal("timed_out"), v.literal("waiting_for_children"), v.literal("waiting_for_finalize"), v.literal("finalize_failed"), v.literal("completed"), v.literal("cancelled") ); var vTaskExecutionOnChildrenFinishedProcessingStatus = v.union( v.literal("idle"), v.literal("processing"), v.literal("processed") ); var vTaskExecutionCloseStatus = v.union( v.literal("idle"), v.literal("ready"), v.literal("closing"), v.literal("closed") ); var vTaskExecutionDBInsertValue = v.object({ shard: v.number(), rootTaskId: v.optional(v.string()), rootExecutionId: v.optional(v.string()), parentTaskId: v.optional(v.string()), parentExecutionId: v.optional(v.string()), parentExecutionDocId: v.optional(v.id("taskExecutions")), indexInParentChildren: v.optional(v.number()), isOnlyChildOfParent: v.optional(v.boolean()), isFinalizeOfParent: v.optional(v.boolean()), taskId: v.string(), executionId: v.string(), isSleepingTask: v.boolean(), sleepingTaskUniqueId: v.optional(v.string()), retryOptions: v.object({ maxAttempts: v.number(), baseDelayMs: v.optional(v.number()), maxDelayMs: v.optional(v.number()), delayMultiplier: v.optional(v.number()) }), sleepMsBeforeRun: v.number(), timeoutMs: v.number(), areChildrenSequential: v.boolean(), input: v.string(), executorId: v.optional(v.string()), status: vTaskExecutionStatus, isFinished: v.boolean(), runOutput: v.optional(v.string()), output: v.optional(v.string()), error: v.optional(vDurableExecutionError), retryAttempts: v.number(), startAt: v.number(), startedAt: v.optional(v.number()), expiresAt: v.optional(v.number()), waitingForChildrenStartedAt: v.optional(v.number()), waitingForFinalizeStartedAt: v.optional(v.number()), finishedAt: v.optional(v.number()), children: v.optional( v.array( v.object({ taskId: v.string(), executionId: v.string() }) ) ), acc: v.number(), ocfpStatus: vTaskExecutionOnChildrenFinishedProcessingStatus, ocfpExpiresAt: v.optional(v.number()), ocfpFinishedAt: v.optional(v.number()), finalize: v.optional( v.object({ taskId: v.string(), executionId: v.string() }) ), closeStatus: vTaskExecutionCloseStatus, closeExpiresAt: v.optional(v.number()), closedAt: v.optional(v.number()), npc: v.boolean(), createdAt: v.number(), updatedAt: v.number() }); var schema_default = defineSchema({ taskExecutions: defineTable(vTaskExecutionDBInsertValue).index("by_executionId", ["executionId"]).index("by_sleepingTaskUniqueId", ["sleepingTaskUniqueId"]).index("by_shard_status_startAt", ["shard", "status", "startAt"]).index("by_shard_status_ocfpStatus_acc_updatedAt", [ "shard", "status", "ocfpStatus", "acc", "updatedAt" ]).index("by_shard_closeStatus_updatedAt", ["shard", "closeStatus", "updatedAt"]).index("by_status_isSleepingTask_expiresAt", ["status", "isSleepingTask", "expiresAt"]).index("by_ocfpExpiresAt", ["ocfpExpiresAt"]).index("by_closeExpiresAt", ["closeExpiresAt"]).index("by_shard_executorId_npc_updatedAt", ["shard", "executorId", "npc", "updatedAt"]).index("by_parentExecutionId_isFinished", ["parentExecutionId", "isFinished"]).index("by_shard_isFinished_closeStatus_updatedAt", [ "shard", "isFinished", "closeStatus", "updatedAt" ]), locks: defineTable({ key: v.string(), expiresAt: v.number() }).index("by_key", ["key"]).index("by_key_expiresAt", ["key", "expiresAt"]) }); // src/common.ts var vTaskExecutionDBUpdateRequest = v2.object({ executorId: v2.optional(v2.string()), status: v2.optional(vTaskExecutionStatus), isFinished: v2.optional(v2.boolean()), runOutput: v2.optional(v2.string()), output: v2.optional(v2.string()), error: v2.optional(vDurableExecutionError), retryAttempts: v2.optional(v2.number()), startAt: v2.optional(v2.number()), startedAt: v2.optional(v2.number()), expiresAt: v2.optional(v2.number()), waitingForChildrenStartedAt: v2.optional(v2.number()), waitingForFinalizeStartedAt: v2.optional(v2.number()), finishedAt: v2.optional(v2.number()), children: v2.optional( v2.array( v2.object({ taskId: v2.string(), executionId: v2.string() }) ) ), acc: v2.optional(v2.number()), ocfpStatus: v2.optional(vTaskExecutionOnChildrenFinishedProcessingStatus), ocfpExpiresAt: v2.optional(v2.number()), ocfpFinishedAt: v2.optional(v2.number()), finalize: v2.optional( v2.object({ taskId: v2.string(), executionId: v2.string() }) ), closeStatus: v2.optional(vTaskExecutionCloseStatus), closeExpiresAt: v2.optional(v2.number()), closedAt: v2.optional(v2.number()), npc: v2.optional(v2.boolean()), updatedAt: v2.number(), unset: v2.optional( v2.object({ executorId: v2.optional(v2.boolean()), runOutput: v2.optional(v2.boolean()), error: v2.optional(v2.boolean()), startedAt: v2.optional(v2.boolean()), expiresAt: v2.optional(v2.boolean()), ocfpExpiresAt: v2.optional(v2.boolean()), closeExpiresAt: v2.optional(v2.boolean()) }) ) }); function taskExecutionStorageUpdateRequestToDBUpdate(update) { const dbUpdate = omitUndefinedValues({ ...update, unset: void 0 }); const updateUnset = update.unset; if (updateUnset) { for (const key in updateUnset) { if (updateUnset[key]) { dbUpdate[key] = void 0; } } } return dbUpdate; } var vTaskExecutionStorageGetByIdFilters = v2.object({ isSleepingTask: v2.optional(v2.boolean()), status: v2.optional(vTaskExecutionStatus), isFinished: v2.optional(v2.boolean()) }); function applyTaskExecutionIdFilters(execution, filters) { if (filters?.isSleepingTask != null && execution.isSleepingTask !== filters.isSleepingTask) { return false; } if (filters?.status != null && execution.status !== filters.status) { return false; } if (filters?.isFinished != null && execution.isFinished !== filters.isFinished) { return false; } return true; } // src/component/_generated/api.js import { anyApi, componentsGeneric } from "convex/server"; var internal = anyApi; var components = componentsGeneric(); // src/component/_generated/server.js import { actionGeneric, httpActionGeneric, queryGeneric, mutationGeneric, internalActionGeneric, internalMutationGeneric, internalQueryGeneric, componentsGeneric as componentsGeneric2 } from "convex/server"; var query = queryGeneric; var mutation = mutationGeneric; var internalMutation = internalMutationGeneric; var action = actionGeneric; // src/component/lib.ts var LOCK_EXPIRATION_MS = 6e4; var acquireLock = mutation({ args: { key: v3.string() }, handler: async (ctx, { key }) => { const lock = await ctx.db.query("locks").withIndex("by_key", (q) => q.eq("key", key)).first(); if (!lock) { const lockId = await ctx.db.insert("locks", { key, expiresAt: Date.now() + LOCK_EXPIRATION_MS }); return lockId; } if (lock.expiresAt <= Date.now()) { await ctx.db.patch(lock._id, { expiresAt: Date.now() + LOCK_EXPIRATION_MS }); return lock._id; } return null; } }); var releaseLock = mutation({ args: { id: v3.id("locks") }, handler: async (ctx, { id }) => { await ctx.db.delete(id); } }); async function insertOneHelper(ctx, execution) { let existingDBValue = await ctx.db.query("taskExecutions").withIndex("by_executionId", (q) => q.eq("executionId", execution.executionId)).first(); if (existingDBValue) { throw new Error("Execution id already exists"); } if (execution.sleepingTaskUniqueId) { existingDBValue = await ctx.db.query("taskExecutions").withIndex( "by_sleepingTaskUniqueId", (q) => q.eq("sleepingTaskUniqueId", execution.sleepingTaskUniqueId) ).first(); if (existingDBValue) { throw new Error("Sleeping task unique id already exists"); } } const now = Date.now() + 50; execution.createdAt = now; execution.updatedAt = now; if (execution.startAt < now) { execution.startAt = now; } if (execution.expiresAt && execution.expiresAt < now) { execution.expiresAt = now + 1e3; } if (execution.ocfpExpiresAt && execution.ocfpExpiresAt < now) { execution.ocfpExpiresAt = now + 1e3; } if (execution.closeExpiresAt && execution.closeExpiresAt < now) { execution.closeExpiresAt = now + 1e3; } await ctx.db.insert("taskExecutions", execution); } async function insertManyHelper(ctx, executions) { for (const execution of executions) { await insertOneHelper(ctx, execution); } } async function updateOneHelper(ctx, dbValue, update) { const now = Date.now() + 50; if (update.updatedAt < now) { update.updatedAt = now; } if (update.startAt && update.startAt < now) { update.startAt = now; } if (update.startedAt && update.startedAt < now) { update.startedAt = now; } if (update.expiresAt && update.expiresAt < now) { update.expiresAt = now + 1e3; } if (update.finishedAt && update.finishedAt < now) { update.finishedAt = now; } if (update.ocfpExpiresAt && update.ocfpExpiresAt < now) { update.ocfpExpiresAt = now + 1e3; } if (update.ocfpFinishedAt && update.ocfpFinishedAt < now) { update.ocfpFinishedAt = now; } if (update.closeExpiresAt && update.closeExpiresAt < now) { update.closeExpiresAt = now + 1e3; } if (update.closedAt && update.closedAt < now) { update.closedAt = now; } await ctx.db.patch(dbValue._id, taskExecutionStorageUpdateRequestToDBUpdate(update)); } async function updateManyHelper(ctx, dbValues, update) { for (const dbValue of dbValues) { await updateOneHelper(ctx, dbValue, update); } } var insertMany = mutation({ args: { executions: v3.array(vTaskExecutionDBInsertValue) }, handler: async (ctx, args) => { await insertManyHelper(ctx, args.executions); } }); var getManyById = query({ args: { requests: v3.array( v3.object({ executionId: v3.string(), filters: v3.optional(vTaskExecutionStorageGetByIdFilters) }) ) }, handler: async (ctx, { requests }) => { const dbValues = []; for (const { executionId, filters } of requests) { const dbValue = await ctx.db.query("taskExecutions").withIndex("by_executionId", (q) => q.eq("executionId", executionId)).first(); if (!dbValue || !applyTaskExecutionIdFilters(dbValue, filters)) { dbValues.push(null); } else { dbValues.push(dbValue); } } return dbValues; } }); var getManyBySleepingTaskUniqueId = query({ args: { requests: v3.array( v3.object({ sleepingTaskUniqueId: v3.string() }) ) }, handler: async (ctx, { requests }) => { const dbValues = []; for (const { sleepingTaskUniqueId } of requests) { const dbValue = await ctx.db.query("taskExecutions").withIndex( "by_sleepingTaskUniqueId", (q) => q.eq("sleepingTaskUniqueId", sleepingTaskUniqueId) ).first(); if (!dbValue) { dbValues.push(null); } else { dbValues.push(dbValue); } } return dbValues; } }); var updateManyById = mutation({ args: { requests: v3.array( v3.object({ executionId: v3.string(), filters: v3.optional(vTaskExecutionStorageGetByIdFilters), update: vTaskExecutionDBUpdateRequest }) ) }, handler: async (ctx, { requests }) => { for (const { executionId, filters, update } of requests) { const dbValue = await ctx.db.query("taskExecutions").withIndex("by_executionId", (q) => q.eq("executionId", executionId)).first(); if (!dbValue || !applyTaskExecutionIdFilters(dbValue, filters)) { continue; } await updateOneHelper(ctx, dbValue, update); } } }); var updateManyByIdAndInsertChildrenIfUpdated = mutation({ args: { requests: v3.array( v3.object({ executionId: v3.string(), filters: v3.optional(vTaskExecutionStorageGetByIdFilters), update: vTaskExecutionDBUpdateRequest, childrenTaskExecutionsToInsertIfAnyUpdated: v3.array(vTaskExecutionDBInsertValue) }) ) }, handler: async (ctx, { requests }) => { for (const { executionId, filters, update, childrenTaskExecutionsToInsertIfAnyUpdated } of requests) { const dbValue = await ctx.db.query("taskExecutions").withIndex("by_executionId", (q) => q.eq("executionId", executionId)).first(); if (!dbValue || !applyTaskExecutionIdFilters(dbValue, filters)) { continue; } await updateOneHelper(ctx, dbValue, update); for (const execution of childrenTaskExecutionsToInsertIfAnyUpdated) { execution.parentExecutionDocId = dbValue._id; } await insertManyHelper(ctx, childrenTaskExecutionsToInsertIfAnyUpdated); } } }); var updateByStatusAndStartAtLessThanAndReturn = mutation( async (ctx, { shard, status, startAtLessThan, update, updateExpiresAtWithStartedAt, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_shard_status_startAt", (q) => q.eq("shard", shard).eq("status", status).lt("startAt", startAtLessThan) ).order("asc").take(limit); for (const dbValue of dbValues) { const finalUpdate = { ...update, expiresAt: updateExpiresAtWithStartedAt + dbValue.timeoutMs }; await updateOneHelper(ctx, dbValue, finalUpdate); } return dbValues; } ); var updateByStatusAndOCFPStatusAndACCZeroAndReturn = mutation( async (ctx, { shard, status, ocfpStatus, update, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_shard_status_ocfpStatus_acc_updatedAt", (q) => q.eq("shard", shard).eq("status", status).eq("ocfpStatus", ocfpStatus).eq("acc", 0).lt("updatedAt", Date.now()) ).order("asc").take(limit); if (dbValues.length === 0) { return []; } await updateManyHelper(ctx, dbValues, update); return dbValues; } ); var updateByCloseStatusAndReturn = mutation( async (ctx, { shard, closeStatus, update, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_shard_closeStatus_updatedAt", (q) => q.eq("shard", shard).eq("closeStatus", closeStatus).lt("updatedAt", Date.now()) ).order("asc").take(limit); if (dbValues.length === 0) { return []; } await updateManyHelper(ctx, dbValues, update); return dbValues; } ); var updateByStatusAndIsSleepingTaskAndExpiresAtLessThan = mutation( async (ctx, { status, isSleepingTask, expiresAtLessThan, update, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_status_isSleepingTask_expiresAt", (q) => q.eq("status", status).eq("isSleepingTask", isSleepingTask).gte("expiresAt", 0).lt("expiresAt", expiresAtLessThan) ).order("asc").take(limit); if (dbValues.length === 0) { return 0; } await updateManyHelper(ctx, dbValues, update); return dbValues.length; } ); var updateByOCFPExpiresAt = mutation( async (ctx, { ocfpExpiresAtLessThan, update, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_ocfpExpiresAt", (q) => q.gte("ocfpExpiresAt", 0).lt("ocfpExpiresAt", ocfpExpiresAtLessThan) ).order("asc").take(limit); if (dbValues.length === 0) { return 0; } await updateManyHelper(ctx, dbValues, update); return dbValues.length; } ); var updateByCloseExpiresAt = mutation( async (ctx, { closeExpiresAtLessThan, update, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_closeExpiresAt", (q) => q.gte("closeExpiresAt", 0).lt("closeExpiresAt", closeExpiresAtLessThan) ).order("asc").take(limit); if (dbValues.length === 0) { return 0; } await updateManyHelper(ctx, dbValues, update); return dbValues.length; } ); var updateByExecutorIdAndNPCAndReturn = mutation( async (ctx, { shard, executorId, npc, update, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_shard_executorId_npc_updatedAt", (q) => q.eq("shard", shard).eq("executorId", executorId).eq("npc", npc).lt("updatedAt", Date.now()) ).order("asc").take(limit); if (dbValues.length === 0) { return []; } await updateManyHelper(ctx, dbValues, update); return dbValues; } ); var getManyByParentExecutionId = query({ args: { requests: v3.array( v3.object({ parentExecutionId: v3.string() }) ) }, handler: async (ctx, { requests }) => { const dbValues = []; for (const { parentExecutionId } of requests) { const currDbValues = await ctx.db.query("taskExecutions").withIndex( "by_parentExecutionId_isFinished", (q) => q.eq("parentExecutionId", parentExecutionId) ).collect(); dbValues.push(currDbValues); } return dbValues; } }); var updateManyByParentExecutionIdAndIsFinished = mutation({ args: { requests: v3.array( v3.object({ parentExecutionId: v3.string(), isFinished: v3.boolean(), update: vTaskExecutionDBUpdateRequest }) ) }, handler: async (ctx, { requests }) => { for (const { parentExecutionId, isFinished, update } of requests) { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_parentExecutionId_isFinished", (q) => q.eq("parentExecutionId", parentExecutionId).eq("isFinished", isFinished) ).collect(); await updateManyHelper(ctx, dbValues, update); } } }); var updateAndDecrementParentACCByIsFinishedAndCloseStatus = mutation( async (ctx, { shard, isFinished, closeStatus, update, limit }) => { const dbValues = await ctx.db.query("taskExecutions").withIndex( "by_shard_isFinished_closeStatus_updatedAt", (q) => q.eq("shard", shard).eq("isFinished", isFinished).eq("closeStatus", closeStatus).lt("updatedAt", Date.now()) ).order("asc").take(limit); if (dbValues.length === 0) { return 0; } await updateManyHelper(ctx, dbValues, update); const parentExecutionDocIdToDecrementValueMap = /* @__PURE__ */ new Map(); const singleChildParentExecutionDocIds = /* @__PURE__ */ new Set(); for (const dbValue of dbValues) { if (dbValue.parentExecutionDocId != null && !dbValue.isFinalizeOfParent) { if (dbValue.isOnlyChildOfParent) { singleChildParentExecutionDocIds.add(dbValue.parentExecutionDocId); } else { parentExecutionDocIdToDecrementValueMap.set( dbValue.parentExecutionDocId, (parentExecutionDocIdToDecrementValueMap.get(dbValue.parentExecutionDocId) ?? 0) + 1 ); } } } if (singleChildParentExecutionDocIds.size > 0) { for (const parentExecutionDocId of singleChildParentExecutionDocIds) { await ctx.db.patch(parentExecutionDocId, { acc: 0, updatedAt: update.updatedAt }); } } if (parentExecutionDocIdToDecrementValueMap.size > 0) { const parentExecutionDocIds = [...parentExecutionDocIdToDecrementValueMap.keys()].sort(); for (const parentDBId of parentExecutionDocIds) { const parentDBValue = await ctx.db.get(parentDBId); if (parentDBValue) { await updateOneHelper(ctx, parentDBValue, { acc: parentDBValue.acc - (parentExecutionDocIdToDecrementValueMap.get(parentDBId) ?? 0), updatedAt: update.updatedAt }); } } } return dbValues.length; } ); var deleteById = mutation({ args: { executionId: v3.string() }, handler: async (ctx, { executionId }) => { const dbValue = await ctx.db.query("taskExecutions").withIndex("by_executionId", (q) => q.eq("executionId", executionId)).first(); if (!dbValue) { return; } await ctx.db.delete(dbValue._id); } }); var deleteMany = internalMutation({ args: { batchSize: v3.number() }, handler: async (ctx, { batchSize }) => { const dbValues = await ctx.db.query("taskExecutions").take(batchSize); if (dbValues.length === 0) { return 0; } for (const dbValue of dbValues) { await ctx.db.delete(dbValue._id); } return dbValues.length; } }); var deleteAll = action({ args: {}, handler: async (ctx) => { while (true) { const deletedCount = await ctx.runMutation(internal.lib.deleteMany, { batchSize: 100 }); if (deletedCount < 100) { break; } } } }); export { acquireLock, deleteAll, deleteById, deleteMany, getManyById, getManyByParentExecutionId, getManyBySleepingTaskUniqueId, insertMany, releaseLock, updateAndDecrementParentACCByIsFinishedAndCloseStatus, updateByCloseExpiresAt, updateByCloseStatusAndReturn, updateByExecutorIdAndNPCAndReturn, updateByOCFPExpiresAt, updateByStatusAndIsSleepingTaskAndExpiresAtLessThan, updateByStatusAndOCFPStatusAndACCZeroAndReturn, updateByStatusAndStartAtLessThanAndReturn, updateManyById, updateManyByIdAndInsertChildrenIfUpdated, updateManyByParentExecutionIdAndIsFinished }; //# sourceMappingURL=lib.js.map