From 03138f8345dabcc0b0b52596fd92c8efa2d0ba40 Mon Sep 17 00:00:00 2001
From: Thales Macedo Garitezi <thalesmg@gmail.com>
Date: Wed, 21 Feb 2024 09:04:54 -0300
Subject: [PATCH] refactor(mgmt): move persistent session counting function to
 mgmt module

Follow up to https://github.com/emqx/emqx/pull/12500#discussion_r1496634087
---
 .../src/emqx_persistent_session_ds_state.erl  | 22 +------------------
 .../src/emqx_mgmt_api_clients.erl             | 18 ++++++++++++++-
 2 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/apps/emqx/src/emqx_persistent_session_ds_state.erl b/apps/emqx/src/emqx_persistent_session_ds_state.erl
index efa03e101e..bea8f25e34 100644
--- a/apps/emqx/src/emqx_persistent_session_ds_state.erl
+++ b/apps/emqx/src/emqx_persistent_session_ds_state.erl
@@ -36,11 +36,7 @@
 -export([get_rank/2, put_rank/3, del_rank/2, fold_ranks/3]).
 -export([get_subscriptions/1, put_subscription/4, del_subscription/3]).
 
--export([
-    make_session_iterator/0,
-    session_iterator_next/2,
-    session_count/0
-]).
+-export([make_session_iterator/0, session_iterator_next/2]).
 
 -export_type([
     t/0,
@@ -369,12 +365,6 @@ del_rank(Key, Rec) ->
 fold_ranks(Fun, Acc, Rec) ->
     gen_fold(ranks, Fun, Acc, Rec).
 
--spec session_count() -> non_neg_integer().
-session_count() ->
-    %% N.B.: this is potentially costly.  Should not be called in hot paths.
-    %% `mnesia:table_info(_, size)' is always zero for rocksdb, so we need to traverse...
-    do_session_count(make_session_iterator(), 0).
-
 -spec make_session_iterator() -> session_iterator().
 make_session_iterator() ->
     mnesia:dirty_first(?session_tab).
@@ -577,16 +567,6 @@ ro_transaction(Fun) ->
 
 %%
 
-do_session_count('$end_of_table', N) ->
-    N;
-do_session_count(Cursor, N) ->
-    case session_iterator_next(Cursor, 1) of
-        {[], _} ->
-            N;
-        {_, NextCursor} ->
-            do_session_count(NextCursor, N + 1)
-    end.
-
 -compile({inline, check_sequence/1}).
 
 -ifdef(CHECK_SEQNO).
diff --git a/apps/emqx_management/src/emqx_mgmt_api_clients.erl b/apps/emqx_management/src/emqx_mgmt_api_clients.erl
index f0e7ce3cb5..72d4a48cba 100644
--- a/apps/emqx_management/src/emqx_mgmt_api_clients.erl
+++ b/apps/emqx_management/src/emqx_mgmt_api_clients.erl
@@ -916,7 +916,7 @@ add_persistent_session_count(QueryState0 = #{total := Totals0}) ->
             %% to traverse the whole table), but also hard to deduplicate live connections
             %% from it...  So this count will possibly overshoot the true count of
             %% sessions.
-            SessionCount = emqx_persistent_session_ds_state:session_count(),
+            SessionCount = persistent_session_count(),
             Totals = Totals0#{undefined => SessionCount},
             QueryState0#{total := Totals};
         false ->
@@ -965,6 +965,22 @@ no_persistent_sessions() ->
             true
     end.
 
+-spec persistent_session_count() -> non_neg_integer().
+persistent_session_count() ->
+    %% N.B.: this is potentially costly.  Should not be called in hot paths.
+    %% `mnesia:table_info(_, size)' is always zero for rocksdb, so we need to traverse...
+    do_persistent_session_count(init_persistent_session_iterator(), 0).
+
+do_persistent_session_count('$end_of_table', N) ->
+    N;
+do_persistent_session_count(Cursor, N) ->
+    case emqx_persistent_session_ds_state:session_iterator_next(Cursor, 1) of
+        {[], _} ->
+            N;
+        {_, NextCursor} ->
+            do_persistent_session_count(NextCursor, N + 1)
+    end.
+
 do_persistent_session_query(ResultAcc, QueryState) ->
     case emqx_persistent_message:is_persistence_enabled() of
         true ->