Erlang/OTP integration of systemd's "notify" start-up type and service watchdog feature

example_app.erl

-module(example_app).
-behaviour(application).
-export([start/2, prep_stop/1, stop/1]).

-spec start(application:start_type(), any()) -> {ok, pid()} | {error, term()}.
start(_StartType, _StartArgs) ->
    case example_sup:start_link() of
        {ok, _PID} = Res ->
            ok = systemd:ready(),
            Res;
        {error, _Reason} = Err ->
            Err
    end.

-spec prep_stop(term()) -> term().
prep_stop(State) ->
    ok = systemd:stopping(),
    State.

-spec stop(term()) -> ok.
stop(_State) ->
    ok.

example_sup.erl

-module(example_sup).
-behaviour(supervisor).
-export([start_link/0, init/1]).

-spec start_link() -> {ok, pid()} | {error, term()}.
start_link() ->
    supervisor:start_link({local, ?MODULE}, ?MODULE, []).

-spec init([]) -> {ok, {supervisor:sup_flags(), [supervisor:child_spec()]}}.
init([]) ->
    SupFlags = #{},
    ChildSpecs = [#{id => systemd, start => {systemd, start_link, []}},
                  #{id => example, start => {example, start_link, []}}],
    {ok, {SupFlags, ChildSpecs}}.

systemd.erl

%%% Copyright (c) 2020 Holger Weiss <[email protected]>.
%%%
%%% Permission to use, copy, modify, and/or distribute this software for any
%%% purpose with or without fee is hereby granted, provided that the above
%%% copyright notice and this permission notice appear in all copies.
%%%
%%% THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
%%% WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
%%% MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
%%% SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
%%% WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
%%% OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
%%% CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

-module(systemd).
-behaviour(gen_server).
-export([start_link/0,
         ready/0,
         reloading/0,
         stopping/0]).
-export([init/1,
         handle_call/3,
         handle_cast/2,
         handle_info/2,
         terminate/2,
         code_change/3]).
-export_type([state/0,
              watchdog_timeout/0]).

-include_lib("kernel/include/logger.hrl").

-record(systemd_state,
        {socket :: gen_udp:socket() | undefined,
         destination :: inet:local_address() | undefined,
         interval :: pos_integer() | undefined,
         last_ping :: integer() | undefined}).

-opaque state() :: #systemd_state{}.
-opaque watchdog_timeout() :: pos_integer() | hibernate.

%% API.

-spec start_link() -> {ok, pid()} | ignore | {error, term()}.
start_link() ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).

-spec ready() -> ok.
ready() ->
    cast_notification(<<"READY=1">>).

-spec reloading() -> ok.
reloading() ->
    cast_notification(<<"RELOADING=1">>).

-spec stopping() -> ok.
stopping() ->
    cast_notification(<<"STOPPING=1">>).

%% Behaviour callbacks.

-spec init(any())
      -> {ok, state()} | {ok, state(), watchdog_timeout()} | {stop, term()}.
