aae_progress_report and other functions needed for Tictac AAE CLI treestatus subcommand (#4)

* aae_schedulenextrebuild/2

* get/setters of various state fields needed for riak admin tictacaae

* touch object_splitfun directly for purposes of get/setting storeheads

this, to avoid introducing dependency on riak_kv, where
riak_object:aae_from_object_binary/1 lives, which actually creates the
object_splitfun function.

* ct test for get/set_rebuild_schedule

* ct for get/set_storeheads

* uncomment, unbreak other tests in basic_SUITE

* remove extra, redundant steps in get_set_storeheads

* avoid exposing internal pids, do produce_aae_progress_report on site

* more expressive function names for aae progress report work

* aae_schedulenextrebuild/2 better named aae_prompt_nextrebuild

* avoid passing entire state record for funs that only need some fields

* prompt_nextrebuild done right, with a test

* callers of produce_report should assume key_store is always a good pid

* splitfun get/set ct test

---------

Co-authored-by: Andriy Zavada <andriy.zavada@tiot.jp>

Author: 2 people

src/aae_controller.erl

@@ -17,40 +17,46 @@
 -include("include/aae.hrl").
 
 -export([init/1,
-            handle_call/3,
-            handle_cast/2,
-            handle_info/2,
-            terminate/2,
-            code_change/3]).
+         handle_call/3,
+         handle_cast/2,
+         handle_info/2,
+         terminate/2,
+         code_change/3]).
 
 -export([aae_start/6,
-            aae_start/7,
-            aae_start/8,
-            aae_nextrebuild/1,
-            aae_put/7,
-            aae_close/1,
-            aae_destroy/1,
-            aae_fetchroot/3,
-            aae_mergeroot/3,
-            aae_fetchbranches/4,
-            aae_mergebranches/4,
-            aae_fetchclocks/5,
-            aae_fetchclocks/7,
-            aae_rebuildtrees/5,
-            aae_rebuildtrees/4,
-            aae_rebuildstore/2,
-            aae_rebuildstore/3,
-            aae_fold/6,
-            aae_fold/8,
-            aae_bucketlist/1,
-            aae_loglevel/2,
-            aae_ping/3,
-            aae_runnerprompt/1
+         aae_start/7,
+         aae_start/8,
+         aae_nextrebuild/1,
+         aae_set_rebuild_schedule/2,
+         aae_get_rebuild_schedule/1,
+         aae_prompt_nextrebuild/2,
+         aae_get_object_splitfun/1,
+         aae_set_object_splitfun/2,
+         aae_put/7,
+         aae_close/1,
+         aae_destroy/1,
+         aae_fetchroot/3,
+         aae_mergeroot/3,
+         aae_fetchbranches/4,
+         aae_mergebranches/4,
+         aae_fetchclocks/5,
+         aae_fetchclocks/7,
+         aae_rebuildtrees/5,
+         aae_rebuildtrees/4,
+         aae_rebuildstore/2,
+         aae_rebuildstore/3,
+         aae_fold/6,
+         aae_fold/8,
+         aae_bucketlist/1,
+         aae_loglevel/2,
+         aae_ping/3,
+         aae_runnerprompt/1,
+         aae_produce_progress_report/1
         ]).
 
 -export([foldobjects_buildtrees/2,
-            hash_clocks/2,
-            wrapped_splitobjfun/1]).
+         hash_clocks/2,
+         wrapped_splitobjfun/1]).
 
 -export([wait_on_sync/5]).
 
@@ -215,6 +221,37 @@ aae_start(
 aae_nextrebuild(Pid) ->
     gen_server:call(Pid, rebuild_time, ?SYNC_TIMEOUT).
 
+-spec aae_prompt_nextrebuild(pid(), non_neg_integer()) -> ok.
+%% @doc
+%% Set the next_rebuild time to now + given delay
+aae_prompt_nextrebuild(Pid, Delay) ->
+    gen_server:call(Pid, {prompt_nextrebuild, Delay}, ?SYNC_TIMEOUT).
+
+-spec aae_get_rebuild_schedule(pid()) -> rebuild_schedule().
+%% @doc
+%% Get rebuild schedule
+aae_get_rebuild_schedule(Pid) ->
+    gen_server:call(Pid, get_rebuild_schedule, ?SYNC_TIMEOUT).
+
+-spec aae_set_rebuild_schedule(pid(), rebuild_schedule()) -> ok.
+%% @doc
+%% Set rebuild schedule
+aae_set_rebuild_schedule(Pid, RS) ->
+    gen_server:call(Pid, {set_rebuild_schedule, RS}, ?SYNC_TIMEOUT).
+
+-spec aae_get_object_splitfun(pid()) -> function().
+%% @doc
+%% Get object_splitfun field. It is only used to infer the value of 'storeheads'.
+aae_get_object_splitfun(Pid) ->
+    gen_server:call(Pid, get_object_splitfun, ?SYNC_TIMEOUT).
+
+-spec aae_set_object_splitfun(pid(), function()) -> ok.
+%% @doc
+%% Set object_splitfun. Used to emulate re-initing the controller
+%% with a different value of 'storeheads'.
+aae_set_object_splitfun(Pid, A) ->
+    gen_server:call(Pid, {set_object_splitfun, A}, ?SYNC_TIMEOUT).
+
 -spec aae_put(
     pid(),
     responsible_preflist(), 
@@ -481,6 +518,15 @@ aae_ping(Pid, RequestTime, From) ->
 aae_runnerprompt(Pid) ->
     gen_server:cast(Pid, runner_prompt).
 
+-spec aae_produce_progress_report(pid()) -> list({atom(), term()}).
+%% @doc
+%% Generate a 'progress report', with items such as status,
+%% last_rebuild, next_rebuild, for the aae controller handling this
+%% partition. This report can be displayed by `riak admin tictacaae
+%% treestatus` command.
+aae_produce_progress_report(Pid) ->
+    gen_server:call(Pid, produce_report).
+
 
 %%%============================================================================
 %%% gen_server callbacks
@@ -585,6 +631,17 @@ init([Opts]) ->
 
 handle_call(rebuild_time, _From, State) ->  
     {reply, State#state.next_rebuild, State};
+handle_call(get_rebuild_schedule, _From, State = #state{rebuild_schedule = RS}) ->
+    {reply, RS, State};
+handle_call({set_rebuild_schedule, RS}, _From, State) ->
+    {reply, ok, State#state{rebuild_schedule = RS}};
+handle_call(get_object_splitfun, _From, State = #state{object_splitfun = A}) ->
+    {reply, A, State};
+handle_call({set_object_splitfun, A}, _From, State) ->
+    {reply, ok, State#state{object_splitfun = A}};
+handle_call({prompt_nextrebuild, SecsFromNow}, _From, State) ->
+    {Mega, Sec, Micros} = os:timestamp(),
+    {reply, ok, State#state{next_rebuild = {Mega, Sec + SecsFromNow, Micros}}};
 handle_call(close, _From, State) ->
     ok = maybe_flush_puts(State#state.key_store, 
                             State#state.objectspecs_queue,
@@ -945,6 +1002,11 @@ handle_call(bucket_list,  _From, State) ->
                             true),
     R = aae_keystore:store_bucketlist(State#state.key_store),
     {reply, R, State};
+handle_call(produce_report, _From, State = #state{key_store = KeyStore,
+                                                  next_rebuild = NextRebuild,
+                                                  tree_caches = TreeCaches}) ->
+    R = produce_report(KeyStore, NextRebuild, TreeCaches),
+    {reply, R, State};
 handle_call({ping, RequestTime}, _From, State) ->
     T = max(0, timer:now_diff(os:timestamp(), RequestTime)),
     aae_util:log(aae15, [T div 1000], State#state.log_levels),
@@ -1154,6 +1216,16 @@ wrapped_splitobjfun(ObjectSplitFun) ->
         end
     end.
 
+produce_report(KeyStore, NextRebuild, TreeCaches) ->
+    TotalDirtySegments =
+        lists:sum(
+          [aae_treecache:cache_segment_count(P) || {_, P} <- TreeCaches]),
+    [{last_rebuild, aae_keystore:store_last_rebuild(KeyStore)},
+     {next_rebuild, NextRebuild},
+     {total_dirty_segments, TotalDirtySegments}
+    ].
+
+
 %%%============================================================================
 %%% Internal functions
 %%%============================================================================

src/aae_keystore.erl

@@ -58,6 +58,7 @@
         store_fold/8,
         store_fetchclock/3,
         store_bucketlist/1,
+        store_last_rebuild/1,
         store_loglevel/2
     ]
 ).
@@ -442,6 +443,12 @@ store_fetchclock(Pid, Bucket, Key) ->
 store_bucketlist(Pid) ->
     gen_fsm:sync_send_all_state_event(Pid, bucket_list, infinity).
 
+-spec store_last_rebuild(pid()) -> erlang:timestamp() | never.
+%% @doc
+%% Get the last rebuild time
+store_last_rebuild(Pid) ->
+    gen_fsm:sync_send_all_state_event(Pid, last_rebuild, infinity).
+
 -spec store_loglevel(pid(), aae_util:log_levels()) -> ok.
 %% @doc
 %% Alter the log level at runtime
@@ -670,6 +677,8 @@ native({prompt, rebuild_complete}, State) ->
 handle_sync_event(bucket_list, _From, StateName, State) ->
     Folder = bucket_list(State#state.store_type, State#state.store),
     {reply, Folder, StateName, State};
+handle_sync_event(last_rebuild, _From, StateName, State = #state{last_rebuild = A}) ->
+    {reply, A, StateName, State};
 handle_sync_event(current_status, _From, StateName, State) ->
     {reply, {StateName, State#state.current_guid}, StateName, State};
 handle_sync_event(ping, _From, StateName, State) ->
@@ -1291,7 +1300,6 @@ disklog_filename(RootPath, GUID) ->
     filename:join(RootPath, GUID ++ ?DISKLOG_EXT).
 
 
-
 %%%============================================================================
 %%% Test
 %%%============================================================================
@@ -1859,4 +1867,4 @@ generate_objectspecs(Op, B, KeyList) ->
 
 
 
--endif.
+-endif.

src/aae_treecache.erl

@@ -17,17 +17,18 @@
             format_status/1]).
 
 -export([cache_open/3,
-            cache_new/3,
-            cache_alter/4,
-            cache_root/1,
-            cache_leaves/2,
-            cache_markdirtysegments/3,
-            cache_replacedirtysegments/3,
-            cache_destroy/1,
-            cache_startload/1,
-            cache_completeload/2,
-            cache_loglevel/2,
-            cache_close/1]).
+         cache_new/3,
+         cache_alter/4,
+         cache_root/1,
+         cache_leaves/2,
+         cache_markdirtysegments/3,
+         cache_replacedirtysegments/3,
+         cache_destroy/1,
+         cache_startload/1,
+         cache_completeload/2,
+         cache_loglevel/2,
+         cache_close/1,
+         cache_segment_count/1]).
 
 -define(PENDING_EXT, ".pnd").
 -define(FINAL_EXT, ".aae").
@@ -85,6 +86,12 @@ cache_new(RootPath, PartitionID, LogLevels) ->
 cache_destroy(AAECache) ->
     gen_server:cast(AAECache, destroy).
 
+-spec cache_segment_count(pid()) -> non_neg_integer().
+%% @doc
+%% Expose dirty_segments length, for aae-progress-report.
+cache_segment_count(AAECache) ->
+    gen_server:call(AAECache, segment_count, ?SYNC_TIMEOUT).
+
 -spec cache_close(pid()) -> ok.
 %% @doc
 %% Close a cache with saving
@@ -193,14 +200,16 @@ init([Opts]) ->
         },
         hibernate
     }.
-    
+
 
 handle_call(is_restored, _From, State) ->
     {reply, State#state.is_restored, State};
 handle_call(fetch_root, _From, State) ->
     {reply, leveled_tictac:fetch_root(State#state.tree), State};
 handle_call({fetch_leaves, BranchIDs}, _From, State) ->
     {reply, leveled_tictac:fetch_leaves(State#state.tree, BranchIDs), State};
+handle_call(segment_count, _From, State = #state{dirty_segments = A}) ->
+    {reply, length(A), State};
 handle_call(close, _From, State) ->
     case State#state.safe_save of
         true ->
@@ -866,4 +875,4 @@ test_setup_funs(InitialKeys) ->
     {AddFun, AlterFun, RemoveFun}.
 
 
--endif.
+-endif.