import assert from "assert" ; import { AsyncLocalStorage } from "async_hooks" ; import { setImmediate as setImmediatePromise } from "timers/promises" ; import { TypedEventTarget , kWrapListener } from "../event" ; import { Awaitable } from "./awaitable" ; const inputGateStorage = new AsyncLocalStorage < InputGate > ( ) ; const outputGateStorage = new AsyncLocalStorage < OutputGate > ( ) ; /**
 * Waits for the context's input gate (if any) to be open before returning.
 * Should be called before returning result of async I/O (e.g. setTimeout, KV).
 */ export function waitForOpenInputGate ( ) { const inputGate = inputGateStorage . getStore ( ) ; return inputGate ?. waitForOpen ( ) ; } /**
 * Runs closure with the context's input gate (if any) closed, unless
 * allowConcurrency is true. Should be called when performing storage
 * operations.
 */ export function runWithInputGateClosed < T > ( closure , allowConcurrency = false ) { if ( allowConcurrency ) return closure ( ) ; const inputGate = inputGateStorage . getStore ( ) ; if ( inputGate === undefined ) return closure ( ) ; return inputGate . runWithClosed ( closure ) ; } /**
 * Waits for the context's output gate (if any) to be open before returning.
 * Should be called before making async I/O requests (e.g. fetch, KV put)
 * which may be using unconfirmed values.
 */ export function waitForOpenOutputGate ( ) { const outputGate = outputGateStorage . getStore ( ) ; return outputGate ?. waitForOpen ( ) ; } /**
 * Registers promise with context's output gate (if any) unless allowUnconfirmed
 * is true. Should be called immediately (before any await) if returning a
 * promise for a storage write operation.
 */ export function waitUntilOnOutputGate < T > ( promise , allowUnconfirmed = false ) { if ( allowUnconfirmed ) return promise ; const outputGate = outputGateStorage . getStore ( ) ; outputGate ?. waitUntil ( promise ) ; return promise ; } export class InputGate { #lockCount = 0 ; readonly #resolveQueue : ( ( ) => void ) [ ] = [ ] ; readonly #parent ? : InputGate ; constructor ( parent ? ) { this . #parent = parent ; } /** Waits for input gate to open, then runs closure under the input gate */ async runWith < T > ( closure ) { await this . waitForOpen ( ) ; return inputGateStorage . run ( this , closure ) ; } /** Waits for input gate to be open (e.g. before returning from async I/O) */ async waitForOpen ( ) { await setImmediatePromise ( ) ; if ( this . #lockCount === 0 ) return ; return new Promise ( ( resolve ) => this . #resolveQueue . push ( resolve ) ) ; } /**
     * Runs a closure with this input gate closed (e.g. performing Durable Object
     * storage operation, blockConcurrencyWhile). Once complete, if the gate is
     * now open, resolves some waiters.
     */ async runWithClosed < T > ( closure ) { this . #lock ( ) ; await Promise . resolve ( ) ; const childInputGate = new InputGate ( this ) ; try { return await inputGateStorage . run ( childInputGate , closure ) ; } finally { setImmediate ( this . #unlock ) ; } } #lock ( ) { this . #lockCount ++ ; if ( this . #parent ) this . #parent . #lock ( ) ; } #unlock = async ( ) : Promise < void > => { assert ( this . #lockCount > 0 ) ; this . #lockCount -- ; while ( this . #lockCount === 0 && this . #resolveQueue . length ) { this . #resolveQueue . shift ( ) ! ( ) ; await setImmediatePromise ( ) ; } if ( this . #parent ) return this . #parent . #unlock ( ) ; } ; } export class OutputGate { readonly #waitUntil : Promise < unknown > [ ] = [ ] ; /** Runs closure under the output gate, then waits for output gate to open */ async runWith < T > ( closure ) { try { return await outputGateStorage . run ( this , closure ) ; } finally { await this . waitForOpen ( ) ; } } /** Waits for promises registered with this gate via waitUntil to resolve */ async waitForOpen ( ) { await Promise . all ( this . #waitUntil ) ; } /**
     * Registers a promise with this output gate. The gate won't open until this
     * promise resolves.
     */ waitUntil ( promise ) { this . #waitUntil . push ( promise ) ; } } export class InputGatedEventTarget < EventMap extends Record < string , Event > > extends TypedEventTarget < EventMap > { protected [ kWrapListener ] < Type extends keyof EventMap > ( listener ) { const inputGate = inputGateStorage . getStore ( ) ; return inputGate ? async ( event ) => { await inputGate . waitForOpen ( ) ; listener ( event ) ; } : listener ; } }