repodb / lib / pool.ml

(** Lock-free connection pool using kcas for thread-safe operations.

    Works with any concurrency library (Eio, Lwt, direct-style) because kcas
    provides lock-free data structures independent of any specific scheduler.

    {[
      let config =
        Pool.
          {
            max_size = 10;
            connect = (fun () -> Repodb_postgresql.connect conninfo);
            close = Repodb_postgresql.close;
            validate = None;
          }
      in
      let pool = Pool.create config in
      Pool.with_connection pool (fun conn -> Repo.all conn ~table:users ~decode)
    ]} *)

open Kcas
open Kcas_data

type pool_error =
  | Pool_empty
  | Pool_closed
  | Pool_timeout
  | Connection_error of string

type 'conn config = {
  max_size : int;
  connect : unit -> ('conn, string) result;
  close : 'conn -> unit;
  validate : ('conn -> bool) option;
}

type 'conn pooled = { conn : 'conn; created_at : float }

type 'conn t = {
  config : 'conn config;
  available : 'conn pooled Queue.t;
  total_count : int Loc.t;
  in_use_count : int Loc.t;
  closed : bool Loc.t;
}

let create config =
  {
    config;
    available = Queue.create ();
    total_count = Loc.make 0;
    in_use_count = Loc.make 0;
    closed = Loc.make false;
  }

let size t = Loc.get t.total_count
let in_use t = Loc.get t.in_use_count
let available t = Queue.length t.available
let is_closed t = Loc.get t.closed

let validate_connection t pooled =
  match t.config.validate with
  | None -> true
  | Some validate -> validate pooled.conn

let close_pooled t pooled =
  t.config.close pooled.conn;
  Loc.decr t.total_count

let try_create_connection t =
  let tx ~xt =
    if Xt.get ~xt t.closed then `Closed
    else
      let total = Xt.get ~xt t.total_count in
      if total >= t.config.max_size then `At_capacity
      else begin
        Xt.set ~xt t.total_count (total + 1);
        Xt.set ~xt t.in_use_count (Xt.get ~xt t.in_use_count + 1);
        `Can_create
      end
  in
  match Xt.commit { tx } with
  | `Closed -> Error Pool_closed
  | `At_capacity -> Error Pool_empty
  | `Can_create -> (
      match t.config.connect () with
      | Ok conn ->
          let pooled = { conn; created_at = Unix.gettimeofday () } in
          Ok pooled
      | Error msg ->
          Loc.decr t.total_count;
          Loc.decr t.in_use_count;
          Error (Connection_error msg))

let rec acquire_from_pool t =
  if Loc.get t.closed then Error Pool_closed
  else
    match Queue.take_opt t.available with
    | None -> Error Pool_empty
    | Some pooled ->
        if validate_connection t pooled then begin
          Loc.incr t.in_use_count;
          Ok pooled.conn
        end
        else begin
          close_pooled t pooled;
          acquire_from_pool t
        end

let acquire t =
  match acquire_from_pool t with
  | Ok conn -> Ok conn
  | Error Pool_empty -> (
      match try_create_connection t with
      | Ok pooled -> Ok pooled.conn
      | Error e -> Error e)
  | Error e -> Error e

let rec acquire_blocking ?timeoutf t =
  if Loc.get t.closed then Error Pool_closed
  else
    match acquire t with
    | Ok conn -> Ok conn
    | Error Pool_empty -> (
        match timeoutf with
        | Some timeout -> (
            match Queue.take_blocking ~timeoutf:timeout t.available with
            | exception Kcas.Timeout.Timeout -> Error Pool_timeout
            | pooled ->
                if validate_connection t pooled then begin
                  Loc.incr t.in_use_count;
                  Ok pooled.conn
                end
                else begin
                  close_pooled t pooled;
                  acquire_blocking ?timeoutf t
                end)
        | None ->
            let pooled = Queue.take_blocking t.available in
            if validate_connection t pooled then begin
              Loc.incr t.in_use_count;
              Ok pooled.conn
            end
            else begin
              close_pooled t pooled;
              acquire_blocking ?timeoutf t
            end)
    | Error e -> Error e

let release t conn =
  if Loc.get t.closed then t.config.close conn
  else begin
    Loc.decr t.in_use_count;
    let pooled = { conn; created_at = Unix.gettimeofday () } in
    Queue.add pooled t.available
  end

let with_connection t f =
  match acquire t with
  | Error e -> Error e
  | Ok conn -> (
      match f conn with
      | result ->
          release t conn;
          Ok result
      | exception exn ->
          release t conn;
          raise exn)

let with_connection_blocking ?timeoutf t f =
  match acquire_blocking ?timeoutf t with
  | Error e -> Error e
  | Ok conn -> (
      match f conn with
      | result ->
          release t conn;
          Ok result
      | exception exn ->
          release t conn;
          raise exn)

