1// Alchemy modifications are licensed under Apache-2.0.
2// This file includes third-party code; see /THIRD_PARTY_LICENSES.md.
3import { ms } from "itty-time";
4import type { Engine } from "../engine.ts";
5import type { WorkflowSleepDuration } from "cloudflare:workers";
6
7export const ENGINE_TIMEOUT = ms("5 minutes" satisfies WorkflowSleepDuration);
8
9let latestGracePeriodTimestamp: number | undefined = undefined;
10
11export type WaitingPromiseType = "pause";
12
13export type GracePeriodCallback = (engine: Engine, timeoutMs: number) => void;
14
15export class GracePeriodSemaphore {
16 #counter: number = 0;
17 readonly callback: GracePeriodCallback;
18 readonly timeoutMs: number;
19 #waitingPromises: Array<{
20 rejectCallback: () => void;
21 resolveCallback: (value: unknown) => void;
22 type: WaitingPromiseType;
23 }> = [];
24 #canInitiateSteps = true;
25 #waitingSteps: Array<{
26 rejectCallback: () => void;
27 resolveCallback: (value: unknown) => void;
28 }> = [];
29
30 constructor(callback: GracePeriodCallback, timeoutMs: number) {
31 this.callback = callback;
32 this.timeoutMs = timeoutMs;
33 }
34
35 // acquire takes engine to be the same as release
36 async acquire(_engine: Engine) {
37 if (!this.#canInitiateSteps) {
38 await new Promise((resolve, reject) => {
39 this.#waitingSteps.push({
40 resolveCallback: resolve,
41 rejectCallback: reject,
42 });
43 });
44 }
45 // when the counter goes from 0 to 1 - we can safely reject the previous grace period
46 if (this.#counter == 0) {
47 latestGracePeriodTimestamp = undefined;
48 }
49 this.#counter += 1;
50 }
51
52 async release(engine: Engine) {
53 this.#counter = Math.max(this.#counter - 1, 0);
54 if (this.#counter == 0) {
55 // Trigger timeout promise, no need to await here,
56 // this can be triggered slightly after it's not time sensitive
57 this.callback(engine, this.timeoutMs);
58 // Resolve any promises waiting for all steps to finish (e.g. pause)
59 for (const promise of this.#waitingPromises) {
60 promise.resolveCallback(undefined);
61 }
62 this.#waitingPromises = [];
63 }
64 }
65
66 async waitUntilNothingIsRunning(
67 type: WaitingPromiseType,
68 callback: () => Promise<void>,
69 ): Promise<void> {
70 this.#canInitiateSteps = false;
71 if (this.#counter > 0) {
72 try {
73 await new Promise((resolve, reject) => {
74 this.#waitingPromises.push({
75 resolveCallback: resolve,
76 rejectCallback: reject,
77 type,
78 });
79 });
80 } catch {
81 // If the promise gets rejected (e.g. resume cancels the pause),
82 // allow steps to run again and unblock any that were waiting
83 for (const promise of this.#waitingSteps) {
84 promise.resolveCallback(undefined);
85 }
86 this.#waitingSteps = [];
87 this.#canInitiateSteps = true;
88 return;
89 }
90 }
91 await callback();
92 // Allow steps to run again and unblock any that were waiting
93 for (const promise of this.#waitingSteps) {
94 promise.resolveCallback(undefined);
95 }
96 this.#waitingSteps = [];
97 this.#canInitiateSteps = true;
98 }
99
100 cancelWaitingPromisesByType(type: WaitingPromiseType) {
101 const sameTypePromises = this.#waitingPromises.filter((val) => val.type === type);
102 if (sameTypePromises.length === 0) {
103 return;
104 }
105
106 for (const promise of sameTypePromises) {
107 promise.rejectCallback();
108 }
109
110 this.#waitingPromises = this.#waitingPromises.filter((val) => val.type !== type);
111
112 this.#canInitiateSteps = true;
113
114 // Unblock any steps that were waiting to acquire while the pause was pending
115 for (const promise of this.#waitingSteps) {
116 promise.resolveCallback(undefined);
117 }
118 this.#waitingSteps = [];
119 }
120
121 dispose() {
122 // Reject all waiting step promises so they stop blocking
123 for (const promise of this.#waitingSteps) {
124 promise.rejectCallback();
125 }
126 this.#waitingSteps = [];
127
128 // Reject all waiting promises
129 for (const promise of this.#waitingPromises) {
130 promise.rejectCallback();
131 }
132 this.#waitingPromises = [];
133
134 this.#canInitiateSteps = false;
135 }
136
137 isRunningStep() {
138 return this.#counter > 0;
139 }
140}
141
142export const startGracePeriod: GracePeriodCallback = async (engine: Engine, timeoutMs: number) => {
143 const gracePeriodHandler = async () => {
144 const thisTimestamp = new Date().valueOf();
145
146 // TODO: Should the grace period be 5 mins every time or 5 mins across lifetimes?
147 // At the moment, it looks like this will reset to 5 mins if a metal crashes
148 // There might possibly waste memory waiting for `timeoutMs` every time
149 // We should eventually strongly persist and respect this value across lifetimes
150
151 // We are starting a new grace period
152 // 1. There should not be one already set
153 // 2. Or if there is, it should be in the past
154 if (!(latestGracePeriodTimestamp === undefined || latestGracePeriodTimestamp < thisTimestamp)) {
155 throw new Error(
156 "Can't start grace period since there is already an active one started on " +
157 latestGracePeriodTimestamp,
158 );
159 }
160
161 latestGracePeriodTimestamp = thisTimestamp;
162 await scheduler.wait(timeoutMs);
163 if (thisTimestamp !== latestGracePeriodTimestamp || engine.timeoutHandler.isRunningStep()) {
164 return;
165 }
166 // priorityQueue is set before the user code runs which implies that a grace period cannot start
167 // before init finishes where it is set
168
169 // Ensure next alarm is set before we abort
170 await engine.priorityQueue?.handleNextAlarm();
171 // await engine.abort(ABORT_REASONS.GRACE_PERIOD_COMPLETE);
172 };
173 void gracePeriodHandler().catch(() => {
174 // Swallow — the engine is shutting down (abort kills the context,
175 // which rejects the scheduler.wait inside gracePeriodHandler)
176 });
177};