WIP: proof-of-concept zstd compression for replication

This requires OTP-28 to build.

Author: the-mikedavis

src/osiris.hrl

@@ -47,6 +47,9 @@
 -define(FILE_OPTS_WRITE, [raw, binary, write, read]).
 -define(INDEX_RECORD_SIZE_B, 29).
 
+%% signature of chunks compressed
+-define(COMPRESSED_MAGIC, 7).
+-define(COMPRESSED_VERSION, 0).
 
 %% chunk types
 -define(CHNK_USER, 0).

src/osiris_log.erl

@@ -420,7 +420,8 @@
          %% that will be included in snapshots written to new segments
          readers_counter_fun = fun(_) -> ok end :: function(),
          shared :: atomics:atomics_ref(),
-         filter_size = ?DEFAULT_FILTER_SIZE :: osiris_bloom:filter_size()
+         filter_size = ?DEFAULT_FILTER_SIZE :: osiris_bloom:filter_size(),
+         compression :: undefined | {zstd, zstd:context()}
          }).
 -record(ra,
         {size = ?HEADER_SIZE_B + ?DEFAULT_FILTER_SIZE :: non_neg_integer(),
@@ -1064,14 +1065,22 @@ init_data_reader_at(ChunkId, FilePos, File,
             Cnt = make_counter(Config),
             counters:put(Cnt, ?C_OFFSET, ChunkId - 1),
             CountersFun(1),
+            Compression = case Config of
+                              #{compression := zstd} ->
+                                 {ok, Ctx} = zstd:context(compress),
+                                 {zstd, Ctx};
+                              _ ->
+                                  undefined
+                          end,
             {ok,
              #?MODULE{cfg =
                       #cfg{directory = Dir,
                            counter = Cnt,
                            counter_id = counter_id(Config),
                            name = Name,
                            readers_counter_fun = CountersFun,
-                           shared = Shared
+                           shared = Shared,
+                           compression = Compression
                           },
                       mode =
                       #read{type = data,
@@ -1716,7 +1725,8 @@ send_file(Sock, State) ->
     {error, term()} |
     {end_of_stream, state()}.
 send_file(Sock,
-          #?MODULE{mode = #read{type = RType,
+          #?MODULE{cfg = #cfg{compression = Compression},
+                   mode = #read{type = RType,
                                 chunk_selector = Selector,
                                 transport = Transport,
                                 filter_size = RaFs}} = State0,
@@ -1755,19 +1765,32 @@ send_file(Sock,
                     FrameHeader = Callback(Header, ToSend + byte_size(HeaderData)),
                     case ra_read(DataPos, ToSend, Ra) of
                         Data when byte_size(Data) == ToSend ->
-                            case send(Transport, Sock, [FrameHeader, HeaderData, Data]) of
+                            case send(Transport, Compression, Sock, [FrameHeader, HeaderData, Data]) of
                                 ok ->
                                     State = State1#?MODULE{mode = Read},
                                     {ok, State};
                                 Err ->
                                     Err
                             end;
+                        _ when Compression =/= undefined ->
+                            case file:pread(Fd, DataPos, ToSend) of
+                                {ok, Data} ->
+                                    case send(Transport, Compression, Sock, [FrameHeader, HeaderData, Data]) of
+                                        ok ->
+                                            State = State1#?MODULE{mode = Read},
+                                            {ok, State};
+                                        Err ->
+                                            Err
+                                    end;
+                                Err ->
+                                    Err
+                            end;
                         _ ->
                             %% the read ahead buffer could not satisfy the
                             %% data read
-                            case send(Transport, Sock, [FrameHeader, HeaderData]) of
+                            case send(Transport, Compression, Sock, [FrameHeader, HeaderData]) of
                                 ok ->
-                                    case sendfile(Transport, Fd, Sock,
+                                    case sendfile(Transport, Compression, Fd, Sock,
                                                   DataPos, ToSend) of
                                         ok ->
                                             State = State1#?MODULE{mode = Read},
@@ -2526,30 +2549,40 @@ max_segment_size_reached(
     CurrentSizeBytes >= MaxSizeBytes orelse
     CurrentSizeChunks >= MaxSizeChunks.
 
-sendfile(_Transport, _Fd, _Sock, _Pos, 0) ->
+sendfile(_Transport, _Compression, _Fd, _Sock, _Pos, 0) ->
     ok;
-sendfile(tcp = Transport, Fd, Sock, Pos, ToSend) ->
+sendfile(tcp = Transport, undefined, Fd, Sock, Pos, ToSend) ->
     case file:sendfile(Fd, Sock, Pos, ToSend, []) of
         {ok, 0} ->
             %% TODO add counter for this?
-            sendfile(Transport, Fd, Sock, Pos, ToSend);
+            sendfile(Transport, undefined, Fd, Sock, Pos, ToSend);
         {ok, BytesSent} ->
-            sendfile(Transport, Fd, Sock, Pos + BytesSent, ToSend - BytesSent);
+            sendfile(Transport, undefined, Fd, Sock, Pos + BytesSent, ToSend - BytesSent);
         {error, _} = Err ->
             Err
     end;
-sendfile(ssl, Fd, Sock, Pos, ToSend) ->
+sendfile(Transport, Compression, Fd, Sock, Pos, ToSend) ->
     case file:pread(Fd, Pos, ToSend) of
         {ok, Data} ->
-            ssl:send(Sock, Data);
+            send(Transport, Compression, Sock, Data);
         {error, _} = Err ->
             Err
     end.
 
-send(tcp, Sock, Data) ->
-    gen_tcp:send(Sock, Data);
-send(ssl, Sock, Data) ->
-    ssl:send(Sock, Data).
+send(tcp, Compression, Sock, Data) ->
+    gen_tcp:send(Sock, compress_data(Compression, Data));
+send(ssl, Compression, Sock, Data) ->
+    ssl:send(Sock, compress_data(Compression, Data)).
+
+compress_data({zstd, Ctx}, Data) ->
+    CompressedData = zstd:compress(Data, Ctx),
+    Header = <<?COMPRESSED_MAGIC:4/unsigned,
+               ?COMPRESSED_VERSION:4/unsigned,
+               3:8/unsigned,
+               (iolist_size(CompressedData)):32/unsigned>>,
+    [Header | CompressedData];
+compress_data(undefined, Data) ->
+    Data.
 
 last_timestamp_in_index_file(IdxFile) ->
     case file:open(IdxFile, [raw, binary, read]) of

src/osiris_replica.erl

@@ -679,6 +679,35 @@ parse_chunk(<<?MAGIC:4/unsigned,
     TotalSize = ?HEADER_SIZE_B + FSize + Size + TSize,
     <<Chunk:TotalSize/binary, _/binary>> = All,
     parse_chunk(Rem, undefined, [{{FirstOffset, Timestamp}, Chunk} | Acc]);
+parse_chunk(<<?COMPRESSED_MAGIC:4/unsigned,
+              ?COMPRESSED_VERSION:4/unsigned,
+              CompressionType:8/unsigned,
+              CompressedLen:32/unsigned,
+              CompressedData:CompressedLen/binary,
+              Rem/binary>>,
+            undefined, Acc) ->
+    Chunk =  case CompressionType of
+                 3 -> %% zstd
+                     zstd_decompress(CompressedData);
+                 1 ->
+                     zlib:gunzip(CompressedData)
+             end,
+    <<?MAGIC:4/unsigned,
+      ?VERSION:4/unsigned,
+      _ChType:8/unsigned,
+      _NumEntries:16/unsigned,
+      _NumRecords:32/unsigned,
+      Timestamp:64/signed,
+      _Epoch:64/unsigned,
+      FirstOffset:64/unsigned,
+      _Rest/binary>> = Chunk,
+    parse_chunk(Rem, undefined, [{{FirstOffset, Timestamp}, Chunk} | Acc]);
+parse_chunk(<<?COMPRESSED_MAGIC:4/unsigned,
+              ?COMPRESSED_VERSION:4/unsigned,
+              _Partial/binary>> = All,
+            undefined, Acc) ->
+    %% This could be more efficient: decompression functions can take iodata.
+    {All, lists:reverse(Acc)};
 parse_chunk(Bin, undefined, Acc)
     when byte_size(Bin) =< ?HEADER_SIZE_B ->
     {Bin, lists:reverse(Acc)};
@@ -714,6 +743,14 @@ parse_chunk(Bin, {FirstOffsetTs, IOData, RemSize}, Acc) ->
     {{FirstOffsetTs, [Bin | IOData], RemSize - byte_size(Bin)},
      lists:reverse(Acc)}.
 
+-if(?OTP_RELEASE >= 28).
+zstd_decompress(Data) ->
+    iolist_to_binary(zstd:decompress(Data)).
+-else.
+zstd_decompress(_Data) ->
+    erlang:error(zstd_not_supported).
+-endif.
+
 notify_offset_listeners(#?MODULE{cfg = #cfg{reference = Ref,
                                             event_formatter = EvtFmt},
                                  committed_chunk_id = CommittedChId,

src/osiris_replica_reader.erl

@@ -123,7 +123,7 @@ init(#{hosts := Hosts,
                    [Host, Port]),
             CntId = {?MODULE, ExtRef, Host, Port},
             CntSpec = {CntId, ?COUNTER_FIELDS},
-            Config = #{counter_spec => CntSpec, transport => Transport},
+            Config = #{counter_spec => CntSpec, transport => Transport, compression => zstd},
             %% send token to replica to complete connection setup
             ok = send(Transport, Sock, Token),
             Ret = osiris_writer:init_data_reader(LeaderPid, TailInfo, Config),