pr feedback

Author: magnetised

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

@@ -197,20 +197,22 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
   end
 
   @impl NimblePool
-  def handle_ping(%__MODULE__{} = state, pool_state) do
-    Logger.debug(fn -> ["Closing idle SQLite ", to_string(state.mode), " connection"] end)
-
+  def handle_ping(%__MODULE__{} = _conn, _pool_state) do
     # 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")
+  @impl NimblePool
+  def terminate_worker(reason, conn, pool_state) do
+    Logger.debug(fn ->
+      ["Closing SQLite ", to_string(conn.mode), " connection for reason: ", inspect(reason)]
+    end)
+
+    _ = close(conn)
+    :ok = Statistics.worker_stop(Keyword.get(pool_state, :stack_id))
+
+    {:ok, pool_state}
   end
 
   @max_recovery_attempts 2

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

@@ -61,7 +61,8 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
        page_size: 0,
        stats: %{},
        connections: 0,
-       use_pool?: Keyword.get(args, :exclusive_mode, false),
+       first_run?: true,
+       exclusive_mode?: Keyword.get(args, :exclusive_mode, false),
        enable_memory_stats?: enable_memory_stats?,
        measurement_period: measurement_period,
        pool_opts: args
@@ -115,7 +116,7 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
     state
     |> open_connection(fn conn ->
       with {:ok, memstat_available?, page_size} <-
-             initialize_connection(conn, state.enable_memory_stats?),
+             initialize_connection(conn, state.first_run?, state.enable_memory_stats?),
            {:ok, stats} <-
              Connection.fetch_all(
                conn,
@@ -130,21 +131,37 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
         %{state | stats: stats}
 
       {:error, reason} ->
-        Logger.warning(["Failed to read statistics: ", inspect(reason)])
+        Logger.warning(["Failed to read SQLite statistics: ", inspect(reason)])
+        state
+
+      :error ->
+        Logger.warning("Failed to read SQLite statistics")
         state
     end
+    |> then(fn
+      %{first_run?: true} = state ->
+        %{state | first_run?: false}
+
+      state ->
+        state
+    end)
     |> tap(fn state ->
       Process.send_after(self(), :read_stats, state.measurement_period)
     end)
   end
 
-  defp open_connection(%{use_pool?: true} = state, fun) do
+  # In exclusive_mode we *must* use a pooled connection because the db maybe
+  # in-memory. This is ok because in this mode we never close the single
+  # connection instance so using it won't prevent closing idle connections.
+  defp open_connection(%{exclusive_mode?: true} = state, fun) do
     Connection.checkout_write!(state.stack_id, :read_stats, fn %{conn: conn} ->
       fun.(conn)
     end)
   end
 
-  defp open_connection(%{use_pool?: false, pool_opts: pool_opts} = _state, fun) do
+  # read the stats over a completely new connection to avoid waking a pool
+  # worker and preventing it from reaching the idle timeout
+  defp open_connection(%{exclusive_mode?: false, pool_opts: pool_opts} = _state, fun) do
     with {:ok, conn} <- Connection.open(pool_opts) do
       try do
         fun.(conn)
@@ -154,11 +171,12 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
     end
   end
 
-  defp initialize_connection(conn, enable_memory_stats?) do
+  defp initialize_connection(conn, first_run?, enable_memory_stats?) do
     memstat_available? =
       if enable_memory_stats? do
         case Connection.enable_extension(conn, "memstat") do
           :ok ->
+            if first_run?, do: Logger.info("SQLite memory statistics enabled")
             true
 
           {:error, reason} ->

packages/sync-service/test/electric/shape_cache/shape_status/shape_db_test.exs

@@ -611,15 +611,13 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDbTest do
   end
 
   describe "pool scaling" do
-    @tag shape_db_opts: [connection_idle_timeout: 10, statistics_collection_period: 10]
+    @tag shape_db_opts: [connection_idle_timeout: 5, statistics_collection_period: 5]
     test "scales to 0 if nothing is active", ctx do
       ShapeDb.Connection.checkout!(ctx.stack_id, :test, fn _conn ->
-        Process.sleep(1)
-        {:ok, %{connections: n}} = ShapeDb.statistics(ctx.stack_id)
-        assert n == 1
+        assert :ok = assert_stats_match(ctx, connections: 1)
       end)
 
-      assert :ok = assert_zero_memory(ctx)
+      assert :ok = assert_stats_match(ctx, connections: 0, total_memory: 0)
     end
 
     @tag shape_db_opts: [
@@ -629,30 +627,26 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDbTest do
          ]
     test "does not scale the pool in exclusive mode", ctx do
       ShapeDb.Connection.checkout!(ctx.stack_id, :test, fn _conn ->
-        Process.sleep(1)
-        {:ok, %{connections: n}} = ShapeDb.statistics(ctx.stack_id)
-        assert n == 1
+        assert :ok = assert_stats_match(ctx, connections: 1)
       end)
 
-      :error = assert_zero_memory(ctx)
+      assert :error = assert_stats_match(ctx, connections: 0, total_memory: 0)
     end
 
-    defp assert_zero_memory(ctx, repeats \\ 10)
+    defp assert_stats_match(ctx, match, repeats \\ 10)
 
-    defp assert_zero_memory(_ctx, 0) do
+    defp assert_stats_match(_ctx, _match, 0) do
       :error
     end
 
-    defp assert_zero_memory(ctx, repeats) do
+    defp assert_stats_match(ctx, match, repeats) do
       {:ok, stats} = ShapeDb.statistics(ctx.stack_id)
 
-      case stats do
-        %{connections: 0, total_memory: 0} ->
-          :ok
-
-        _ ->
-          Process.sleep(10)
-          assert_zero_memory(ctx, repeats - 1)
+      if Enum.all?(match, fn {k, v} -> stats[k] == v end) do
+        :ok
+      else
+        Process.sleep(10)
+        assert_stats_match(ctx, match, repeats - 1)
       end
     end
   end