From a31bd89e2a36cf4d9dc9acd5178ac7e760a50d60 Mon Sep 17 00:00:00 2001 From: Snehil Shah Date: Mon, 27 Jul 2026 20:47:50 +0530 Subject: [PATCH] feat: add client_idle_in_transaction_timeout Signed-off-by: Snehil Shah --- docs/configuration/tenants.md | 2 + lib/supavisor.ex | 11 ++- lib/supavisor/client_handler.ex | 37 +++++++++- lib/supavisor/client_handler/data.ex | 1 + lib/supavisor/db_handler.ex | 7 +- lib/supavisor/manager.ex | 6 +- lib/supavisor/protocol/server.ex | 9 +++ lib/supavisor/tenants/tenant.ex | 2 + ...add_client_idle_in_transaction_timeout.exs | 9 +++ .../idle_in_transaction_timeout_test.exs | 69 +++++++++++++++++++ test/supavisor/client_handler_test.exs | 38 +++++++++- test/supavisor/db_handler_test.exs | 1 + 12 files changed, 183 insertions(+), 9 deletions(-) create mode 100644 priv/repo/migrations/20260727145257_add_client_idle_in_transaction_timeout.exs create mode 100644 test/integration/idle_in_transaction_timeout_test.exs diff --git a/docs/configuration/tenants.md b/docs/configuration/tenants.md index 4e1560b2d..4dbdc042e 100644 --- a/docs/configuration/tenants.md +++ b/docs/configuration/tenants.md @@ -44,4 +44,6 @@ connection `client_idle_timeout` - the maximum duration of an idle client connection +`client_idle_in_transaction_timeout` - the maximum duration a client may sit idle inside an open transaction + `allow_list` - a list of CIDR ranges which are allowed to connect diff --git a/lib/supavisor.ex b/lib/supavisor.ex index d4ca1aceb..34f8d5e35 100644 --- a/lib/supavisor.ex +++ b/lib/supavisor.ex @@ -120,8 +120,15 @@ defmodule Supavisor do | {:error, Supavisor.Errors.WorkerNotFoundError.t()} def subscribe(id, pid \\ self()) do with {:ok, workers} <- get_local_workers(id), - {:ok, ps, idle_timeout} <- Manager.subscribe(workers.manager, pid) do - {:ok, %{workers: workers, ps: ps, idle_timeout: idle_timeout}} + {:ok, ps, idle_timeout, idle_in_transaction_timeout} <- + Manager.subscribe(workers.manager, pid) do + {:ok, + %{ + workers: workers, + ps: ps, + idle_timeout: idle_timeout, + idle_in_transaction_timeout: idle_in_transaction_timeout + }} end end diff --git a/lib/supavisor/client_handler.ex b/lib/supavisor/client_handler.ex index 42d15f3f2..0b4b677eb 100644 --- a/lib/supavisor/client_handler.ex +++ b/lib/supavisor/client_handler.ex @@ -112,6 +112,7 @@ defmodule Supavisor.ClientHandler do stream_state: MessageStreamer.new_stream_state(FrontendMessageHandler), stats: %{}, idle_timeout: 0, + idle_in_transaction_timeout: 0, heartbeat_interval: 0, connection_start: now, state_entered_at: now, @@ -333,7 +334,8 @@ defmodule Supavisor.ClientHandler do data | manager: manager_ref, db_connection: db_connection, - idle_timeout: opts.idle_timeout + idle_timeout: opts.idle_timeout, + idle_in_transaction_timeout: opts.idle_in_transaction_timeout } Registry.register(@clients_registry, data.id, @@ -545,6 +547,25 @@ defmodule Supavisor.ClientHandler do :keep_state_and_data end + # The backend is idle inside an open transaction and holding the connection. + def handle_event(:cast, {:db_status, :idle_in_transaction}, :busy, data) do + {:keep_state_and_data, arm_idle_in_transaction(data)} + end + + def handle_event({:timeout, :idle_in_transaction}, :idle_in_transaction_terminate, _state, data) do + Logger.warning( + "ClientHandler: Terminate an idle-in-transaction connection by " <> + "#{data.idle_in_transaction_timeout} timeout" + ) + + HandlerHelpers.sock_send( + data.sock, + Server.encode_error_message(Server.idle_in_transaction_timeout()) + ) + + {:stop, :normal} + end + def handle_event(:cast, {:send_error_and_terminate, error_message}, _state, data) do HandlerHelpers.sock_send(data.sock, error_message) {:stop, :normal} @@ -694,7 +715,7 @@ defmodule Supavisor.ClientHandler do Logger.debug("ClientHandler: Receive sync") :ok = sock_send(msg, data) - {:keep_state, data, handle_actions(data)} + {:keep_state, data, handle_actions(data) ++ cancel_idle_in_transaction(data)} end # Any message when idle - checkout and send to db @@ -714,7 +735,7 @@ defmodule Supavisor.ClientHandler do def handle_event(_kind, {proto, _, msg}, :busy, data) when proto in @proto do case handle_data(msg, data) do {:ok, updated_data} -> - {:keep_state, updated_data} + {:keep_state, updated_data, cancel_idle_in_transaction(updated_data)} {:error, exception} -> Error.terminate_with_error(data, exception, :authenticated) @@ -931,6 +952,16 @@ defmodule Supavisor.ClientHandler do defp idle_timeout_action(_data), do: [] + defp arm_idle_in_transaction(%{idle_in_transaction_timeout: timeout}) when timeout > 0, + do: [{{:timeout, :idle_in_transaction}, timeout, :idle_in_transaction_terminate}] + + defp arm_idle_in_transaction(_data), do: [] + + defp cancel_idle_in_transaction(%{idle_in_transaction_timeout: timeout}) when timeout > 0, + do: [{{:timeout, :idle_in_transaction}, :cancel}] + + defp cancel_idle_in_transaction(_data), do: [] + defp record_state_duration(old_state, new_state, data) do now = System.monotonic_time() diff --git a/lib/supavisor/client_handler/data.ex b/lib/supavisor/client_handler/data.ex index abce20f08..0f85421e9 100644 --- a/lib/supavisor/client_handler/data.ex +++ b/lib/supavisor/client_handler/data.ex @@ -35,6 +35,7 @@ defmodule Supavisor.ClientHandler.Data do :stream_state, :stats, :idle_timeout, + :idle_in_transaction_timeout, :heartbeat_interval, :connection_start, :state_entered_at, diff --git a/lib/supavisor/db_handler.ex b/lib/supavisor/db_handler.ex index b26d96eb7..8e56e2f40 100644 --- a/lib/supavisor/db_handler.ex +++ b/lib/supavisor/db_handler.ex @@ -461,13 +461,18 @@ defmodule Supavisor.DbHandler do # messages and the last status is idle and not mid-transaction. outstanding = data.expected_rfq - count batch_done? = outstanding <= 0 and last_status == ?I + idle_in_transaction? = outstanding <= 0 and last_status in [?T, ?E] data = %{data | expected_rfq: max(outstanding, 0)} # db_status must be enqueued in the ClientHandler's mailbox before the final # ReadyForQuery reaches the client socket: the client sends its next query as # soon as it reads ReadyForQuery, and if that arrives while the ClientHandler # is still :busy it gets forwarded to a connection the client no longer owns. - if batch_done?, do: ClientHandler.db_status(data.caller, :ready_for_query) + cond do + batch_done? -> ClientHandler.db_status(data.caller, :ready_for_query) + idle_in_transaction? -> ClientHandler.db_status(data.caller, :idle_in_transaction) + true -> :ok + end send_result = if to_send == [], do: :ok, else: client_send(data, to_send) diff --git a/lib/supavisor/manager.ex b/lib/supavisor/manager.ex index 1d6c4c0b0..dd9447b1f 100644 --- a/lib/supavisor/manager.ex +++ b/lib/supavisor/manager.ex @@ -193,6 +193,7 @@ defmodule Supavisor.Manager do default_pool_size: default_pool_size, default_max_clients: default_max_clients, client_idle_timeout: client_idle_timeout, + client_idle_in_transaction_timeout: client_idle_in_transaction_timeout, sni_hostname: sni_hostname, feature_flags: feature_flags } = tenant_record @@ -225,6 +226,7 @@ defmodule Supavisor.Manager do default_parameter_status: ps, max_clients: max_clients, idle_timeout: client_idle_timeout, + idle_in_transaction_timeout: client_idle_in_transaction_timeout, connection_params: connection_params, mode: mode, replica_type: replica_type, @@ -267,7 +269,9 @@ defmodule Supavisor.Manager do state end - {:reply, {:ok, state.parameter_status, state.idle_timeout}, new_state} + {:reply, + {:ok, state.parameter_status, state.idle_timeout, state.idle_in_transaction_timeout}, + new_state} {:error, _} = error -> {:reply, error, state} diff --git a/lib/supavisor/protocol/server.ex b/lib/supavisor/protocol/server.ex index 99a2b3e99..b7a332f4c 100644 --- a/lib/supavisor/protocol/server.ex +++ b/lib/supavisor/protocol/server.ex @@ -189,6 +189,15 @@ defmodule Supavisor.Protocol.Server do } end + def idle_in_transaction_timeout do + %{ + "S" => "FATAL", + "V" => "FATAL", + "C" => "25P03", + "M" => "terminating connection due to idle-in-transaction timeout" + } + end + @spec scram_request :: iodata() def scram_request, do: @scram_request diff --git a/lib/supavisor/tenants/tenant.ex b/lib/supavisor/tenants/tenant.ex index 86776f2d6..8239eb2d6 100644 --- a/lib/supavisor/tenants/tenant.ex +++ b/lib/supavisor/tenants/tenant.ex @@ -29,6 +29,7 @@ defmodule Supavisor.Tenants.Tenant do field(:sni_hostname, :string) field(:default_max_clients, :integer, default: 1000) field(:client_idle_timeout, :integer, default: 0) + field(:client_idle_in_transaction_timeout, :integer, default: 0) field(:client_heartbeat_interval, :integer, default: 60) field(:allow_list, {:array, :string}, default: ["0.0.0.0/0", "::/0"]) field(:availability_zone, :string) @@ -69,6 +70,7 @@ defmodule Supavisor.Tenants.Tenant do :sni_hostname, :default_max_clients, :client_idle_timeout, + :client_idle_in_transaction_timeout, :client_heartbeat_interval, :allow_list, :availability_zone, diff --git a/priv/repo/migrations/20260727145257_add_client_idle_in_transaction_timeout.exs b/priv/repo/migrations/20260727145257_add_client_idle_in_transaction_timeout.exs new file mode 100644 index 000000000..12196704e --- /dev/null +++ b/priv/repo/migrations/20260727145257_add_client_idle_in_transaction_timeout.exs @@ -0,0 +1,9 @@ +defmodule Supavisor.Repo.Migrations.AddClientIdleInTransactionTimeout do + use Ecto.Migration + + def change do + alter table("tenants", prefix: "_supavisor") do + add(:client_idle_in_transaction_timeout, :integer, null: false, default: 0) + end + end +end diff --git a/test/integration/idle_in_transaction_timeout_test.exs b/test/integration/idle_in_transaction_timeout_test.exs new file mode 100644 index 000000000..6aae6a042 --- /dev/null +++ b/test/integration/idle_in_transaction_timeout_test.exs @@ -0,0 +1,69 @@ +defmodule Supavisor.Integration.IdleInTransactionTimeoutTest do + use Supavisor.DataCase, async: false + + alias Supavisor.Support.ProtocolClient + + @moduletag :integration + + @timeout_ms 1000 + + setup do + %{db_conf: Application.get_env(:supavisor, Supavisor.Repo)} + end + + test "server disconnects a client left idle inside a transaction", %{db_conf: db_conf} do + sock = db_conf |> idle_in_transaction_tenant() |> connect() + + # Open a transaction and then go idle. + :ok = :gen_tcp.send(sock, :pgo_protocol.encode_query_message("BEGIN")) + :ok = ProtocolClient.recv_until_ready_for_query(sock, "") + + assert {:ok, err} = :gen_tcp.recv(sock, 0, @timeout_ms * 5) + assert err =~ "idle-in-transaction" + assert err =~ "25P03" + assert {:error, :closed} = :gen_tcp.recv(sock, 0, 1000) + end + + defp connect(%{tenant: tenant, port: port, user: user, password: password}) do + {:ok, sock} = :gen_tcp.connect(~c"127.0.0.1", port, [:binary, active: false]) + ProtocolClient.authenticate(sock, "#{user}.#{tenant}", password) + sock + end + + # A require_user transaction-mode tenant with client_idle_in_transaction_timeout set. + defp idle_in_transaction_tenant(db_conf) do + suffix = :crypto.strong_rand_bytes(6) |> Base.encode16(case: :lower) + tenant_id = "idle_in_transaction_#{System.unique_integer([:positive])}_#{suffix}" + + {:ok, _} = + Supavisor.Tenants.create_tenant(%{ + db_host: db_conf[:hostname], + db_port: db_conf[:port], + db_database: db_conf[:database], + external_id: tenant_id, + require_user: true, + default_parameter_status: %{"server_version" => "15.0"}, + client_idle_in_transaction_timeout: @timeout_ms, + client_heartbeat_interval: 0, + users: [ + %{ + "db_user" => db_conf[:username], + "db_password" => db_conf[:password], + "pool_size" => 2, + "max_clients" => 10, + "mode_type" => "transaction", + "is_manager" => true + } + ] + }) + + on_exit(fn -> Supavisor.Tenants.delete_tenant_by_external_id(tenant_id) end) + + %{ + tenant: tenant_id, + port: Application.get_env(:supavisor, :proxy_port_transaction), + user: db_conf[:username], + password: db_conf[:password] + } + end +end diff --git a/test/supavisor/client_handler_test.exs b/test/supavisor/client_handler_test.exs index 2366db9f9..2b75b3e4a 100644 --- a/test/supavisor/client_handler_test.exs +++ b/test/supavisor/client_handler_test.exs @@ -187,7 +187,7 @@ defmodule Supavisor.ClientHandlerTest do batch = <> <> <> <> <> - assert {:keep_state, _data} = + assert {:keep_state, _data, _actions} = @subject.handle_event(:info, {:tcp, :sock, batch}, :busy, data) assert_received {:"$gen_cast", {:expect_ready_for_query, 3}} @@ -205,10 +205,44 @@ defmodule Supavisor.ClientHandlerTest do batch = <> <> <> - assert {:keep_state, _data} = + assert {:keep_state, _data, _actions} = @subject.handle_event(:info, {:tcp, :sock, batch}, :busy, data) refute_received {:"$gen_cast", {:expect_ready_for_query, _count}} end end + + describe "idle-in-transaction timeout" do + test "arms the timeout when the backend goes idle-in-transaction" do + data = %{idle_in_transaction_timeout: 5000} + + assert {:keep_state_and_data, + [{{:timeout, :idle_in_transaction}, 5000, :idle_in_transaction_terminate}]} = + @subject.handle_event(:cast, {:db_status, :idle_in_transaction}, :busy, data) + end + + test "does not arm when the timeout is disabled" do + data = %{idle_in_transaction_timeout: 0} + + assert {:keep_state_and_data, []} = + @subject.handle_event(:cast, {:db_status, :idle_in_transaction}, :busy, data) + end + + test "terminates the client with a FATAL 25P03 error when the timeout fires" do + {sock, recv} = sockpair() + data = %{sock: {:gen_tcp, sock}, idle_in_transaction_timeout: 5000} + + assert {:stop, :normal} = + @subject.handle_event( + {:timeout, :idle_in_transaction}, + :idle_in_transaction_terminate, + :busy, + data + ) + + assert {:ok, bin} = :gen_tcp.recv(recv, 0, 1000) + assert bin =~ "idle-in-transaction" + assert bin =~ "25P03" + end + end end diff --git a/test/supavisor/db_handler_test.exs b/test/supavisor/db_handler_test.exs index b61ae3a4f..496797b64 100644 --- a/test/supavisor/db_handler_test.exs +++ b/test/supavisor/db_handler_test.exs @@ -1172,6 +1172,7 @@ defmodule Supavisor.DbHandlerTest do Db.handle_event(:info, {:tcp, :sock, in_transaction}, :busy, data) refute_received {:"$gen_cast", {:db_status, :ready_for_query}} + assert_received {:"$gen_cast", {:db_status, :idle_in_transaction}} assert data.expected_rfq == 0 end