1// Alchemy modifications are licensed under Apache-2.0.
2// This file includes third-party code; see /THIRD_PARTY_LICENSES.md.
3import Heap from "heap-js";
4import type { InstanceMetadata, WakerPriorityEntry, WakerPriorityType } from "../instance.ts";
5
6const wakerPriorityEntryComparator = (a: WakerPriorityEntry, b: WakerPriorityEntry) => {
7 return a.targetTimestamp - b.targetTimestamp;
8};
9
10const enum SQLiteBoolean {
11 FALSE = 0,
12 TRUE = 1,
13}
14
15const enum EntryType {
16 RETRY = 0,
17 SLEEP = 1,
18 TIMEOUT = 2,
19}
20
21type PriorityQueueDBEntry = {
22 id: number;
23 created_on: string;
24 target_timestamp: number;
25 action: SQLiteBoolean;
26 entryType: number;
27 hash: string;
28};
29
30export class TimePriorityQueue {
31 #heap: Heap<WakerPriorityEntry> = new Heap(wakerPriorityEntryComparator);
32 #ctx: DurableObjectState;
33
34 constructor(ctx: DurableObjectState, _instanceMetadata: InstanceMetadata) {
35 this.#ctx = ctx;
36
37 this.#heap.init(this.getEntries());
38 }
39
40 popPastEntries(): Array<WakerPriorityEntry> | undefined {
41 // early return if there is nothing in the queue
42 if (this.#heap.length === 0) {
43 return;
44 }
45 const res: Array<WakerPriorityEntry> = [];
46 const currentTimestamp = new Date().valueOf();
47 // heap-js does not have a ordered iterator that doesn't consume the input so we
48 // peek the first one, and pop if it's old until it's empty or in the future
49 while (true) {
50 const element = this.#heap.peek();
51 if (element === undefined) {
52 break;
53 }
54 if (element.targetTimestamp > currentTimestamp) {
55 break;
56 }
57 // at this point, targetTimestamp is older so we can pop this node because it's no
58 // longer relevant
59 res.push(element);
60 this.#heap.pop();
61 }
62 this.#ctx.storage.transactionSync(() => {
63 for (const entry of res) {
64 this.removeEntryDB(entry);
65 }
66 });
67 return res;
68 }
69
70 /**
71 * `add` is ran using a transaction so it's race condition free, if it's ran atomically
72 * @param entry
73 */
74 async add(entry: WakerPriorityEntry) {
75 await this.#ctx.storage.transaction(async () => {
76 // TODO: Handle this
77 // const waker = this.#env.WAKERS.idFromName(
78 // this.#instanceMetadata.instance.id
79 // );
80 // const wakerStub = this.#env.WAKERS.get(waker);
81 // We can optimise this by only calling it if time is sooner than any other
82 // await wakerStub.wake(
83 // new Date(entry.targetTimestamp),
84 // this.#instanceMetadata
85 // );
86
87 this.#heap.add(entry);
88 this.addEntryDB(entry);
89 });
90 }
91
92 /**
93 * `remove` is ran using a transaction so it's race condition free, if it's ran atomically
94 * @param entry
95 */
96 remove(entry: Omit<WakerPriorityEntry, "targetTimestamp">) {
97 this.#ctx.storage.transactionSync(() => {
98 this.removeFirst((e) => {
99 if (e.hash === entry.hash && e.type === entry.type) {
100 return true;
101 }
102 return false;
103 });
104 });
105 }
106
107 offsetAll(offset: number) {
108 // Clear the entire PQ table and re-insert only the offset entries.
109 // We can't use the append-only add/remove pattern here because the
110 // UNIQUE (action, entryType, hash) constraint would conflict with
111 // the original action=1 rows still in the table.
112 this.#ctx.storage.transactionSync(() => {
113 const entries = this.#heap.toArray();
114
115 // Wipe the table — removes all historical add/remove rows
116 this.#ctx.storage.sql.exec("DELETE FROM priority_queue");
117
118 const newEntries = entries.map((value) => ({
119 ...value,
120 targetTimestamp: value.targetTimestamp + offset,
121 }));
122 for (const entry of newEntries) {
123 this.addEntryDB(entry);
124 }
125 // re-init in-memory heap
126 this.#heap = new Heap(wakerPriorityEntryComparator);
127 this.#heap.init(newEntries);
128 });
129 }
130
131 popTypeAll(entryType: WakerPriorityType) {
132 this.#ctx.storage.transactionSync(() => {
133 this.filter((e) => e.type !== entryType);
134 });
135 }
136
137 // Idempotent, perhaps name should suggest so
138 async handleNextAlarm() {
139 const nextWakeCall = this.#heap.peek();
140 if (nextWakeCall === undefined) {
141 return;
142 }
143 // TODO: Handle this
144 // const waker = this.#env.WAKERS.idFromName(
145 // this.#instanceMetadata.instance.id
146 // );
147 // const wakerStub = this.#env.WAKERS.get(waker);
148 // await wakerStub.wake(
149 // new Date(nextWakeCall.targetTimestamp),
150 // this.#instanceMetadata
151 // );
152 }
153
154 getFirst(callbackFn: (a: WakerPriorityEntry) => boolean): WakerPriorityEntry | undefined {
155 // clone it so that people cant just modify the entry on the PQ
156 return structuredClone(this.#heap.toArray().find(callbackFn));
157 }
158
159 private removeFirst(callbackFn: (a: WakerPriorityEntry) => boolean) {
160 const elements = this.#heap.toArray();
161 const index = elements.findIndex(callbackFn);
162 if (index === -1) {
163 return;
164 }
165
166 const removedEntry = elements.splice(index, 1)[0];
167 this.removeEntryDB(removedEntry);
168
169 this.#heap = new Heap(wakerPriorityEntryComparator);
170 this.#heap.init(elements);
171 }
172
173 private filter(callbackFn: (a: WakerPriorityEntry) => boolean) {
174 const filteredElements = this.#heap.toArray().filter(callbackFn);
175 const removedElements = this.#heap.toArray().filter((a) => !callbackFn(a));
176 this.#ctx.storage.transactionSync(() => {
177 for (const entry of removedElements) {
178 this.removeEntryDB(entry);
179 }
180 });
181 this.#heap = new Heap(wakerPriorityEntryComparator);
182 this.#heap.init(filteredElements);
183 }
184
185 length() {
186 return this.#heap.length;
187 }
188
189 private getEntries() {
190 const entries = [
191 ...this.#ctx.storage.sql.exec("SELECT * FROM priority_queue ORDER BY id"),
192 ] as Array<PriorityQueueDBEntry>;
193
194 const activeEntries: Array<WakerPriorityEntry> = [];
195
196 entries.forEach((val) => {
197 const entryType = toWakerPriorityType(val.entryType);
198 // 0 - removed
199 if (val.action == 0) {
200 const index = activeEntries.findIndex(
201 (activeVal) => val.hash == activeVal.hash && entryType == activeVal.type,
202 );
203 // if it's found remove it from the active list
204 if (index !== -1) {
205 activeEntries.splice(index, 1);
206 }
207 } else {
208 // 1 - added
209 const index = activeEntries.findIndex(
210 (activeVal) => val.hash == activeVal.hash && entryType == activeVal.type,
211 );
212 // if it's found remove it from the active list
213 if (index === -1) {
214 activeEntries.push({
215 hash: val.hash,
216 targetTimestamp: val.target_timestamp,
217 type: entryType,
218 });
219 }
220 }
221 });
222 return activeEntries;
223 }
224
225 private removeEntryDB(entry: WakerPriorityEntry) {
226 this.#ctx.storage.sql.exec(
227 `
228 INSERT INTO priority_queue (target_timestamp, action, entryType, hash)
229 VALUES (?, ?, ? ,?)
230 `,
231 entry.targetTimestamp,
232 SQLiteBoolean.FALSE,
233 fromWakerPriorityType(entry.type),
234 entry.hash,
235 );
236 }
237
238 checkIfExistedInPast(entry: Omit<WakerPriorityEntry, "targetTimestamp">) {
239 return (
240 this.#ctx.storage.sql
241 .exec(
242 "SELECT * FROM priority_queue WHERE entryType = ? AND hash = ? AND action = ?",
243 fromWakerPriorityType(entry.type),
244 entry.hash,
245 0,
246 )
247 .toArray().length >= 1
248 );
249 }
250
251 private addEntryDB(entry: WakerPriorityEntry) {
252 this.#ctx.storage.sql.exec(
253 `
254 INSERT INTO priority_queue (target_timestamp, action, entryType, hash)
255 VALUES (?, ?, ? ,?)
256 `,
257 entry.targetTimestamp,
258 SQLiteBoolean.TRUE,
259 fromWakerPriorityType(entry.type),
260 entry.hash,
261 );
262 }
263}
264
265const toWakerPriorityType = (entryType: EntryType): WakerPriorityType => {
266 switch (entryType) {
267 case EntryType.RETRY:
268 return "retry";
269 case EntryType.SLEEP:
270 return "sleep";
271 case EntryType.TIMEOUT:
272 return "timeout";
273 }
274};
275
276const fromWakerPriorityType = (entryType: WakerPriorityType): EntryType => {
277 switch (entryType) {
278 case "retry":
279 return EntryType.RETRY;
280 case "sleep":
281 return EntryType.SLEEP;
282 case "timeout":
283 return EntryType.TIMEOUT;
284 default:
285 throw new Error(`WakerPriorityType "${entryType}" has not been handled`);
286 }
287};