Permalink
Fetching contributors…
Cannot retrieve contributors at this time
282 lines (259 sloc) 9.48 KB
%% -------------------------------------------------------------------
%%
%% riaknostic - automated diagnostic tools for Riak
%%
%% Copyright (c) 2011 Basho Technologies, Inc. All Rights Reserved.
%%
%% This file is provided to you 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 Provides convenient access to Riak configuration values. When
%% the {@link riaknostic. riaknostic} module calls {@link
%% prepare/0. prepare/0}, Riak's <code>app.config</code> and
%% <code>vm.args</code> files will be parsed and memoized, and lager
%% will be started on the console at the configured severity level.
%% @end
-module(riaknostic_config).
-export([prepare/0,
data_directories/0,
get_app_env/1,
get_app_env/2,
get_vm_env/1,
base_dir/0,
etc_dir/0,
node_name/0,
cookie/0,
user/0]).
%% @doc Prepares appropriate configuration so the riaknostic script
%% can run. This is called by the riaknostic module and you do
%% not need to invoke it.
-spec prepare() -> ok | {error, iodata()}.
prepare() ->
prepare([fun start_lager/0, fun load_app_config/0, fun load_vm_args/0]).
prepare([]) ->
ok;
prepare([Fun|T]) ->
case Fun() of
{error, Reason} ->
{error, Reason};
_ ->
prepare(T)
end.
%% @doc Determines where Riak is configured to store data. Returns a
%% list of paths to directories defined by storage backends.
-spec data_directories() -> [ file:filename() ].
data_directories() ->
KVBackend = get_app_env([riak_kv, storage_backend]),
SearchBackend = get_app_env([riak_search, storage_backend], merge_index_backend),
Dirs = case get_app_env([riak_search, enabled]) of
true ->
data_directory(KVBackend) ++ data_directory(SearchBackend);
_ ->
data_directory(KVBackend)
end,
[ filename:absname(Dir, base_dir()) || Dir <- Dirs, Dir =/= undefined ].
%% @doc Get a key out of the app.config file, or if it doesn't exist,
%% return the Default. You specify a nested key using a list of atoms,
%% e.g. <code>[riak_kv, storage_backend]</code>.
%% @see get_app_env/1
-spec get_app_env([atom()], term()) -> term().
get_app_env(Keys, Default) ->
case get_app_env(Keys) of
undefined ->
Default;
Value ->
Value
end.
%% @doc Get a key out of the app.config file. You specify a nested
%% key using a list of atoms, e.g. <code>[riak_kv, storage_backend]</code>.
%% @equiv get_app_env(Keys, undefined)
-spec get_app_env([atom()]) -> undefined | term().
get_app_env(Keys) ->
{ok, Env} = application:get_env(riaknostic, app_config),
find_nested_key(Keys, Env).
%% @doc Get an -env flag out of the vm.args file.
-spec get_vm_env(string()) -> string() | undefined.
get_vm_env(Key) ->
case application:get_env(riaknostic, vm_env) of
undefined ->
undefined;
{ok, PList} ->
proplists:get_value(Key, PList)
end.
%% @doc Determines the user/uid that the installed Riak runs as.
-spec user() -> string().
user() ->
case application:get_env(riaknostic, user) of
{ok, Value} ->
Value;
_ -> undefined
end.
%% @doc The base directory from which the Riak script runs.
-spec base_dir() -> file:filename().
base_dir() ->
case application:get_env(riaknostic, base) of
undefined ->
filename:absname(".");
{ok, Path} ->
filename:absname(Path)
end.
%% @doc The Riak configuration directory.
-spec etc_dir() -> file:filename().
etc_dir() ->
case application:get_env(riaknostic, etc) of
undefined ->
filename:absname("./etc", base_dir());
{ok, Path} ->
filename:absname(Path, base_dir())
end.
%% @doc The local Riak node name. Includes whether the node uses short
%% or long nodenames for distributed Erlang.
-spec node_name() -> {shortnames | longnames, Name::string()}.
node_name() ->
case application:get_env(riaknostic, node_name) of
undefined ->
undefined;
{ok, Node} ->
Node
end.
%% @doc The Riak node's distributed Erlang cookie.
-spec cookie() -> atom().
cookie() ->
case application:get_env(riaknostic, cookie) of
undefined ->
undefined;
{ok, Cookie} ->
list_to_atom(Cookie)
end.
%% Private functions
start_lager() ->
application:load(lager),
case application:get_env(riaknostic, log_level) of
undefined ->
{error, "Log level not set!"};
{ok, Level} ->
application:set_env(lager, crash_log, undefined),
application:set_env(lager, handlers, [{lager_console_backend, Level}]),
lager:start()
end.
load_app_config() ->
case filelib:is_dir(etc_dir()) of
false ->
{error, "Cannot find Riak configuration directory!"};
true ->
AppConfig = filename:join([etc_dir(), "app.config"]),
case file:consult(AppConfig) of
{ok, [Config]} ->
application:set_env(riaknostic, app_config, Config);
_ ->
{error, io_lib:format("Riak config file ~s is malformed!", [AppConfig])}
end
end.
load_vm_args() ->
case filelib:is_dir(etc_dir()) of
false ->
{error, "Cannot find Riak configuration directory!"};
true ->
VmArgs = filename:join([etc_dir(), "vm.args"]),
case file:read_file(VmArgs) of
{error, Reason} ->
{error, io_lib:format("Could not read ~s, received error ~w!", [VmArgs, Reason])};
{ok, Binary} ->
load_vm_args(Binary)
end
end.
load_vm_args(Bin) when is_binary(Bin) ->
load_vm_args(re:split(Bin, "\s*\r?\n\s*", [{return, list}, trim]));
load_vm_args([]) ->
ok;
load_vm_args([[$#|_]|T]) ->
load_vm_args(T);
load_vm_args([""|T]) ->
load_vm_args(T);
load_vm_args(["-sname " ++ NodeName|T]) ->
application:set_env(riaknostic, node_name, {shortnames, string:strip(NodeName)}),
load_vm_args(T);
load_vm_args(["-name " ++ NodeName|T]) ->
application:set_env(riaknostic, node_name, {longnames, string:strip(NodeName)}),
load_vm_args(T);
load_vm_args(["-setcookie " ++ Cookie|T]) ->
application:set_env(riaknostic, cookie, string:strip(Cookie)),
load_vm_args(T);
load_vm_args(["-env " ++ Env|T]) ->
[Key, Value] = re:split(Env, "\s+", [{return, list}, trim]),
add_or_insert_env(vm_env, {Key, Value}),
load_vm_args(T);
load_vm_args([[$+|EmuFlags]|T]) ->
[Flag|Rest] = re:split(EmuFlags, "\s+", [{return,list}, trim]),
add_or_insert_env(emu_flags, {[$+|Flag], Rest}),
load_vm_args(T);
load_vm_args([[$-|InitFlags]|T]) ->
[Flag|Rest] = re:split(InitFlags, "\s+", [{return,list}, trim]),
add_or_insert_env(init_flags, {[$-|Flag], Rest}),
load_vm_args(T);
load_vm_args([Line|_]) ->
{error, io_lib:format("Erroneous line in vm.args: ~s", [Line])}.
add_or_insert_env(Key, Value) ->
case application:get_env(riaknostic, Key) of
undefined ->
application:set_env(riaknostic, Key, [Value]);
{ok, List} ->
application:set_env(riaknostic, Key, [Value|List])
end.
find_nested_key(_, undefined) ->
undefined;
find_nested_key([], Val) ->
Val;
find_nested_key([Key|T], PList) ->
find_nested_key(T, proplists:get_value(Key, PList)).
%% Determine the data directory(ies) for the configured storage backend
-spec data_directory(atom()) -> [ file:filename() ].
data_directory(riak_kv_bitcask_backend) ->
[ get_app_env([bitcask, data_root]) ];
data_directory(riak_kv_eleveldb_backend) ->
[ get_app_env([eleveldb, data_root]) ];
data_directory(merge_index_backend) ->
[ get_app_env([merge_index, data_root]) ];
data_directory(riak_kv_innostore_backend) ->
[ get_app_env([innostore, data_home_dir]),
get_app_env([innostore, log_group_home_dir]) ];
data_directory(riak_kv_multi_backend) ->
[ multi_data_directory(Backend) ||
Backend <- get_app_env([riak_kv, multi_backend]),
element(2, Backend) =/= riak_kv_memory_backend ];
data_directory(_) -> %% Memory or unknown backend
[].
%% Extracts data paths from multi_backend config
multi_data_directory({_, riak_kv_bitcask_backend, Props}) ->
case proplists:get_value(data_root, Props) of
undefined ->
get_app_env([bitcask, data_root]);
Path when is_list(Path) ->
Path
end;
multi_data_directory({_, riak_kv_eleveldb_backend, Props}) ->
case proplists:get_value(data_root, Props) of
undefined ->
get_app_env([eleveldb, data_root]);
Path when is_list(Path) ->
Path
end;
multi_data_directory({_, riak_kv_innostore_backend, Props}) ->
case proplists:get_value(data_home_dir, Props) of
undefined ->
get_app_env([innostore, data_home_dir]);
Path when is_list(Path) ->
Path
end.