%% Copyright (c) 2018 EMQ Technologies Co., Ltd. All Rights Reserved. %% %% 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. %% A per-topic gen_server whihch manages a number of per-partition wolff_producer workers. -module(wolff_producers). %% APIs -export([start_link/3]). -export([start_linked_producers/3, stop_linked/1]). -export([start_supervised/3, stop_supervised/1, stop_supervised/3]). -export([pick_producer/2, lookup_producer/2]). %% gen_server callbacks -export([code_change/3, handle_call/3, handle_cast/2, handle_info/2, init/1, terminate/2]). -export_type([producers/0]). -include("wolff.hrl"). -opaque producers() :: #{workers := #{partition() => pid()} | ets:tab(), partitioner := partitioner(), client_id => wolff:client_id(), topic => kpro:topic() }. -type topic() :: kpro:topic(). -type partition() :: kpro:partition(). -type config() :: wolff_producer:config(). -type partitioner() :: random %% default | roundrobin | first_key_dispatch | fun((PartitionCount :: pos_integer(), [wolff:msg()]) -> partition()) | partition(). -define(down(Reason), {down, Reason}). -define(rediscover_client, rediscover_client). -define(rediscover_client_tref, rediscover_client_tref). -define(rediscover_client_delay, 1000). -define(init_producers, init_producers). -define(init_producers_delay, 1000). -define(not_initialized, not_initialized). -define(partition_count_refresh_interval_seconds, 300). -define(refresh_partition_count, refresh_partition_count). %% @doc start wolff_producdrs gen_server start_link(ClientId, Topic, Config) -> Name = get_name(Config), gen_server:start_link({local, Name}, ?MODULE, {ClientId, Topic, Config}, []). %% @doc start wolff_producdrs gen_server -spec start_linked_producers(wolff:client_id() | pid(), topic(), config()) -> {ok, producers()} | {error, any()}. start_linked_producers(ClientId, Topic, ProducerCfg) when is_binary(ClientId) -> {ok, ClientPid} = wolff_client_sup:find_client(ClientId), start_linked_producers(ClientId, ClientPid, Topic, ProducerCfg); start_linked_producers(ClientPid, Topic, ProducerCfg) when is_pid(ClientPid) -> ClientId = wolff_client:get_id(ClientPid), start_linked_producers(ClientId, ClientPid, Topic, ProducerCfg). start_linked_producers(ClientId, ClientPid, Topic, ProducerCfg) -> case wolff_client:get_leader_connections(ClientPid, Topic) of {ok, Connections} -> Workers = start_link_producers(ClientId, Topic, Connections, ProducerCfg), ok = put_partition_cnt(ClientId, Topic, maps:size(Workers)), Partitioner = maps:get(partitioner, ProducerCfg, random), {ok, #{client_id => ClientId, topic => Topic, workers => Workers, partitioner => Partitioner}}; {error, Reason} -> {error, Reason} end. stop_linked(#{workers := Workers}) when is_map(Workers) -> lists:foreach(fun({_, Pid}) -> wolff_producer:stop(Pid) end, maps:to_list(Workers)). %% @doc Start supervised producers. -spec start_supervised(wolff:client_id(), topic(), config()) -> {ok, producers()} | {error, any()}. start_supervised(ClientId, Topic, ProducerCfg) -> {ok, Pid} = wolff_producers_sup:ensure_present(ClientId, Topic, ProducerCfg), case gen_server:call(Pid, get_workers, infinity) of ?not_initialized -> %% This means wolff_client failed to fetch metadata %% for this topic, wolff_producers will keep retrying, %% but we have to return an error here {error, failed_to_initialize_producers_in_time}; Ets -> {ok, #{client_id => ClientId, topic => Topic, workers => Ets, partitioner => maps:get(partitioner, ProducerCfg, random) }} end. %% @doc Ensure workers and clean up meta data. -spec stop_supervised(producers()) -> ok. stop_supervised(#{client_id := ClientId, workers := NamedEts, topic := Topic}) -> stop_supervised(ClientId, Topic, NamedEts). %% @doc Ensure workers and clean up meta data. -spec stop_supervised(wolff:client_id(), topic(), wolff:name()) -> ok. stop_supervised(ClientId, Topic, NamedEts) -> wolff_producers_sup:ensure_absence(ClientId, NamedEts), case wolff_client_sup:find_client(ClientId) of {ok, Pid} -> ok = wolff_client:delete_producers_metadata(Pid, Topic); {error, _} -> %% not running ok end. %% @doc Retrieve the per-partition producer pid. -spec pick_producer(producers(), [wolff:msg()]) -> {partition(), pid()}. pick_producer(#{workers := Workers, partitioner := Partitioner, client_id := ClientId, topic := Topic }, Batch) -> Count = partition_cnt(ClientId, Topic), Partition = pick_partition(Count, Partitioner, Batch), do_pick_producer(Partitioner, Partition, Count, Workers). do_pick_producer(Partitioner, Partition, Count, Workers) -> Pid = lookup_producer(Workers, Partition), case is_pid(Pid) andalso is_process_alive(Pid) of true -> {Partition, Pid}; false when Partitioner =:= random -> pick_next_alive(Workers, Partition, Count); false when Partitioner =:= roundrobin -> R = {Partition, Pid} = pick_next_alive(Workers, Partition, Count), _ = put(wolff_roundrobin, (Partition + 1) rem Count), R; false -> erlang:error({producer_down, Pid}) end. pick_next_alive(Workers, Partition, Count) -> pick_next_alive(Workers, (Partition + 1) rem Count, Count, _Tried = 1). pick_next_alive(_Workers, _Partition, Count, Count) -> erlang:error(all_producers_down); pick_next_alive(Workers, Partition, Count, Tried) -> Pid = lookup_producer(Workers, Partition), case is_alive(Pid) of true -> {Partition, Pid}; false -> pick_next_alive(Workers, (Partition + 1) rem Count, Count, Tried + 1) end. is_alive(Pid) -> is_pid(Pid) andalso is_process_alive(Pid). lookup_producer(#{workers := Workers}, Partition) -> lookup_producer(Workers, Partition); lookup_producer(Workers, Partition) when is_map(Workers) -> maps:get(Partition, Workers); lookup_producer(Workers, Partition) -> [{Partition, Pid}] = ets:lookup(Workers, Partition), Pid. pick_partition(_Count, Partition, _) when is_integer(Partition) -> Partition; pick_partition(Count, F, Batch) when is_function(F) -> F(Count, Batch); pick_partition(Count, Partitioner, _) when not is_integer(Count); Count =< 0 -> error({invalid_partition_count, Count, Partitioner}); pick_partition(Count, random, _) -> rand:uniform(Count) - 1; pick_partition(Count, roundrobin, _) -> Partition = case get(wolff_roundrobin) of undefined -> 0; Number -> Number end, _ = put(wolff_roundrobin, (Partition + 1) rem Count), Partition; pick_partition(Count, first_key_dispatch, [#{key := Key} | _]) -> erlang:phash2(Key) rem Count. init({ClientId, Topic, Config}) -> erlang:process_flag(trap_exit, true), self() ! ?rediscover_client, {ok, #{client_id => ClientId, client_pid => false, topic => Topic, config => Config, ets => ?not_initialized, refresh_tref => start_partition_refresh_timer(Config) }}. handle_info(?refresh_partition_count, #{refresh_tref := Tref, config := Config} = St0) -> %% this message can be sent from anywhere, %% so we should ensure the timer is cancelled before starting a new one ok = ensure_timer_cancelled(Tref), St = refresh_partition_count(St0), {noreply, St#{refresh_tref := start_partition_refresh_timer(Config)}}; handle_info(?rediscover_client, #{client_id := ClientId, client_pid := false } = St0) -> St1 = St0#{?rediscover_client_tref => false}, case wolff_client_sup:find_client(ClientId) of {ok, Pid} -> _ = erlang:monitor(process, Pid), St2 = St1#{client_pid := Pid}, St3 = maybe_init_producers(St2), St = maybe_restart_producers(St3), {noreply, St}; {error, Reason} -> log_error("failed_to_discover_client", #{reason => Reason, client_id => ClientId}), {noreply, ensure_rediscover_client_timer(St1)} end; handle_info(?init_producers, St) -> %% this is a retry of last failure when initializing producer procs {noreply, maybe_init_producers(St)}; handle_info({'DOWN', _, process, Pid, Reason}, #{client_id := ClientId, client_pid := Pid } = St) -> log_error("client_pid_down", #{client_id => ClientId, client_pid => Pid, reason => Reason}), %% client down, try to discover it after a delay %% producers should all monitor client pid, %% expect their 'EXIT' signals soon {noreply, ensure_rediscover_client_timer(St#{client_pid := false})}; handle_info({'EXIT', Pid, Reason}, #{ets := Ets, topic := Topic, client_id := ClientId, client_pid := ClientPid, config := Config } = St) -> case ets:match(Ets, {'$1', Pid}) of [] -> %% this should not happen, hence error level log_error("unknown_EXIT_message", #{pid => Pid, reason => Reason}); [[Partition]] -> case is_alive(ClientPid) of true -> %% wolff_producer is not designed to crash & restart %% if this happens, it's likely a bug in wolff_producer module log_error("producer_down", #{topic => Topic, partition => Partition, partition_worker => Pid, reason => Reason}), ok = start_producer_and_insert_pid(Ets, ClientId, Topic, Partition, Config); false -> %% no client, restart will be triggered when client connection is back. ets:insert(Ets, {Partition, ?down(Reason)}) end end, {noreply, St}; handle_info(Info, St) -> log_error("unknown_info", #{info => Info}), {noreply, St}. handle_call(get_workers, _From, #{ets := Ets} = St) -> {reply, Ets, St}; handle_call(Call, From, St) -> log_error("unknown_call", #{call => Call, from => From}), {reply, {error, unknown_call}, St}. handle_cast(Cast, St) -> log_error("unknown_cast", #{cast => Cast}), {noreply, St}. code_change(_OldVsn, St, _Extra) -> {ok, St}. terminate(_, _St) -> ok. ensure_rediscover_client_timer(#{?rediscover_client_tref := false} = St) -> Tref = erlang:send_after(?rediscover_client_delay, self(), ?rediscover_client), St#{?rediscover_client_tref := Tref}. log(Level, Msg, Args) -> logger:log(Level, Args#{msg => Msg}). log_error(Msg, Args) -> log(error, Msg, Args). log_warning(Msg, Args) -> log(warning, Msg, Args). log_info(Msg, Args) -> log(info, Msg, Args). start_link_producers(ClientId, Topic, Connections, Config) -> lists:foldl( fun({Partition, MaybeConnPid}, Acc) -> {ok, WorkerPid} = wolff_producer:start_link(ClientId, Topic, Partition, MaybeConnPid, Config), Acc#{Partition => WorkerPid} end, #{}, Connections). maybe_init_producers(#{ets := ?not_initialized, topic := Topic, client_id := ClientId, config := Config } = St) -> case start_linked_producers(ClientId, Topic, Config) of {ok, #{workers := Workers}} -> Ets = ets:new(get_name(Config), [protected, named_table, {read_concurrency, true}]), true = ets:insert(Ets, maps:to_list(Workers)), St#{ets := Ets}; {error, Reason} -> log_error("failed_to_init_producers", #{topic => Topic, reason => Reason}), erlang:send_after(?init_producers_delay, self(), ?init_producers), St end; maybe_init_producers(St) -> St. maybe_restart_producers(#{ets := ?not_initialized} = St) -> St; maybe_restart_producers(#{ets := Ets, client_id := ClientId, topic := Topic, config := Config } = St) -> lists:foreach( fun({Partition, Pid}) -> case is_alive(Pid) of true -> ok; false -> start_producer_and_insert_pid(Ets, ClientId, Topic, Partition, Config) end end, ets:tab2list(Ets)), St. get_name(Config) -> maps:get(name, Config, ?MODULE). start_producer_and_insert_pid(Ets, ClientId, Topic, Partition, Config) -> {ok, Pid} = wolff_producer:start_link(ClientId, Topic, Partition, ?conn_down(to_be_discovered), Config), ets:insert(Ets, {Partition, Pid}), ok. %% Config is not used so far. start_partition_refresh_timer(Config) -> IntervalSeconds = maps:get(partition_count_refresh_interval_seconds, Config, ?partition_count_refresh_interval_seconds), case IntervalSeconds of 0 -> undefined; _ -> Interval = timer:seconds(IntervalSeconds), erlang:send_after(Interval, self(), ?refresh_partition_count) end. refresh_partition_count(#{client_pid := Pid} = St) when not is_pid(Pid) -> %% client is to be (re)discovered St; refresh_partition_count(#{ets := ?not_initialized} = St) -> %% to be initialized St; refresh_partition_count(#{client_pid := Pid, topic := Topic} = St) -> case wolff_client:get_leader_connections(Pid, Topic) of {ok, Connections} -> start_new_producers(St, Connections); {error, Reason} -> log_warning("failed_to_refresh_partition_count_will_retry", #{topic => Topic, reason => Reason}), St end. start_new_producers(#{client_id := ClientId, topic := Topic, config := Config, ets := Ets } = St, Connections0) -> %% process only the newly discovered connections F = fun({Partition, _MaybeConnPid}) -> [] =:= ets:lookup(Ets, Partition) end, Connections = lists:filter(F, Connections0), Workers = start_link_producers(ClientId, Topic, Connections, Config), true = ets:insert(Ets, maps:to_list(Workers)), true = (length(Connections0) >= partition_cnt(ClientId, Topic)), %% assert case maps:size(Workers) > 0 of true -> log_info("started_producers_for_newly_discovered_partitions", #{workers => Workers}), ok = put_partition_cnt(ClientId, Topic, length(Connections0)); false -> ok end, St. partition_cnt(ClientId, Topic) -> persistent_term:get({?MODULE, ClientId, Topic}). put_partition_cnt(ClientId, Topic, Count) -> persistent_term:put({?MODULE, ClientId, Topic}, Count). ensure_timer_cancelled(Tref) when is_reference(Tref) -> _ = erlang:cancel_timer(Tref), ok; ensure_timer_cancelled(_) -> ok.