gen_durable

defmodule GenDurable.Migration do
  @moduledoc """
  Library-owned schema migration (Oban-style).

  The host application does not copy the DDL; it writes a one-line Ecto
  migration that delegates here:

      defmodule MyApp.Repo.Migrations.SetupGenDurable do
        use Ecto.Migration

        def up,   do: GenDurable.Migration.up()
        def down, do: GenDurable.Migration.down()
      end

  ## Options

    * `:prefix`  — Postgres schema the tables live in (default `"public"`).
      The runtime queries reference the tables **unqualified**, so a non-default
      prefix also requires the schema on the repo's `search_path`
      (e.g. `after_connect: {Postgrex, :query!, ["SET search_path TO my_schema, public", []]}`).
    * `:version` — schema version to migrate to (default: latest for `up/1`,
      `0` for `down/1`).

  The installed schema version is recorded in `COMMENT ON TABLE gen_durable`,
  so `up/1` only applies the increments that are missing. This keeps the
  host-facing call stable as the schema evolves across library releases.
  """

  use Ecto.Migration

  @latest_version 2

  @doc "Migrate the schema up to `:version` (default: latest)."
  def up(opts \\ []) do
    prefix = prefix(opts)
    target = Keyword.get(opts, :version, @latest_version)
    current = migrated_version(prefix)

    if prefix != "public" do
      execute("CREATE SCHEMA IF NOT EXISTS #{prefix}")
    end

    if current < target do
      for v <- (current + 1)..target, do: change(v, :up, prefix)
      record_version(prefix, target)
    end

    :ok
  end

  @doc "Migrate the schema down to `:version` (default: 0, i.e. drop everything)."
  def down(opts \\ []) do
    prefix = prefix(opts)
    target = Keyword.get(opts, :version, 0)
    current = migrated_version(prefix)

    if current > target do
      for v <- current..(target + 1)//-1, do: change(v, :down, prefix)
      if target > 0, do: record_version(prefix, target)
    end

    :ok
  end

  # --- version 1: DDL ------------------------------------------------

  defp change(1, :up, p) do
    execute("""
    CREATE TYPE #{p}.durable_status AS ENUM
      ('runnable', 'executing', 'awaiting_signal', 'awaiting_children', 'done', 'failed')
    """)

    execute("""
    CREATE TABLE #{p}.gen_durable (
      id            bigint generated always as identity primary key,
      fsm           text not null,
      fsm_version   int  not null default 1,
      step          text not null,
      status        #{p}.durable_status not null default 'runnable',
      state         jsonb not null default '{}',
      result        jsonb,
      awaits        text[],

      queue         text     not null default 'default',
      priority      smallint not null default 0,

      -- concurrency_key: a semaphore of size K on the key. K comes from the
      -- `concurrency_limits:` config by name (split_part(key, ':', 1)); an
      -- unconfigured key defaults to K = 1 (mutual exclusion). concurrency_shard
      -- is set at claim time for CONFIGURED keys only (which bucket shard the
      -- slot was drawn from — the release credits it back addressed); its
      -- NULLness is also the tier discriminator: unconfigured claims keep it
      -- NULL and are fenced by the unique index below, configured claims are
      -- fenced by the bucket CHECK instead.
      concurrency_key   text,
      concurrency_shard smallint,

      eligible_at   timestamptz not null default now(),
      attempt       int  not null default 0,
      last_error    text,

      -- await timeout: set by the park when the step passed `timeout:`; the
      -- reaper sweeps expired parks back to runnable (a wake, not a failure).
      await_deadline timestamptz,

      -- rate limiting: rate_limit is the bucket key for the CURRENT step
      -- (NULL ⇒ not rate-limited), weight is how many budget units this step's execution
      -- consumes (default 1). Both rewritten on every transition; kept on :retry.
      rate_limit    text,
      weight        double precision not null default 1,

      locked_by        text,
      lease_expires_at timestamptz,

      parent_id        bigint references #{p}.gen_durable(id) on delete set null,
      children_pending int not null default 0,

      -- correlation_key: the instance's business identity — both the signal address
      -- and the uniqueness guard. correlation_scope is the set of statuses in
      -- which the key is "occupied", supplied by the caller (defaults to the non-terminal
      -- statuses ⇒ unique among live, freed on termination). The guard equals the key
      -- while the status is occupied, else NULL (drops out of the unique/address index).
      -- scope is durable_status[] (not text[]) so the generated column stays IMMUTABLE
      -- — the enum->text cast (enum_out) is only STABLE.
      correlation_key   text,
      correlation_scope #{p}.durable_status[] not null default '{}',
      correlation_guard text generated always as (
        case when correlation_key is not null and status = any(correlation_scope)
             then correlation_key end
      ) stored,

      inserted_at   timestamptz not null default now(),
      updated_at    timestamptz not null default now()
    )
    """)

    execute("""
    CREATE INDEX gen_durable_pick ON #{p}.gen_durable (queue, priority, eligible_at)
      WHERE status = 'runnable'
    """)

    execute("""
    CREATE INDEX gen_durable_lease ON #{p}.gen_durable (lease_expires_at)
      WHERE status = 'executing'
    """)

    # Supports the await-timeout sweep: parked rows with an armed deadline.
    execute("""
    CREATE INDEX gen_durable_await_deadline ON #{p}.gen_durable (await_deadline)
      WHERE status = 'awaiting_signal' AND await_deadline IS NOT NULL
    """)

    # The K = 1 tier of concurrency_key (unconfigured keys), enforced by the
    # database: at most ONE executing row per key can ever be committed. The
    # picker's NOT EXISTS guard reads it as an optimization; the UNIQUE arbiter
    # is the correctness — a cross-node claim race ends in a unique violation
    # and the losing pick retries. The "lock" is the row's own executing status,
    # so it spans exactly the step window and releases with any outcome (or the
    # reaper, on a crash). Configured (gated) keys claim with a non-null
    # concurrency_shard and drop out of this index — their cap is the bucket
    # CHECK below. Non-keyed claims (the common case) never write to it.
    execute("""
    CREATE UNIQUE INDEX gen_durable_concurrency_active ON #{p}.gen_durable (concurrency_key)
      WHERE status = 'executing' AND concurrency_key IS NOT NULL
        AND concurrency_shard IS NULL
    """)

    # correlation_key: one partial unique index does double duty — it enforces
    # uniqueness among "occupied" statuses (per the :unique policy / scope) AND backs
    # the address lookup in deliver_signal (`correlation_guard = $1` resolves to the
    # single occupied instance). A row whose status leaves the scope drops out, freeing the key.
    execute("""
    CREATE UNIQUE INDEX gen_durable_correlation ON #{p}.gen_durable (correlation_guard)
      WHERE correlation_guard IS NOT NULL
    """)

    execute("""
    CREATE INDEX gen_durable_parent ON #{p}.gen_durable (parent_id)
      WHERE parent_id IS NOT NULL
    """)

    # Supports the GC sweep: terminal rows ordered by termination instant.
    execute("""
    CREATE INDEX gen_durable_gc ON #{p}.gen_durable (updated_at)
      WHERE status IN ('done', 'failed')
    """)

    execute("""
    CREATE TABLE #{p}.signals (
      id          bigint generated always as identity primary key,
      target_id   bigint not null references #{p}.gen_durable(id) on delete cascade,
      name        text   not null,
      payload     jsonb  not null default '{}',
      dedup_key   text,
      inserted_at timestamptz not null default now(),
      unique (target_id, dedup_key)
    )
    """)

    execute("CREATE INDEX signals_target ON #{p}.signals (target_id, name)")

    # rate limiting. Configs are seeded at engine start from the `rate_limits:`
    # option; the picker joins them. Buckets are token-bucket counters, one row
    # per distinct key, minted pre-debited by the first pick that grants from
    # the key (see the pick's r_mint CTE) and swept by gc_buckets when idle.
    execute("""
    CREATE TABLE #{p}.gen_durable_rate_configs (
      name  text primary key,
      rate  double precision not null,   -- tokens per second (allowed / period)
      burst double precision not null    -- bucket capacity
    )
    """)

    # The primary key doubles as the cold-mint arbiter: two picks minting the
    # same cold key cannot both commit — the loser aborts and retries against
    # the winner's row (see pick/6's rescue).
    execute("""
    CREATE TABLE #{p}.gen_durable_rate_buckets (
      key         text primary key,
      tokens      double precision not null,
      last_refill timestamptz not null default clock_timestamp()
    )
    """)

    # concurrency gates: `cap` is the whole-key limit, split across `shards`
    # bucket rows to spread the release chains (completions of one key
    # serialize on their shard's row lock, so shards multiply the per-key
    # completion ceiling). Configs are seeded at engine start.
    execute("""
    CREATE TABLE #{p}.gen_durable_concurrency_configs (
      name   text primary key,
      cap    int not null,
      shards int not null
    )
    """)

    # One row per (gate key, shard). `available` is the free-slot count for the
    # shard; claims debit it (batched, in the pick), releases credit it back
    # addressed (a rider in the outcome). The CHECK is the hard cap: an
    # over-admission (or a double release) is uncommittable. Crash paths leak
    # conservatively (available too LOW — under-admission), healed by the GC
    # reconciler from the executing-rows truth.
    execute("""
    CREATE TABLE #{p}.gen_durable_concurrency_buckets (
      key       text not null,
      shard     smallint not null,
      cap       int not null,
      available int not null,
      PRIMARY KEY (key, shard),
      CHECK (available >= 0 AND available <= cap)
    )
    """)
  end

  defp change(1, :down, p) do
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_concurrency_buckets")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_concurrency_configs")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_rate_buckets")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_rate_configs")
    execute("DROP TABLE IF EXISTS #{p}.signals")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable")
    execute("DROP TYPE IF EXISTS #{p}.durable_status")
  end

  # --- version 2: unified, sharded rate + concurrency buckets --------

  # The four limiter tables collapse into two. Rate joins concurrency in one
  # sharded shape so the pick can lock a key's shards with SKIP LOCKED —
  # concurrent (cross-node) pickers grab disjoint shard subsets instead of
  # blocking on one bucket row across the whole claim.
  #
  # Config and bucket rows are ephemeral: configs are re-seeded from the
  # `rate_limits:`/`concurrency_limits:` options at every engine start, and
  # buckets are minted on demand by the pick and swept by the GC. So the
  # upgrade is a plain drop-and-recreate with nothing to preserve — an
  # existing install runs only this increment; a fresh one runs 1 then 2.
  defp change(2, :up, p) do
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_concurrency_buckets")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_concurrency_configs")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_rate_buckets")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_rate_configs")

    # `kind` is part of the primary key so a rate limit and a concurrency gate
    # may share a name (and a bucket key) without colliding — they lived in
    # separate tables before, so this keeps the two namespaces independent.
    execute("""
    CREATE TABLE #{p}.gen_durable_bucket_configs (
      kind     text not null,
      name     text not null,
      rate     double precision,           -- tokens/sec; NULL for concurrency
      capacity double precision not null,   -- burst (rate) or cap (concurrency)
      shards   smallint not null,
      PRIMARY KEY (kind, name),
      CHECK (kind IN ('rate', 'conc'))
    )
    """)

    # One row per (kind, key, shard). `available` is tokens (rate) or free
    # slots (conc); `capacity` is the shard's ceiling (burst/shards or a
    # slot split of cap). The CHECK is the hard cap: over-admission or a
    # double credit is uncommittable. `last_refill` drives the rate refill
    # and is NULL for concurrency (which refills by completion credit).
    execute("""
    CREATE TABLE #{p}.gen_durable_buckets (
      kind        text not null,
      key         text not null,
      shard       smallint not null,
      capacity    double precision not null,
      available   double precision not null,
      last_refill timestamptz,
      PRIMARY KEY (kind, key, shard),
      CHECK (available >= 0 AND available <= capacity)
    )
    """)
  end

  defp change(2, :down, p) do
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_buckets")
    execute("DROP TABLE IF EXISTS #{p}.gen_durable_bucket_configs")

    execute("""
    CREATE TABLE #{p}.gen_durable_rate_configs (
      name  text primary key,
      rate  double precision not null,
      burst double precision not null
    )
    """)

    execute("""
    CREATE TABLE #{p}.gen_durable_rate_buckets (
      key         text primary key,
      tokens      double precision not null,
      last_refill timestamptz not null default clock_timestamp()
    )
    """)

    execute("""
    CREATE TABLE #{p}.gen_durable_concurrency_configs (
      name   text primary key,
      cap    int not null,
      shards int not null
    )
    """)

    execute("""
    CREATE TABLE #{p}.gen_durable_concurrency_buckets (
      key       text not null,
      shard     smallint not null,
      cap       int not null,
      available int not null,
      PRIMARY KEY (key, shard),
      CHECK (available >= 0 AND available <= cap)
    )
    """)
  end

  # --- helpers ---------------------------------------------------------------

  defp prefix(opts), do: Keyword.get(opts, :prefix, "public")

  defp record_version(p, version) do
    execute("COMMENT ON TABLE #{p}.gen_durable IS '#{version}'")
  end

  # Current installed version from the table comment; 0 if the table is absent.
  defp migrated_version(p) do
    query = """
    SELECT pg_catalog.obj_description(c.oid, 'pg_class')
    FROM pg_catalog.pg_class c
    JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
    WHERE c.relname = 'gen_durable' AND n.nspname = $1
    """

    case repo().query(query, [p]) do
      {:ok, %{rows: [[version]]}} when is_binary(version) -> String.to_integer(version)
      _ -> 0
    end
  end
end