Add socket protocol support and trim the encode path

The shackle protocol is now configurable with the protocol app env,
so pools can run on shackle_socket instead of gen_tcp. The setup
path (startup, auth, use keyspace) picks gen_tcp or the socket
module based on the socket type.

encode_long_string no longer copies the query into a fresh binary,
and the query flags are computed with direct map lookups instead of
four query_opts round trips.

Against a local Scylla (64 callers, pool of 16), shackle_socket cuts
client CPU from ~24.8 to ~19.8 us per request; throughput is server
bound and unchanged. Requires a shackle release with shackle_socket
when the protocol env is set; the default is unchanged.

Author: lpgauth

src/marina_client.erl

@@ -30,7 +30,7 @@ init(_Opts) ->
         frame_flags = marina_utils:frame_flags()
     }}.
 
--spec setup(inet:socket(), state()) ->
+-spec setup(shackle:socket(), state()) ->
     {ok, state()} |
     {error, atom(), state()}.
 

src/marina_pool.erl

@@ -105,6 +105,7 @@ start(<<A, B, C, D>> = RpcAddress) ->
     PoolSize = ?GET_ENV(pool_size, ?DEFAULT_POOL_SIZE),
     PoolStrategy = ?GET_ENV(pool_strategy, ?DEFAULT_POOL_STRATEGY),
     Port = ?GET_ENV(port, ?DEFAULT_PORT),
+    Protocol = ?GET_ENV(protocol, shackle_tcp),
     Reconnect = ?GET_ENV(reconnect, ?DEFAULT_RECONNECT),
     ReconnectTimeMax = ?GET_ENV(reconnect_time_max,
         ?DEFAULT_RECONNECT_MAX),
@@ -115,6 +116,7 @@ start(<<A, B, C, D>> = RpcAddress) ->
     ClientOptions = [
         {ip, Ip},
         {port, Port},
+        {protocol, Protocol},
         {reconnect, Reconnect},
         {reconnect_time_max, ReconnectTimeMax},
         {reconnect_time_min, ReconnectTimeMin},

src/marina_request.erl

@@ -158,44 +158,34 @@ encode_batch_values(Values) ->
      [marina_types:encode_bytes(V) || V <- Values]].
 
 flags(QueryOpts) ->
-    {Mask1, Values} = values_flag(QueryOpts),
-    Mask2 = skip_metadata(QueryOpts),
-    {Mask3, PageSize} = page_size_flag(QueryOpts),
-    {Mask4, PagingState} = paging_state(QueryOpts),
+    {Mask1, Values} = values_flag(maps:get(values, QueryOpts, undefined)),
+    Mask2 = skip_metadata(maps:get(skip_metadata, QueryOpts, false)),
+    {Mask3, PageSize} =
+        page_size_flag(maps:get(page_size, QueryOpts, undefined)),
+    {Mask4, PagingState} =
+        paging_state(maps:get(paging_state, QueryOpts, undefined)),
     Flags = Mask1 + Mask2 + Mask3 + Mask4,
     [Flags, Values, PageSize, PagingState].
 
-page_size_flag(QueryOpts) ->
-    case marina_utils:query_opts(page_size, QueryOpts) of
-        undefined ->
-            {0, []};
-        PageSize ->
-            {4, marina_types:encode_int(PageSize)}
-    end.
-
-paging_state(QueryOpts) ->
-    case marina_utils:query_opts(paging_state, QueryOpts) of
-        undefined ->
-            {0, []};
-        PagingState ->
-            {8, marina_types:encode_bytes(PagingState)}
-    end.
-
-skip_metadata(QueryOpts) ->
-    case marina_utils:query_opts(skip_metadata, QueryOpts) of
-        false -> 0;
-        true -> 2
-    end.
-
-values_flag(QueryOpts) ->
-    case marina_utils:query_opts(values, QueryOpts) of
-        undefined ->
-            {0, []};
-        Values ->
-            ValuesCount = length(Values),
-            EncodedValues = [marina_types:encode_bytes(Value) ||
-                Value <- Values],
-            Values2 = [marina_types:encode_short(ValuesCount),
-                EncodedValues],
-            {1, Values2}
-    end.
+page_size_flag(undefined) ->
+    {0, []};
+page_size_flag(PageSize) ->
+    {4, marina_types:encode_int(PageSize)}.
+
+paging_state(undefined) ->
+    {0, []};
+paging_state(PagingState) ->
+    {8, marina_types:encode_bytes(PagingState)}.
+
+skip_metadata(false) ->
+    0;
+skip_metadata(true) ->
+    2.
+
+values_flag(undefined) ->
+    {0, []};
+values_flag(Values) ->
+    ValuesCount = length(Values),
+    EncodedValues = [marina_types:encode_bytes(Value) || Value <- Values],
+    Values2 = [marina_types:encode_short(ValuesCount), EncodedValues],
+    {1, Values2}.

