%%%
%%% 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.
%%%
%%%=============================================================================
%%% @doc This is a utility module to help force commit offsets to kafka.
%%% @end
%%%=============================================================================
-module(brod_cg_commits).
-behaviour(gen_server).
-behaviour(brod_group_member).
-export([ run/2
]).
-export([ start_link/2
, stop/1
, sync/1
]).
%% callbacks for brod_group_coordinator
-export([ get_committed_offsets/2
, assignments_received/4
, assignments_revoked/1
, assign_partitions/3
]).
-export([ code_change/3
, handle_call/3
, handle_cast/2
, handle_info/2
, init/1
, terminate/2
]).
-include("brod_int.hrl").
-type topic() :: brod:topic().
-type partition() :: brod:partition().
-type offset() :: brod:offset().
-type group_id() :: brod:group_id().
-type member_id() :: brod:group_member_id().
-type retention() :: integer(). %% -1 to use whatever configured in kafka
-type offsets() :: latest | earliest | [{partition(), offset()}].
-type prop_key() :: id | topic | retention | protocol | offsets.
-type prop_val() :: group_id() | topic() | retention() | offsets()
| brod_group_coordinator:protocol_name().
-type group_input() :: [{prop_key(), prop_val()}].
-if(?OTP_RELEASE >= 25).
-type pending_sync() :: ?undef | gen_server:from().
-else.
-type pending_sync() :: ?undef | {pid(), reference()}.
-endif.
-record(state,
{ client :: brod:client()
, groupId :: brod:group_id()
, memberId :: ?undef | member_id()
, generationId :: ?undef | brod:group_generation_id()
, coordinator :: pid()
, topic :: ?undef | topic()
, offsets :: ?undef | offsets()
, is_elected = false :: boolean()
, pending_sync :: pending_sync()
, is_done = false :: boolean()
}).
-type state() :: #state{}.
%%%_* APIs =====================================================================
%% @doc Force commit offsets.
run(ClientId, GroupInput) ->
{ok, Pid} = start_link(ClientId, GroupInput),
ok = sync(Pid),
ok = stop(Pid).
%% @doc Start (link) a group member.
%% The member will try to join the consumer group and
%% get assignments for the given topic-partitions,
%% then commit given offsets to kafka.
%% In case not all given partitions are assigned to it,
%% it will terminate with an exit exception
-spec start_link(brod:client(), group_input()) -> {ok, pid()} | {error, any()}.
start_link(Client, GroupInput) ->
gen_server:start_link(?MODULE, {Client, GroupInput}, []).
%% @doc Stop the process.
-spec stop(pid()) -> ok.
stop(Pid) ->
Mref = erlang:monitor(process, Pid),
ok = gen_server:cast(Pid, stop),
receive
{'DOWN', Mref, process, Pid, _Reason} ->
ok
end.
%% @doc Make a call to the resetter process, the call will be blocked
%% until offsets are committed.
-spec sync(pid()) -> ok.
sync(Pid) ->
ok = gen_server:call(Pid, sync, infinity).
%%%_* APIs for group coordinator ===============================================
%% @doc Called by group coordinator when there is new assignment received.
-spec assignments_received(pid(), member_id(), integer(),
brod:received_assignments()) -> ok.
assignments_received(Pid, MemberId, GenerationId, TopicAssignments) ->
gen_server:cast(Pid, {new_assignments, MemberId,
GenerationId, TopicAssignments}).
%% @doc Called by group coordinator before re-joinning the consumer group.
-spec assignments_revoked(pid()) -> ok.
assignments_revoked(Pid) ->
gen_server:call(Pid, unsubscribe_all_partitions, infinity).
%% @doc This function is called only when `partition_assignment_strategy'
%% is set for `callback_implemented' in group config.
-spec assign_partitions(pid(), [brod:group_member()],
[{brod:topic(), brod:partition()}]) ->
[{member_id(), [brod:partition_assignment()]}].
assign_partitions(Pid, Members, TopicPartitionList) ->
Call = {assign_partitions, Members, TopicPartitionList},
gen_server:call(Pid, Call, infinity).
%% @doc Called by group coordinator when initializing the assignments
%% for subscriber.
%% NOTE: this function is called only when it is DISABLED to commit offsets
%% to kafka. i.e. offset_commit_policy is set to consumer_managed
-spec get_committed_offsets(pid(), [{brod:topic(), brod:partition()}]) ->
{ok, [{{brod:topic(), brod:partition()}, brod:offset()}]}.
get_committed_offsets(_Pid, _TopicPartitions) -> {ok, []}.
%%%_* gen_server callbacks =====================================================
init({Client, GroupInput}) ->
ok = brod_utils:assert_client(Client),
GroupId = proplists:get_value(id, GroupInput),
ok = brod_utils:assert_group_id(GroupId),
Topic = proplists:get_value(topic, GroupInput),
ProtocolName = proplists:get_value(protocol, GroupInput),
Retention = proplists:get_value(retention, GroupInput),
Offsets = proplists:get_value(offsets, GroupInput),
%% use callback_implemented strategy so I know I am elected leader
%% when `assign_partitions' callback is called.
Config = [ {partition_assignment_strategy, callback_implemented}
, {offset_retention_seconds, Retention}
, {protocol_name, ProtocolName}
, {rejoin_delay_seconds, 2}
],
{ok, Pid} = brod_group_coordinator:start_link(Client, GroupId, [Topic],
Config, ?MODULE, self()),
State = #state{ client = Client
, groupId = GroupId
, coordinator = Pid
, topic = Topic
, offsets = Offsets
},
{ok, State}.
handle_info(Info, State) ->
log(State, info, "Info discarded:~p", [Info]),
{noreply, State}.
handle_call(sync, From, State0) ->
State1 = State0#state{pending_sync = From},
State = maybe_reply_sync(State1),
{noreply, State};
handle_call({assign_partitions, Members, TopicPartitions}, _From,
#state{topic = MyTopic,
offsets = Offsets} = State) ->
log(State, info, "Assigning all topic partitions to self", []),
MyTP = [{MyTopic, P} || {P, _} <- Offsets],
%% Assert that my topic partitions are included in
%% subscriptions collected from ALL members
Pred = fun(TP) -> not lists:member(TP, TopicPartitions) end,
case lists:filter(Pred, MyTP) of
[] ->
%% all of the give partitions are valid
ok;
BadPartitions ->
PartitionNumbers = [P || {_T, P} <- BadPartitions],
log(State, error, "Nonexisting partitions in input: ~p",
[PartitionNumbers]),
erlang:exit({non_existing_partitions, PartitionNumbers})
end,
%% To keep it simple, assign all my topic-partitions to self
%% but discard all other topic-partitions.
%% After all, I will leave group as soon as offsets are committed
Result = assign_all_to_self(Members, MyTP),
{reply, Result, State#state{is_elected = true}};
handle_call(unsubscribe_all_partitions, _From, #state{} = State) ->
%% nothing to do, because I do not subscribe to any partition
{reply, ok, State};
handle_call(Call, _From, State) ->
{reply, {error, {unknown_call, Call}}, State}.
handle_cast({new_assignments, _MemberId, GenerationId, Assignments},
#state{ is_elected = IsLeader
, offsets = OffsetsToCommit
, coordinator = Pid
, topic = MyTopic
} = State) ->
%% Write a log if I am not a leader,
%% hope the desired partitions are all assigned to me
IsLeader orelse log(State, info, "Not elected", []),
Groupped0 =
brod_utils:group_per_key(
fun(#brod_received_assignment{ topic = Topic
, partition = Partition
, begin_offset = Offset
}) ->
{Topic, {Partition, Offset}}
end, Assignments),
%% Discard other topics if for whatever reason the group leader assigns
%% irrelevant topic-partitions to me
Groupped = lists:filter(fun({Topic, _}) -> Topic =:= MyTopic end, Groupped0),
log(State, info, "current offsets:\n~p", [Groupped]),
%% Assert all desired partitions are in assignment
case Groupped of
[] ->
log(State, error, "Topic ~s is not received in assignment", [MyTopic]),
erlang:exit({bad_topic_assignment, Groupped0});
[{MyTopic, PartitionOffsetList}] ->
MyPartitions = [P || {P, _O} <- OffsetsToCommit],
ReceivedPartitions = [P || {P, _O} <- PartitionOffsetList],
case MyPartitions -- ReceivedPartitions of
[] ->
ok;
Left ->
log(State, error,
"Partitions ~p are not received in assignment, "
"There is probably another active group member subscribing "
"to topic ~s, stop it and retry\n", [MyTopic, Left]),
erlang:exit({unexpected_assignments, Left})
end
end,
%% Stage all offsets in coordinator process
lists:foreach(
fun({Partition, Offset}) ->
%% -1 here, because brod_group_coordinator +1 to commit
OffsetToCommit = Offset - 1,
brod_group_coordinator:ack(Pid, GenerationId, MyTopic,
Partition, OffsetToCommit)
end, OffsetsToCommit),
%% Now force it to commit
case brod_group_coordinator:commit_offsets(Pid) of
ok -> ok;
{error, Reason} ->
log(State, error, "Failed to commit, reason:\n~p", [Reason]),
erlang:exit(commit_failed)
end,
{noreply, set_done(State)};
handle_cast(stop, State) ->
{stop, normal, State};
handle_cast(_Cast, State) ->
{noreply, State}.
code_change(_OldVsn, State, _Extra) ->
{ok, State}.
terminate(_Reason, #state{}) ->
ok.
%%%_* Internal Functions =======================================================
set_done(State) ->
maybe_reply_sync(State#state{is_done = true}).
maybe_reply_sync(#state{is_done = false} = State) ->
State;
maybe_reply_sync(#state{pending_sync = ?undef} = State) ->
State;
maybe_reply_sync(#state{pending_sync = From} = State) ->
gen_server:reply(From, ok),
log(State, info, "done\n", []),
State#state{pending_sync = ?undef}.
%% I am the current leader because I am assigning partitions.
%% My member ID should be positioned at the head of the member list.
-spec assign_all_to_self([brod:group_member()], [{topic(), partition()}]) ->
[{member_id(), [brod:partition_assignment()]}].
assign_all_to_self([{MyMemberId, _} | Members], TopicPartitions) ->
Groupped = brod_utils:group_per_key(TopicPartitions),
[ {MyMemberId, Groupped}
| [{Id, []} || {Id, _MemberMeta} <- Members]
].
-spec log(state(), logger:level(), io:format(), [term()]) ->
ok.
log(#state{groupId = GroupId}, Level, Fmt, Args) ->
?BROD_LOG(Level,
"Group member (~s,coor=~p):\n" ++ Fmt,
[GroupId, self() | Args]).
%%%_* Emacs ====================================================================
%%% Local Variables:
%%% allout-layout: t
%%% erlang-indent-level: 2
%%% End:
