From 3d4f4ee910db4fc6455370671dd8ef1b8e71c3c4 Mon Sep 17 00:00:00 2001
From: Stas <sts@abc3.dev>
Date: Wed, 6 Nov 2024 15:00:43 +0100
Subject: [PATCH] fix: better handling of resetting active_counter

---
 lib/supavisor/client_handler.ex | 25 +++++++++++++++++--------
 lib/supavisor/db_handler.ex     | 12 ++++++------
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/lib/supavisor/client_handler.ex b/lib/supavisor/client_handler.ex
index ab49ee38..b79a6d65 100644
--- a/lib/supavisor/client_handler.ex
+++ b/lib/supavisor/client_handler.ex
@@ -519,21 +519,21 @@ defmodule Supavisor.ClientHandler do
       do: :ok = sock_send_maybe_active_once(msg, data),
       else: :ok = HandlerHelpers.sock_send(data.sock, Server.ready_for_query())
 
-    {:keep_state, %{data | active_count: 0}, handle_actions(data)}
+    {:keep_state, %{data | active_count: reset_active_count(data)}, handle_actions(data)}
   end
 
   def handle_event(:info, {proto, _, <<?S, 4::32, _::binary>> = msg}, _, data)
       when proto in @proto do
     Logger.debug("ClientHandler: Receive sync while not idle")
     :ok = sock_send_maybe_active_once(msg, data)
-    {:keep_state, %{data | active_count: 0}, handle_actions(data)}
+    {:keep_state, %{data | active_count: reset_active_count(data)}, handle_actions(data)}
   end
 
   def handle_event(:info, {proto, _, <<?H, 4::32, _::binary>> = msg}, _, data)
       when proto in @proto do
     Logger.debug("ClientHandler: Receive flush while not idle")
     :ok = sock_send_maybe_active_once(msg, data)
-    {:keep_state, %{data | active_count: 0}, handle_actions(data)}
+    {:keep_state, %{data | active_count: reset_active_count(data)}, handle_actions(data)}
   end
 
   # incoming query with a single pool
@@ -645,17 +645,16 @@ defmodule Supavisor.ClientHandler do
       :ready_for_query ->
         Logger.debug("ClientHandler: Client is ready")
 
-        HandlerHelpers.sock_send(data.sock, bin)
+        :ok = sock_send_maybe_active_once(bin, data)
+
         db_pid = handle_db_pid(data.mode, data.pool, data.db_pid)
 
         {_, stats} = Telem.network_usage(:client, data.sock, data.id, data.stats)
 
         Telem.client_query_time(data.query_start, data.id)
 
-        if data.active_count > @switch_active_count,
-          do: HandlerHelpers.activate(data.sock)
-
-        {:next_state, :idle, %{data | db_pid: db_pid, stats: stats, active_count: 0},
+        {:next_state, :idle,
+         %{data | db_pid: db_pid, stats: stats, active_count: reset_active_count(data)},
          handle_actions(data)}
 
       :read_sql_error ->
@@ -1123,4 +1122,14 @@ defmodule Supavisor.ClientHandler do
       {:stop, {:shutdown, :subscribe_retries}}
     end
   end
+
+  @spec reset_active_count(map()) :: 0
+  def reset_active_count(data) do
+    if data.active_count >= @switch_active_count do
+      Logger.debug("ClientHandler: Activate socket #{inspect(data.active_count)}")
+      HandlerHelpers.activate(data.sock)
+    end
+
+    0
+  end
 end
diff --git a/lib/supavisor/db_handler.ex b/lib/supavisor/db_handler.ex
index e747d4d5..1cc2a8bd 100644
--- a/lib/supavisor/db_handler.ex
+++ b/lib/supavisor/db_handler.ex
@@ -294,10 +294,10 @@ defmodule Supavisor.DbHandler do
       when is_pid(caller) and proto in @proto do
     Logger.debug("DbHandler: Got write replica message  #{inspect(bin)}")
 
-    if data.active_count > @switch_active_count,
-      do: HandlerHelpers.active_once(data.sock)
-
     if String.ends_with?(bin, Server.ready_for_query()) do
+      if data.active_count >= @switch_active_count,
+        do: HandlerHelpers.activate(data.sock)
+
       {_, stats} = Telem.network_usage(:db, data.sock, data.id, data.stats)
 
       # in transaction mode, we need to notify the client when the transaction is finished,
@@ -311,11 +311,11 @@ defmodule Supavisor.DbHandler do
           %{data | stats: stats, active_count: 0}
         end
 
-      if data.active_count > @switch_active_count,
-        do: HandlerHelpers.activate(data.sock)
-
       {:next_state, :idle, data, {:next_event, :internal, :check_anon_buffer}}
     else
+      if data.active_count > @switch_active_count,
+        do: HandlerHelpers.active_once(data.sock)
+
       HandlerHelpers.sock_send(data.client_sock, bin)
       {:keep_state, %{data | active_count: data.active_count + 1}}
     end