import { resolver } from "@rocicorp/resolver";
import { assert } from "../../../shared/src/asserts.js";
import { Queue } from "../../../shared/src/queue.js";
import { promiseVoid } from "../../../shared/src/resolved-promises.js";
import "../types/pg.js";
import { runTx } from "./run-transaction.js";
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, "initialWorkers must be positive");
    assert(
      maxWorkers >= initialWorkers,
      "maxWorkers must be >= initialWorkers"
    );
    this.#lc = lc;
    this.#mode = mode;
    this.#init = init ? this.#stmtRunner(init) : void 0;
    this.#cleanup = cleanup ? this.#stmtRunner(cleanup) : void 0;
    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: (key2, value2) => this.addLoggingContext(key2, value2)
    };
  }
  /**
   * 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) {
      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) => {
      const start = performance.now();
      try {
        lc.debug?.("started transaction");
        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);
            task = await this.#tasks.dequeue(timeoutTask, timeoutMs);
          }
        } finally {
          if (this.#cleanup) {
            await executeTask(this.#cleanup);
          }
        }
        const elapsed = performance.now() - start;
        lc.debug?.(`closing transaction (${elapsed.toFixed(3)} ms)`);
        return last;
      } catch (e) {
        if (e !== this.#failure) {
          this.fail(e);
        }
        throw e;
      }
    };
    this.#workers.push(
      runTx(db, worker, { mode: this.#mode }).catch((e) => {
        if (e instanceof RollbackSignal) {
          lc.debug?.("aborted transaction");
        } else {
          throw e;
        }
      }).finally(() => this.#numWorkers--)
    );
    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 };
        }
        const last = stmts.reduce(
          (_, stmt) => stmt.execute().then(() => {
            if (++this.#stmts % 1e3 === 0) {
              const log = this.#stmts % 1e4 === 0 ? "info" : "debug";
              const q = stmt;
              lc[log]?.(
                `executed ${this.#stmts}th statement (${(performance.now() - this.#start).toFixed(3)} ms)`,
                { statement: q.string }
              );
            }
          }).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);
    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 !== void 0 && !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,
      () => `Cannot unref ${count} when refCount is ${this.#refCount}`
    );
    this.#refCount -= count;
    if (this.#refCount === 0) {
      this.setDone();
    }
  }
  isRunning() {
    return this.#db !== void 0 && !this.#done && this.#failure === void 0;
  }
  /**
   * Signals all workers to fail their transactions with the given {@link err}.
   */
  fail(err) {
    if (!this.#failure) {
      this.#failure = ensureError(err);
      const level = this.#failure instanceof ControlFlowError ? "debug" : "info";
      this.#lc[level]?.(this.#failure);
      for (let i = 0; i < this.#numWorkers; i++) {
        this.#tasks.enqueue(this.#failure);
      }
    }
  }
}
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
  };
}
class ControlFlowError extends Error {
  constructor(cause) {
    super();
    this.cause = cause;
  }
}
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;
}
const IDLE_TIMEOUT_MS = 5e3;
const KEEPALIVE_TIMEOUT_MS = 6e4;
const KEEPALIVE_TASK = (tx) => [tx`SELECT 1`.simple()];
const TIMEOUT_TASKS = {
  forInitialWorkers: {
    timeoutMs: KEEPALIVE_TIMEOUT_MS,
    task: KEEPALIVE_TASK
  },
  forExtraWorkers: {
    timeoutMs: IDLE_TIMEOUT_MS,
    task: "done"
  }
};
export {
  ControlFlowError,
  TIMEOUT_TASKS,
  TransactionPool,
  importSnapshot
};
//# sourceMappingURL=transaction-pool.js.map
