-module(grisp_updater_loader). -behavior(gen_server). -behaviour(grisp_updater_sink). %--- Includes ------------------------------------------------------------------ -include_lib("kernel/include/logger.hrl"). -include("grisp_updater.hrl"). %--- Exports ------------------------------------------------------------------- % API -export([start_link/1]). -export([schedule_load/3]). -export([cancel_load/1]). -export([abort/0]). % Behaviour grisp_updater_sink Callbacks -export([sink_error/3]). -export([sink_data/3]). -export([sink_done/3]). % Behaviour gen_server Callbacks -export([init/1]). -export([handle_call/3]). -export([handle_cast/2]). -export([handle_info/2]). %--- Records ------------------------------------------------------------------- -record(pending, { url :: binary(), block :: #block{}, target :: #target{}, stream :: undefined | reference() }). -record(stream, { id :: non_neg_integer(), ref :: reference(), block_left :: undefined | non_neg_integer(), data_left :: non_neg_integer(), data_offset :: non_neg_integer(), block_hash :: undefined | term(), data_hash :: term(), inflater :: term() }). -record(state, { concurrency :: pos_integer(), pending = #{} :: #{non_neg_integer() => #pending{}}, streams = #{} :: #{reference() => #stream{}}, schedule = queue:new() :: queue:queue() }). %--- Macros -------------------------------------------------------------------- -define(DEFAULT_CONCURRENCY, 1). %--- API Functions ------------------------------------------------------------- start_link(Opts) -> gen_server:start_link({local, ?MODULE}, ?MODULE, Opts, []). schedule_load(Url, Block, Target) -> gen_server:cast(?MODULE, {schedule, Url, Block, Target}). cancel_load(BlockId) -> gen_server:call(?MODULE, {cancel, BlockId}). abort() -> gen_server:call(?MODULE, abort). %--- Behaviour grisp_updater_sink Callbacks ------------------------------------ sink_error(StreamRef, Params, Reason) -> gen_server:cast(?MODULE, {sink_error, StreamRef, Params, Reason}). sink_data(StreamRef, Params, Data) -> gen_server:cast(?MODULE, {sink_data, StreamRef, Params, Data}). sink_done(StreamRef, Params, Data) -> gen_server:cast(?MODULE, {sink_done, StreamRef, Params, Data}). %--- Behaviour gen_server Callbacks -------------------------------------------- init(Opts) -> ?LOG_INFO("Starting GRiSP updater's block loader ..."), {ok, #state{ concurrency = maps:get(concurrency, Opts, ?DEFAULT_CONCURRENCY) }}. handle_call(abort, _From, State) -> {reply, ok, do_abort(State)}; handle_call({cancel, BlockId}, _From, State) -> {reply, ok, do_cancel(State, BlockId)}; handle_call(Request, From, State) -> ?LOG_WARNING("Unexpected call from ~p: ~p", [From, Request]), {reply, {error, unexpected_call}, State}. handle_cast({schedule, Url, Block, Target}, State) -> {noreply, do_schedule(State, Url, Block, Target)}; handle_cast({sink_error, StreamRef, Params, Reason}, State) -> {noreply, got_sink_error(State, StreamRef, Params, Reason)}; handle_cast({sink_data, StreamRef, Params, Data}, State) -> {noreply, got_sink_data(State, StreamRef, Params, Data)}; handle_cast({sink_done, StreamRef, Params, Data}, State) -> {noreply, got_sink_done(State, StreamRef, Params, Data)}; handle_cast(Request, State) -> ?LOG_WARNING("Unexpected cast: ~p", [Request]), {noreply, State}. handle_info(Info, State) -> ?LOG_WARNING("Unexpected messagge: ~p", [Info]), {noreply, State}. %--- Internal ------------------------------------------------------------------ block_path(#block{encoding = #raw_encoding{block_path = Path}}) -> Path; block_path(#block{encoding = #gzip_encoding{block_path = Path}}) -> Path. do_abort(State) -> cancel_pendings(State#state{schedule = queue:new()}). do_cancel(State, BlockId) -> cancel_pending(State, BlockId). do_schedule(#state{pending = PendMap, schedule = Sched} = State, Url, #block{id = Id} = Block, Target) -> ?LOG_DEBUG("Scheduling block ~b for streaming from ~s/~s", [Id, Url, block_path(Block)]), Pending = #pending{ url = Url, block = Block, target = Target, stream = undefined }, error = maps:find(Id, PendMap), PendMap2 = PendMap#{Id => Pending}, Sched2 = queue:in(Id, Sched), start_streams(State#state{pending = PendMap2, schedule = Sched2}). start_streams(#state{concurrency = Concurrency, streams = StreamMap, schedule = Sched} = State) -> case maps:size(StreamMap) < Concurrency of false -> State; true -> case queue:out(Sched) of {empty, Sched2} -> State#state{schedule = Sched2}; {{value, Id}, Sched2} -> State2 = State#state{schedule = Sched2}, start_streams(start_stream(State2, Id)) end end. block_ref(BaseUrl, Block) -> Path = block_path(Block), case uri_string:parse(Path) of #{scheme := _} -> {Path, <<>>}; #{} -> {BaseUrl, Path} end. start_stream(#state{pending = PendMap, streams = StreamMap} = State, Id) -> #{Id := #pending{url = Url, block = Block} = Pending} = PendMap, {BlockUrl, BlockPath} = block_ref(Url, Block), ?LOG_DEBUG("Start streaming block ~b from ~s/~s", [Id, BlockUrl, BlockPath]), case grisp_updater_source:stream(BlockUrl, BlockPath, ?MODULE, Id) of {error, Reason} -> grisp_updater_manager:loader_failed(Id, Reason), State#state{pending = maps:remove(Id, PendMap)}; {ok, StreamRef} -> Pending2 = Pending#pending{stream = StreamRef}, Stream = init_stream(Block, StreamRef), StreamMap2 = StreamMap#{StreamRef => Stream}, PendMap2 = PendMap#{Id => Pending2}, State#state{pending = PendMap2, streams = StreamMap2} end. got_sink_error(State, _StreamRef, Id, {http_error, Code} = Reason) when 400 =< Code, Code < 500 -> % all 4xx http error codes (client errors) % imply that the request should not be repeated grisp_updater_manager:loader_error(Id, Reason), State; got_sink_error(#state{pending = PendMap, streams = StreamMap} = State, StreamRef, Id, Reason) -> %TODO: Figure out which errors are fatal and which are recoverable case maps:find(StreamRef, StreamMap) of error -> State; {ok, #stream{id = Id}} -> #{Id := #pending{stream = StreamRef}} = PendMap, grisp_updater_manager:loader_failed(Id, Reason), start_streams(cancel_pending(State, Id)) end. got_sink_data(#state{pending = PendMap, streams = StreamMap} = State, StreamRef, Id, Data) -> case maps:find(StreamRef, StreamMap) of error -> State; {ok, #stream{id = Id}} -> #{Id := #pending{stream = StreamRef}} = PendMap, case got_data(State, StreamRef, Id, Data) of {error, Reason} -> %TODO: Figure out which errors are fatal and which are recoverable grisp_updater_manager:loader_failed(Id, Reason), start_streams(cancel_pending(State, Id)); {ok, State2} -> State2 end end. got_sink_done(#state{pending = PendMap, streams = StreamMap} = State, StreamRef, Id, Data) -> case maps:find(StreamRef, StreamMap) of error -> State; {ok, #stream{id = Id}} -> #{Id := #pending{stream = StreamRef}} = PendMap, case got_data(State, StreamRef, Id, Data) of {error, Reason} -> %TODO: Figure out which errors are fatal and which are recoverable grisp_updater_manager:loader_failed(Id, Reason), start_streams(cancel_pending(State, Id)); {ok, State2} -> start_streams(stream_terminated(State2, StreamRef)) end end. got_data(#state{pending = PendMap, streams = StreamMap} = State, StreamRef, Id, BlockData) -> #{Id := #pending{stream = StreamRef} = Pending} = PendMap, #{StreamRef := #stream{id = Id} = Stream} = StreamMap, #pending{ block = #block{id = Id, encoding = Encoding}, stream = StreamRef } = Pending, case got_data(State, Pending, Stream, Encoding, BlockData) of {ok, Stream2, State2} -> StreamMap2 = StreamMap#{StreamRef := Stream2}, {ok, State2#state{streams = StreamMap2}}; {error, _Reason} = Error -> Error end. got_data(State, Pending, Stream, #raw_encoding{}, BlockData) -> #pending{ block = #block{data_offset = DataOffset}, target = #target{device = Device, offset = DeviceOffset} } = Pending, #stream{ data_left = DataLeft, data_offset = WriteOffset, data_hash = DataHash } = Stream, ChunkSize = byte_size(BlockData), case ChunkSize =< DataLeft of false -> {error, block_too_large}; true -> DataHash2 = crypto:hash_update(DataHash, BlockData), Seek = DeviceOffset + DataOffset + WriteOffset, case grisp_updater_storage:write(Device, Seek, BlockData) of {error, Reason} -> {error, {write_error, Reason}}; ok -> Stream2 = Stream#stream{ data_left = DataLeft - ChunkSize, data_offset = WriteOffset + ChunkSize, data_hash = DataHash2 }, {ok, Stream2, State} end end; got_data(State, Pending, Stream, #gzip_encoding{}, BlockData) -> #pending{ block = #block{data_offset = DataOffset}, target = #target{device = Device, offset = DeviceOffset} } = Pending, #stream{ block_left = BlockLeft, data_left = DataLeft, data_offset = WriteOffset, data_hash = DataHash, block_hash = BlockHash, inflater = Z } = Stream, ChunkSize = byte_size(BlockData), case ChunkSize =< BlockLeft of false -> {error, encoded_block_too_large}; true -> BlockHash2 = crypto:hash_update(BlockHash, BlockData), try zlib:inflate(Z, BlockData) of Data -> InflatedChuckSize = iolist_size(Data), case InflatedChuckSize =< DataLeft of false -> {error, decoded_block_too_large}; true -> DataHash2 = crypto:hash_update(DataHash, Data), Seek = DeviceOffset + DataOffset + WriteOffset, case grisp_updater_storage:write(Device, Seek, Data) of {error, Reason} -> {error, {write_error, Reason}}; ok -> Stream2 = Stream#stream{ block_left = BlockLeft - ChunkSize, data_left = DataLeft - InflatedChuckSize, data_offset = WriteOffset + InflatedChuckSize, data_hash = DataHash2, block_hash = BlockHash2 }, {ok, Stream2, State} end end catch error:Reason -> {error, {deflate_error, Reason}} end end. stream_terminated(#state{pending = PendMap, streams = StreamMap} = State, StreamRef) -> #{StreamRef := #stream{id = Id} = Stream} = StreamMap, #{Id := #pending{stream = StreamRef, block = Block}} = PendMap, #block{id = Id, encoding = Encoding} = Block, State2 = stream_terminated(State, Block, Stream, Encoding), cleanup_stream(Stream), State2#state{pending = maps:remove(Id, PendMap), streams = maps:remove(StreamRef, StreamMap)}. stream_terminated(State, Block, Stream, #raw_encoding{}) -> #block{id = Id, data_hash_data = ExpDataHash} = Block, #stream{data_left = DataLeft, data_hash = DataHash} = Stream, GotDataHash = crypto:hash_final(DataHash), case {DataLeft, GotDataHash} of {0, ExpDataHash} -> grisp_updater_manager:loader_done(Id); {V, _} when V =:= 0 -> grisp_updater_manager:loader_failed(Id, data_size_mismatch); {_, H} when H =/= ExpDataHash -> grisp_updater_manager:loader_failed(Id, data_hash_mismatch) end, State; stream_terminated(State, Block, Stream, #gzip_encoding{}) -> #block{ id = Id, data_hash_data = ExpDataHash, encoding = #gzip_encoding{ block_hash_data = ExpBlockHash } } = Block, #stream{ block_left = BlockLeft, data_left = DataLeft, data_hash = DataHash, block_hash = BlockHash } = Stream, GotDataHash = crypto:hash_final(DataHash), GotBlockHash = crypto:hash_final(BlockHash), case {BlockLeft, DataLeft, GotBlockHash, GotDataHash} of {0, 0, ExpBlockHash, ExpDataHash} -> grisp_updater_manager:loader_done(Id); {V, _, _, _} when V =:= 0 -> grisp_updater_manager:loader_failed(Id, block_size_mismatch); {_, V, _, _} when V =:= 0 -> grisp_updater_manager:loader_failed(Id, data_size_mismatch); {_, _, H, _} when H =/= ExpBlockHash -> grisp_updater_manager:loader_failed(Id, block_hash_mismatch); {_, _, _, H} when H =/= ExpDataHash -> grisp_updater_manager:loader_failed(Id, data_hash_mismatch) end, State. init_stream(#block{encoding = #raw_encoding{}} = Block, StreamRef) -> #block{ id = Id, data_size = DataSize, data_hash_type = DataHashType } = Block, #stream{ id = Id, ref = StreamRef, data_left = DataSize, data_offset = 0, data_hash = crypto:hash_init(DataHashType) }; init_stream(#block{encoding = #gzip_encoding{}} = Block, StreamRef) -> #block{ id = Id, data_size = DataSize, data_hash_type = DataHashType, encoding = #gzip_encoding{ block_size = BlockSize, block_hash_type = BlockHashType } } = Block, Z = zlib:open(), zlib:inflateInit(Z, 16 + 15, reset), #stream{ id = Id, ref = StreamRef, block_left = BlockSize, data_left = DataSize, data_offset = 0, data_hash = crypto:hash_init(DataHashType), block_hash = crypto:hash_init(BlockHashType), inflater = Z }. cancel_pendings(#state{pending = PendMap} = State) -> lists:foldl(fun (#pending{stream = undefined}, S) -> S; (#pending{stream = StreamRef}, S) -> cancel_stream(S, StreamRef) end, State#state{pending = #{}}, maps:values(PendMap)). cancel_pending(#state{pending = PendMap} = State, BlockId) -> case maps:take(BlockId, PendMap) of error -> State; {#pending{stream = undefined}, PendMap2} -> State#state{pending = PendMap2}; {#pending{stream = StreamRef}, PendMap2} -> cancel_stream(State#state{pending = PendMap2}, StreamRef) end. cancel_stream(#state{streams = StreamMap} = State, StreamRef) -> case maps:take(StreamRef, StreamMap) of error -> State; {Stream, StreamMap2} -> grisp_updater_source:cancel(StreamRef), cleanup_stream(Stream), State#state{streams = StreamMap2} end. cleanup_stream(#stream{inflater = undefined}) -> ok; cleanup_stream(#stream{inflater = Z}) -> catch zlib:inflateEnd(Z), catch zlib:close(Z), ok.