1// Alchemy modifications are licensed under Apache-2.0.
2// This file includes third-party code; see /THIRD_PARTY_LICENSES.md.
3import { InstanceEvent } from "../instance.ts";
4import { MODIFIER_KEYS } from "../modifier.ts";
5import type { RestartFromStep } from "../binding.ts";
6import type { RawInstanceLog } from "../instance.ts";
7
8const EVENT_MAP_PREFIX = "EVENT_MAP";
9const RESTART_FROM_STEP_KEY = "RESTART_FROM_STEP";
10
11const KV_STEP_SUFFIXES = [
12 "-value",
13 "-error",
14 "-config",
15 "-metadata",
16 "-log-written",
17 "-pending",
18 "-value-stream-meta",
19] as const;
20
21export function resolveGroupKeysToWipe(
22 sql: DurableObjectStorage["sql"],
23 param: RestartFromStep,
24): Set<string> | null {
25 const stepTypeToEvent: Record<string, InstanceEvent> = {
26 do: InstanceEvent.STEP_START,
27 sleep: InstanceEvent.SLEEP_START,
28 waitForEvent: InstanceEvent.WAIT_START,
29 };
30 const targetEvent = stepTypeToEvent[param.type ?? "do"];
31 const targetCount = param.count ?? 1;
32
33 const cursor = sql.exec<RawInstanceLog>(
34 "SELECT event, target, groupKey FROM states WHERE event IN (?, ?, ?) ORDER BY id",
35 InstanceEvent.STEP_START,
36 InstanceEvent.SLEEP_START,
37 InstanceEvent.WAIT_START,
38 );
39
40 let nameOccurrence = 0;
41 let found = false;
42 const groupKeys = new Set<string>();
43
44 for (const row of cursor) {
45 if (row.groupKey === null) {
46 continue;
47 }
48 const groupKey = String(row.groupKey);
49
50 if (found) {
51 groupKeys.add(groupKey);
52 continue;
53 }
54
55 if (row.target === null) {
56 continue;
57 }
58 const rawStepName = String(row.target).replace(/-\d+$/, "");
59 if (rawStepName !== param.name) {
60 continue;
61 }
62 if (row.event !== targetEvent) {
63 continue;
64 }
65
66 nameOccurrence++;
67 if (nameOccurrence === targetCount) {
68 found = true;
69 groupKeys.add(groupKey);
70 }
71 }
72
73 return found ? groupKeys : null;
74}
75
76function getMockedEventMapKeys(allKeys: Map<string, unknown>): Set<string> {
77 const mockEventTypes = new Set<string>();
78 for (const key of allKeys.keys()) {
79 if (key.startsWith(MODIFIER_KEYS.MOCK_EVENT)) {
80 mockEventTypes.add(key.slice(MODIFIER_KEYS.MOCK_EVENT.length));
81 }
82 }
83
84 if (mockEventTypes.size === 0) {
85 return new Set();
86 }
87
88 const preserved = new Set<string>();
89 for (const key of allKeys.keys()) {
90 if (key.startsWith(`${EVENT_MAP_PREFIX}\n`)) {
91 const eventType = key.split("\n")[1];
92 if (eventType !== undefined && mockEventTypes.has(eventType)) {
93 preserved.add(key);
94 }
95 }
96 }
97
98 return preserved;
99}
100
101async function deleteGroupKeyBatch(
102 storage: DurableObjectStorage,
103 batch: Array<string>,
104): Promise<void> {
105 const kvKeys: Array<string> = [];
106 for (const groupKey of batch) {
107 for (const suffix of KV_STEP_SUFFIXES) {
108 kvKeys.push(`${groupKey}${suffix}`);
109 }
110 storage.sql.exec("DELETE FROM states WHERE groupKey = ?", groupKey);
111 storage.sql.exec(
112 "DELETE FROM streaming_step_chunks WHERE attempt != 0 AND cache_key = ?",
113 groupKey,
114 );
115 }
116 await storage.delete(kvKeys);
117}
118
119export async function wipeRestartState(
120 storage: DurableObjectStorage,
121 engineStatusKey: string,
122 pauseDatetimeKey: string,
123 groupKeysToWipe: Set<string> | null,
124): Promise<void> {
125 if (groupKeysToWipe) {
126 await deleteGroupKeyBatch(storage, Array.from(groupKeysToWipe));
127 storage.sql.exec(
128 "DELETE FROM states WHERE groupKey IS NULL AND event NOT IN (?, ?)",
129 InstanceEvent.WORKFLOW_START,
130 InstanceEvent.WORKFLOW_QUEUED,
131 );
132 } else {
133 const cursor = storage.sql.exec<{ groupKey: string }>(
134 "SELECT DISTINCT groupKey FROM states WHERE groupKey IS NOT NULL",
135 );
136 const allGroupKeys = [...cursor].map((r) => r.groupKey);
137 if (allGroupKeys.length > 0) {
138 await deleteGroupKeyBatch(storage, allGroupKeys);
139 }
140 storage.sql.exec("DELETE FROM states");
141 }
142
143 const keysToDelete: Array<string> = [engineStatusKey, pauseDatetimeKey];
144
145 const allKeys = await storage.list();
146 const preservedEventMapKeys = getMockedEventMapKeys(allKeys);
147 for (const key of allKeys.keys()) {
148 if (key.startsWith(`${EVENT_MAP_PREFIX}\n`) && !preservedEventMapKeys.has(key)) {
149 keysToDelete.push(key);
150 }
151 }
152 await storage.delete(keysToDelete);
153
154 storage.sql.exec("DELETE FROM priority_queue");
155}
156
157export async function readAndClearRestartFromStep(
158 storage: DurableObjectStorage,
159): Promise<RestartFromStep | undefined> {
160 const value = await storage.get<RestartFromStep>(RESTART_FROM_STEP_KEY);
161 await storage.delete(RESTART_FROM_STEP_KEY);
162 return value;
163}
164
165export async function storeRestartFromStep(
166 storage: DurableObjectStorage,
167 from: RestartFromStep,
168): Promise<void> {
169 await storage.put(RESTART_FROM_STEP_KEY, from);
170}