let drain t =
  let rec loop () =
    match Queue.take_opt t.available with
    | None -> ()
    | Some pooled ->
        close_pooled t pooled;
        loop ()
  in
  loop ()

let shutdown t =
  Loc.set t.closed true;
  drain t

type stats = { total : int; available : int; in_use : int; closed : bool }

let stats t =
  {
    total = size t;
    available = available t;
    in_use = in_use t;
    closed = is_closed t;
  }

let error_to_string = function
  | Pool_empty -> "Pool empty: no connections available"
  | Pool_closed -> "Pool closed"
  | Pool_timeout -> "Timeout waiting for connection"
  | Connection_error msg -> Printf.sprintf "Connection error: %s" msg

(** Multi-server pool with round-robin load balancing.

    Distributes connections across multiple database servers without requiring
    an external load balancer. Each server maintains its own connection pool.

    {[
      let config =
        Pool.
          {
            servers = [ "host1"; "host2"; "host3" ];
            max_size_per_server = 5;
            connect = (fun conninfo -> Repodb_postgresql.connect conninfo);
            close = Repodb_postgresql.close;
            validate = None;
          }
      in
      let multi = Pool.create_multi config in
      Pool.with_connection_multi multi (fun conn ->
          Repo.all conn ~table:users ~decode)
    ]} *)

