Re-abstract the stats timer thing in reader.

2 files changed, 27 insertions, 13 deletions

diff --git a/src/rabbit_event.erl b/src/rabbit_event.erl
index d784a6aa..e2d0094d 100644
--- a/src/rabbit_event.erl
+++ b/src/rabbit_event.erl
@@ -34,6 +34,7 @@
 -include("rabbit.hrl").
 
 -export([init_stats_timer/0, ensure_stats_timer/3, stop_stats_timer/2]).
+-export([ensure_stats_timer_after/2, reset_stats_timer_after/1]).
 -export([stats_level/1]).
 -export([notify/2]).
 
@@ -43,6 +44,8 @@
 -spec(ensure_stats_timer/3 ::
         (state(), fun (() -> 'ok'), fun (() -> 'ok')) -> state()).
 -spec(stop_stats_timer/2 :: (state(), fun (() -> 'ok')) -> state()).
+-spec(ensure_stats_timer_after/2 :: (state(), fun (() -> 'ok')) -> state()).
+-spec(reset_stats_timer_after/1 :: (state()) -> 'ok').
 -spec(stats_level/1 :: (state()) -> atom()).
 -spec(notify/2 :: (atom(), term()) -> 'ok').
 
@@ -74,6 +77,18 @@ stop_stats_timer(State = #state{timer = TRef}, NowFun) ->
     NowFun(),
     State#state{timer = undefined}.
 
+ensure_stats_timer_after(State = #state{level = none}, _TimerFun) ->
+    State;
+ensure_stats_timer_after(State = #state{timer = undefined}, TimerFun) ->
+    {ok, TRef} = timer:apply_after(?STATS_INTERVAL,
+                                   erlang, apply, [TimerFun, []]),
+    State#state{timer = TRef};
+ensure_stats_timer_after(State, _TimerFun) ->
+    State.
+
+reset_stats_timer_after(State) ->
+    State#state{timer = undefined}.
+
 stats_level(#state{level = Level}) ->
     Level.
 
diff --git a/src/rabbit_reader.erl b/src/rabbit_reader.erl
index ffd40bed..7a45922e 100644
--- a/src/rabbit_reader.erl
+++ b/src/rabbit_reader.erl
@@ -60,7 +60,7 @@
 %---------------------------------------------------------------------------
 
 -record(v1, {sock, connection, callback, recv_ref, connection_state,
-             queue_collector, stats_timer_ref, stats_level}).
+             queue_collector, stats_timer}).
 
 -define(STATISTICS_KEYS, [pid, recv_oct, recv_cnt, send_oct, send_cnt,
                           send_pend, state, channels]).
@@ -251,7 +251,6 @@ start_connection(Parent, Deb, Sock, SockTransform) ->
                       handshake_timeout),
     ProfilingValue = setup_profiling(),
     {ok, Collector} = rabbit_queue_collector:start_link(),
-    {ok, StatsLevel} = application:get_env(rabbit, collect_statistics),
     try
         mainloop(Parent, Deb, switch_callback(
                                 #v1{sock = ClientSock,
@@ -265,8 +264,8 @@ start_connection(Parent, Deb, Sock, SockTransform) ->
                                     recv_ref = none,
                                     connection_state = pre_init,
                                     queue_collector = Collector,
-                                    stats_timer_ref = undefined,
-                                    stats_level = StatsLevel},
+                                    stats_timer =
+                                        rabbit_event:init_stats_timer()},
                                 handshake, 8))
     catch
         Ex -> (if Ex == connection_closed_abruptly ->
@@ -355,7 +354,10 @@ mainloop(Parent, Deb, State = #v1{sock= Sock, recv_ref = Ref}) ->
             mainloop(Parent, Deb, State);
         {'$gen_cast', emit_stats} ->
             internal_emit_stats(State),
-            mainloop(Parent, Deb, State#v1{stats_timer_ref = undefined});
+            mainloop(Parent, Deb,
+                     State#v1{stats_timer =
+                                  rabbit_event:reset_stats_timer_after(
+                                    State#v1.stats_timer)});
         {system, From, Request} ->
             sys:handle_system_msg(Request, From,
                                   Parent, ?MODULE, Deb, State);
@@ -603,14 +605,11 @@ check_version(ClientVersion, ServerVersion) ->
           (ClientMajor == ServerMajor andalso
            ClientMinor >= ServerMinor).
 
-ensure_stats_timer(State = #v1{stats_level = none}) ->
-    State;
-ensure_stats_timer(State = #v1{stats_timer_ref = undefined}) ->
-    {ok, TRef} = timer:apply_after(?STATS_INTERVAL,
-                                   rabbit_reader, emit_stats, [self()]),
-    State#v1{stats_timer_ref = TRef};
-ensure_stats_timer(State) ->
-    State.
+ensure_stats_timer(State = #v1{stats_timer = StatsTimer}) ->
+    Self = self(),
+    State#v1{stats_timer = rabbit_event:ensure_stats_timer_after(
+                             StatsTimer,
+                             fun() -> emit_stats(Self) end)}.
 
 %%--------------------------------------------------------------------------