scale sqlite pool to 0

Author: magnetised

.changeset/unlucky-files-protect.md

@@ -0,0 +1,5 @@
+---
+'@core/sync-service': patch
+---
+
+Add SQLite connection pool scaling to minimize memory usage in quiet instances

packages/sync-service/lib/electric/shape_cache/shape_status/shape_db/connection.ex

@@ -4,6 +4,7 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
   alias Exqlite.Sqlite3
   alias Electric.ShapeCache.ShapeStatus.ShapeDb.Query
   alias Electric.ShapeCache.ShapeStatus.ShapeDb.PoolRegistry
+  alias Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics
   alias Electric.Telemetry.OpenTelemetry
 
   require Logger
@@ -65,7 +66,7 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
 
   defguardp is_raw_connection(conn) when is_reference(conn)
 
-  defstruct [:conn, :stmts]
+  defstruct [:conn, :mode, :stmts]
 
   def migrate(conn, opts) when is_raw_connection(conn) do
     # because we embed the storage version into the db path
@@ -119,7 +120,7 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
   end
 
   def optimize(conn) when is_raw_connection(conn) do
-    execute_all(conn, ["PRAGMA optimize=0x10002"])
+    execute_all(conn, ["PRAGMA optimize=0x10002", "PRAGMA shrink_memory"])
   end
 
   def enable_extension(conn, extension) when is_raw_connection(conn) do
@@ -131,17 +132,25 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
 
   @impl NimblePool
   def init_worker(pool_state) do
-    if Keyword.get(pool_state, :exclusive_mode, false) do
+    with {:ok, conn} <- init_worker_for_pool(pool_state) do
+      :ok = Statistics.worker_start(Keyword.get(pool_state, :stack_id))
+      {:ok, conn, pool_state}
+    end
+  end
+
+  defp init_worker_for_pool(pool_state) do
+    if(Keyword.get(pool_state, :exclusive_mode, false)) do
       init_worker_exclusive(pool_state)
     else
       init_worker_pooled(pool_state)
     end
   end
 
   defp init_worker_pooled(pool_state) do
-    with {:ok, conn} <- open(pool_state),
-         stmts <- Query.prepare!(conn, pool_state) do
-      {:ok, %__MODULE__{conn: conn, stmts: stmts}, pool_state}
+    with mode = Keyword.get(pool_state, :mode, :readwrite),
+         {:ok, conn} <- open(pool_state),
+         stmts <- Query.prepare!(conn, mode) do
+      {:ok, %__MODULE__{conn: conn, mode: mode, stmts: stmts}}
     end
   end
 
@@ -151,10 +160,11 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
   # over the same, single, connection since every connection is a separate db
   # in in-memory mode.
   defp init_worker_exclusive(pool_state) do
-    with {:ok, conn} <- open(pool_state, integrity_check: true),
+    with mode = :readwrite,
+         {:ok, conn} <- open(pool_state, integrity_check: true),
          {:ok, _version} <- migrate(conn, pool_state),
-         stmts <- Query.prepare!(conn, Keyword.put(pool_state, :mode, :readwrite)) do
-      {:ok, %__MODULE__{conn: conn, stmts: stmts}, pool_state}
+         stmts <- Query.prepare!(conn, mode) do
+      {:ok, %__MODULE__{conn: conn, mode: mode, stmts: stmts}}
     end
   end
 
@@ -186,6 +196,23 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
     {:ok, client_state, pool_state}
   end
 
+  @impl NimblePool
+  def handle_ping(%__MODULE__{} = state, pool_state) do
+    Logger.debug(fn -> ["Closing idle SQLite ", to_string(state.mode), " connection"] end)
+
+    # the idle timeout is only enabled in non-exclusive mode, so we're free
+    # to close all the connections, including write.
+
+    :ok = close(state)
+    :ok = Statistics.worker_stop(Keyword.get(pool_state, :stack_id))
+
+    {:remove, :idle}
+  end
+
+  def shrink_memory(conn) when is_raw_connection(conn) do
+    execute(conn, "PRAGMA shrink_memory")
+  end
+
   @max_recovery_attempts 2
 
   def open(pool_state, opts \\ []) do
@@ -258,6 +285,26 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
     end
   end
 