init(_Opts) ->
    process_flag(trap_exit, true),
    case os:getenv("NOTIFY_SOCKET") of
        [$@ | _Abstract] ->
            ?LOG_CRITICAL("Abstract NOTIFY_SOCKET not supported"),
            {stop, esocktnosupport};
        Path when is_list(Path), length(Path) > 0 ->
            ?LOG_DEBUG("Got NOTIFY_SOCKET: ~s", [Path]),
            Destination = {local, Path},
            case gen_udp:open(0, [local]) of
                {ok, Socket} ->
                    Interval = get_watchdog_interval(),
                    State = #systemd_state{socket = Socket,
                                           destination = Destination,
                                           interval = Interval},
                    if is_integer(Interval), Interval > 0 ->
                            ?LOG_INFO("Watchdog notifications enabled"),
                            {ok, set_last_ping(State), Interval};
                       true ->
                            ?LOG_INFO("Watchdog notifications disabled"),
                            {ok, State}
                    end;
                {error, Reason} ->
                    ?LOG_CRITICAL("Cannot open IPC socket: ~p", [Reason]),
                    {stop, Reason}
            end;
        _ ->
            ?LOG_INFO("Got no NOTIFY_SOCKET, notifications disabled"),
            {ok, #systemd_state{}}
    end.

-spec handle_call(term(), {pid(), term()}, state())
      -> {reply, {error, badarg}, state(), watchdog_timeout()}.
handle_call(Request, From, State) ->
    ?LOG_ERROR("Got unexpected request from ~p: ~p", [From, Request]),
    {reply, {error, badarg}, State, get_timeout(State)}.

-spec handle_cast({notify, binary()} | term(), state())
      -> {noreply, state(), watchdog_timeout()}.
handle_cast({notify, Notification},
            #systemd_state{destination = undefined} = State) ->
    ?LOG_DEBUG("No NOTIFY_SOCKET, dropping ~s notification", [Notification]),
    {noreply, State, get_timeout(State)};
handle_cast({notify, Notification}, State) ->
    try notify(State, Notification)
    catch _:Err ->
            ?LOG_ERROR("Cannot send ~s notification: ~p", [Notification, Err])
    end,
    {noreply, State, get_timeout(State)};
handle_cast(Msg, State) ->
    ?LOG_ERROR("Got unexpected message: ~p", [Msg]),
    {noreply, State, get_timeout(State)}.

-spec handle_info(timeout | term(), state())
      -> {noreply, state(), watchdog_timeout()}.
handle_info(timeout, #systemd_state{interval = Interval} = State)
  when is_integer(Interval), Interval > 0 ->
    try notify(State, <<"WATCHDOG=1">>)
    catch _:Err ->
            ?LOG_ERROR("Cannot ping watchdog: ~p", [Err])
    end,
    {noreply, set_last_ping(State), Interval};
handle_info(Info, State) ->
    ?LOG_ERROR("Got unexpected info: ~p", [Info]),
    {noreply, State, get_timeout(State)}.

-spec terminate(normal | shutdown | {shutdown, term()} | term(), state()) -> ok.
terminate(Reason, #systemd_state{socket = undefined}) ->
    ?LOG_DEBUG("Terminating ~s (~p)", [?MODULE, Reason]),
    ok;
terminate(Reason, #systemd_state{socket = Socket}) ->
    ?LOG_DEBUG("Closing socket and terminating ~s (~p)", [?MODULE, Reason]),
    ok = gen_udp:close(Socket).

-spec code_change({down, term()} | term(), state(), term()) -> {ok, state()}.
code_change(_OldVsn, State, _Extra) ->
    ?LOG_INFO("Got code change request"),
    {ok, State}.

%% Internal functions.

-spec get_watchdog_interval() -> integer() | undefined.
get_watchdog_interval() ->
    case os:getenv("WATCHDOG_USEC") of
        WatchdogUSec when is_list(WatchdogUSec), length(WatchdogUSec) > 0 ->
            Interval = round(0.5 * list_to_integer(WatchdogUSec)),
            ?LOG_DEBUG("Watchdog interval: ~B microseconds", [Interval]),
            erlang:convert_time_unit(Interval, microsecond, millisecond);
        _ ->
            undefined
    end.

-spec get_timeout(state()) -> watchdog_timeout().
get_timeout(#systemd_state{interval = undefined}) ->
    ?LOG_DEBUG("Watchdog interval is undefined, hibernating"),
    hibernate;
get_timeout(#systemd_state{interval = Interval, last_ping = LastPing}) ->
    case Interval - (erlang:monotonic_time(millisecond) - LastPing) of
        Timeout when Timeout > 0 ->
            ?LOG_DEBUG("Calculated new timeout value: ~B", [Timeout]),
            Timeout;
        _ ->
            ?LOG_DEBUG("Calculated new timeout value: 1"),
            1
    end.

-spec set_last_ping(state()) -> state().
set_last_ping(State) ->
    LastPing = erlang:monotonic_time(millisecond),
    State#systemd_state{last_ping = LastPing}.

-spec notify(state(), binary()) -> ok.
notify(#systemd_state{socket = Socket, destination = Destination},
       Notification) ->
    ?LOG_DEBUG("Notifying systemd: ~s", [Notification]),
    ok = gen_udp:send(Socket, Destination, 0, Notification).

-spec cast_notification(binary()) -> ok.
cast_notification(Notification) ->
    ok = gen_server:cast(?MODULE, {notify, Notification}).