Ensure AAE tree wait fails test if timeout

Some refactoring to make sure errors get reported properly when things
go wrong. Wait until functions need to assert on the return value, and
this return value is more useful if it contains errors from nodes timing
out, Bad RPCs, etc.

Author: engelsanchez

src/rt.erl

@@ -828,61 +828,77 @@ wait_until_nodes_agree_about_ownership(Nodes) ->
 %% AAE support
 wait_until_aae_trees_built(Nodes) ->
     lager:info("Wait until AAE builds all partition trees across ~p", [Nodes]),
-    %% Wait until all nodes report no undefined trees
-    wait_until(fun() ->
-                       lists:foldl(aae_tree_built_fun(), true, Nodes)
-               end).
+    BuiltFun = fun() -> lists:foldl(aae_tree_built_fun(), true, Nodes) end,
+    ?assertEqual(ok, wait_until(BuiltFun)),
+    ok.
 
 aae_tree_built_fun() ->
-    fun(_, false) ->
-            false;
-       (Node, _AllBuilt = true) ->
-            InfoRes = get_aae_tree_info(Node),
-            case all_trees_have_build_times(InfoRes) of
-                true ->
-                    {ok, Info} = InfoRes,
-                    Partitions = [I || {I, _} <- Info],
-                    lager:debug("Check if really built by locking"),
-
-                    %% Try to lock each partition. If you get not_built,
-                    %% the manager has not detected the built process has
-                    %% died yet.
-                    %% Notice that the process locking is spawned by the
-                    %% pmap. That's important! as it should die eventually
-                    %% so the test can lock on the tree.
-                    AllBuilt = lists:all(fun(V) -> V == true end,
-                                         rt:pmap(index_built_fun(Node), Partitions)),
-                    lager:debug("For node ~p all built = ~p", [Node, AllBuilt]),
-                    AllBuilt;
-                false ->
-                    false
-            end
+    fun(Node, _AllBuilt = true) ->
+            case get_aae_tree_info(Node) of
+                {ok, TreeInfos} ->
+                    case all_trees_have_build_times(TreeInfos) of
+                        true ->
+                            Partitions = [I || {I, _} <- TreeInfos],
+                            all_aae_trees_built(Node, Partitions);
+                        false ->
+                            some_trees_not_built
+                    end;
+                Err ->
+                    Err
+            end;
+       (_Node, Err) ->
+            Err
+    end.
+
+% It is unlikely but possible to get a tree built time from compute_tree_info
+% but an attempt to use the tree returns not_built. This is because the build
+% process has finished, but the lock on the tree won't be released until it
+% dies and the manager detects it. Yes, this is super freaking paranoid.
+all_aae_trees_built(Node, Partitions) ->
+    %% Notice that the process locking is spawned by the
+    %% pmap. That's important! as it should die eventually
+    %% so the lock is released and the test can lock the tree.
+    IndexBuilts = rt:pmap(index_built_fun(Node), Partitions),
+    BadOnes = [R || R <- IndexBuilts, R /= true],
+    case BadOnes of
+        [] ->
+            true;
+        _ ->
+            BadOnes
     end.
 
 get_aae_tree_info(Node) ->
     case rpc:call(Node, riak_kv_entropy_info, compute_tree_info, []) of
         {badrpc, _} ->
-            {error, badrpc};
+            {error, {badrpc, Node}};
         Info  ->
             lager:debug("Entropy table on node ~p : ~p", [Node, Info]),
             {ok, Info}
     end.
 
-all_trees_have_build_times({badrpc, _}) ->
-    false;
-all_trees_have_build_times({ok, Info}) ->
+all_trees_have_build_times(Info) ->
     not lists:keymember(undefined, 2, Info).
 
 index_built_fun(Node) ->
     fun(Idx) ->
-            {ok, TreePid} = rpc:call(Node, riak_kv_vnode,
-                                     hashtree_pid, [Idx]),
-            TreeLocked =
-                rpc:call(Node, riak_kv_index_hashtree, get_lock,
-                         [TreePid, for_riak_test]),
-            lager:debug("Partition ~p : ~p", [Idx, TreeLocked]),
-            TreeLocked == ok
-                orelse TreeLocked == already_locked
+            case rpc:call(Node, riak_kv_vnode,
+                                     hashtree_pid, [Idx]) of
+                {ok, TreePid} ->
+                    case rpc:call(Node, riak_kv_index_hashtree,
+                                  get_lock, [TreePid, for_riak_test]) of
+                        {badrpc, _} ->
+                            {error, {badrpc, Node}};
+                        TreeLocked when TreeLocked == ok;
+                                        TreeLocked == already_locked ->
+                            true;
+                        Err ->
+                            % Either not_built or some unhandled result,
+                            % in which case update this case please!
+                            {error, {index_not_built, Node, Idx, Err}}
+                    end;
+                {badrpc, _} ->
+                    {error, {badrpc, Node}}
+            end
     end.
 
 %%%===================================================================