+  def close(%__MODULE__{conn: conn, stmts: stmts}) do
+    # Need to release all the prepared statements on the connection or the
+    # close doesn't actually release the resources.
+    stmts
+    |> Query.active_stmts()
+    |> Enum.each(fn {name, stmt} ->
+      case Sqlite3.release(conn, stmt) do
+        :ok ->
+          :ok
+
+        {:error, reason} ->
+          Logger.warning(
+            "Failed to release prepared statement #{inspect(name)}: #{inspect(reason)}"
+          )
+      end
+    end)
+
+    close(conn)
+  end
+
   def close(conn) when is_raw_connection(conn) do
     Sqlite3.close(conn)
   end

packages/sync-service/lib/electric/shape_cache/shape_status/shape_db/migrator.ex

@@ -20,44 +20,38 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Migrator do
     Logger.metadata(stack_id: stack_id)
     Electric.Telemetry.Sentry.set_tags_context(stack_id: stack_id)
 
-    with {:ok, conn} <- apply_migration(stack_id, args, exclusive_mode) do
-      {:ok, schedule_optimize(stack_id, conn), :hibernate}
+    with :ok <- apply_migration(stack_id, args, exclusive_mode) do
+      {:ok, schedule_optimize(stack_id), :hibernate}
     end
   end
 
   defp apply_migration(_stack_id, _opts, true = _exclusive?) do
     # In exclusive  mode we *must* apply the migrations within the pool
     # connection initialization because we might be using a memory db.
     # We return nil to trigger checkout-mode.
-    {:ok, nil}
+    :ok
   end
 
   defp apply_migration(_stack_id, opts, false = _exclusive?) do
     with {:ok, conn} <- ShapeDb.Connection.open(opts, integrity_check: true),
          {:ok, _version} <- ShapeDb.Connection.migrate(conn, opts),
-         :ok = ShapeDb.Connection.optimize(conn) do
-      {:ok, conn}
+         :ok = ShapeDb.Connection.optimize(conn),
+         :ok = ShapeDb.Connection.close(conn) do
+      :ok
     end
   end
 
   @impl GenServer
-  def handle_info(:optimize, {stack_id, nil}) do
+  def handle_info(:optimize, stack_id) do
     ShapeDb.Connection.checkout_write!(stack_id, :optimize, fn %{conn: conn} ->
       :ok = ShapeDb.Connection.optimize(conn)
     end)
 
-    {:noreply, schedule_optimize(stack_id, nil), :hibernate}
+    {:noreply, schedule_optimize(stack_id), :hibernate}
   end
 
-  def handle_info(:optimize, {stack_id, conn}) do
-    Logger.notice("Optimizing shape db tables")
-    :ok = ShapeDb.Connection.optimize(conn)
-
-    {:noreply, schedule_optimize(stack_id, conn), :hibernate}
-  end
-
-  defp schedule_optimize(stack_id, conn) do
+  defp schedule_optimize(stack_id) do
     Process.send_after(self(), :optimize, @optimization_period)
-    {stack_id, conn}
+    stack_id
   end
 end

packages/sync-service/lib/electric/shape_cache/shape_status/shape_db/query.ex

@@ -52,7 +52,9 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Query do
                    Keyword.take(@read_queries, [:handle_lookup])
                  )
 
-  defstruct Enum.uniq(Keyword.keys(@read_queries) ++ Keyword.keys(@write_queries))
+  @stmt_names Enum.uniq(Keyword.keys(@read_queries) ++ Keyword.keys(@write_queries))
+
+  defstruct @stmt_names
 
   alias Electric.ShapeCache.ShapeStatus.ShapeDb.Connection, as: Conn
 
@@ -68,8 +70,18 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Query do
       stream_query: 3
     ]
 
-  def prepare!(conn, opts) do
-    case Keyword.get(opts, :mode, :readwrite) do
+  def active_stmts(nil) do
+    []
+  end
+
+  def active_stmts(%__MODULE__{} = query) do
+    @stmt_names
+    |> Enum.map(&{&1, Map.fetch!(query, &1)})
+    |> Enum.reject(&is_nil(elem(&1, 1)))
+  end
+
+  def prepare!(conn, mode) do
+    case mode do
       :readwrite ->
         struct(__MODULE__, prepare_stmts!(conn, @read_queries ++ @write_queries))