import { resolver } from '@rocicorp/resolver';
import { assert } from "../../../shared/src/asserts.js";
import { stringify } from "../../../shared/src/bigint-json.js";
import { Queue } from "../../../shared/src/queue.js";
import { promiseVoid } from "../../../shared/src/resolved-promises.js";
import { disableStatementTimeout, } from "../types/pg.js";
import * as Mode from "./mode-enum.js";
/**
 * A TransactionPool is a pool of one or more {@link postgres.TransactionSql}
 * objects that participate in processing a dynamic queue of tasks.
 *
 * This can be used for serializing a set of tasks that arrive asynchronously
 * to a single transaction (for writing) or performing parallel reads across
 * multiple connections at the same snapshot (e.g. read only snapshot transactions).
 */
export class TransactionPool {
    #lc;
    #mode;
    #init;
    #cleanup;
    #tasks = new Queue();
    #workers = [];
    #initialWorkers;
    #maxWorkers;
    #timeoutTask;
    #numWorkers;
    #numWorking = 0;
    #db; // set when running. stored to allow adaptive pool sizing.
    #refCount = 1;
    #done = false;
    #failure;
    /**
     * @param init A {@link Task} that is run in each Transaction before it begins
     *             processing general tasks. This can be used to to set the transaction
     *             mode, export/set snapshots, etc. This will be run even if
     *             {@link fail} has been called on the pool.
     * @param cleanup A {@link Task} that is run in each Transaction before it closes.
     *                This will be run even if {@link fail} has been called, or if a
     *                preceding Task threw an Error.
     * @param initialWorkers The initial number of transaction workers to process tasks.
     *                       This is the steady state number of workers that will be kept
     *                       alive if the TransactionPool is long lived.
     *                       This must be greater than 0. Defaults to 1.
     * @param maxWorkers When specified, allows the pool to grow to `maxWorkers`. This
     *                   must be greater than or equal to `initialWorkers`. On-demand
     *                   workers will be shut down after an idle timeout of 5 seconds.
     */
    constructor(lc, mode, init, cleanup, initialWorkers = 1, maxWorkers = initialWorkers, timeoutTasks = TIMEOUT_TASKS) {
        assert(initialWorkers > 0);
        assert(maxWorkers >= initialWorkers);
        this.#lc = lc;
        this.#mode = mode;
        this.#init = init ? this.#stmtRunner(init) : undefined;
        this.#cleanup = cleanup ? this.#stmtRunner(cleanup) : undefined;
        this.#initialWorkers = initialWorkers;
        this.#numWorkers = initialWorkers;
        this.#maxWorkers = maxWorkers;
        this.#timeoutTask = timeoutTasks;
    }
    /**
     * Starts the pool of workers to process Tasks with transactions opened from the
     * specified {@link db}.
     */
    run(db) {
        assert(!this.#db, 'already running');
        this.#db = db;
        for (let i = 0; i < this.#numWorkers; i++) {
            this.#addWorker(db);
        }
        return this;
    }
    /**
     * Adds context parameters to internal LogContext. This is useful for context values that
     * are not known when the TransactionPool is constructed (e.g. determined after a database
     * call when the pool is running).
     *
     * Returns an object that can be used to add more parameters.
     */
    addLoggingContext(key, value) {
        this.#lc = this.#lc.withContext(key, value);
        return {
            addLoggingContext: (key, value) => this.addLoggingContext(key, value),
        };
    }
    /**
     * Returns a promise that:
     *
     * * resolves after {@link setDone} has been called (or the the pool as been {@link unref}ed
     *   to a 0 ref count), once all added tasks have been processed and all transactions have been
     *   committed or closed.
     *
     * * rejects if processing was aborted with {@link fail} or if processing any of
     *   the tasks resulted in an error. All uncommitted transactions will have been
     *   rolled back.
     *
     * Note that partial failures are possible if processing writes with multiple workers
     * (e.g. `setDone` is called, allowing some workers to commit, after which other
     *  workers encounter errors). Using a TransactionPool in this manner does not make
     * sense in terms of transactional semantics, and is thus not recommended.
     *
     * For reads, however, multiple workers is useful for performing parallel reads
     * at the same snapshot. See {@link synchronizedSnapshots} for an example.
     * Resolves or rejects when all workers are done or failed.
     */
    async done() {
        const numWorkers = this.#workers.length;
        await Promise.all(this.#workers);
        if (numWorkers < this.#workers.length) {
            // If workers were added after the initial set, they must be awaited to ensure
            // that the results (i.e. rejections) of all workers are accounted for. This only
            // needs to be re-done once, because the fact that the first `await` completed
            // guarantees that the pool is in a terminal state and no new workers can be added.
            await Promise.all(this.#workers);
        }
        this.#lc.debug?.('transaction pool done');
    }
    #addWorker(db) {
        const id = this.#workers.length + 1;
        const lc = this.#lc.withContext('tx', id);
        const tt = this.#workers.length < this.#initialWorkers
            ? this.#timeoutTask.forInitialWorkers
            : this.#timeoutTask.forExtraWorkers;
        const { timeoutMs } = tt;
        const timeoutTask = tt.task === 'done' ? 'done' : this.#stmtRunner(tt.task);
        const worker = async (tx) => {
            try {
                lc.debug?.('started transaction');
                disableStatementTimeout(tx);
                let last = promiseVoid;
                const executeTask = async (runner) => {
                    runner !== this.#init && this.#numWorking++;
                    const { pending } = await runner.run(tx, lc, () => {
                        runner !== this.#init && this.#numWorking--;
                    });
                    last = pending ?? last;
                };
                let task = this.#init ?? (await this.#tasks.dequeue(timeoutTask, timeoutMs));
                try {
                    while (task !== 'done') {
                        if (task instanceof Error ||
                            (task !== this.#init && this.#failure)) {
                            throw this.#failure ?? task;
                        }
                        await executeTask(task);
                        // await the next task.
                        task = await this.#tasks.dequeue(timeoutTask, timeoutMs);
                    }
                }
                finally {
                    // Execute the cleanup task even on failure.
                    if (this.#cleanup) {
                        await executeTask(this.#cleanup);
                    }
                }
                lc.debug?.('closing transaction');
                // Given the semantics of a Postgres transaction, the last statement
                // will only succeed if all of the preceding statements succeeded.
                return last;
            }
            catch (e) {
                if (e !== this.#failure) {
                    lc.error?.('error from worker', e);
                    this.fail(e); // A failure in any worker should fail the pool.
                }
                throw e;
            }
        };
        this.#workers.push(db
            .begin(this.#mode, worker)
            .catch(e => {
            if (e instanceof RollbackSignal) {
                // A RollbackSignal is used to gracefully rollback the postgres.js
                // transaction block. It should not be thrown up to the application.
                lc.debug?.('aborted transaction');
            }
            else {
                throw e;
            }
        })
            .finally(() => this.#numWorkers--));
        // After adding the worker, enqueue a terminal signal if we are in either of the
        // terminal states (both of which prevent more tasks from being enqueued), to ensure
        // that the added worker eventually exits.
        if (this.#done) {
            this.#tasks.enqueue('done');
        }
        if (this.#failure) {
            this.#tasks.enqueue(this.#failure);
        }
    }
    /**
     * Processes the statements produced by the specified {@link Task},
     * returning a Promise that resolves when the statements are either processed
     * by the database or rejected.
     *
     * Note that statement failures will result in failing the entire
     * TransactionPool (per transaction semantics). However, the returned Promise
     * itself will resolve rather than reject. As such, it is fine to ignore
     * returned Promises in order to pipeline requests to the database. It is
     * recommended to occasionally await them (e.g. after some threshold) in
     * order to avoid memory blowup in the case of database slowness.
     */
    process(task) {
        const r = resolver();
        this.#process(this.#stmtRunner(task, r));
        return r.promise;
    }
    #start = performance.now();
    #stmts = 0;
    /**
     * Implements the semantics specified in {@link process()}.
     *
     * Specifically:
     * * `freeWorker()` is called as soon as the statements are produced,
     *   allowing them to be pipelined to the database.
     * * Statement errors result in failing the transaction pool.
     * * The client-supplied Resolver resolves on success or failure;
     *   it is never rejected.
     */
    #stmtRunner(task, r = resolver()) {
        return {
            run: async (tx, lc, freeWorker) => {
                let stmts;
                try {
                    stmts = await task(tx, lc);
                }
                catch (e) {
                    r.resolve();
                    throw e;
                }
                finally {
                    freeWorker();
                }
                if (stmts.length === 0) {
                    r.resolve();
                    return { pending: null };
                }
                // Execute the statements (i.e. send to the db) immediately.
                // The last result is returned for the worker to await before
                // closing the transaction.
                const last = stmts.reduce((_, stmt) => stmt
                    .execute()
                    .then(() => {
                    if (++this.#stmts % 1000 === 0) {
                        const q = stmt;
                        lc.debug?.(`executed ${this.#stmts}th statement (${(performance.now() - this.#start).toFixed(3)} ms)`, {
                            statement: q.string,
                            params: stringify(q.parameters),
                        });
                    }
                })
                    .catch(e => this.fail(e)), promiseVoid);
                return { pending: last.then(r.resolve) };
            },
            rejected: r.resolve,
        };
    }
    /**
     * Processes and returns the result of executing the {@link ReadTask} from
     * within the transaction. An error thrown by the task will result in
     * rejecting the returned Promise, but will not affect the transaction pool
     * itself.
     */
    processReadTask(readTask) {
        const r = resolver();
        this.#process(this.#readRunner(readTask, r));
        return r.promise;
    }
    /**
     * Implements the semantics specified in {@link processReadTask()}.
     *
     * Specifically:
     * * `freeWorker()` is called as soon as the result is produced,
     *   before resolving the client-supplied Resolver.
     * * Errors result in rejecting the client-supplied Resolver but
     *   do not affect transaction pool.
     */
    #readRunner(readTask, r) {
        return {
            run: async (tx, lc, freeWorker) => {
                let result;
                try {
                    result = await readTask(tx, lc);
                    freeWorker();
                    r.resolve(result);
                }
                catch (e) {
                    freeWorker();
                    r.reject(e);
                }
                return { pending: null };
            },
            rejected: r.reject,
        };
    }
    #process(runner) {
        assert(!this.#done, 'already set done');
        if (this.#failure) {
            runner.rejected(this.#failure);
            return;
        }
        this.#tasks.enqueue(runner);
        // Check if the pool size can and should be increased.
        if (this.#numWorkers < this.#maxWorkers) {
            const outstanding = this.#tasks.size();
            if (outstanding > this.#numWorkers - this.#numWorking) {
                this.#db && this.#addWorker(this.#db);
                this.#numWorkers++;
                this.#lc.debug?.(`Increased pool size to ${this.#numWorkers}`);
            }
        }
    }
    /**
     * Ends all workers with a ROLLBACK. Throws if the pool is already done
     * or aborted.
     */
    abort() {
        this.fail(new RollbackSignal());
    }
    /**
     * Signals to all workers to end their transaction once all pending tasks have
     * been completed. Throws if the pool is already done or aborted.
     */
    setDone() {
        assert(!this.#done, 'already set done');
        this.#done = true;
        for (let i = 0; i < this.#numWorkers; i++) {
            this.#tasks.enqueue('done');
        }
    }
    /**
     * An alternative to explicitly calling {@link setDone}, `ref()` increments an internal reference
     * count, and {@link unref} decrements it. When the reference count reaches 0, {@link setDone} is
     * automatically called. A TransactionPool is initialized with a reference count of 1.
     *
     * `ref()` should be called before sharing the pool with another component, and only after the
     * pool has been started with {@link run()}. It must not be called on a TransactionPool that is
     * already done (either via {@link unref()} or {@link setDone()}. (Doing so indicates a logical
     * error in the code.)
     *
     * It follows that:
     * * The creator of the TransactionPool is responsible for running it.
     * * The TransactionPool should be ref'ed before being sharing.
     * * The receiver of the TransactionPool is only responsible for unref'ing it.
     *
     * On the other hand, a transaction pool that fails with a runtime error can still be ref'ed;
     * attempts to use the pool will result in the runtime error as expected.
     */
    // TODO: Get rid of the ref-counting stuff. It's no longer needed.
    ref(count = 1) {
        assert(this.#db !== undefined && !this.#done, `Cannot ref() a TransactionPool that is not running`);
        this.#refCount += count;
    }
    /**
     * Decrements the internal reference count, automatically invoking {@link setDone} when it reaches 0.
     */
    unref(count = 1) {
        assert(count <= this.#refCount);
        this.#refCount -= count;
        if (this.#refCount === 0) {
            this.setDone();
        }
    }
    isRunning() {
        return this.#db !== undefined && !this.#done && this.#failure === undefined;
    }
    /**
     * Signals all workers to fail their transactions with the given {@link err}.
     */
    fail(err) {
        if (!this.#failure) {
            this.#failure = ensureError(err); // Fail fast: this is checked in the worker loop.
            if (this.#failure instanceof ControlFlowError) {
                this.#lc.debug?.(this.#failure);
            }
            else {
                this.#lc.error?.(this.#failure);
            }
            for (let i = 0; i < this.#numWorkers; i++) {
                // Enqueue the Error to terminate any workers waiting for tasks.
                this.#tasks.enqueue(this.#failure);
            }
        }
    }
}
/**
 * Init Tasks for Postgres snapshot synchronization across transactions.
 *
 * https://www.postgresql.org/docs/9.3/functions-admin.html#:~:text=Snapshot%20Synchronization%20Functions,identical%20content%20in%20the%20database.
 */
export function synchronizedSnapshots() {
    const { promise: snapshotExported, resolve: exportSnapshot, reject: failExport, } = resolver();
    const { promise: snapshotCaptured, resolve: captureSnapshot, reject: failCapture, } = resolver();
    // Set by the first worker to run its initTask, who becomes responsible for
    // exporting the snapshot. TODO: Plumb the workerNum and use that instead.
    let firstWorkerRun = false;
    // Note: Neither init task should `await`, as processing in each pool can proceed
    //       as soon as the statements have been sent to the db. However, the `cleanupExport`
    //       task must `await` the result of `setSnapshot` to ensure that exporting transaction
    //       does not close before the snapshot has been captured.
    return {
        exportSnapshot: tx => {
            if (!firstWorkerRun) {
                firstWorkerRun = true;
                const stmt = tx `SELECT pg_export_snapshot() AS snapshot; SET TRANSACTION READ ONLY;`.simple();
                // Intercept the promise to propagate the information to `snapshotExported`.
                stmt.then(result => exportSnapshot(result[0].snapshot), failExport);
                return [stmt]; // Also return the stmt so that it gets awaited (and errors handled).
            }
            return snapshotExported.then(snapshotID => [
                tx.unsafe(`SET TRANSACTION SNAPSHOT '${snapshotID}'`),
                tx `SET TRANSACTION READ ONLY`.simple(),
            ]);
        },
        setSnapshot: tx => snapshotExported.then(snapshotID => {
            const stmt = tx.unsafe(`SET TRANSACTION SNAPSHOT '${snapshotID}'`);
            // Intercept the promise to propagate the information to `cleanupExport`.
            stmt.then(captureSnapshot, failCapture);
            return [stmt];
        }),
        cleanupExport: async () => {
            await snapshotCaptured;
            return [];
        },
        snapshotID: snapshotExported,
    };
}
/**
 * Returns `init` and `cleanup` {@link Task}s for a TransactionPool that ensure its workers
 * share a single view of the database. This is used for View Notifier and View Syncer logic
 * that allows multiple entities to perform parallel reads on the same snapshot of the database.
 */
export function sharedSnapshot() {
    const { promise: snapshotExported, resolve: exportSnapshot, reject: failExport, } = resolver();
    // Set by the first worker to run its initTask, who becomes responsible for
    // exporting the snapshot.
    let firstWorkerRun = false;
    // Set when any worker is done, signalling that all non-sentinel Tasks have been
    // dequeued, and thus any subsequently spawned workers should skip their initTask
    // since the snapshot is no longer needed (and soon to become invalid).
    let firstWorkerDone = false;
    return {
        init: (tx, lc) => {
            if (!firstWorkerRun) {
                firstWorkerRun = true;
                const stmt = tx `SELECT pg_export_snapshot() AS snapshot;`.simple();
                // Intercept the promise to propagate the information to `snapshotExported`.
                stmt.then(result => exportSnapshot(result[0].snapshot), failExport);
                return [stmt]; // Also return the stmt so that it gets awaited (and errors handled).
            }
            if (!firstWorkerDone) {
                return snapshotExported.then(snapshotID => [
                    tx.unsafe(`SET TRANSACTION SNAPSHOT '${snapshotID}'`),
                ]);
            }
            lc.debug?.('All work is done. No need to set snapshot');
            return [];
        },
        cleanup: () => {
            firstWorkerDone = true;
            return [];
        },
        snapshotID: snapshotExported,
    };
}
/**
 * @returns An `init` Task for importing a snapshot from another transaction.
 */
export function importSnapshot(snapshotID) {
    const { promise: imported, resolve, reject } = resolver();
    return {
        init: tx => {
            const stmt = tx.unsafe(`SET TRANSACTION SNAPSHOT '${snapshotID}'`);
            stmt.then(() => resolve(), reject);
            return [stmt];
        },
        imported,
    };
}
/**
 * A superclass of Errors used for control flow that is needed to handle
 * another Error but does not constitute an error condition itself (e.g.
 * aborting transactions after a previous one fails). Subclassing this Error
 * will result in lowering the log level from `error` to `debug`.
 */
export class ControlFlowError extends Error {
    constructor(cause) {
        super();
        this.cause = cause;
    }
}
/**
 * Internal error used to rollback the worker transaction. This is used
 * instead of executing a `ROLLBACK` statement because the postgres.js
 * library will otherwise try to execute an extraneous `COMMIT`, which
 * results in outputting a "no transaction in progress" warning to the
 * database logs.
 *
 * Throwing an exception, on the other hand, executes the postgres.js
 * codepath that calls `ROLLBACK` instead.
 */
class RollbackSignal extends ControlFlowError {
    name = 'RollbackSignal';
    message = 'rolling back transaction';
}
function ensureError(err) {
    if (err instanceof Error) {
        return err;
    }
    const error = new Error();
    error.cause = err;
    return error;
}
// TODO: Get rid of the timeout stuff. It's no longer needed.
const IDLE_TIMEOUT_MS = 5_000;
const KEEPALIVE_TIMEOUT_MS = 60_000;
const KEEPALIVE_TASK = tx => [tx `SELECT 1`.simple()];
// Production timeout tasks. Overridden in tests.
export const TIMEOUT_TASKS = {
    forInitialWorkers: {
        timeoutMs: KEEPALIVE_TIMEOUT_MS,
        task: KEEPALIVE_TASK,
    },
    forExtraWorkers: {
        timeoutMs: IDLE_TIMEOUT_MS,
        task: 'done',
    },
};
//# sourceMappingURL=transaction-pool.js.map