Merge pull request #105 from kafka4beam/251024-stream-line-message-encoding

251024 streamline message encoding

Author: zmstone

changelog.md

@@ -3,6 +3,7 @@
   - Added producer process label `{wolff_producer, KafkaClientId, Topic, Partition}` [#102](https://github.com/kafka4beam/wolff/pull/102)
   - Upgrade to `kafka_protocol-4.3.0` for better CRC32C performance.
   - Add client config `allow_auto_topic_creation` (default = false).
+  - Streamline batch encoding.
 
 * 4.0.13 (merge 1.5.19)
   - Handle `record_list_too_large` error returned from Kafka.

src/wolff_producer.erl

@@ -127,6 +127,7 @@
                   , inflight_calls := non_neg_integer()
                   , topic := topic()
                   , calls := ?EMPTY | calls()
+                  , sender => pid()
                   }.
 
 %% @doc Start a per-partition producer worker.
@@ -231,7 +232,8 @@ send_sync(Pid, Batch0, Timeout) ->
 init(#{client_id := ClientId, topic := Topic, partition := Partition} = St) ->
   erlang:process_flag(trap_exit, true),
   ok = set_process_label(ClientId, Topic, Partition),
-  {ok, St, {continue, do_init}}.
+  {ok, Sender} = wolff_sender:start_link(self(), ClientId, Topic, Partition),
+  {ok, St#{sender => Sender}, {continue, do_init}}.
 
 do_init(#{client_id := ClientId,
           topic := Topic,
@@ -423,10 +425,12 @@ use_defaults(Config, [{K, V} | Rest]) ->
     false -> use_defaults(Config#{K => V}, Rest)
   end.
 
-resend_requests(#{q_items := Items} = Sent, St, ?reconnect) ->
+resend_requests(#{q_items := Items, req_ref := {alias, OldRef}} = Sent, St, ?reconnect) ->
+  erlang:unalias(OldRef),
   {ok, Ref} = send_request(Items, St),
   [Sent#{req_ref := Ref}];
-resend_requests(#{q_items := Items} = Sent, St, ?message_too_large) ->
+resend_requests(#{q_items := Items, req_ref := {alias, OldRef}} = Sent, St, ?message_too_large) ->
+  erlang:unalias(OldRef),
   Refs = lists:map(fun(Item) -> {ok, Ref} = send_request([Item], St), Ref end, Items),
   #{attempts := Attempts, q_ack_ref := Q_AckRef} = Sent,
   make_sent_items_list(Items, Refs, Attempts, Q_AckRef).
@@ -445,25 +449,25 @@ make_sent_items_list([Item | Items], [Ref | Refs], Attempts, Q_AckRef) ->
      attempts => Attempts
     } | make_sent_items_list(Items, Refs, Attempts, Q_AckRef)].
 
-send_request(QueueItems, #{conn := Conn} = St) ->
-  #kpro_req{ref = Ref} = Req = make_request(QueueItems, St),
-  ok = request_async(Conn, Req),
-  {ok, Ref}.
-
-make_request(QueueItems,
+send_request(QueueItems,
              #{config := #{required_acks := RequiredAcks,
                            ack_timeout := AckTimeout,
                            compression := Compression
                           },
                produce_api_vsn := Vsn,
                topic := Topic,
-               partition := Partition}) ->
+               partition := Partition,
+               conn := Conn,
+               sender := Sender
+              }) ->
   Batch = get_batch_from_queue_items(QueueItems),
-  kpro_req_lib:produce(Vsn, Topic, Partition, Batch,
-                       #{ack_timeout => AckTimeout,
-                         required_acks => RequiredAcks,
-                         compression => Compression
-                        }).
+  Opts = #{ack_timeout => AckTimeout,
+           required_acks => RequiredAcks,
+           compression => Compression
+           },
+  Ref = {alias, erlang:alias([reply])},
+  ok = wolff_sender:do(Sender, Conn, Ref, Vsn, Topic, Partition, Batch, Opts),
+  {ok, Ref}.
 
 resend_sent_reqs(#{sent_reqs := SentReqs,
                    config := Config
@@ -504,7 +508,6 @@ send_to_kafka(#{sent_reqs := SentReqs,
                 replayq := Q,
                 config := #{required_acks := RequiredAcks,
                             max_batch_bytes := BytesLimit} = Config,
-                conn := Conn,
                 pending_acks := PendingAcks
                } = St0) ->
   {NewQ, QAckRef, Items} =
@@ -518,7 +521,7 @@ send_to_kafka(#{sent_reqs := SentReqs,
   NewInflightCalls = InflightCalls + NrOfCalls,
   _ = wolff_metrics:inflight_set(Config, NewInflightCalls),
   NoAck = (RequiredAcks =:= none),
-  #kpro_req{ref = Ref} = Req = make_request(Items, St0),
+  {ok, Ref} = send_request(Items, St0),
   Sent = #{req_ref => Ref,
            q_items => Items,
            q_ack_ref => QAckRef,
@@ -530,7 +533,6 @@ send_to_kafka(#{sent_reqs := SentReqs,
              inflight_calls := NewInflightCalls,
              pending_acks := NewPendingAcks
             },
-  ok = request_async(Conn, Req),
   St2 = maybe_fake_kafka_ack(NoAck, St1),
   maybe_send_to_kafka(St2).
 
@@ -876,9 +878,6 @@ varint_bytes(N) -> vb(zz(N)).
 -compile({inline, zz/1}).
 zz(I) -> (I bsl 1) bxor (I bsr 63).
 
-request_async(Conn, Req) when is_pid(Conn) ->
-  ok = kpro:send(Conn, Req).
-
 %% collect send calls which are already sent to mailbox,
 %% the collection is size-limited by the max_linger_bytes config.
 collect_send_calls(Call, Bytes, ?EMPTY, Max) ->

src/wolff_sender.erl

@@ -0,0 +1,74 @@
+%% Copyright (c) 2025 EMQ Technologies Co., Ltd. All Rights Reserved.
+%%
+%% Licensed under the Apache License, Version 2.0 (the "License");
+%% you may not use this file except in compliance with the License.
+%% You may obtain a copy of the License at
+%%
+%%     http://www.apache.org/licenses/LICENSE-2.0
+%%
+%% Unless required by applicable law or agreed to in writing, software
+%% distributed under the License is distributed on an "AS IS" BASIS,
+%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+%% See the License for the specific language governing permissions and
+%% limitations under the License.
+
+%% This module encodes input batch into Kafka produce request and send
+%% the encoded batch to kpro_connection process.
+%% It offloads the heavy lifting work:
+%% - Batch encoding
+%% - CRC32C
+%% - Compression
+%% From wolff_producer process.
+-module(wolff_sender).
+
+-behaviour(gen_server).
+
+-export([start_link/4, do/8]).
+
+-include_lib("kafka_protocol/include/kpro.hrl").
+
+%% gen_server callbacks
+-export([code_change/3, handle_call/3, handle_cast/2, handle_info/2, init/1, terminate/2]).
+
+-spec do(pid(), pid(), {alias, reference()}, kpro:vsn(),
+         kpro:topic(), kpro:partition(), [kpro:msg_input()], kpro:produce_opts()) -> ok.
+do(SenderPid, ConnPid, Ref, Vsn, Topic, Partition, Batch, Opts) ->
+    gen_server:cast(SenderPid, {send, ConnPid, Ref, Vsn, Topic, Partition, Batch, Opts}).
+
+%% @doc Start the wolff_sender gen_server.
+-spec start_link(pid(), wolff:client_id(), kpro:topic(), kpro:partition()) -> {ok, pid()} | {error, any()}.
+start_link(Owner, ClientId, Topic, Partition) ->
+    gen_server:start_link(?MODULE, {Owner, ClientId, Topic, Partition}, []).
+
+%% gen_server callbacks
+
+init({Owner, ClientId, Topic, Partition}) ->
+    ok = set_process_label(ClientId, Topic, Partition),
+    {ok, #{owner => Owner, client_id => ClientId, topic => Topic, partition => Partition}}.
+
+handle_call(_Request, _From, State) ->
+    {reply, {error, unknown_request}, State}.
+
+handle_cast({send, ConnPid, Ref, Vsn, Topic, Partition, Batch, Opts}, State) ->
+    Req = kpro_req_lib:produce(Vsn, Topic, Partition, Batch, Opts, Ref),
+    ok = kpro:send(ConnPid, Req),
+    {noreply, State};
+handle_cast(_Msg, State) ->
+    {noreply, State}.
+
+handle_info(_Info, State) ->
+    {noreply, State}.
+
+terminate(_Reason, _State) ->
+    ok.
+
+code_change(_OldVsn, State, _Extra) ->
+    {ok, State}.
+
+-if(?OTP_RELEASE >= 27).
+set_process_label(ClientId, Topic, Partition) ->
+    proc_lib:set_label({?MODULE, ClientId, Topic, Partition}).
+-else.
+set_process_label(_ClientId, _Topic, _Partition) ->
+    ok.
+-endif.