From d32bc49af73c2270fefc8801ae8d6b95954f245e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Jan=20Niemier?= <lukasz@supabase.io>
Date: Mon, 25 Nov 2024 13:08:14 +0100
Subject: [PATCH] feat: change `Protocol.Server.cancel_message/2` function a
 macro

This allows us to use this piece of code in function head in pattern
match. This reduces places where the magic value is required and make
the code a little bit neater and cleaner.
---
 lib/supavisor/client_handler.ex  |  6 +++---
 lib/supavisor/handler_helpers.ex |  3 ++-
 lib/supavisor/protocol/server.ex | 11 ++++++-----
 test/supavisor/protocol_test.exs | 21 ++++++++++++---------
 4 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/lib/supavisor/client_handler.ex b/lib/supavisor/client_handler.ex
index f4864395..35244a51 100644
--- a/lib/supavisor/client_handler.ex
+++ b/lib/supavisor/client_handler.ex
@@ -11,7 +11,6 @@ defmodule Supavisor.ClientHandler do
   @behaviour :ranch_protocol
   @behaviour :gen_statem
   @proto [:tcp, :ssl]
-  @cancel_query_msg <<16::32, 1234::16, 5678::16>>
   @switch_active_count Application.compile_env(:supavisor, :switch_active_count)
   @subscribe_retries Application.compile_env(:supavisor, :subscribe_retries)
   @timeout_subscribe 500
@@ -22,10 +21,11 @@ defmodule Supavisor.ClientHandler do
     Helpers,
     Monitoring.Telem,
     Protocol.Client,
-    Protocol.Server,
     Tenants
   }
 
+  require Supavisor.Protocol.Server, as: Server
+
   @impl true
   def start_link(ref, transport, opts) do
     pid = :proc_lib.spawn_link(__MODULE__, :init, [ref, transport, opts])
@@ -111,7 +111,7 @@ defmodule Supavisor.ClientHandler do
   end
 
   # cancel request
-  def handle_event(:info, {_, _, <<@cancel_query_msg, pid::32, key::32>>}, _, _) do
+  def handle_event(:info, {_, _, Server.cancel_message(pid, key)}, _, _) do
     Logger.debug("ClientHandler: Got cancel query for #{inspect({pid, key})}")
     :ok = HandlerHelpers.send_cancel_query(pid, key)
     {:stop, {:shutdown, :cancel_query}}
diff --git a/lib/supavisor/handler_helpers.ex b/lib/supavisor/handler_helpers.ex
index 91f8055e..f46afee9 100644
--- a/lib/supavisor/handler_helpers.ex
+++ b/lib/supavisor/handler_helpers.ex
@@ -2,7 +2,8 @@ defmodule Supavisor.HandlerHelpers do
   @moduledoc false
 
   alias Phoenix.PubSub
-  alias Supavisor.Protocol.Server
+
+  require Supavisor.Protocol.Server, as: Server
 
   @spec sock_send(Supavisor.sock(), iodata()) :: :ok | {:error, term()}
   def sock_send({mod, sock}, data) do
diff --git a/lib/supavisor/protocol/server.ex b/lib/supavisor/protocol/server.ex
index 9a865134..e545830f 100644
--- a/lib/supavisor/protocol/server.ex
+++ b/lib/supavisor/protocol/server.ex
@@ -28,6 +28,12 @@ defmodule Supavisor.Protocol.Server do
           }
   end
 
+  defmacro cancel_message(pid, key) do
+    quote do
+      <<unquote(@msg_cancel_header)::binary, unquote(pid)::32, unquote(key)::32>>
+    end
+  end
+
   @spec decode(iodata()) :: [Pkt.t()] | []
   def decode(data) do
     decode(data, [])
@@ -456,11 +462,6 @@ defmodule Supavisor.Protocol.Server do
     <<byte_size(bin) + 9::32, 0, 3, 0, 0, bin::binary, 0>>
   end
 
-  @spec cancel_message(non_neg_integer, non_neg_integer) :: iodata
-  def cancel_message(pid, key) do
-    [@msg_cancel_header, <<pid::32, key::32>>]
-  end
-
   @spec has_read_only_error?(list) :: boolean
   def has_read_only_error?(pkts) do
     Enum.any?(pkts, fn
diff --git a/test/supavisor/protocol_test.exs b/test/supavisor/protocol_test.exs
index 1fc7c00d..1c87a980 100644
--- a/test/supavisor/protocol_test.exs
+++ b/test/supavisor/protocol_test.exs
@@ -1,6 +1,9 @@
 defmodule Supavisor.ProtocolTest do
   use ExUnit.Case, async: true
-  alias Supavisor.Protocol.Server, as: S
+
+  @subject Supavisor.Protocol.Server
+
+  require Supavisor.Protocol.Server
 
   @initial_data %{
     "DateStyle" => "ISO, MDY",
@@ -27,12 +30,12 @@ defmodule Supavisor.ProtocolTest do
                     97, 105, 108, 101, 100, 0, 0>>
 
   test "encode_parameter_status/1" do
-    result = S.encode_parameter_status(@initial_data)
+    result = @subject.encode_parameter_status(@initial_data)
 
     for {key, value} <- @initial_data do
       assert :erlang.is_binary(key)
       assert :erlang.is_binary(value)
-      encoded = S.encode_pkt(:parameter_status, key, value)
+      encoded = @subject.encode_pkt(:parameter_status, key, value)
       assert Enum.member?(result, encoded)
     end
   end
@@ -40,29 +43,29 @@ defmodule Supavisor.ProtocolTest do
   test "encode_pkt/3" do
     key = "TimeZone"
     value = "UTC"
-    result = S.encode_pkt(:parameter_status, key, value)
+    result = @subject.encode_pkt(:parameter_status, key, value)
 
     assert result == [<<?S, 17::32>>, [key, <<0>>, value, <<0>>]]
   end
 
   test "backend_key_data/0" do
-    {header, payload} = S.backend_key_data()
+    {header, payload} = @subject.backend_key_data()
     len = byte_size(payload) + 4
 
     assert [
-             %S.Pkt{
+             %@subject.Pkt{
                tag: :backend_key_data,
                len: 13,
                payload: %{pid: _, key: _}
              }
-           ] = S.decode([header, payload] |> IO.iodata_to_binary())
+           ] = @subject.decode([header, payload] |> IO.iodata_to_binary())
 
     assert header == <<?K, len::32>>
     assert byte_size(payload) == 8
   end
 
   test "decode_payload for error_response" do
-    assert S.decode(@auth_bin_error) == [
+    assert @subject.decode(@auth_bin_error) == [
              %Supavisor.Protocol.Server.Pkt{
                tag: :error_response,
                len: 112,
@@ -84,7 +87,7 @@ defmodule Supavisor.ProtocolTest do
     key = 123_456
     expected = <<0, 0, 0, 16, 4, 210, 22, 46, 0, 0, 0, 123, 0, 1, 226, 64>>
 
-    assert S.cancel_message(pid, key)
+    assert @subject.cancel_message(pid, key)
            |> IO.iodata_to_binary() == expected
   end
 end