erlflake

%%%---------------------------------------------------------------------------
%%% Copyright (C) 2021, Fix Solution All Rights Reserved
%%% Unauthorized copy of this file is through any medium strictly not allowed.
%%%
%%% Authors : Alpha Shaw <shawalpha5@gmail.com>
%%% Created :  20 Apr 2021  by <shawalpha5@gmail.com>
%%% 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.