From bc20e0cb80092fd24df3bb11abe31c24226b4677 Mon Sep 17 00:00:00 2001
From: Kamil Waz <kamil@kamilwaz.com>
Date: Tue, 22 Nov 2022 12:57:57 +0100
Subject: [PATCH] Improve error handling in stat

---
 .../admin/mongoose_graphql_stats_domain.erl   | 10 +++-
 .../admin/mongoose_graphql_stats_global.erl   | 18 ++++---
 src/stats_api.erl                             | 54 +++++++++----------
 3 files changed, 44 insertions(+), 38 deletions(-)

diff --git a/src/graphql/admin/mongoose_graphql_stats_domain.erl b/src/graphql/admin/mongoose_graphql_stats_domain.erl
index ea7616175aa..13e67dc6fad 100644
--- a/src/graphql/admin/mongoose_graphql_stats_domain.erl
+++ b/src/graphql/admin/mongoose_graphql_stats_domain.erl
@@ -12,6 +12,12 @@
 -include("jlib.hrl").
 
 execute(_Ctx, Domain, <<"registeredUsers">>, _Args) ->
-    {ok, stats_api:stats(<<"registeredusers">>, Domain)};
+    globalStats(<<"registeredusers">>, Domain);
 execute(_Ctx, Domain, <<"onlineUsers">>, _Args) ->
-    {ok, stats_api:stats(<<"onlineusers">>, Domain)}.
+    globalStats(<<"onlineusers">>, Domain).
+
+globalStats(Name, Domain) ->
+    case stats_api:stats(Name, Domain) of
+        {ok, _} = Result -> Result;
+        Error -> make_error(Error, #{domain => Domain})
+    end.
diff --git a/src/graphql/admin/mongoose_graphql_stats_global.erl b/src/graphql/admin/mongoose_graphql_stats_global.erl
index 0d2500e2ced..9a884f86284 100644
--- a/src/graphql/admin/mongoose_graphql_stats_global.erl
+++ b/src/graphql/admin/mongoose_graphql_stats_global.erl
@@ -12,14 +12,20 @@
 -include("jlib.hrl").
 
 execute(_Ctx, globalStats, <<"uptimeSeconds">>, _Args) ->
-    {ok, stats_api:stats(<<"uptimeseconds">>)};
+    globalStats(<<"uptimeseconds">>);
 execute(_Ctx, globalStats, <<"registeredUsers">>, _Args) ->
-    {ok, stats_api:stats(<<"registeredusers">>)};
+    globalStats(<<"registeredusers">>);
 execute(_Ctx, globalStats, <<"onlineUsersNode">>, _Args) ->
-    {ok, stats_api:stats(<<"onlineusersnode">>)};
+    globalStats(<<"onlineusersnode">>);
 execute(_Ctx, globalStats, <<"onlineUsers">>, _Args) ->
-    {ok, stats_api:stats(<<"onlineusers">>)};
+    globalStats(<<"onlineusers">>);
 execute(_Ctx, globalStats, <<"incomingS2S">>, _Args) ->
-    {ok, stats_api:incoming_s2s_number()};
+    stats_api:incoming_s2s_number();
 execute(_Ctx, globalStats, <<"outgoingS2S">>, _Args) ->
-    {ok, stats_api:outgoing_s2s_number()}.
+    stats_api:outgoing_s2s_number().
+
+globalStats(Name) ->
+    case stats_api:stats(Name) of
+        {ok, _} = Result -> Result;
+        Error -> make_error(Error, #{})
+    end.
diff --git a/src/stats_api.erl b/src/stats_api.erl
index c23196374ab..b003be432e8 100644
--- a/src/stats_api.erl
+++ b/src/stats_api.erl
@@ -5,38 +5,32 @@
 -include("mongoose.hrl").
 -include("ejabberd_commands.hrl").
 
--spec incoming_s2s_number() -> non_neg_integer().
+-spec incoming_s2s_number() -> {ok, non_neg_integer()}.
 incoming_s2s_number() ->
-    length(supervisor:which_children(ejabberd_s2s_in_sup)).
+    {ok, length(supervisor:which_children(ejabberd_s2s_in_sup))}.
 
--spec outgoing_s2s_number() -> non_neg_integer().
+-spec outgoing_s2s_number() -> {ok, non_neg_integer()}.
 outgoing_s2s_number() ->
-    length(supervisor:which_children(ejabberd_s2s_out_sup)).
+    {ok, length(supervisor:which_children(ejabberd_s2s_out_sup))}.
 
--spec stats(binary()) -> integer() | {error, string()}.
-stats(Name) ->
-    case Name of
-        <<"uptimeseconds">> ->
-            trunc(element(1, erlang:statistics(wall_clock))/1000);
-        <<"registeredusers">> ->
-            Domains = lists:flatmap(fun mongoose_domain_api:get_domains_by_host_type/1,
-                                    ?ALL_HOST_TYPES),
-            lists:sum([ejabberd_auth:get_vh_registered_users_number(Domain) || Domain <- Domains]);
-        <<"onlineusersnode">> ->
-            ejabberd_sm:get_node_sessions_number();
-        <<"onlineusers">> ->
-            ejabberd_sm:get_total_sessions_number();
-        _ ->
-            {error, "Wrong command name."}
-    end.
+-spec stats(binary()) -> {ok, integer()} | {not_found, string()}.
+stats(<<"uptimeseconds">>) ->
+    {ok, trunc(element(1, erlang:statistics(wall_clock)) / 1000)};
+stats(<<"registeredusers">>) ->
+    Domains = lists:flatmap(fun mongoose_domain_api:get_domains_by_host_type/1,
+                            ?ALL_HOST_TYPES),
+    {ok, lists:sum([ejabberd_auth:get_vh_registered_users_number(Domain) || Domain <- Domains])};
+stats(<<"onlineusersnode">>) ->
+    {ok, ejabberd_sm:get_node_sessions_number()};
+stats(<<"onlineusers">>) ->
+    {ok, ejabberd_sm:get_total_sessions_number()};
+stats(_Name) ->
+    {not_found, "Stats not found"}.
 
--spec stats(binary(), jid:server()) -> integer() | {error, string()}.
-stats(Name, Host) ->
-    case Name of
-        <<"registeredusers">> ->
-            ejabberd_auth:get_vh_registered_users_number(Host);
-        <<"onlineusers">> ->
-            ejabberd_sm:get_vh_session_number(Host);
-        _ ->
-            {error, "Wrong command name."}
-    end.
+-spec stats(binary(), jid:server()) -> {ok, integer()} | {not_found, string()}.
+stats(<<"registeredusers">>, Host) ->
+    {ok, ejabberd_auth:get_vh_registered_users_number(Host)};
+stats(<<"onlineusers">>, Host) ->
+    {ok, ejabberd_sm:get_vh_session_number(Host)};
+stats(_Name, _Host) ->
+    {not_found, "Stats not found"}.