346 lines
10 KiB
OCaml
346 lines
10 KiB
OCaml
(* Copyright (C) 2014--2025 Petter A. Urkedal <paurkedal@gmail.com>
|
|
*
|
|
* This library is free software; you can redistribute it and/or modify it
|
|
* under the terms of the GNU Lesser General Public License as published by
|
|
* the Free Software Foundation, either version 3 of the License, or (at your
|
|
* option) any later version, with the LGPL-3.0 Linking Exception.
|
|
*
|
|
* This library is distributed in the hope that it will be useful, but WITHOUT
|
|
* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
|
|
* FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
|
|
* License for more details.
|
|
*
|
|
* You should have received a copy of the GNU Lesser General Public License
|
|
* and the LGPL-3.0 Linking Exception along with this library. If not, see
|
|
* <http://www.gnu.org/licenses/> and <https://spdx.org>, respectively.
|
|
*)
|
|
|
|
module Config = Caqti_pool_config
|
|
|
|
let default_max_size =
|
|
try int_of_string (Sys.getenv "CAQTI_POOL_MAX_SIZE") with Not_found -> 8
|
|
|
|
let default_log_src = Logs.Src.create "Caqti_platform.Pool"
|
|
|
|
module type ALARM = sig
|
|
type switch
|
|
type stdenv
|
|
|
|
type t
|
|
|
|
val schedule :
|
|
sw: switch ->
|
|
stdenv: stdenv ->
|
|
Mtime.t -> (unit -> unit) -> t
|
|
|
|
val unschedule : t -> unit
|
|
end
|
|
|
|
module type S = sig
|
|
type switch
|
|
type stdenv
|
|
|
|
include Caqti_pool_sig.S
|
|
|
|
val create :
|
|
?config: Caqti_pool_config.t ->
|
|
?check: ('a -> (bool -> unit) -> unit) ->
|
|
?validate: ('a -> bool fiber) ->
|
|
?log_src: Logs.Src.t ->
|
|
sw: switch ->
|
|
stdenv: stdenv ->
|
|
(unit -> ('a, 'e) result fiber) -> ('a -> unit fiber) ->
|
|
('a, 'e) t
|
|
end
|
|
|
|
module Make
|
|
(System : System_sig.CORE)
|
|
(Alarm : ALARM
|
|
with type stdenv := System.stdenv
|
|
and type switch := System.Switch.t) =
|
|
struct
|
|
open System
|
|
open System.Fiber.Infix
|
|
|
|
type reaper_state =
|
|
| Idle
|
|
| Working
|
|
| Waiting of Alarm.t
|
|
|
|
let (>>=?) m f =
|
|
m >>= function Ok x -> f x | Error e -> Fiber.return (Error e)
|
|
|
|
module Task = struct
|
|
type t = {priority: float; condition: Condition.t}
|
|
let wake {condition; _} = Condition.signal condition
|
|
let compare {priority = pA; _} {priority = pB; _} = Float.compare pB pA
|
|
end
|
|
|
|
module Taskq = Heap.Make (Task)
|
|
|
|
type 'a entry = {
|
|
resource: 'a;
|
|
mutable used_count: int;
|
|
mutable used_latest: Mtime.t;
|
|
}
|
|
|
|
type ('a, +'e) t = {
|
|
stdenv: stdenv;
|
|
switch: Switch.t;
|
|
create: unit -> ('a, 'e) result Fiber.t;
|
|
free: 'a -> unit Fiber.t;
|
|
check: 'a -> (bool -> unit) -> unit;
|
|
validate: 'a -> bool Fiber.t;
|
|
log_src: Logs.Src.t;
|
|
max_idle_size: int;
|
|
max_idle_age: Mtime.Span.t option;
|
|
max_size: int;
|
|
max_use_count: int option;
|
|
|
|
(* Mutable *)
|
|
mutex: Mutex.t;
|
|
mutable cur_size: int;
|
|
queue: 'a entry Queue.t;
|
|
mutable waiting: Taskq.t;
|
|
mutable reaper_state: reaper_state;
|
|
}
|
|
|
|
(*
|
|
let configure c pool =
|
|
Option.iter (fun x -> pool.max_size <- x) Config.(get max_size c);
|
|
Option.iter (fun x -> pool.max_idle_size <- x) Config.(get max_idle_size c);
|
|
Option.iter (fun x -> pool.max_idle_age <- x) Config.(get max_idle_age c);
|
|
Option.iter (fun x -> pool.max_use_count <- x) Config.(get max_use_count c)
|
|
*)
|
|
|
|
let create
|
|
?(config = Caqti_pool_config.default)
|
|
?(check = fun _ f -> f true)
|
|
?(validate = fun _ -> Fiber.return true)
|
|
?(log_src = default_log_src)
|
|
~sw
|
|
~stdenv
|
|
create free =
|
|
let max_size =
|
|
Config.(get max_size) config |> Option.value ~default:default_max_size in
|
|
let max_idle_size =
|
|
Config.(get max_idle_size) config |> Option.value ~default:max_size in
|
|
let max_idle_age =
|
|
Config.(get max_idle_age) config |> Option.value ~default:None in
|
|
let max_use_count =
|
|
Config.(get max_use_count) config |> Option.value ~default:(Some 100) in
|
|
assert (max_size > 0);
|
|
assert (max_size >= max_idle_size);
|
|
assert (Option.fold ~none:true ~some:(fun n -> n > 0) max_use_count);
|
|
{
|
|
stdenv; switch = sw;
|
|
create; free; check; validate; log_src;
|
|
max_idle_size; max_size; max_use_count; max_idle_age;
|
|
cur_size = 0;
|
|
queue = Queue.create ();
|
|
waiting = Taskq.empty;
|
|
reaper_state = Idle;
|
|
mutex = Mutex.create ();
|
|
}
|
|
|
|
let size pool = pool.cur_size (* TODO: atomic *)
|
|
|
|
let wait_lck ~priority pool =
|
|
let condition = Condition.create () in
|
|
pool.waiting <- Taskq.push Task.({priority; condition}) pool.waiting;
|
|
Condition.wait condition pool.mutex
|
|
|
|
let schedule_lck pool =
|
|
if not (Taskq.is_empty pool.waiting) then begin
|
|
let task, taskq = Taskq.pop_e pool.waiting in
|
|
pool.waiting <- taskq;
|
|
Task.wake task
|
|
end
|
|
|
|
let realloc pool =
|
|
let on_error () =
|
|
Mutex.lock pool.mutex >|= fun () ->
|
|
pool.cur_size <- pool.cur_size - 1;
|
|
schedule_lck pool;
|
|
Mutex.unlock pool.mutex
|
|
in
|
|
Fiber.cleanup
|
|
(fun () ->
|
|
pool.create () >>=
|
|
(function
|
|
| Ok resource ->
|
|
Fiber.return @@
|
|
Ok {resource; used_count = 0; used_latest = Mtime_clock.now ()}
|
|
| Error err ->
|
|
on_error () >|= fun () ->
|
|
Error err))
|
|
on_error
|
|
|
|
let rec acquire ~priority pool =
|
|
Mutex.lock pool.mutex >>= fun () ->
|
|
acquire_and_unlock ~priority pool
|
|
and acquire_and_unlock ~priority pool =
|
|
if Queue.is_empty pool.queue then begin
|
|
if pool.cur_size < pool.max_size then
|
|
begin
|
|
pool.cur_size <- pool.cur_size + 1;
|
|
Mutex.unlock pool.mutex;
|
|
realloc pool
|
|
end
|
|
else
|
|
begin
|
|
wait_lck ~priority pool >>= fun () ->
|
|
acquire_and_unlock ~priority pool
|
|
end
|
|
end else begin
|
|
let entry = Queue.take pool.queue in
|
|
Mutex.unlock pool.mutex;
|
|
pool.validate entry.resource >>= fun ok ->
|
|
if ok then
|
|
Fiber.return (Ok entry)
|
|
else
|
|
begin
|
|
Log.warn ~src:pool.log_src (fun f ->
|
|
f "Dropped pooled connection due to invalidation.") >>= fun () ->
|
|
realloc pool
|
|
end
|
|
end
|
|
|
|
let can_reuse_lck pool entry =
|
|
pool.cur_size <= pool.max_idle_size
|
|
&& Option.fold ~none:true ~some:(fun n -> entry.used_count < n)
|
|
pool.max_use_count
|
|
|
|
let dispose_expiring_lck pool =
|
|
let rec process_queue_lck () =
|
|
(* We hold the lock and only release it while freeing resources, which
|
|
* means that some other operations may have been processed right before
|
|
* each recursive call. *)
|
|
(match Queue.peek_opt pool.queue, pool.max_idle_age with
|
|
| None, _ | _, None -> Fiber.return ()
|
|
| Some entry, Some max_idle_age ->
|
|
let now = Mtime_clock.now () in
|
|
(match Mtime.add_span entry.used_latest max_idle_age with
|
|
| None ->
|
|
Logs.warn ~src:pool.log_src (fun f -> f
|
|
"Cannot schedule pool expiration check due to \
|
|
Mtime overflow.");
|
|
pool.reaper_state <- Idle;
|
|
Fiber.return ()
|
|
| Some expiry when Mtime.compare now expiry < 0 ->
|
|
let alarm =
|
|
Alarm.schedule
|
|
~sw:pool.switch ~stdenv:pool.stdenv
|
|
expiry on_alarm
|
|
in
|
|
pool.reaper_state <- Waiting alarm;
|
|
Fiber.return ()
|
|
| Some _ ->
|
|
let entry = Queue.take pool.queue in
|
|
pool.cur_size <- pool.cur_size - 1;
|
|
Mutex.unlock pool.mutex;
|
|
pool.free entry.resource >>= fun () ->
|
|
Mutex.lock pool.mutex >>= fun () ->
|
|
assert (pool.reaper_state = Working);
|
|
process_queue_lck ()))
|
|
and on_alarm () =
|
|
async ~sw:pool.switch begin fun () ->
|
|
Mutex.lock pool.mutex >>= fun () ->
|
|
pool.reaper_state <- Working;
|
|
process_queue_lck () >|= fun () ->
|
|
Mutex.unlock pool.mutex
|
|
end
|
|
in
|
|
(match pool.reaper_state, pool.max_idle_age with
|
|
| Idle, None | Working, _ | Waiting _, Some _ ->
|
|
Fiber.return ()
|
|
| Idle, Some _ ->
|
|
pool.reaper_state <- Working;
|
|
process_queue_lck ()
|
|
| Waiting alarm, None ->
|
|
(* Not reachable unless live reconfiguration is implemented. *)
|
|
Alarm.unschedule alarm;
|
|
pool.reaper_state <- Idle;
|
|
Fiber.return ())
|
|
|
|
let release pool entry =
|
|
Mutex.lock pool.mutex >>= fun () ->
|
|
entry.used_count <- entry.used_count + 1;
|
|
if not (can_reuse_lck pool entry) then
|
|
begin
|
|
pool.cur_size <- pool.cur_size - 1;
|
|
Mutex.unlock pool.mutex;
|
|
pool.free entry.resource >>= fun () ->
|
|
Mutex.lock pool.mutex >|= fun () ->
|
|
schedule_lck pool;
|
|
Mutex.unlock pool.mutex
|
|
end
|
|
else
|
|
begin
|
|
Mutex.unlock pool.mutex;
|
|
(* TODO: Consider changing the signature of check to return a bool
|
|
* Fiber.t to avoid the async call. *)
|
|
pool.check entry.resource begin fun ok ->
|
|
async ~sw:pool.switch @@ fun () ->
|
|
Mutex.lock pool.mutex >>= fun () ->
|
|
begin
|
|
if ok then
|
|
begin
|
|
entry.used_latest <- Mtime_clock.now ();
|
|
Queue.add entry pool.queue;
|
|
dispose_expiring_lck pool
|
|
end
|
|
else
|
|
begin
|
|
Logs.warn ~src:pool.log_src (fun f ->
|
|
f "Will not repool connection due to invalidation.");
|
|
pool.cur_size <- pool.cur_size - 1;
|
|
Fiber.return ()
|
|
end
|
|
end >|= fun () ->
|
|
schedule_lck pool;
|
|
Mutex.unlock pool.mutex
|
|
end;
|
|
Fiber.return ()
|
|
end
|
|
|
|
let use ?(priority = 0.0) f pool =
|
|
acquire ~priority pool >>=? fun entry ->
|
|
Fiber.finally
|
|
(fun () -> f entry.resource)
|
|
(fun () -> release pool entry)
|
|
|
|
let rec drain pool =
|
|
Mutex.lock pool.mutex >>= fun () ->
|
|
drain_and_unlock pool
|
|
and drain_and_unlock pool =
|
|
if pool.cur_size = 0 then
|
|
begin
|
|
(match pool.reaper_state with
|
|
| Idle | Working -> ()
|
|
| Waiting alarm ->
|
|
Alarm.unschedule alarm;
|
|
pool.reaper_state <- Idle);
|
|
Mutex.unlock pool.mutex;
|
|
Fiber.return ()
|
|
end
|
|
else
|
|
(match Queue.take_opt pool.queue with
|
|
| None ->
|
|
wait_lck ~priority:0.0 pool >>= fun () ->
|
|
drain_and_unlock pool
|
|
| Some entry ->
|
|
pool.cur_size <- pool.cur_size - 1;
|
|
Mutex.unlock pool.mutex;
|
|
pool.free entry.resource >>= fun () ->
|
|
drain pool)
|
|
|
|
end
|
|
|
|
module No_alarm = struct
|
|
type t = unit
|
|
let schedule ~sw:_ ~stdenv:_ _ _ = ()
|
|
let unschedule _ = ()
|
|
end
|
|
|
|
module Make_without_alarm (System : System_sig.CORE) = Make (System) (No_alarm)
|