type 'conn multi_config = {
  servers : string list;
  max_size_per_server : int;
  connect : string -> ('conn, string) result;
  close : 'conn -> unit;
  validate : ('conn -> bool) option;
}

type 'conn multi_t = {
  pools : 'conn t array;
  pool_conninfos : string array;
  conn_to_pool : ('conn, int) Stdlib.Hashtbl.t;
  healthy : bool Loc.t array;
  next_index : int Loc.t;
  closed : bool Loc.t;
}

let create_multi config =
  if config.servers = [] then
    invalid_arg "Pool.create_multi: servers list cannot be empty";
  let n = List.length config.servers in
  let pools =
    Array.of_list
      (List.map
         (fun conninfo ->
           let pool_config =
             {
               max_size = config.max_size_per_server;
               connect = (fun () -> config.connect conninfo);
               close = config.close;
               validate = config.validate;
             }
           in
           create pool_config)
         config.servers)
  in
  let healthy = Array.init n (fun _ -> Loc.make true) in
  {
    pools;
    pool_conninfos = Array.of_list config.servers;
    conn_to_pool = Stdlib.Hashtbl.create (n * config.max_size_per_server);
    healthy;
    next_index = Loc.make 0;
    closed = Loc.make false;
  }

let multi_size multi =
  Array.fold_left (fun acc pool -> acc + size pool) 0 multi.pools

let multi_in_use multi =
  Array.fold_left (fun acc pool -> acc + in_use pool) 0 multi.pools

let multi_available multi =
  Array.fold_left (fun acc pool -> acc + available pool) 0 multi.pools

let multi_is_closed multi = Loc.get multi.closed
let multi_server_count multi = Array.length multi.pools

let multi_is_healthy multi idx =
  if idx < 0 || idx >= Array.length multi.healthy then false
  else Loc.get multi.healthy.(idx)

let mark_unhealthy multi idx =
  if idx >= 0 && idx < Array.length multi.healthy then
    Loc.set multi.healthy.(idx) false

let mark_healthy multi idx =
  if idx >= 0 && idx < Array.length multi.healthy then
    Loc.set multi.healthy.(idx) true

let count_healthy multi =
  Array.fold_left
    (fun acc h -> if Loc.get h then acc + 1 else acc)
    0 multi.healthy

(** Find next healthy server using round-robin *)
let next_healthy_index multi =
  let n = Array.length multi.pools in
  let start = Loc.fetch_and_add multi.next_index 1 mod n in
  let rec find_healthy attempts idx =
    if attempts >= n then None
    else if Loc.get multi.healthy.(idx) then Some idx
    else find_healthy (attempts + 1) ((idx + 1) mod n)
  in
  find_healthy 0 start

let acquire_multi multi =
  if Loc.get multi.closed then Error Pool_closed
  else
    match next_healthy_index multi with
    | None -> Error Pool_empty (* All servers unhealthy *)
    | Some idx -> (
        match acquire multi.pools.(idx) with
        | Ok conn ->
            Stdlib.Hashtbl.replace multi.conn_to_pool conn idx;
            Ok conn
        | Error Pool_empty ->
            (* This server's pool is exhausted, try others *)
            let n = Array.length multi.pools in
            let rec try_others attempts current_idx =
              if attempts >= n then Error Pool_empty
              else if current_idx <> idx && Loc.get multi.healthy.(current_idx)
              then
                match acquire multi.pools.(current_idx) with
                | Ok conn ->
                    Stdlib.Hashtbl.replace multi.conn_to_pool conn current_idx;
                    Ok conn
                | Error Pool_empty ->
                    try_others (attempts + 1) ((current_idx + 1) mod n)
                | Error e -> Error e
              else try_others (attempts + 1) ((current_idx + 1) mod n)
            in
            try_others 1 ((idx + 1) mod n)
        | Error e -> Error e)

let release_multi multi conn =
  match Stdlib.Hashtbl.find_opt multi.conn_to_pool conn with
  | Some idx ->
      Stdlib.Hashtbl.remove multi.conn_to_pool conn;
      if Loc.get multi.closed then multi.pools.(idx).config.close conn
      else release multi.pools.(idx) conn
  | None ->
      (* Connection not tracked, close it directly using first pool's config *)
      if Array.length multi.pools > 0 then multi.pools.(0).config.close conn

let with_connection_multi multi f =
  match acquire_multi multi with
  | Error e -> Error e
  | Ok conn -> (
      match f conn with
      | result ->
          release_multi multi conn;
          Ok result
      | exception exn ->
          release_multi multi conn;
          raise exn)

let rec acquire_multi_blocking ?timeoutf multi =
  if Loc.get multi.closed then Error Pool_closed
  else
    match acquire_multi multi with
    | Ok conn -> Ok conn
    | Error Pool_empty -> (
        (* All pools exhausted, wait on a healthy one *)
        match next_healthy_index multi with
        | None -> Error Pool_empty
        | Some idx -> (
            match acquire_blocking ?timeoutf multi.pools.(idx) with
            | Ok conn ->
                Stdlib.Hashtbl.replace multi.conn_to_pool conn idx;
                Ok conn
            | Error Pool_timeout -> Error Pool_timeout
            | Error Pool_closed -> Error Pool_closed
            | Error Pool_empty -> acquire_multi_blocking ?timeoutf multi
            | Error e -> Error e))
    | Error e -> Error e

let with_connection_multi_blocking ?timeoutf multi f =
  match acquire_multi_blocking ?timeoutf multi with
  | Error e -> Error e
  | Ok conn -> (
      match f conn with
      | result ->
          release_multi multi conn;
          Ok result
      | exception exn ->
          release_multi multi conn;
          raise exn)

let drain_multi multi = Array.iter drain multi.pools

let shutdown_multi multi =
  Loc.set multi.closed true;
  Array.iter shutdown multi.pools;
  Stdlib.Hashtbl.clear multi.conn_to_pool

type multi_stats = {
  total_servers : int;
  healthy_servers : int;
  per_server : (string * stats) list;
  aggregate : stats;
}

let stats_multi multi =
  let per_server =
    Array.to_list
      (Array.mapi
         (fun i pool -> (multi.pool_conninfos.(i), stats pool))
         multi.pools)
  in
  let aggregate =
    {
      total = multi_size multi;
      available = multi_available multi;
      in_use = multi_in_use multi;
      closed = multi_is_closed multi;
    }
  in
  {
    total_servers = Array.length multi.pools;
    healthy_servers = count_healthy multi;
    per_server;
    aggregate;
  }

module Make (D : Driver.S) = struct
  type nonrec t = D.connection t
  type conn = D.connection

  let create ~max_size ~conninfo ?validate () =
    let config =
      {
        max_size;
        connect =
          (fun () -> D.connect conninfo |> Result.map_error D.error_message);
        close = D.close;
        validate;
      }
    in
    create config

  let acquire = acquire
  let acquire_blocking = acquire_blocking
  let release = release
  let with_connection = with_connection
  let with_connection_blocking = with_connection_blocking
  let drain = drain
  let shutdown = shutdown
  let stats = stats
  let size = size
  let available = available
  let in_use = in_use
  let is_closed = is_closed

  module Multi = struct
    type t = D.connection multi_t

    let create ~servers ?validate () =
      let cfg =
        {
          servers;
          max_size_per_server = 10;
          connect =
            (fun conninfo ->
              D.connect conninfo |> Result.map_error D.error_message);
          close = D.close;
          validate;
        }
      in
      create_multi cfg

    let create_sized ~servers ~max_size_per_server ?validate () =
      let cfg =
        {
          servers;
          max_size_per_server;
          connect =
            (fun conninfo ->
              D.connect conninfo |> Result.map_error D.error_message);
          close = D.close;
          validate;
        }
      in
      create_multi cfg

    let acquire = acquire_multi
    let acquire_blocking = acquire_multi_blocking
    let release = release_multi
    let with_connection = with_connection_multi
    let with_connection_blocking = with_connection_multi_blocking
    let drain = drain_multi
    let shutdown = shutdown_multi
    let stats = stats_multi
    let size = multi_size
    let available = multi_available
    let in_use = multi_in_use
    let is_closed = multi_is_closed
    let server_count = multi_server_count
    let is_healthy = multi_is_healthy
    let mark_healthy = mark_healthy
    let mark_unhealthy = mark_unhealthy
  end
end