src/marina_types.erl

@@ -160,10 +160,10 @@ encode_list(Values) ->
 encode_long(Value) ->
     <<Value:64/signed>>.
 
--spec encode_long_string(binary()) -> binary().
+-spec encode_long_string(binary()) -> iodata().
 
 encode_long_string(Value) ->
-    <<(encode_int(size(Value)))/binary, Value/binary>>.
+    [encode_int(size(Value)), Value].
 
 -spec encode_short(integer()) -> binary().
 encode_short(Value) ->

src/marina_utils.erl

@@ -17,7 +17,7 @@
 ]).
 
 %% public
--spec authenticate(inet:socket()) ->
+-spec authenticate(shackle:socket()) ->
     ok | {error, atom()}.
 
 authenticate(Socket) ->
@@ -103,18 +103,18 @@ server_to_pool(Node) ->
     PoolBin = erlang:iolist_to_binary(lists:join(<<"_">>, PoolSplit)),
     erlang:binary_to_atom(PoolBin).
 
--spec sync_msg(inet:socket(), iodata()) ->
+-spec sync_msg(shackle:socket(), iodata()) ->
     {ok, term()} | {error, term()}.
 
 sync_msg(Socket, Msg) ->
-    case gen_tcp:send(Socket, Msg) of
+    case sock_send(Socket, Msg) of
         ok ->
             rcv_buf(Socket, <<>>);
         {error, Reason} ->
             {error, Reason}
     end.
 
--spec startup(inet:socket()) ->
+-spec startup(shackle:socket()) ->
     {ok, binary() | undefined} | {error, atom()}.
 
 startup(Socket) ->
@@ -140,7 +140,7 @@ timeout(Timeout, Timestamp) ->
     Diff = timer:now_diff(os:timestamp(), Timestamp) div 1000,
     Timeout - Diff.
 
--spec use_keyspace(inet:socket()) ->
+-spec use_keyspace(shackle:socket()) ->
     ok | {error, atom()}.
 
 use_keyspace(Socket) ->
@@ -164,7 +164,7 @@ authenticate(Username, Password, Socket) when is_binary(Username),
     end.
 
 rcv_buf(Socket, Buffer) ->
-    case gen_tcp:recv(Socket, 0, ?DEFAULT_RECV_TIMEOUT) of
+    case sock_recv(Socket) of
         {ok, Msg} ->
             Buffer2 = <<Buffer/binary, Msg/binary>>,
             case marina_frame:decode(Buffer2) of
@@ -177,6 +177,16 @@ rcv_buf(Socket, Buffer) ->
             {error, Reason}
     end.
 
+sock_recv(Socket) when is_port(Socket) ->
+    gen_tcp:recv(Socket, 0, ?DEFAULT_RECV_TIMEOUT);
+sock_recv(Socket) ->
+    socket:recv(Socket, 0, ?DEFAULT_RECV_TIMEOUT).
+
+sock_send(Socket, Msg) when is_port(Socket) ->
+    gen_tcp:send(Socket, Msg);
+sock_send(Socket, Msg) ->
+    socket:send(Socket, Msg).
+
 use_keyspace(undefined, _Socket) ->
     ok;
 use_keyspace(Keyspace, Socket) when is_binary(Keyspace)->