make stats reading async

so stats retrieval is not blocked by concurrent reading of stats

Author: magnetised

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

@@ -197,10 +197,16 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Connection do
   end
 
   @impl NimblePool
-  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.
-    {:remove, :idle}
+  def handle_ping(%__MODULE__{} = conn, pool_state) do
+    if Keyword.get(pool_state, :exclusive_mode, false) do
+      # keep the write connection alive in exclusive mode — closing it
+      # would destroy an in-memory database
+      {:ok, conn}
+    else
+      # the idle timeout is only enabled in non-exclusive mode, so we're free
+      # to close all the connections, including write.
+      {:remove, :idle}
+    end
   end
 
   @impl NimblePool

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

@@ -65,6 +65,7 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
        exclusive_mode?: Keyword.get(args, :exclusive_mode, false),
        enable_memory_stats?: enable_memory_stats?,
        measurement_period: measurement_period,
+       task: nil,
        pool_opts: args
      }}
   end
@@ -74,6 +75,41 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
     {:noreply, read_stats(state), :hibernate}
   end
 
+  def handle_info({ref, read_stats_result}, %{task: %{ref: ref}} = state) do
+    state =
+      case read_stats_result do
+        {:ok, stats} ->
+          %{state | stats: stats}
+
+        {:error, 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)
+
+    {:noreply, state}
+  end
+
+  def handle_info({:DOWN, ref, :process, _pid, _reason}, %{task: %{ref: ref}} = state) do
+    Process.send_after(self(), :read_stats, state.measurement_period)
+    {:noreply, %{state | task: nil}}
+  end
+
+  def handle_info(msg, state) do
+    Logger.warning(["Received unexpected message: ", inspect(msg)])
+    {:noreply, state}
+  end
+
   @impl GenServer
   def handle_call(:statistics, _from, state) do
     {:reply, {:ok, Map.put(state.stats, :connections, state.connections)}, state}
@@ -113,41 +149,25 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
   end
 
   defp do_read_stats(state, include_memory?) do
-    state
-    |> open_connection(fn conn ->
-      with {:ok, memstat_available?, page_size} <-
-             initialize_connection(conn, state.first_run?, state.enable_memory_stats?),
-           {:ok, stats} <-
-             Connection.fetch_all(
-               conn,
-               stats_query(memstat_available? && include_memory?),
-               []
-             ) do
-        {:ok, analyze_stats(stats, page_size)}
-      end
-    end)
-    |> case do
-      {:ok, stats} ->
-        %{state | stats: stats}
-
-      {:error, 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}
+    # Read the stats in an async task so that we don't block reading the stats
+    # and get spurious timeout errors
+    task =
+      Task.async(fn ->
+        open_connection(state, fn conn ->
+          with {:ok, memstat_available?, page_size} <-
+                 initialize_connection(conn, state.first_run?, state.enable_memory_stats?),
+               {:ok, stats} <-
+                 Connection.fetch_all(
+                   conn,
+                   stats_query(memstat_available? && include_memory?),
+                   []
+                 ) do
+            {:ok, analyze_stats(stats, page_size)}
+          end
+        end)
+      end)
 
-      state ->
-        state
-    end)
-    |> tap(fn state ->
-      Process.send_after(self(), :read_stats, state.measurement_period)
-    end)
+    %{state | task: task}
   end
 
   # In exclusive_mode we *must* use a pooled connection because the db maybe
@@ -176,7 +196,7 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Statistics do
       if enable_memory_stats? do
         case Connection.enable_extension(conn, "memstat") do
           :ok ->
-            if first_run?, do: Logger.info("SQLite memory statistics enabled")
+            if first_run?, do: Logger.notice("SQLite memory statistics enabled")
             true
 
           {:error, reason} ->

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

@@ -22,6 +22,7 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDb.Supervisor do
     exclusive_mode = Keyword.get(opts, :exclusive_mode, false)
     idle_timeout = Keyword.get(opts, :connection_idle_timeout, @default_connection_idle_timeout)
     # don't close the write connection in exclusive mode
+    # NimblePool treats `worker_idle_timeout: nil` as no idle timeout
     write_pool_idle_timeout = if(exclusive_mode, do: nil, else: idle_timeout)
 
     read_pool_spec =

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

@@ -548,17 +548,31 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDbTest do
   describe "statistics" do
     @tag shape_db_opts: [enable_memory_stats?: true]
     test "export memory and disk usage when enabled", ctx do
-      assert {:ok, %{total_memory: memory, disk_size: disk_size}} =
-               ShapeDb.statistics(ctx.stack_id)
+      assert {:ok, %{total_memory: memory, disk_size: disk_size}} = wait_statistics(ctx)
 
       assert memory > 0
       assert disk_size > 0
     end
 
     test "only exports disk usage by default", ctx do
-      assert {:ok, %{total_memory: 0, disk_size: disk_size}} = ShapeDb.statistics(ctx.stack_id)
+      assert {:ok, %{total_memory: 0, disk_size: disk_size}} = wait_statistics(ctx)
       assert disk_size > 0
     end
+
+    defp wait_statistics(ctx, attempts \\ 10)
+
+    defp wait_statistics(_ctx, 0), do: :error
+
+    defp wait_statistics(ctx, remaining_attempts) do
+      case ShapeDb.statistics(ctx.stack_id) do
+        {:ok, %{updated_at: %DateTime{}} = stats} ->
+          {:ok, stats}
+
+        {:ok, _} ->
+          Process.sleep(10)
+          wait_statistics(ctx, remaining_attempts - 1)
+      end
+    end
   end
 
   describe "recovery" do
@@ -623,14 +637,14 @@ defmodule Electric.ShapeCache.ShapeStatus.ShapeDbTest do
     @tag shape_db_opts: [
            exclusive_mode: true,
            connection_idle_timeout: 10,
-           statistics_collection_period: 10
+           statistics_collection_period: 1000
          ]
     test "does not scale the pool in exclusive mode", ctx do
       ShapeDb.Connection.checkout!(ctx.stack_id, :test, fn _conn ->
         assert :ok = assert_stats_match(ctx, connections: 1)
       end)
 
-      assert :error = assert_stats_match(ctx, connections: 0, total_memory: 0)
+      assert :error = assert_stats_match(ctx, connections: 0)
     end
 
     defp assert_stats_match(ctx, match, repeats \\ 10)