%%%--------------------------------------------------------------------------- %%% Copyright (C) 2021, Fix Solution All Rights Reserved %%% Unauthorized copy of this file is through any medium strictly not allowed. %%% %%% Authors : Alpha Shaw %%% Created : 20 Apr 2021 by %%% Purpose : %%% %%% %%% 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(worker). -author("Alpha Umaru Shaw"). -include_lib("erlflake.hrl"). -include_lib("erlwater/include/logger.hrl"). -behaviour(gen_server). %% API -export([start_link/0, reserve/0]). %% gen_server callbacks -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]). -define(NODE, "erlflake"). -define(SERVER, ?MODULE). -record(state, { conn_pid, worker_id }). %%%=================================================================== %%% API %%%=================================================================== %%-------------------------------------------------------------------- %% @doc %% Starts the server %% %% @end %%-------------------------------------------------------------------- -spec(start_link() -> {ok, Pid :: pid()} | ignore | {error, Reason :: term()}). start_link() -> gen_server:start_link({local, ?SERVER}, ?MODULE, [], []). reserve() -> gen_server:call(?SERVER, reserve_worker). %%%=================================================================== %%% gen_server callbacks %%%=================================================================== init(_) -> case get_local_worker_id() of Id when is_integer(Id) -> ?LOG_INFO("Found a local worker ID: ~p~n", [Id]), {ok, #state{worker_id = Id}}; _ -> ?LOG_INFO( "No explicit local worker ID is detected. " "Fallbacking to zookeeper for dynamic provisioning", []), case erlzk:start() of ok -> case zk_connect() of {ok, Pid} -> case erlzk:create(Pid, "/" ++ ?NODE) of {error, Reason} when Reason /= node_exists -> {stop, Reason}; _ -> {ok, #state{conn_pid = Pid}} end; {error, Reason} -> {stop, Reason} end; {error, Reason} -> {stop, Reason} end end. handle_call(reserve_worker, _From, State) -> if State#state.worker_id == ?Undef -> {Id, State2} = reserve_worker(State), {reply, Id, State2#state{worker_id = Id}}; true -> {reply, {ok, State#state.worker_id}, State} end; handle_call(_Request, _From, State) -> {reply, ok, State}. handle_cast(_Request, State) -> {noreply, State}. handle_info(_Info, State) -> {noreply, State}. terminate(_Reason, _State) -> ok. code_change(_OldVsn, State, _Extra) -> {ok, State}. %%%=================================================================== %%% Internal functions %%%=================================================================== reserve_worker(#state{conn_pid = Pid} = State) -> case reserve_worker_id(Pid) of Id when is_integer(Id) -> {{ok, format_worker(Id)}, State}; {error, Reason} = Error -> ?LOG_INFO("** Unable to reserv a worker ID **~n" "Reason: ~p~n", [Reason]), {Error, State} end. format_worker(Worker) when Worker =< ?MAX_SEQUENCE -> Worker rem (?MAX_WORKER + 1); format_worker(Worker) -> Worker. reserve_worker_id(Pid) -> create_path(Pid, "/" ++ ?NODE, undefined, 4). create_path(_, _, Res, 0) -> Res; create_path(Pid, Path, _, Retries) -> case erlzk:create(Pid, Path, ephemeral_sequential) of {ok, NewPath} -> Seq = string:sub_string(NewPath, length(Path) + 1), erlang:abs(erlwater:to_integer(Seq)); {error, _} = Err -> create_path(Pid, Path, Err, Retries - 1) end. zk_connect() -> HostPort = zookeeper_host_port(), ?LOG_INFO("Connecting to zookeeper at '~s'~n", [ew_net:ip_to_binary(HostPort)]), attempt_connect(HostPort, undefined, 4). attempt_connect(_, Res, 0) -> Res; attempt_connect(HostPort, _, Retries) -> case erlzk:connect([HostPort], 30000) of {ok, Pid} -> {ok, Pid}; {error, _} = Err -> attempt_connect(HostPort, Err, Retries - 1) end. zookeeper_host_port() -> Host = case ew_os:get_env(?ZK_HOST_ENV_VAR, ?Undef) of ?Undef -> erlwater:get_binary_app_env(erlflake, zk_host, ?DEFAULT_ZK_HOST); Host0 -> iolist_to_binary(Host0) end, Port = case ew_os:get_env(?ZK_PORT_ENV_VAR) of ?Undef -> erlwater:get_int_app_env(erlflake, zk_port, ?DEFAULT_ZK_PORT); Port0 -> list_to_integer(Port0) end, {binary_to_list(Host), Port}. get_local_worker_id() -> case erlwater:get_int_env(?WORKER_ID_ENV_VAR) of ?Undef -> erlwater:get_int_app_env(erlflake, worker_id); Id -> Id end.