Current section

Files

Jump to
kafka_protocol src kpro_compress.erl
Raw

src/kpro_compress.erl

%%% Copyright (c) 2018-2021, Klarna Bank AB (publ)
%%%
%%% 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.
%%%
-module(kpro_compress).
-export([ compress/2
, decompress/2
, codec_to_method/1
, method_to_codec/1
, provide/1
]).
-include("kpro_private.hrl").
%% @doc Set snappy or lz4 compression modules.
%% This should override the default usage of `snappyer' and `lz4b_frame'.
-spec provide([{snappy | lz4, module()}]) -> ok.
provide(Libs) ->
lists:foreach(fun({Name, Module}) ->
persistent_term:put({?MODULE, Name}, Module)
end, Libs).
%% @doc Translate codec in kafka batch attributes to compression method.
-spec codec_to_method(byte()) -> kpro:compress_option().
codec_to_method(A) when ?KPRO_IS_GZIP_ATTR(A) -> ?gzip;
codec_to_method(A) when ?KPRO_IS_SNAPPY_ATTR(A) -> ?snappy;
codec_to_method(A) when ?KPRO_IS_LZ4_ATTR(A) -> ?lz4;
codec_to_method(_) -> ?no_compression.
%% @doc Translate compression method to bits for kafka batch attributes.
method_to_codec(?gzip) -> ?KPRO_COMPRESS_GZIP;
method_to_codec(?snappy) -> ?KPRO_COMPRESS_SNAPPY;
method_to_codec(?lz4) -> ?KPRO_COMPRESS_LZ4;
method_to_codec(?no_compression) -> ?KPRO_COMPRESS_NONE.
%% @doc Compress encoded batch.
-spec compress(kpro:compress_option(), iodata()) -> iodata().
compress(?no_compression, IoData) -> IoData;
compress(?gzip, IoData) -> zlib:gzip(IoData);
compress(Name, IoData) -> do_compress(Name, IoData).
%% @doc Decompress batch.
-spec decompress(kpro:compress_option(), binary()) -> binary().
decompress(?no_compression, Bin) -> Bin;
decompress(?gzip, Bin) -> zlib:gunzip(Bin);
decompress(?snappy, Bin) -> java_snappy_unpack(Bin);
decompress(?lz4, Bin) -> do_decompress(?lz4, Bin).
%%%_* Internals ================================================================
%% Java snappy implementation has its own non-standard
%% magic header, see org/xerial/snappy/SnappyCodec.java
java_snappy_unpack(<<130, "SNAPPY", 0,
_Version:32, _MinCompatibleV:32, Chunks/binary>>) ->
java_snappy_unpack_chunks(Chunks, []);
java_snappy_unpack(Bin) ->
do_decompress(snappy, Bin).
java_snappy_unpack_chunks(<<>>, Acc) ->
iolist_to_binary(Acc);
java_snappy_unpack_chunks(Chunks, Acc) ->
<<Len:32/unsigned-integer, Rest/binary>> = Chunks,
case Len =:= 0 of
true ->
Rest =:= <<>> orelse erlang:error({Len, Rest}), %% assert
Acc;
false ->
<<Data:Len/binary, Tail/binary>> = Rest,
Decompressed = do_decompress(?snappy, Data),
java_snappy_unpack_chunks(Tail, [Acc, Decompressed])
end.
do_compress(Name, IoData) ->
Module = get_module(Name),
iodata(Module:compress(IoData)).
do_decompress(Name, Bin) ->
Module = get_module(Name),
iodata(Module:decompress(Bin)).
get_module(snappy) ->
get_module(snappy, snappyer);
get_module(lz4) ->
get_module(lz4, lz4b_frame).
get_module(Name, Default) ->
persistent_term:get({?MODULE, Name}, Default).
iodata({ok, IoData}) -> IoData;
iodata({error, Reason}) -> error(Reason);
iodata(IoData) when is_list(IoData) orelse is_binary(IoData) -> IoData.
%%%_* Emacs ====================================================================
%%% Local Variables:
%%% allout-layout: t
%%% erlang-indent-level: 2
%%% End: