From e17424598c4bb8293c25dddda35cd7a528a791c1 Mon Sep 17 00:00:00 2001 From: Jonathan Freedman Date: Tue, 27 Mar 2012 13:37:24 -0700 Subject: [PATCH] * bits and pieces of usability improvement. better errors, tweaked command line args. ability to whitelist apps for thunderbeam --- include/magicbeam.hrl | 2 +- include/magicbeam_config.hrl | 3 +- src/hotbeam.erl | 13 ++++---- src/magicbeam.erl | 61 ++++++++++++++++++++++-------------- src/magicbeam_app.erl | 12 ++++--- src/magicbeam_util.erl | 24 +++++++++++--- src/shellbeam.erl | 4 +-- src/thunderbeam.erl | 37 ++++++++++++---------- 8 files changed, 98 insertions(+), 58 deletions(-) diff --git a/include/magicbeam.hrl b/include/magicbeam.hrl index f07ab42..e82bf42 100644 --- a/include/magicbeam.hrl +++ b/include/magicbeam.hrl @@ -6,7 +6,7 @@ -record(hotbeam, {mod, beam, src, beam_time=0, src_time=0, last=0}). -record(hotbeam_state, {enable, src, apps=[], beams=[], tref=undefined, scantime=0, hotload_count=0, compile_count=0}). --record(thunderbeam_state, {enabled, tref=undefined, killed=0, base, variable, immune_app=[], immune_proc=[]}). +-record(thunderbeam_state, {enabled, tref=undefined, killed=0, base, variable, immune_app=[], immune_proc=[], force_apps=[]}). -define(enow(), calendar:datetime_to_gregorian_seconds(calendar:universal_time())). diff --git a/include/magicbeam_config.hrl b/include/magicbeam_config.hrl index 837a827..dd2a041 100644 --- a/include/magicbeam_config.hrl +++ b/include/magicbeam_config.hrl @@ -6,6 +6,7 @@ -define(THUNDERBEAM_IMMUNE_PROC, ?THUNDERBEAM_NEVER_KILL_PROC ++ magicbeam_srv:cfgget(thunderbeam_immune_proc, [])). -define(THUNDERBEAM_IMMUNE_APP, magicbeam_srv:cfgget(thunderbeam_immune_app, [stdlib,kernel,mnesia,sasl,inets])). -define(THUNDERBEAM_FORCE_KILL, magicbeam_srv:cfgget(thunderbeam_force_kill, false)). +-define(THUNDERBEAM_FORCE_APPS, magicbeam_srv:cfgget(thunderbeam_force_app, [])). -define(HOTBEAM_ENABLED, magicbeam_srv:cfgget(hotbeam_enabled, false)). -define(HOTBEAM_COMPILE, magicbeam_srv:cfgget(hotbeam_compile, false)). @@ -13,5 +14,5 @@ -define(HOTBEAM_LOOP, 1000). -define(SHELLBEAM_ANSI, magicbeam_srv:cfgget(shellbeam_ansi, false)). --define(SHELLBEAM_PROMPT, magicbeam_srv:cfgget(shellbeam_prompt, atom_to_list(node()) ++ "OTP 4 LYFE")). +-define(SHELLBEAM_PROMPT, magicbeam_srv:cfgget(shellbeam_prompt, atom_to_list(node()) ++ " OTP 4 LYFE")). -define(SHELLBEAM_MODULES, magicbeam_srv:cfgget(shellbeam_modules, [magicbeam_shell])). diff --git a/src/hotbeam.erl b/src/hotbeam.erl index 2615f0f..e28d2ab 100644 --- a/src/hotbeam.erl +++ b/src/hotbeam.erl @@ -234,13 +234,14 @@ p_rescan([App | Apps], #hotbeam_state{} = State) -> end. p_filetime(File) -> - case {info, file:read_file_info(File)} of - {info, {ok, #file_info{mtime = MTime}}} -> - case {dst, calendar:local_time_to_universal_time_dst(MTime)} of - {dst, [UTC]} -> UTC; - {dst, [_, UTC]} -> UTC + case file:read_file_info(File) of + {ok, #file_info{mtime = MTime}} -> + case calendar:local_time_to_universal_time_dst(MTime) of + [UTC] -> UTC; + [_, UTC] -> UTC end; - {info, {error, enoent}} -> error + {error, enoent} -> error; + {error, enotdir} -> error end. p_lazy_load(#hotbeam{mod = Mod} = HB, #hotbeam_state{compile_count = CC, beams = Beams} = State) -> diff --git a/src/magicbeam.erl b/src/magicbeam.erl index 2801502..9c02d77 100644 --- a/src/magicbeam.erl +++ b/src/magicbeam.erl @@ -3,8 +3,11 @@ -export([main/1, behaviour_info/1, rehash/0, start/0, stop/0, info/0, rpc_shell/1, remote_shell/1]). -start() -> application:start(magicbeam). -stop() -> application:stop(magicbeam). +start() -> + ok = magicbeam_util:start_deps(), + ok = application:start(magicbeam). + +stop() -> ok = application:stop(magicbeam). info() -> [ @@ -48,10 +51,10 @@ options() -> {unload, $u, "unload", undefined, "Unload application from remote node."}, {help, $h, "help", undefined, "Receive Help."}, {node, $n, "node", string, "Remote Node."}, - {short, $s, "short", undefined, "Use short names."}, - {cookie, $c, "cookie", string, "Cookie to Use."}, - {callback, $m, "module", string, "Callback module to register with."}, - {shell, $r, "shell", undefined, "Start shell on remote node."} + {short, undefined, "short", undefined, "Use short names."}, + {setcookie, $c, "cookie", string, "Cookie to Use."}, + {callback, undefined, "module", string, "Callback module to register with."}, + {shell, $s, "shell", undefined, "Start shell on remote node."} ]. maybe_work(Opts) -> @@ -64,13 +67,14 @@ maybe_work(Opts) -> case {lists:member(load, Opts), lists:member(unload, Opts)} of {true, false} -> load(Node, Opts); {false, true} -> unload(Node, Opts); - {_, _} -> - case lists:member(shell, Opts) of - true -> shell(Node, Opts); - false -> - help(), - erlang:halt(1) - end + {_, false} -> maybe_shell(Node, Opts) + end. + +maybe_shell(Node, Opts) -> + case lists:member(shell, Opts) of + true -> shell(Node, Opts); + false -> + erlang:halt(0) end. init(Node, Opts) -> @@ -80,19 +84,27 @@ init(Node, Opts) -> true -> [shortnames] end, ok = case net_kernel:start(NKSO) of - {ok, _PID} -> ok + {ok, _PID} -> ok; + {error, {already_started, _PID}} -> ok end, - case proplists:get_value(cookie, Opts) of - undefined -> ok; + case proplists:get_value(setcookie, Opts) of + undefined -> connect(Node); Cookie -> true = erlang:set_cookie(Node, list_to_atom(Cookie)), - case net_kernel:hidden_connect(Node) of - true -> ok; - false -> - magicbeam_util:error_out("Unable to establish connection with " ++ atom_to_list(Node)) - end + connect(Node) end. +connect(Node) -> + EF = fun() -> magicbeam_util:error_out("Unable to establish connection with " ++ atom_to_list(Node)) end, + case net_kernel:hidden_connect(Node) of + true -> + case net_adm:ping(Node) of + pong -> ok; + pang -> EF() + end; + false -> + EF() + end. load(Node, Opts) -> ok = init(Node, Opts), @@ -102,7 +114,7 @@ load(Node, Opts) -> end, ok = magicbeam_util:inject(Node, CB), io:format("Injected magicbeam.~n"), - erlang:halt(0). + maybe_shell(Node, Opts). unload(Node, Opts) -> ok = init(Node, Opts), @@ -112,7 +124,10 @@ unload(Node, Opts) -> shell(Node, Opts) -> ok = init(Node, Opts), - remote_shell(Node). + case magicbeam_util:loaded(Node) of + true -> remote_shell(Node); + false -> magicbeam_util:error_out("Application not loaded on " ++ atom_to_list(Node)) + end. remote_shell(Node) -> block_until_done(user_drv:start(['tty_sl -c -e', {magicbeam, rpc_shell, [Node]}])), diff --git a/src/magicbeam_app.erl b/src/magicbeam_app.erl index a1a737b..0b3c712 100644 --- a/src/magicbeam_app.erl +++ b/src/magicbeam_app.erl @@ -7,9 +7,11 @@ -record(state, {ssh}). rpc_start(AppSpec) -> - lists:foreach(fun(A) -> application:start(A) end, proplists:get_value(applications, AppSpec)), ok = application:load(AppSpec), - ok = application:start(magicbeam). + case magicbeam_util:start_deps() of + ok -> ok = application:start(magicbeam); + {error, _} = E -> E + end. rpc_stop() -> application:stop(magicbeam), @@ -84,9 +86,9 @@ ssh_verify_path(Path) -> case {F("authorized_keys"), F("ssh_host_dsa_key"), F("ssh_host_rsa_key")} of {true, true, true} -> Path; - {false, true, true} -> - ?warn("missing file ~s", [Path ++ "authorized_keys"]), - Path; + {false, true, true} -> + ?warn("missing file ~s", [Path ++ "authorized_keys"]), + Path; {A, B, C} -> ?error("missing magicbeam ssh files in path ~p authorized_keys:~p ssh_host_dsa_key:~p ssh_host_rsa_key:~p", [Path, A, B, C]), undefined diff --git a/src/magicbeam_util.erl b/src/magicbeam_util.erl index 6c2dbf0..2bfcc5d 100644 --- a/src/magicbeam_util.erl +++ b/src/magicbeam_util.erl @@ -2,7 +2,7 @@ -include("magicbeam.hrl"). -author('jonafree@gmail.com'). --export([appenv/2, inject/2, remove/1, error_out/1, loaded/1]). +-export([appenv/2, inject/2, remove/1, error_out/1, loaded/1, start_deps/0]). appenv(Key, Default) -> case application:get_env(magicbeam, Key) of @@ -22,7 +22,8 @@ inject(Node, Mod) -> inject(Node, Mod, []) -> case rpc(Node, magicbeam_app, rpc_start, [app_spec(Mod)]) of ok -> ok; - error -> error_out("Unable to start application on " ++ atom_to_list(Node)) + {error, {magicbeam, {start_dep, App, _}}} -> error_out("Unable to start dependency " ++ atom_to_list(App) ++ " on " ++ atom_to_list(Node)); + E when E == error ; is_tuple(E) -> error_out("Unable to start application on " ++ atom_to_list(Node)) end; inject(Node, Mod, [M | Tail]) when is_atom(Mod) -> {M, Bin, FName} = code:get_object_code(M), @@ -57,7 +58,8 @@ remove(Node, [Mod | Mods]) when is_atom(Mod) -> rpc(Node, M, F, A) -> case rpc:call(Node, M, F, A) of {badrpc, _E} -> error; - {error, _E} -> error; + {error, {magicbeam, _}} = E -> E; + {error, _} -> error; Value -> Value end. @@ -74,5 +76,19 @@ loaded(Node) -> false -> false end; - error -> false + error -> false + end. + +start_deps() -> + case application:get_key(magicbeam, applications) of + undefined -> {error, {magicbeam, not_loaded}}; + {ok, Deps} -> start_dep(Deps) + end. + +start_dep([]) -> ok; +start_dep([App | Deps]) -> + case application:start(App) of + ok -> start_dep(Deps); + {error,{already_started,App}} -> start_dep(Deps); + {error, R} -> {error, {magicbeam, {start_dep, App, R}}} end. diff --git a/src/shellbeam.erl b/src/shellbeam.erl index f7ed3c2..2173938 100644 --- a/src/shellbeam.erl +++ b/src/shellbeam.erl @@ -131,10 +131,10 @@ command_match([{_, _, optional} | MT], [_H | _] = TT, Ar) -> command_match([{_, auto} | MT], [H | TT], Ar) -> case catch list_to_integer(H) of I when is_integer(I) -> command_match(MT, TT, Ar ++ [I]); - {'EXIT',{badarg,[{erlang,list_to_integer,[H]} | _]}} -> + {'EXIT',{badarg, _}} -> case catch list_to_existing_atom(H) of A when is_atom(A) -> command_match(MT, TT, Ar ++ [A]); - {'EXIT',{badarg,[{erlang,list_to_existing_atom,[H]} | _]}} -> command_match(MT, TT, Ar ++ [H]) + {'EXIT',{badarg,_}} -> command_match(MT, TT, Ar ++ [H]) end end; command_match(_, _, _) -> false. diff --git a/src/thunderbeam.erl b/src/thunderbeam.erl index 1d43d24..955bb95 100644 --- a/src/thunderbeam.erl +++ b/src/thunderbeam.erl @@ -2,6 +2,10 @@ -behaviour(gen_server). -author('jonafree@gmail.com'). +-ifdef(TEST). +-compile(export_all). +-endif. + -include("magicbeam.hrl"). -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]). @@ -29,24 +33,25 @@ handle_cast(rehash, State) -> {noreply, p_rehash(State)}. handle_info(loop, State) -> {noreply, p_timer(p_kill(State))}; handle_info(_, State) -> {noreply, State}. -terminate(_Reason, _State) -> +terminate(_Reason, _State) -> ok. code_change(_Old, State, _Extra) -> {ok, State}. p_rehash(State) -> p_timer(State#thunderbeam_state{ - enabled = ?THUNDERBEAM_ENABLED, - base = ?THUNDERBEAM_WAIT_BASE, - variable = ?THUNDERBEAM_WAIT_VARIABLE, - immune_proc = ?THUNDERBEAM_IMMUNE_PROC, - immune_app = ?THUNDERBEAM_IMMUNE_APP - }). + enabled = ?THUNDERBEAM_ENABLED, + base = ?THUNDERBEAM_WAIT_BASE, + variable = ?THUNDERBEAM_WAIT_VARIABLE, + immune_proc = ?THUNDERBEAM_IMMUNE_PROC, + immune_app = ?THUNDERBEAM_IMMUNE_APP, + force_apps = ?THUNDERBEAM_FORCE_APPS + }). p_info(#thunderbeam_state{enabled=Enabled, killed=Killed}) -> [ - {enabled, Enabled}, - {killed, Killed} + {enabled, Enabled}, + {killed, Killed} ]. p_timer(#thunderbeam_state{tref=TRef, enabled = true} = State) when is_tuple(TRef) -> @@ -72,7 +77,7 @@ p_kill(State, Attempts, [PID| PIDS], Count) when ( length(PIDS) + 1 ) == Count - p_kill1(State, Attempts, PID); p_kill(State, Attempts, [_PID | PIDS], Count) -> p_kill(State, Attempts, PIDS, Count). -p_kill1(State, Attempts, PID) -> p_kill1(State, Attempts, PID, erlang:process_info(PID, [registered_name])). +p_kill1(State, Attempts, PID) when is_pid(PID) -> p_kill1(State, Attempts, PID, erlang:process_info(PID, [registered_name])). p_kill1(State, Attempts, PID, [{registered_name, []}]) -> p_kill2(State, Attempts, PID, "N/A"); p_kill1(#thunderbeam_state{immune_proc = IP} = State, Attempts, PID, [{registered_name, Name}]) when is_atom(Name) -> case lists:member(Name, IP) of @@ -83,18 +88,18 @@ p_kill1(#thunderbeam_state{immune_proc = IP} = State, Attempts, PID, [{registere p_kill2(State, Attempts, PID, Name) -> p_kill2(application:get_application(PID), Attempts, State, PID, Name). p_kill2(undefined, _Attempts, State, PID, Name) -> p_kill3(State, PID, Name); -p_kill2({ok, Application}, Attempts, #thunderbeam_state{immune_app = IA} = State, PID, Name) -> - case lists:member(Application, IA) of - false -> p_kill3(State, PID, Name); - true -> - p_kill(State, Attempts - 1) +p_kill2({ok, Application}, Attempts, #thunderbeam_state{immune_app = IA, force_apps = FA} = State, PID, Name) -> + case {lists:member(Application, FA), lists:member(Application, IA)} of + {true, true} -> p_kill3(State, PID, Name); + {false, false} -> p_kill3(State, PID, Name); + _ -> p_kill(State, Attempts - 1) end. p_kill3(#thunderbeam_state{killed = Killed} = State, PID, Name) -> p_kill_warn(PID, Name), case process_info(PID, trap_exit) of {trap_exit, true} -> p_kill_trap_exit(PID, Name, State); - {trap_exit, false} -> + {trap_exit, false} -> erlang:exit(PID, thunderbeam), p_kill_event(PID, Name), State#thunderbeam_state{killed = Killed + 1}