Add pool

Author: Schultzer

lib/pool.ex

@@ -0,0 +1,94 @@
+defmodule SQL.Pool do
+  @moduledoc false
+  use GenServer
+
+  def start_link(opts) do
+    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
+  end
+
+  @impl true
+  def init(state) do
+    {:ok, initialize(state)}
+  end
+
+  def checkout(%{id: id}=sql, pool \\ :default) do
+    start_time = System.monotonic_time()
+    scheduler_id = :erlang.system_info(:scheduler_id)
+    state = :persistent_term.get({__MODULE__, pool})
+    {_size, connections, activations, recent_activations, sockets, schedulers, _indexes, _health} = state
+    {n, workers} = Map.get(schedulers, scheduler_id)
+    case checkout(state, n, workers) do
+      :none=error ->
+        :telemetry.execute([:sql, :checkout], %{pool: pool, duration: System.monotonic_time()-start_time}, %{id: id})
+        error
+      {idx, _load} ->
+        case :atomics.compare_exchange(connections, idx, 0, 1) do
+          :ok ->
+            :counters.add(activations, idx, 1)
+            :counters.add(recent_activations, idx, 1)
+            result = {idx, elem(sockets, idx-1)}
+            :telemetry.execute([:sql, :checkout], %{pool: pool, duration: System.monotonic_time()-start_time}, %{id: id})
+            result
+          _ ->
+            checkout(sql, pool)
+        end
+    end
+  end
+
+  defp checkout(_state, _n, []), do: :none
+  defp checkout({size, _connections, activations, recent_activations, _sockets, _schedulers, indexes, health}=state, n, workers) do
+    # Power-of-three-choices with load weighting
+    workers
+    |> Enum.take_random(min(3, n))
+    |> Enum.map(fn id ->
+      active = :counters.get(activations, id)
+      recent = :counters.get(recent_activations, id)
+      health = :atomics.get(health, id)
+      {id, active + recent / 1000, health} # Weight recent activations
+    end)
+    |> Enum.filter(fn {_, _, health} -> health == 1 end)
+    |> Enum.min_by(fn {_, load, _} -> load end, fn -> nil end)
+    |> case do
+      {id, load, _} -> {id, trunc(load)}
+      _ -> checkout(state, size-n, indexes--workers)
+    end
+  end
+
+  def checkin(idx, pool) do
+    {_size, connections, activations, recent_activations, _sockets, _schedulers, _indexes, _health} = :persistent_term.get({__MODULE__, pool})
+    :ok = :atomics.put(connections, idx, 0)
+    :ok = :counters.sub(activations, idx, 1)
+    :ok = :counters.sub(recent_activations, idx, 1)
+  end
+
+  defp initialize(state) do
+    pool = build_pool(state)
+    :persistent_term.put({__MODULE__, state.name}, pool)
+    Map.put(state, :pool, pool)
+  end
+
+  defp build_pool(state) do
+    schedulers = :erlang.system_info(:schedulers_online)
+    size = state.size
+    connections = :atomics.new(size, signed: false)
+    health = :atomics.new(size, signed: false)
+    indexes = Enum.to_list(1..size)
+    schedulers = indexes
+    |> Enum.reduce(%{}, fn id, acc ->
+      Map.update(acc, rem(id, schedulers)+1, [id], &([id | &1]))
+    end)
+    |> Map.new(fn {k, v} -> {k, {length(v), v}} end)
+    {size, connections, :counters.new(size, [:write_concurrency]), :counters.new(size, [:write_concurrency]), init_sockets(size, connections, health, state.protocol), schedulers, indexes, health}
+  end
+
+  defp init_sockets(size, connections, health, _protocol) do
+    # queries = :persistent_term.get({__MODULE__, :queries})
+    for idx <- 1..size do
+      :atomics.put(connections, idx, 0)
+      :atomics.put(health, idx, 1)
+      # :socket.open(:inet, :stream, protocol)
+      self()
+    end
+    |> List.to_tuple()
+  end
+end