ipregistry

%% Copyright 2019 Ipregistry (https://ipregistry.co).
%%
%% 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 Optional in-memory cache for Ipregistry lookups, with time-based
%% expiration and a bounded size using insertion-order (FIFO) eviction.
%%
%% Start one under your supervision tree with {@link child_spec/1} or
%% standalone with {@link start_link/1}, then reference its name in the
%% `cache' option of {@link ipregistry:new/2}:
%%
%% ```
%% {ok, _} = ipregistry_cache:start_link(#{name => my_ip_cache, ttl => 600000}),
%% Client = ipregistry:new(<<"YOUR_API_KEY">>, #{cache => my_ip_cache}).
%% '''
%%
%% Reads are served directly from a protected ETS table owned by the cache
%% process, so cache hits never block on the server. Writes go through the
%% server and are best-effort: if the cache process is down, lookups keep
%% working without caching rather than failing.
-module(ipregistry_cache).

-behaviour(gen_server).

-export([
    start_link/0, start_link/1,
    child_spec/1,
    stop/1,
    get/2,
    put/3,
    invalidate/2,
    invalidate_all/1,
    size/1
]).

-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2]).

-export_type([cache/0, options/0]).

-define(DEFAULT_NAME, ?MODULE).
-define(DEFAULT_MAX_SIZE, 4096).
%% 10 minutes.
-define(DEFAULT_TTL, 600000).
%% 1 minute.
-define(DEFAULT_SWEEP_INTERVAL, 60000).

-type cache() :: atom() | pid().
%% A cache identified by its registered name (preferred) or its pid. Direct
%% ETS reads require the registered name; when a pid is used, reads go
%% through the server process.

-type options() :: #{
    name => atom(),
    max_size => pos_integer(),
    ttl => pos_integer(),
    sweep_interval => pos_integer()
}.
%% `name' registers the cache process and its ETS table (default
%% `ipregistry_cache'). `ttl' is how long entries stay valid, in
%% milliseconds (default 10 minutes). `max_size' bounds the number of
%% entries (default 4096); when full, the oldest entry is evicted.
%% `sweep_interval' controls how often expired entries are purged, in
%% milliseconds (default 1 minute).

-record(state, {
    table :: ets:table(),
    order :: ets:table(),
    max_size :: pos_integer(),
    ttl :: pos_integer(),
    sweep_interval :: pos_integer(),
    seq = 0 :: non_neg_integer()
}).

%%--------------------------------------------------------------------
%% API
%%--------------------------------------------------------------------

%% @doc Starts a cache with default options.
%% @equiv start_link(#{})
-spec start_link() -> {ok, pid()} | {error, term()}.
start_link() ->
    start_link(#{}).

%% @doc Starts a cache registered under the configured name and links it to
%% the calling process.
-spec start_link(options()) -> {ok, pid()} | {error, term()}.
start_link(Options) when is_map(Options) ->
    Name = maps:get(name, Options, ?DEFAULT_NAME),
    gen_server:start_link({local, Name}, ?MODULE, Options#{name => Name}, []).

%% @doc Returns a supervisor child spec for embedding the cache in your
%% supervision tree.
-spec child_spec(options()) -> supervisor:child_spec().
child_spec(Options) when is_map(Options) ->
    Name = maps:get(name, Options, ?DEFAULT_NAME),
    #{
        id => Name,
        start => {?MODULE, start_link, [Options#{name => Name}]},
        restart => permanent,
        shutdown => 5000,
        type => worker,
        modules => [?MODULE]
    }.

%% @doc Stops the cache process.
-spec stop(cache()) -> ok.
stop(Cache) ->
    gen_server:stop(Cache).

%% @doc Returns the cached value for `Key' when present and unexpired.
-spec get(cache(), term()) -> {ok, term()} | miss.
get(Cache, Key) when is_atom(Cache) ->
    try ets:lookup(Cache, Key) of
        [{_Key, Value, ExpiresAt, _Seq}] ->
            case erlang:monotonic_time(millisecond) =< ExpiresAt of
                true -> {ok, Value};
                false -> miss
            end;
        [] ->
            miss
    catch
        %% The cache process (and its table) is down: degrade to a miss.
        error:badarg -> miss
    end;
get(Cache, Key) when is_pid(Cache) ->
    safe_call(Cache, {get, Key}, miss).

%% @doc Stores `Value' under `Key', refreshing its expiration. Best-effort:
%% returns `ok' even when the cache process is unavailable.
-spec put(cache(), term(), term()) -> ok.
put(Cache, Key, Value) ->
    safe_call(Cache, {put, Key, Value}, ok).

%% @doc Removes the entry for `Key', if present.
-spec invalidate(cache(), term()) -> ok.
invalidate(Cache, Key) ->
    safe_call(Cache, {invalidate, Key}, ok).

%% @doc Removes every entry.
-spec invalidate_all(cache()) -> ok.
invalidate_all(Cache) ->
    safe_call(Cache, invalidate_all, ok).

%% @doc Returns the current number of entries, including expired entries
%% not yet swept. Primarily useful in tests.
-spec size(cache()) -> non_neg_integer().
size(Cache) ->
    safe_call(Cache, size, 0).

safe_call(Cache, Message, Default) ->
    try
        gen_server:call(Cache, Message)
    catch
        exit:{noproc, _} -> Default;
        exit:{normal, _} -> Default;
        exit:{shutdown, _} -> Default
    end.

%%--------------------------------------------------------------------
%% gen_server callbacks
%%--------------------------------------------------------------------

%% @private
init(Options) ->
    Name = maps:get(name, Options),
    Table = ets:new(Name, [named_table, protected, set, {read_concurrency, true}]),
    Order = ets:new(order_table_name(Name), [ordered_set, private]),
    State = #state{
        table = Table,
        order = Order,
        max_size = maps:get(max_size, Options, ?DEFAULT_MAX_SIZE),
        ttl = maps:get(ttl, Options, ?DEFAULT_TTL),
        sweep_interval = maps:get(sweep_interval, Options, ?DEFAULT_SWEEP_INTERVAL)
    },
    schedule_sweep(State),
    {ok, State}.

%% @private
handle_call({put, Key, Value}, _From, State) ->
    {reply, ok, insert(Key, Value, State)};
handle_call({get, Key}, _From, State = #state{table = Table}) ->
    Reply =
        case ets:lookup(Table, Key) of
            [{_Key, Value, ExpiresAt, _Seq}] ->
                case erlang:monotonic_time(millisecond) =< ExpiresAt of
                    true -> {ok, Value};
                    false -> miss
                end;
            [] ->
                miss
        end,
    {reply, Reply, State};
handle_call({invalidate, Key}, _From, State = #state{table = Table}) ->
    true = ets:delete(Table, Key),
    {reply, ok, State};
handle_call(invalidate_all, _From, State = #state{table = Table, order = Order}) ->
    true = ets:delete_all_objects(Table),
    true = ets:delete_all_objects(Order),
    {reply, ok, State};
handle_call(size, _From, State = #state{table = Table}) ->
    {reply, ets:info(Table, size), State};
handle_call(_Request, _From, State) ->
    {reply, {error, unknown_request}, State}.

%% @private
handle_cast(_Message, State) ->
    {noreply, State}.

%% @private
handle_info(sweep, State) ->
    sweep(State),
    schedule_sweep(State),
    {noreply, State};
handle_info(_Message, State) ->
    {noreply, State}.

%% @private
terminate(_Reason, _State) ->
    ok.

%%--------------------------------------------------------------------
%% Internal functions
%%--------------------------------------------------------------------

order_table_name(Name) ->
    list_to_atom(atom_to_list(Name) ++ "_order").

insert(Key, Value, State = #state{table = Table, order = Order, ttl = Ttl, seq = Seq}) ->
    State1 = maybe_evict(Key, State),
    ExpiresAt = erlang:monotonic_time(millisecond) + Ttl,
    true = ets:insert(Table, {Key, Value, ExpiresAt, Seq}),
    true = ets:insert(Order, {Seq, Key}),
    State1#state{seq = Seq + 1}.

%% Evicts the oldest entry when inserting a new key into a full cache.
%% Overwrites of an existing key never need eviction.
maybe_evict(Key, State = #state{table = Table, max_size = MaxSize}) ->
    case ets:member(Table, Key) orelse ets:info(Table, size) < MaxSize of
        true -> State;
        false -> evict_oldest(State)
    end.

evict_oldest(State = #state{table = Table, order = Order}) ->
    case ets:first(Order) of
        '$end_of_table' ->
            State;
        Seq ->
            [{_Seq, Key}] = ets:lookup(Order, Seq),
            true = ets:delete(Order, Seq),
            case ets:lookup(Table, Key) of
                %% Only delete when this order entry is not stale, i.e. the
                %% key was not overwritten with a newer sequence number since.
                [{_Key, _Value, _ExpiresAt, Seq}] ->
                    true = ets:delete(Table, Key),
                    State;
                _ ->
                    %% Stale order entry (key overwritten or already gone):
                    %% keep evicting.
                    evict_oldest(State)
            end
    end.

sweep(#state{table = Table, order = Order}) ->
    Now = erlang:monotonic_time(millisecond),
    ExpiredSpec = [{{'_', '_', '$1', '_'}, [{'<', '$1', Now}], [true]}],
    _ = ets:select_delete(Table, ExpiredSpec),
    %% Drop order entries whose main entry is gone or superseded.
    prune_order(Table, Order, ets:first(Order)).

prune_order(_Table, _Order, '$end_of_table') ->
    ok;
prune_order(Table, Order, Seq) ->
    Next = ets:next(Order, Seq),
    [{_Seq, Key}] = ets:lookup(Order, Seq),
    case ets:lookup(Table, Key) of
        [{_Key, _Value, _ExpiresAt, Seq}] -> ok;
        _ -> true = ets:delete(Order, Seq)
    end,
    prune_order(Table, Order, Next).

schedule_sweep(#state{sweep_interval = Interval}) ->
    _ = erlang:send_after(Interval, self(), sweep),
    ok.