16470769502

Committed 23 Jul 2025 12:34PM UTC coverage: 57.067% (+1.7%) from 55.355%

Build # 16470769502

Build Type

Pull #694

github

Committed by

web-flow

Commit Message

Merge 78a9c0b2c into deaa48192

Pull Request Pull Request #694: feat: improved named prepared statements support

Run Details

175 of 217 new or added lines in 11 files covered. (80.65%)

16 existing lines in 4 files now uncovered.

1292 of 2264 relevant lines covered (57.07%)

1126.08 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

68.52

/lib/supavisor/db_handler.ex

defmodule Supavisor.DbHandler do
  @moduledoc """
  This module contains functions to start a link with the database, send requests to the database, and handle incoming messages from clients.
  It uses the Supavisor.Protocol.Server module to decode messages from the database and sends messages to clients Supavisor.ClientHandler.
  """

  @behaviour :gen_statem

  require Logger
  require Supavisor.Protocol.Server, as: Server

  alias Supavisor.Protocol.PreparedStatements

  alias Supavisor.{
    ClientHandler,
    HandlerHelpers,
    Helpers,
    Monitoring.Telem,
    Protocol.Server
  }

  @type state :: :connect | :authentication | :idle | :busy

  @reconnect_timeout 2_500
  @reconnect_timeout_proxy 500
  @sock_closed [:tcp_closed, :ssl_closed]
  @proto [:tcp, :ssl]
  @switch_active_count Application.compile_env(:supavisor, :switch_active_count)

  def start_link(config),
    do: :gen_statem.start_link(__MODULE__, config, hibernate_after: 5_000)

  def checkout(pid, sock, caller, timeout \\ 15_000),
    do: :gen_statem.call(pid, {:checkout, sock, caller}, timeout)

  @spec checkin(pid()) :: :ok
  def checkin(pid), do: :gen_statem.cast(pid, :checkin)

  @spec handle_prepared_statement_pkts(pid, [PreparedStatements.handled_pkt()]) :: :ok
  def handle_prepared_statement_pkts(pid, pkts) do
    :gen_statem.call(pid, {:handle_ps_pkts, pkts}, 15_000)
  end

  @spec get_state_and_mode(pid()) :: {:ok, {state, Supavisor.mode()}} | {:error, term()}
  def get_state_and_mode(pid) do
    {:ok, :gen_statem.call(pid, :get_state_and_mode, 5_000)}
  catch
    error, reason -> {:error, {error, reason}}
  end

  @spec stop(pid()) :: :ok
  def stop(pid) do
    Logger.debug("DbHandler: Stop pid #{inspect(pid)}")
    :gen_statem.stop(pid, {:shutdown, :client_termination}, 5_000)
  end

  @impl true
  def init(args) do
    Process.flag(:trap_exit, true)
    Helpers.set_log_level(args.log_level)
    Helpers.set_max_heap_size(90)

    {_, tenant} = args.tenant
    Logger.metadata(project: tenant, user: args.user, mode: args.mode)

    data =
      %{
        id: args.id,
        sock: nil,
        sent: false,
        auth: args.auth,
        user: args.user,
        tenant: args.tenant,
        buffer: [],
        anon_buffer: [],
        db_state: nil,
        parameter_status: %{},
        nonce: nil,
        messages: "",
        server_proof: nil,
        stats: %{},
        client_stats: %{},
        prepared_statements: MapSet.new(),
        pending_bin: "",
        action_queue: :queue.new(),
        mode: args.mode,
        replica_type: args.replica_type,
        reply: nil,
        caller: args[:caller] || nil,
        client_sock: args[:client_sock] || nil,
        proxy: args[:proxy] || false,
        active_count: 0,
        reconnect_retries: 0
      }

    Telem.handler_action(:db_handler, :started, args.id)
    {:ok, :connect, data, {:next_event, :internal, :connect}}
  end

  @impl true
  def callback_mode, do: [:handle_event_function]

  @impl true
  def handle_event(:internal, _, :connect, %{auth: auth} = data) do
    Logger.debug("DbHandler: Try to connect to DB")

    sock_opts =
      [
        auth.ip_version,
        mode: :binary,
        packet: :raw,
        # recbuf: 8192,
        # sndbuf: 8192,
        # backlog: 2048,
        # send_timeout: 120,
        # keepalive: true,
        # nopush: true,
        nodelay: true,
        active: false
      ]

    Telem.handler_action(:db_handler, :db_connection, data.id)

    case :gen_tcp.connect(auth.host, auth.port, sock_opts) do
      {:ok, sock} ->
        Logger.debug("DbHandler: auth #{inspect(auth, pretty: true)}")

        case try_ssl_handshake({:gen_tcp, sock}, auth) do
          {:ok, sock} ->
            tenant = if data.proxy, do: Supavisor.tenant(data.id)
            search_path = Supavisor.search_path(data.id)

            case send_startup(sock, auth, tenant, search_path) do
              :ok ->
                :ok = activate(sock)
                {:next_state, :authentication, %{data | sock: sock}}

              {:error, reason} ->
                Logger.error("DbHandler: Send startup error #{inspect(reason)}")
                maybe_reconnect(reason, data)
            end

          {:error, reason} ->
            Logger.error("DbHandler: Handshake error #{inspect(reason)}")
            maybe_reconnect(reason, data)
        end

      other ->
        Logger.error(
          "DbHandler: Connection failed #{inspect(other)} to #{inspect(auth.host)}:#{inspect(auth.port)}"
        )

        maybe_reconnect(other, data)
    end
  end

  def handle_event(:state_timeout, :connect, _state, data) do
    retry = data.reconnect_retries
    Logger.warning("DbHandler: Reconnect #{retry} to DB")

    {:keep_state, %{data | reconnect_retries: retry + 1}, {:next_event, :internal, :connect}}
  end

  def handle_event(:info, {proto, _, bin}, :authentication, data) when proto in @proto do
    {:ok, dec_pkt, _} = Server.decode(bin)
    Logger.debug("DbHandler: dec_pkt, #{inspect(dec_pkt, pretty: true)}")

    resp = Enum.reduce(dec_pkt, %{}, &handle_auth_pkts(&1, &2, data))

    case resp do
      {:authentication_sasl, nonce} ->
        {:keep_state, %{data | nonce: nonce}}

      {:authentication_server_first_message, server_proof} ->
        {:keep_state, %{data | server_proof: server_proof}}

      %{authentication_server_final_message: _server_final} ->
        :keep_state_and_data

      %{authentication_ok: true} ->
        :keep_state_and_data

      :authentication ->
        :keep_state_and_data

      :authentication_md5 ->
        {:keep_state, data}

      :authentication_cleartext ->
        {:keep_state, data}

      {:error_response, ["SFATAL", "VFATAL", "C28P01", reason, _, _, _]} ->
        handle_authentication_error(data, reason)
        Logger.error("DbHandler: Auth error #{inspect(reason)}")
        {:stop, :invalid_password, data}

      {:error_response, error} ->
        Logger.error("DbHandler: Error auth response #{inspect(error)}")
        {:stop, {:encode_and_forward, error}}

      {:ready_for_query, acc} ->
        ps = acc.ps

        Logger.debug(
          "DbHandler: DB ready_for_query: #{inspect(acc.db_state)} #{inspect(ps, pretty: true)}"
        )

        if data.proxy do
          bin_ps = Server.encode_parameter_status(ps)
          send(data.caller, {:parameter_status, bin_ps})
        else
          Supavisor.set_parameter_status(data.id, ps)
        end

        {:next_state, :idle, %{data | parameter_status: ps, reconnect_retries: 0},
         {:next_event, :internal, :check_buffer}}

      other ->
        Logger.error("DbHandler: Undefined auth response #{inspect(other)}")
        {:stop, :auth_error, data}
    end
  end

  def handle_event(:internal, :check_buffer, :idle, %{reply: from} = data) when from != nil do
    Logger.debug("DbHandler: Check buffer")
    {:next_state, :busy, %{data | reply: nil}, {:reply, from, data.sock}}
  end

  def handle_event(:internal, :check_buffer, :idle, %{buffer: buff, caller: caller} = data)
      when is_pid(caller) do
    if buff != [] do
      Logger.debug("DbHandler: Buffer is not empty, try to send #{IO.iodata_length(buff)} bytes")
      buff = Enum.reverse(buff)
      :ok = HandlerHelpers.sock_send(data.sock, buff)
    end

    {:next_state, :busy, %{data | buffer: []}}
  end

  # check if it needs to apply queries from the anon buffer
  def handle_event(:internal, :check_anon_buffer, _, %{anon_buffer: buff, caller: nil} = data) do
    Logger.debug("DbHandler: Check anon buffer")

    if buff != [] do
      Logger.debug(
        "DbHandler: Anon buffer is not empty, try to send #{IO.iodata_length(buff)} bytes"
      )

      buff = Enum.reverse(buff)
      :ok = HandlerHelpers.sock_send(data.sock, buff)
    end

    {:keep_state, %{data | anon_buffer: []}}
  end

  def handle_event(:internal, :check_anon_buffer, _, _) do
    Logger.debug("DbHandler: Anon buffer is empty")
    :keep_state_and_data
  end

  # the process received message from db without linked caller
  def handle_event(:info, {proto, _, bin}, _, %{caller: nil}) when proto in @proto do
    Logger.debug("DbHandler: Got db response #{inspect(bin)} when caller was nil")
    :keep_state_and_data
  end

  def handle_event(:info, {proto, _, bin}, _, %{replica_type: :read} = data)
      when proto in @proto do
    Logger.debug("DbHandler: Got read replica message #{inspect(bin)}")
    {:ok, pkts, rest} = Server.decode(<<data.pending_bin::binary, bin::binary>>)
    data = %{data | pending_bin: rest}

    resp =
      cond do
        Server.has_read_only_error?(pkts) ->
          Logger.error("DbHandler: read only error")

          with [_] <- pkts do
            # need to flush ready_for_query if it's not in same packet
            :ok = receive_ready_for_query()
          end

          :read_sql_error

        List.last(pkts).tag == :ready_for_query ->
          :ready_for_query

        true ->
          :continue
      end

    if resp != :continue do
      HandlerHelpers.sock_send(data.client_sock, bin)
      :ok = ClientHandler.db_status(data.caller, resp)
      {_, stats} = Telem.network_usage(:db, data.sock, data.id, data.stats)
      {:keep_state, %{data | stats: stats, caller: handler_caller(data)}}
    else
      {:keep_state, data}
    end
  end

  # forward the message to the client
  def handle_event(:info, {proto, _, bin}, _, %{caller: caller, reply: nil} = data)
      when is_pid(caller) and proto in @proto do
    Logger.debug("DbHandler: Got write replica message  #{inspect(bin)}")
    {pkts, rest} = Supavisor.Protocol.split_pkts(<<data.pending_bin::binary, bin::binary>>)
    data = %{data | pending_bin: rest}

    last_packet = List.last(pkts)

    if transaction_complete_pkt?(last_packet) do
      HandlerHelpers.activate(data.sock)

      {_, stats} =
        if data.proxy,
          do: {nil, data.stats},
          else: Telem.network_usage(:db, data.sock, data.id, data.stats)

      # in transaction mode, we need to notify the client when the transaction is finished,
      # after which it will unlink the direct db connection process from itself.
      data =
        if data.mode == :transaction do
          # for some reason, if we send the data **before** doing the `ready_for_query` cast,
          # we get race condition that sometimes causes msgs to be sent to the wrong socket
          ClientHandler.db_status(data.caller, :ready_for_query)
          data = handle_server_messages(pkts, data)
          %{data | stats: stats, caller: nil, client_sock: nil, active_count: 0}
        else
          data = handle_server_messages(pkts, data)

          {_, client_stats} =
            if data.proxy,
              do: {nil, data.client_stats},
              else: Telem.network_usage(:client, data.client_sock, data.id, data.client_stats)

          %{data | stats: stats, active_count: 0, client_stats: client_stats}
        end

      {:next_state, :idle, data, {:next_event, :internal, :check_anon_buffer}}
    else
      if data.active_count > @switch_active_count,
        do: HandlerHelpers.active_once(data.sock)

      data = handle_server_messages(pkts, data)
      {:keep_state, %{data | active_count: data.active_count + 1}}
    end
  end

  def handle_event({:call, from}, {:handle_ps_pkts, pkts}, _state, data) do
    {iodata, data} = Enum.reduce(pkts, {[], data}, &handle_prepared_statement_pkt/2)

    {close_pkts, prepared_statements} = evict_exceeding(data)

    :ok = HandlerHelpers.sock_send(data.sock, Enum.reverse([close_pkts | iodata]))

    data = %{
      data
      | action_queue:
          Enum.reduce(close_pkts, data.action_queue, fn _, queue ->
            :queue.in({:intercept, :close}, queue)
          end),
        prepared_statements: prepared_statements
    }

    {:keep_state, data, {:reply, from, :ok}}
  end

  def handle_event({:call, from}, {:checkout, sock, caller}, state, data) do
    Logger.debug("DbHandler: checkout call when state was #{state}")

    # store the reply ref and send it when the state is idle
    if state in [:idle, :busy],
      do: {:keep_state, %{data | client_sock: sock, caller: caller}, {:reply, from, data.sock}},
      else: {:keep_state, %{data | client_sock: sock, caller: caller, reply: from}}
  end

  def handle_event({:call, from}, :ps, _, data) do
    Logger.debug("DbHandler: get parameter status")
    {:keep_state_and_data, {:reply, from, data.parameter_status}}
  end

  def handle_event(_, {closed, _}, :busy, data) when closed in @sock_closed do
    {:stop, {:shutdown, :db_termination}, data}
  end

  def handle_event(_, {closed, _}, state, data) when closed in @sock_closed do
    Logger.error("DbHandler: Connection closed when state was #{state}")

    if Application.get_env(:supavisor, :reconnect_on_db_close),
      do: {:next_state, :connect, data, {:state_timeout, reconnect_timeout(data), :connect}},
      else: {:stop, {:shutdown, :db_termination}, data}
  end

  # linked client_handler went down
  def handle_event(_, {:EXIT, pid, reason}, _state, data) do
    if reason != :normal do
      Logger.error(
        "DbHandler: ClientHandler #{inspect(pid)} went down with reason #{inspect(reason)}"
      )
    end

    HandlerHelpers.sock_send(data.sock, Server.terminate_message())
    HandlerHelpers.sock_close(data.sock)
    {:stop, {:client_handler_down, data.mode}}
  end

  def handle_event({:call, from}, :get_state_and_mode, state, data) do
    {:keep_state_and_data, {:reply, from, {state, data.mode}}}
  end

  def handle_event(type, content, state, data) do
    msg = [
      {"type", type},
      {"content", content},
      {"state", state},
      {"data", data}
    ]

    Logger.debug("DbHandler: Undefined msg: #{inspect(msg, pretty: true)}")

    :keep_state_and_data
  end

  @impl true
  def terminate(:shutdown, _state, data) do
    Telem.handler_action(:db_handler, :stopped, data.id)
    :ok
  end

  def terminate(reason, state, data) do
    Telem.handler_action(:db_handler, :stopped, data.id)

    if data.client_sock != nil do
      message =
        case reason do
          {:encode_and_forward, msg} -> Server.encode_error_message(msg)
          _ -> Server.error_message("XX000", inspect(reason))
        end

      HandlerHelpers.sock_send(data.client_sock, message)
    end

    Logger.error(
      "DbHandler: Terminating with reason #{inspect(reason)} when state was #{inspect(state)}"
    )
  end

  @spec try_ssl_handshake(Supavisor.tcp_sock(), map) :: {:ok, Supavisor.sock()} | {:error, term()}
  defp try_ssl_handshake(sock, %{upstream_ssl: true} = auth) do
    case HandlerHelpers.sock_send(sock, Server.ssl_request()) do
      :ok -> ssl_recv(sock, auth)
      error -> error
    end
  end

  defp try_ssl_handshake(sock, _), do: {:ok, sock}

  @spec ssl_recv(Supavisor.tcp_sock(), map) :: {:ok, Supavisor.ssl_sock()} | {:error, term}
  defp ssl_recv({:gen_tcp, sock} = s, auth) do
    case :gen_tcp.recv(sock, 1, 15_000) do
      {:ok, <<?S>>} -> ssl_connect(s, auth)
      {:ok, <<?N>>} -> {:error, :ssl_not_available}
      {:error, _} = error -> error
    end
  end

  @spec ssl_connect(Supavisor.tcp_sock(), map, pos_integer) ::
          {:ok, Supavisor.ssl_sock()} | {:error, term}
  defp ssl_connect({:gen_tcp, sock}, auth, timeout \\ 5000) do
    opts =
      case auth.upstream_verify do
        :peer ->
          [
            verify: :verify_peer,
            cacerts: [auth.upstream_tls_ca],
            # unclear behavior on pg14
            server_name_indication: auth.sni_hostname || auth.host,
            customize_hostname_check: [{:match_fun, fn _, _ -> true end}]
          ]

        :none ->
          [verify: :verify_none]
      end

    case :ssl.connect(sock, opts, timeout) do
      {:ok, ssl_sock} ->
        {:ok, {:ssl, ssl_sock}}

      {:error, reason} ->
        {:error, reason}
    end
  end

  @spec send_startup(Supavisor.sock(), map(), String.t() | nil, String.t() | nil) ::
          :ok | {:error, term}
  def send_startup(sock, auth, tenant, search_path) do
    user =
      if is_nil(tenant), do: get_user(auth), else: "#{get_user(auth)}.#{tenant}"

    msg =
      :pgo_protocol.encode_startup_message(
        [
          {"user", user},
          {"database", auth.database},
          {"application_name", auth.application_name}
        ] ++ if(search_path, do: [{"options", "--search_path=#{search_path}"}], else: [])
      )

    HandlerHelpers.sock_send(sock, msg)
  end

  @spec activate(Supavisor.sock()) :: :ok | {:error, term}
  defp activate({:gen_tcp, sock}) do
    :inet.setopts(sock, active: true)
  end

  defp activate({:ssl, sock}) do
    :ssl.setopts(sock, active: true)
  end

  defp get_user(auth) do
    if auth.require_user do
      auth.secrets.().db_user
    else
      auth.secrets.().user
    end
  end

  @spec receive_ready_for_query() :: :ok | :timeout_error
  defp receive_ready_for_query do
    receive do
      {_proto, _socket, <<?Z, 5::32, ?I>>} ->
        :ok
    after
      15_000 -> :timeout_error
    end
  end

  @spec handler_caller(map()) :: pid() | nil
  defp handler_caller(%{mode: :session} = data), do: data.caller
  defp handler_caller(_), do: nil

  @spec check_ready(binary()) ::
          {:ready_for_query, :idle | :transaction_block | :failed_transaction_block} | :continue
  def check_ready(bin) do
    bin_size = byte_size(bin)

    case bin do
      <<_::binary-size(bin_size - 6), 90, 0, 0, 0, 5, status_indicator::binary>> ->
        indicator =
          case status_indicator do
            <<?I>> -> :idle
            <<?T>> -> :transaction_block
            <<?E>> -> :failed_transaction_block
            _ -> :continue
          end

        {:ready_for_query, indicator}

      _ ->
        :continue
    end
  end

  @spec handle_auth_pkts(map(), map(), map()) :: any()
  defp handle_auth_pkts(%{tag: :parameter_status, payload: {k, v}}, acc, _),
    do: update_in(acc, [:ps], fn ps -> Map.put(ps || %{}, k, v) end)

  defp handle_auth_pkts(%{tag: :ready_for_query, payload: db_state}, acc, _),
    do: {:ready_for_query, Map.put(acc, :db_state, db_state)}

  defp handle_auth_pkts(%{tag: :backend_key_data, payload: payload}, acc, data) do
    key = self()
    conn = %{host: data.auth.host, port: data.auth.port, ip_ver: data.auth.ip_version}
    Registry.register(Supavisor.Registry.PoolPids, key, Map.merge(payload, conn))
    Logger.debug("DbHandler: Backend #{inspect(key)} data: #{inspect(payload)}")
    Map.put(acc, :backend_key_data, payload)
  end

  defp handle_auth_pkts(%{payload: {:authentication_sasl_password, methods_b}}, _, data) do
    nonce =
      case Server.decode_string(methods_b) do
        {:ok, req_method, _} ->
          Logger.debug("DbHandler: SASL method #{inspect(req_method)}")
          nonce = :pgo_scram.get_nonce(16)
          user = get_user(data.auth)
          client_first = :pgo_scram.get_client_first(user, nonce)
          client_first_size = IO.iodata_length(client_first)

          sasl_initial_response = [
            "SCRAM-SHA-256",
            0,
            <<client_first_size::32-integer>>,
            client_first
          ]

          bin = :pgo_protocol.encode_scram_response_message(sasl_initial_response)
          :ok = HandlerHelpers.sock_send(data.sock, bin)
          nonce

        other ->
          Logger.error("DbHandler: Undefined sasl method #{inspect(other)}")
          nil
      end

    {:authentication_sasl, nonce}
  end

  defp handle_auth_pkts(
         %{payload: {:authentication_server_first_message, server_first}},
         _,
         data
       )
       when data.auth.require_user == false do
    nonce = data.nonce
    server_first_parts = Helpers.parse_server_first(server_first, nonce)

    {client_final_message, server_proof} =
      Helpers.get_client_final(
        :auth_query,
        data.auth.secrets.(),
        server_first_parts,
        nonce,
        data.auth.secrets.().user,
        "biws"
      )

    bin = :pgo_protocol.encode_scram_response_message(client_final_message)
    :ok = HandlerHelpers.sock_send(data.sock, bin)

    {:authentication_server_first_message, server_proof}
  end

  defp handle_auth_pkts(
         %{payload: {:authentication_server_first_message, server_first}},
         _,
         data
       ) do
    nonce = data.nonce
    server_first_parts = :pgo_scram.parse_server_first(server_first, nonce)

    {client_final_message, server_proof} =
      :pgo_scram.get_client_final(
        server_first_parts,
        nonce,
        data.auth.user,
        data.auth.secrets.().password
      )

    bin = :pgo_protocol.encode_scram_response_message(client_final_message)
    :ok = HandlerHelpers.sock_send(data.sock, bin)

    {:authentication_server_first_message, server_proof}
  end

  defp handle_auth_pkts(
         %{payload: {:authentication_server_final_message, server_final}},
         acc,
         _data
       ),
       do: Map.put(acc, :authentication_server_final_message, server_final)

  defp handle_auth_pkts(
         %{payload: :authentication_ok},
         acc,
         _data
       ),
       do: Map.put(acc, :authentication_ok, true)

  defp handle_auth_pkts(%{payload: {:authentication_md5_password, salt}} = dec_pkt, _, data) do
    Logger.debug("DbHandler: dec_pkt, #{inspect(dec_pkt, pretty: true)}")

    digest =
      if data.auth.method == :password do
        Helpers.md5([data.auth.password.(), data.auth.user])
      else
        data.auth.secrets.().secret
      end

    payload = ["md5", Helpers.md5([digest, salt]), 0]
    bin = [?p, <<IO.iodata_length(payload) + 4::signed-32>>, payload]
    :ok = HandlerHelpers.sock_send(data.sock, bin)
    :authentication_md5
  end

  defp handle_auth_pkts(%{payload: :authentication_cleartext_password} = dec_pkt, _, data) do
    Logger.debug("DbHandler: dec_pkt, #{inspect(dec_pkt, pretty: true)}")

    payload = <<data.auth.password.()::binary, 0>>
    bin = [?p, <<IO.iodata_length(payload) + 4::signed-32>>, payload]
    :ok = HandlerHelpers.sock_send(data.sock, bin)
    :authentication_cleartext
  end

  defp handle_auth_pkts(%{tag: :error_response, payload: error}, _acc, _data),
    do: {:error_response, error}

  defp handle_auth_pkts(_e, acc, _data), do: acc

  @spec handle_authentication_error(map(), String.t()) :: any()
  defp handle_authentication_error(%{proxy: false} = data, reason) do
    tenant = Supavisor.tenant(data.id)

    :erpc.multicast([node() | Node.list()], fn ->
      Cachex.del(Supavisor.Cache, {:secrets, tenant, data.user})
      Cachex.del(Supavisor.Cache, {:secrets_check, tenant, data.user})

      Registry.dispatch(Supavisor.Registry.TenantClients, data.id, fn entries ->
        for {client_handler, _meta} <- entries,
            do: send(client_handler, {:disconnect, reason})
      end)
    end)

    Supavisor.stop(data.id)
  end

  defp handle_authentication_error(%{proxy: true}, _reason), do: :ok

  @spec reconnect_timeout(map()) :: pos_integer()
  def reconnect_timeout(%{proxy: true}),
    do: @reconnect_timeout_proxy

  def reconnect_timeout(_),
    do: @reconnect_timeout

  @spec handle_server_messages([binary()], map()) :: map()
  defp handle_server_messages(pkts, data) do
    {packets_to_send, updated_data} = process_messages(pkts, [], data)

    if packets_to_send != [] do
      HandlerHelpers.sock_send(data.client_sock, Enum.reverse(packets_to_send))
    end

    updated_data
  end

  defp process_messages([msg | rest], acc_bins, data) do
    {acc_bins, data} = maybe_inject(data, acc_bins)

    case msg do
      Server.parse_complete_message() ->
        {acc_bins, data} = maybe_filter_message(msg, acc_bins, data, :parse)
        process_messages(rest, acc_bins, data)

      Server.close_complete_message() ->
        {acc_bins, data} = maybe_filter_message(msg, acc_bins, data, :close)
        process_messages(rest, acc_bins, data)

      _other ->
        process_messages(rest, [msg | acc_bins], data)
    end
  end

  defp process_messages([], acc_bins, data) do
    maybe_inject(data, acc_bins)
  end

  defp maybe_filter_message(msg, acc_bins, data, action_type) do
    case :queue.out(data.action_queue) do
      {{:value, {:forward, ^action_type}}, updated_queue} ->
        {[msg | acc_bins], %{data | action_queue: updated_queue}}

      {{:value, {:intercept, ^action_type}}, updated_queue} ->
        {acc_bins, %{data | action_queue: updated_queue}}

      _other ->
        {[msg | acc_bins], data}
    end
  end

  defp maybe_inject(data, acc_bins) do
    case :queue.out(data.action_queue) do
      {{:value, {:inject, :parse}}, q2} ->
        {[Server.parse_complete_message() | acc_bins], %{data | action_queue: q2}}

      _other ->
        {acc_bins, data}
    end
  end

  defp transaction_complete_pkt?(<<?Z, 5::32, ?I>>), do: true
  defp transaction_complete_pkt?(_), do: false

  # If the prepared statement exists for us, it exists for the server, so we just send the
  # bind to the socket. If it doesn't, we must send the parse pkt first.
  #
  # If we received a bind without a parse, we need to intercept the parse response, otherwise,
  # the client will receive an unexpected message.
  defp handle_prepared_statement_pkt({:bind_pkt, stmt_name, pkt, parse_pkt}, {iodata, data}) do
    if stmt_name in data.prepared_statements do
      {[pkt | iodata], data}
    else
      new_data = %{
        data
        | action_queue: :queue.in({:intercept, :parse}, data.action_queue),
          prepared_statements: MapSet.put(data.prepared_statements, stmt_name)
      }

      {[[parse_pkt, pkt] | iodata], new_data}
    end
  end

  defp handle_prepared_statement_pkt({:close_pkt, stmt_name, pkt}, {iodata, data}) do
    {[pkt | iodata],
     %{
       data
       | prepared_statements: MapSet.delete(data.prepared_statements, stmt_name),
         action_queue: :queue.in({:forward, :close}, data.action_queue)
     }}
  end

  defp handle_prepared_statement_pkt({:describe_pkt, _stmt_name, pkt}, {iodata, data}) do
    {[pkt | iodata], data}
  end

  # If we stop generating unique id per statement, and instead do deterministic ids,
  # we need to potentially drop parse pkts and return a parse response
  defp handle_prepared_statement_pkt({:parse_pkt, stmt_name, pkt}, {iodata, data}) do
    if stmt_name in data.prepared_statements do
      {iodata, %{data | action_queue: :queue.in({:inject, :parse}, data.action_queue)}}
    else
      prepared_statements = MapSet.put(data.prepared_statements, stmt_name)

      {[pkt | iodata],
       %{
         data
         | prepared_statements: prepared_statements,
           action_queue: :queue.in({:forward, :parse}, data.action_queue)
       }}
    end
  end

  defp evict_exceeding(%{prepared_statements: prepared_statements, id: id}) do
    limit = PreparedStatements.backend_limit()

    if MapSet.size(prepared_statements) >= limit do
      count = div(limit, 5)
      to_remove = Enum.take_random(prepared_statements, count) |> MapSet.new()
      close_pkts = Enum.map(to_remove, &PreparedStatements.build_close_pkt/1)
      prepared_statements = MapSet.difference(prepared_statements, to_remove)
      Telem.prepared_statements_evicted(count, id)

      {close_pkts, prepared_statements}
    else
      {[], prepared_statements}
    end
  end

  defp maybe_reconnect(reason, data) do
    max_reconnect_retries = Application.get_env(:supavisor, :reconnect_retries)

    if data.reconnect_retries > max_reconnect_retries and data.client_sock != nil do
      {:stop, {:failed_to_connect, reason}}
    else
      {:keep_state_and_data, {:state_timeout, reconnect_timeout(data), :connect}}
    end
  end
end

supabase / supavisor / 16470769502

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous