%%% %%% Copyright (c) 2017-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. %%% %% @private The input source of brod-cli pipe command. %% This module implements a process that reads off the bytes %% from the data source (either stdin or a file) %% and sends the bytes to parent process. %% Messages sent to parent process: %% %% ```{pipe, self(), [{Ts :: integer(), Key :: binary(), Val :: binary()}]}''' %% @end -module(brod_cli_pipe). -ifdef(build_brod_cli). -behaviour(gen_server). -export([ start_link/1 , stop/1 ]). -export([ code_change/3 , handle_call/3 , handle_cast/2 , handle_info/2 , init/1 , terminate/2 ]). -include("brod_int.hrl"). -type arg_name() :: source | kv_deli | msg_deli | prompt | tail | no_exit | blk_size | retry_delay. -type arg_value() :: term(). -define(LINE_BREAK, <<"\n">>). -define(STDIN, standard_io). -define(EOF_RETRY_DELAY_MS, 100). -define(NOT_APPLICABLE, 'N/A'). -define(CONTINUE_MSG, continue). -define(PARENT_BUSY_MSG_QUEUE_LEN_THRESHOLD, 100). -type delimiter() :: binary(). -type epoch_ms() :: integer(). -type read_fun() :: fun((?STDIN | file:io_device(), [binary()]) -> eof | {[{epoch_ms(), Key :: binary(), Val :: binary()}], [binary()]}). -record(state, { parent :: pid() , source :: ?STDIN | {file, string()} , read_fun :: read_fun() , is_eof_exit :: boolean() , is_tail :: boolean() , io_device :: ?undef | ?STDIN | file:io_device() , acc_bytes = [] :: [binary()] , retry_delay :: timeout() }). %% @doc Args explained: %% %% source: 'standard_io' | {file, "path/to/source"} %% kv_deli: 'none' | binary(). %% Delimiter bytes for message key and value %% msg_deli: binary(). Delimiter between kafka messages %% NOTE: eof is always considered a message delimiter %% prompt: boolean(). Applicable when source is standard_io AND %% when kv_deli and msg_deli are both '\n' %% prompts 'key> ' for key input and 'val> ' for value input %% tail: boolean(). Applicable when source is a file %% tell brod-cli to start reading from EOF %% no_exit: boolean(). Do not exit when reaching EOF %% blk_size: Read block size -spec start_link([{arg_name(), arg_value()}]) -> {ok, pid()}. start_link(Args) -> Parent = self(), Arg = fun(Name) -> {_, V} = lists:keyfind(Name, 1, Args), V end, KvDeli = Arg(kv_deli), MsgDeli = Arg(msg_deli), Source = Arg(source), IsLineMode = MsgDeli =:= ?LINE_BREAK, BlkSize = Arg(blk_size), IsPrompt = Arg(prompt), IsTail = Arg(tail), IsNoExit = Arg(no_exit), IsEofExit = not (IsTail orelse IsNoExit), ReadFun = case IsLineMode of true when IsPrompt andalso Source =:= ?STDIN -> make_prompt_line_reader(KvDeli); true -> make_line_reader(KvDeli, _PromptStr = ""); false -> make_stream_reader(KvDeli, MsgDeli, BlkSize, IsEofExit) end, State = #state{ parent = Parent , source = Source , read_fun = ReadFun , is_tail = IsTail , is_eof_exit = IsEofExit , retry_delay = Arg(retry_delay) }, gen_server:start_link({local, ?MODULE}, ?MODULE, State, []). %% @doc Stop gen_server. stop(Pid) -> gen_server:cast(Pid, stop). %% @doc Tell reader to continue. continue() -> self() ! ?CONTINUE_MSG. %% @hidden init(#state{source = Source, is_tail = IsTail} = State0) -> IoDevice = case Source of ?STDIN -> ?STDIN; {file, File} -> {ok, Fd} = file:open(File, [read, binary]), IsTail andalso file:position(Fd, eof), Fd end, State = State0#state{io_device = IoDevice}, _ = continue(), {ok, State}. %% @hidden handle_info(?CONTINUE_MSG, #state{parent = Parent} = State) -> case erlang:process_info(Parent, message_queue_len) of {_, Len} when Len >= ?PARENT_BUSY_MSG_QUEUE_LEN_THRESHOLD -> ok = delay_continue(State), {noreply, State}; _ -> handle_read(State) end; handle_info(_Info, State) -> {noreply, State}. %% @hidden handle_cast(stop, State) -> {stop, normal, State}; handle_cast(_Cast, State) -> {noreply, State}. %% @hidden handle_call(Call, _From, State) -> {reply, {error, {unknown_call, Call}}, State}. %% @hidden code_change(_OldVsn, State, _Extra) -> {ok, State}. %% @hidden terminate(_Reason, _State) -> ok. %%%_* Privates ================================================================= send_to_parent(Parent, Msgs0) -> FilterF = fun(?TKV(_T, K, V)) -> K =/= <<>> orelse V =/= <<>> end, Msgs = lists:filter(FilterF, Msgs0), Msgs =/= [] andalso erlang:send(Parent, {pipe, self(), Msgs}), ok. handle_read(#state{ read_fun = ReadFun , acc_bytes = Acc0 , io_device = IoDevice , parent = Parent } = State0) -> case ReadFun(IoDevice, Acc0) of eof -> handle_eof(State0); {Msgs, Acc} -> _ = continue(), %% continue next trunk ok = send_to_parent(Parent, Msgs), State = State0#state{acc_bytes = Acc}, {noreply, State} end. handle_eof(#state{io_device = ?STDIN} = State) -> %% standard_io pipe closed {stop, normal, State}; handle_eof(#state{is_eof_exit = true} = State) -> {stop, normal, State}; handle_eof(#state{io_device = Fd} = State) -> %% Get current position {ok, LastPos} = file:position(Fd, {cur, 0}), %% Try set position to EOF, %% see if it is the current position case file:position(Fd, eof) of {ok, NewPos} when NewPos < LastPos -> %% File has been truncated. %% Don't know what to do because %% we can not assume the file is truncated to empty {stop, pipe_source_truncated, State}; {ok, _Pos} -> _ = file:position(Fd, LastPos), ok = delay_continue(State), {noreply, State} end. delay_continue(#state{retry_delay = Delay}) -> _ = erlang:send_after(Delay, self(), ?CONTINUE_MSG), ok. -spec make_prompt_line_reader(none | delimiter()) -> read_fun(). make_prompt_line_reader(_KvDeli = none) -> %% Read only value, no key fun(?STDIN, _Acc) -> case read_line(?STDIN, "VAL> ") of eof -> eof; Value -> {[{_Key = <<>>, Value}], []} end end; make_prompt_line_reader(_KvDeli = ?LINE_BREAK) -> fun(?STDIN, _Acc) -> case read_line(?STDIN, "KEY> ") of eof -> eof; Key -> case read_line(?STDIN, "VAL> ") of eof -> {[make_msg(Key, <<>>)], []}; Value -> {[make_msg(Key, Value)], []} end end end; make_prompt_line_reader(KvDeli) -> Prompt = "KEY" ++ binary_to_list(KvDeli) ++ "VAL> ", make_line_reader(KvDeli, Prompt). -spec make_line_reader(none | binary(), string()) -> read_fun(). make_line_reader(KvDeli, Prompt) -> fun(IoDevice, _Acc) -> case read_line(IoDevice, Prompt) of eof -> eof; Key when KvDeli =:= <<"\n">> -> case read_line(IoDevice, Prompt) of eof -> eof; Val -> {[make_msg(Key, Val)], []} end; Val when KvDeli =:= none -> {[make_msg(<<>>, Val)], []}; Line -> [Key, Value] = binary:split(Line, bin(KvDeli)), {[make_msg(Key, Value)], []} end end. -spec make_stream_reader(none | delimiter(), delimiter(), pos_integer(), boolean()) -> read_fun(). make_stream_reader(KvDeli, MsgDeli, BlkSize, IsEofExit) -> IsSameDeli = MsgDeli =:= KvDeli, KvDeliCp = case is_binary(KvDeli) of true -> binary:compile_pattern(KvDeli); false -> none end, MsgDeliCp = binary:compile_pattern(MsgDeli), fun(IoDevice, Acc) -> case file:read(IoDevice, BlkSize) of eof -> case IsEofExit of true when Acc =:= [] -> %% Reached EOF eof; true -> %% Configured to exit when reaching EOF %% try split kv-pairs NOW LastMsg = bin(lists:reverse(Acc)), KvPairs = split_kv_pairs([LastMsg], KvDeliCp, IsSameDeli), {KvPairs, []}; false -> %% Keep looping for the next message delimiter {[], Acc} end; {ok, Bytes} -> Acc1 = add_acc(size(MsgDeli), Bytes, Acc), {Messages, NewAcc} = split_messages(MsgDeliCp, Acc1), KvPairs = split_kv_pairs(Messages, KvDeliCp, IsSameDeli), {KvPairs, NewAcc} end end. -spec add_acc(pos_integer(), binary(), [binary()]) -> [binary()]. add_acc(_DeliSize = 1, Bytes, Acc) -> %% Delimiter is only one byte, in no way could it be cut in half [Bytes | Acc]; add_acc(_DeliSize, Bytes, []) -> [Bytes]; add_acc(DeliSize, Bytes, [Tail | Header]) -> Size = size(Tail) - DeliSize, case Size =< 0 of true -> [<> | Header]; false -> %% cut a DeliSize tail from acc and prepend as current head %% to make sure we will not cut delimiter into two chunks <> = Tail, %% cut NewTail = <>, %% new tail [NewTail, TailH | Header] end. -spec split_messages(binary:cp(), [binary()]) -> {[binary()], [binary()]}. split_messages(MsgDeliCp, [Tail | Header]) -> case binary:split(Tail, MsgDeliCp, [global]) of [_] -> %% no delimiter found {[], [Tail | Header]}; [First0 | More] -> First = bin([lists:reverse(Header), First0]), case lists:reverse(More) of [<<>> | Msgs] -> {[First | lists:reverse(Msgs)], []}; [NewTail | Msgs] -> {[First | lists:reverse(Msgs)], [NewTail]} end end. -spec split_kv_pairs([binary()], none | delimiter(), boolean()) -> brod:value(). split_kv_pairs(Msgs, none, _IsSameDeli) -> lists:map(fun(Msg) -> make_msg(<<>>, Msg) end, Msgs); split_kv_pairs(Msgs, _KvDeliCp, _IsSameDeli = true) -> make_msgs(Msgs); split_kv_pairs(Msgs, KvDeliCp, _IsSameDeli = false) -> lists:map(fun(Msg) -> [K, V] = binary:split(Msg, KvDeliCp), make_msg(K, V) end, Msgs). make_msgs([]) -> []; make_msgs([K, V | Rest]) -> [make_msg(K, V) | make_msgs(Rest)]. make_msg(K, V) -> CreateTs = brod_utils:epoch_ms(), ?TKV(CreateTs, K, V). -spec read_line(?STDIN | file:io_device(), string()) -> eof | binary(). read_line(IoDevice, Prompt) -> case io:get_line(IoDevice, Prompt) of eof -> eof; Line -> Chars = unicode:characters_to_list(Line), unicode:characters_to_binary(rstrip(Chars, "\n")) end. -spec rstrip(string(), string()) -> string(). rstrip(Str, CharSet) -> lists:reverse(lstrip(lists:reverse(Str), CharSet)). -spec lstrip(string(), string()) -> string(). lstrip([], _) -> []; lstrip([C | Rest] = Str, CharSet) -> case lists:member(C, CharSet) of true -> lstrip(Rest, CharSet); false -> Str end. -spec bin(iodata()) -> binary(). bin(X) -> iolist_to_binary(X). -endif. %%%_* Emacs ==================================================================== %%% Local Variables: %%% allout-layout: t %%% erlang-indent-level: 2 %%% End: