%% The process half of albedo/harness/ssh: what each remote host looked like
%% when albedo last probed it, and the probes in flight. A probe is one run of
%% priv/python/albedo_ssh.py, the same ssh layer the model's remote plugin
%% uses, so both ride one ControlMaster per host.
%%
%% A ready answer is kept for a minute and a failure for five seconds; asking
%% again while a probe runs joins it instead of starting another.
-module(albedo_ssh).
-behaviour(gen_server).
-export([probe/2, peek/1, cached/1, step/1, forget/1, commands/2, exec/5, config_hosts/0]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2]).
-define(READY_MS, 60000).
-define(FAILED_MS, 5000).
-define(HELPER_MS, 170000). %% the helper's own probe and staging deadlines, and some
%% A fresh answer, or the one a probe started (or joined) now gives within
%% WaitMs. {error, nil} means it is still warming.
probe(Target, WaitMs) ->
try gen_server:call(ensure(), {probe, Target}, WaitMs)
catch exit:{timeout, _} -> {error, nil}
end.
%% The cached answer, without waiting: a stale or missing one starts a probe
%% and reads as still warming.
peek(Target) -> gen_server:call(ensure(), {peek, Target}).
%% A fresh cached answer, never starting a probe.
cached(Target) -> gen_server:call(ensure(), {cached, Target}).
%% What a running probe is doing past connecting: <<"staging">> while it
%% copies the bundle over, <<>> otherwise.
step(Target) -> gen_server:call(ensure(), {step, Target}).
%% The Host names of ~/.ssh/config and its includes, as a JSON list.
config_hosts() ->
case helper([<<"hosts">>], 5000) of
{ok, Json} -> {ok, Json};
{error, _} -> {error, nil}
end.
forget(Target) -> gen_server:cast(ensure(), {forget, Target}), nil.
%% The commands for a host whose home is known, built locally (no network).
commands(Target, Home) ->
case helper([<<"commands">>, Target, Home], 10000) of
{ok, Json} -> {ok, Json};
{error, _} -> {error, nil}
end.
%% One remote command over ssh (Argv ends with the target), Input written to
%% its stdin: its stdout once it exits 0 within TimeoutMs, else why not.
exec(Argv, Command, AuthSock, Input, TimeoutMs) ->
[Ssh | Options] = Argv,
Env = case AuthSock of
{some, Sock} -> [{"SSH_AUTH_SOCK", binary_to_list(Sock)} | albedo_python:clean_environment()];
none -> albedo_python:clean_environment()
end,
case os:find_executable(binary_to_list(Ssh)) of
false -> {error, <<"ssh not found on PATH">>};
Exe ->
try open_port({spawn_executable, Exe},
[binary, exit_status, use_stdio, hide,
{args, [binary_to_list(A) || A <- Options ++ [Command]]},
{env, Env}]) of
Port ->
_ = (catch port_command(Port, Input)),
case collect(Port, [], erlang:monotonic_time(millisecond) + TimeoutMs) of
{ok, Out} -> {ok, Out};
{error, Why} -> {error, unicode:characters_to_binary(Why)}
end
catch _:Reason -> {error, unicode:characters_to_binary(io_lib:format("~p", [Reason]))}
end
end.
ensure() ->
case gen_server:start({local, ?MODULE}, ?MODULE, [], []) of
{ok, Pid} -> Pid;
{error, {already_started, Pid}} -> Pid
end.
init([]) -> {ok, #{hosts => #{}, running => #{}, steps => #{}}}.
handle_call({probe, Target}, From, S) ->
case fresh(Target, S) of
{ok, Json} -> {reply, {ok, Json}, S};
none -> {noreply, start(Target, From, S)}
end;
handle_call({cached, Target}, _, S) ->
case fresh(Target, S) of
{ok, Json} -> {reply, {ok, Json}, S};
none -> {reply, {error, nil}, S}
end;
handle_call({step, Target}, _, S = #{steps := Steps}) ->
{reply, maps:get(Target, Steps, <<>>), S};
handle_call({peek, Target}, _, S) ->
case fresh(Target, S) of
{ok, Json} -> {reply, {ok, Json}, S};
none -> {reply, {error, nil}, start(Target, none, S)}
end.
handle_cast({forget, Target}, S = #{hosts := Hosts}) ->
{noreply, S#{hosts => maps:remove(Target, Hosts)}}.
handle_info({probed, Target, Json}, S = #{hosts := Hosts, running := Running, steps := Steps}) ->
Waiters = maps:get(Target, Running, []),
[gen_server:reply(From, {ok, Json}) || From <- Waiters],
Now = erlang:monotonic_time(millisecond),
{noreply, S#{hosts => Hosts#{Target => {Json, Now + ttl(Json)}},
running => maps:remove(Target, Running),
steps => maps:remove(Target, Steps)}};
handle_info({stepped, Target, Step}, S = #{running := Running, steps := Steps}) ->
case maps:is_key(Target, Running) of
true -> {noreply, S#{steps => Steps#{Target => Step}}};
false -> {noreply, S}
end;
handle_info(_, S) -> {noreply, S}.
fresh(Target, #{hosts := Hosts}) ->
Now = erlang:monotonic_time(millisecond),
case maps:get(Target, Hosts, none) of
{Json, Until} when Until > Now -> {ok, Json};
_ -> none
end.
start(Target, From, S = #{running := Running}) ->
Waiters = [W || W <- [From], W =/= none],
case maps:find(Target, Running) of
{ok, Earlier} -> S#{running => Running#{Target => Waiters ++ Earlier}};
error ->
Self = self(),
spawn(fun() -> Self ! {probed, Target, run(Target, Self)} end),
S#{running => Running#{Target => Waiters}}
end.
ttl(Json) ->
case catch json:decode(Json) of
#{<<"state">> := <<"ready">>} -> ?READY_MS;
_ -> ?FAILED_MS
end.
%% The probe's answer is its last line; a step line before it is passed on
%% to the server as it arrives.
run(Target, Server) ->
Stepped = fun(Line) ->
case catch json:decode(Line) of
#{<<"step">> := Step} when is_binary(Step) -> Server ! {stepped, Target, Step};
_ -> ok
end
end,
case helper([<<"probe">>, Target], ?HELPER_MS, Stepped) of
{ok, Out} ->
case [L || L <- binary:split(Out, <<"\n">>, [global]), L =/= <<>>] of
[] -> failed("albedo_ssh.py answered nothing");
Lines -> lists:last(Lines)
end;
{error, Why} -> failed(Why)
end.
%% One run of albedo_ssh.py: its stdout, or why there is none. Given an
%% OnLine, each line goes to it as soon as it is complete.
helper(Args, Wait) -> helper(Args, Wait, none).
helper(Args, Wait, OnLine) ->
case albedo_python:local_paths() of
{ok, {Python, Kernel}} ->
Helper = filename:join(filename:dirname(Kernel), <<"albedo_ssh.py">>),
try open_port({spawn_executable, binary_to_list(Python)},
[binary, exit_status, use_stdio, hide,
{args, ["-u", binary_to_list(Helper) | [binary_to_list(A) || A <- Args]]},
{env, albedo_python:clean_environment()}]) of
Port -> collect(Port, [], erlang:monotonic_time(millisecond) + Wait, OnLine, <<>>)
catch _:Reason -> {error, io_lib:format("~p", [Reason])}
end;
{error, {unavailable, Why}} -> {error, Why}
end.
collect(Port, Acc, Deadline) -> collect(Port, Acc, Deadline, none, <<>>).
%% Without OnLine nothing is split, so a large one-line answer (a gather)
%% is never copied again per chunk.
collect(Port, Acc, Deadline, OnLine, Partial) ->
receive
{Port, {data, Data}} when OnLine =:= none ->
collect(Port, [Data | Acc], Deadline, none, Partial);
{Port, {data, Data}} ->
[Rest | Lines] = lists:reverse(binary:split(<<Partial/binary, Data/binary>>, <<"\n">>, [global])),
lists:foreach(OnLine, lists:reverse(Lines)),
collect(Port, [Data | Acc], Deadline, OnLine, Rest);
{Port, {exit_status, 0}} -> {ok, iolist_to_binary(lists:reverse(Acc))};
{Port, {exit_status, Status}} -> {error, io_lib:format("albedo_ssh.py exited ~p", [Status])}
after max(0, Deadline - erlang:monotonic_time(millisecond)) ->
_ = try port_close(Port) catch _:_ -> ok end,
{error, "albedo_ssh.py timed out"}
end.
failed(Why) ->
iolist_to_binary(json:encode(#{state => <<"unreachable">>,
detail => unicode:characters_to_binary(Why)})).