From e0c72c7bdca028975968eda3f5702560a1eb78e5 Mon Sep 17 00:00:00 2001
From: William Yang <mscame@gmail.com>
Date: Wed, 24 Jan 2024 13:29:52 +0100
Subject: [PATCH] feat(listener): count connections under listener

---
 src/quicer_listener.erl        | 19 +++++++++-
 test/quicer_listener_SUITE.erl | 66 ++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/src/quicer_listener.erl b/src/quicer_listener.erl
index 0c74f600..39abb770 100644
--- a/src/quicer_listener.erl
+++ b/src/quicer_listener.erl
@@ -26,7 +26,8 @@
     unlock/2,
     reload/2,
     reload/3,
-    get_handle/2
+    get_handle/2,
+    count_conns/1
 ]).
 
 %% gen_server callbacks
@@ -107,6 +108,11 @@ reload(Pid, ListenOn, NewConf) ->
 get_handle(Pid, Timeout) ->
     gen_server:call(Pid, get_handle, Timeout).
 
+%% @doc count accepted but not yet closed connections
+-spec count_conns(pid()) -> non_neg_integer().
+count_conns(Pid) ->
+    gen_server:call(Pid, count_conns, infinity).
+
 %%%===================================================================
 %%% gen_server callbacks
 %%%===================================================================
@@ -173,6 +179,17 @@ handle_call({reload, NewConf}, _From, State) ->
 handle_call({reload, NewListenOn, NewConf}, _From, State) ->
     {Res, NewState} = do_reload(NewListenOn, NewConf, State),
     {reply, Res, NewState};
+handle_call(
+    count_conns,
+    _From,
+    #state{
+        conn_sup = ConnSup,
+        opts = #{conn_acceptors := NoAcceptors}
+    } =
+        State
+) ->
+    ConnPids = supervisor:which_children(ConnSup),
+    {reply, length(ConnPids) - NoAcceptors, State};
 handle_call(Request, _From, State) ->
     Reply = {error, {unimpl, Request}},
     {reply, Reply, State}.
diff --git a/test/quicer_listener_SUITE.erl b/test/quicer_listener_SUITE.erl
index f3c1dc13..26f1b72a 100644
--- a/test/quicer_listener_SUITE.erl
+++ b/test/quicer_listener_SUITE.erl
@@ -953,6 +953,72 @@ tc_get_listener_owner(Config) ->
     ?assertEqual({ok, self()}, quicer:get_listener_owner(L)),
     quicer:close_listener(L).
 
+tc_count_conns(Config) ->
+    Port0 = select_port(),
+    Port1 = select_port(),
+    ServerConnCallback = example_server_connection,
+    ServerStreamCallback = example_server_stream,
+    ListenerOpts = [
+        {conn_acceptors, 32},
+        {peer_bidi_stream_count, 0},
+        {peer_unidi_stream_count, 2}
+        | default_listen_opts(Config)
+    ],
+    ConnectionOpts = [
+        {conn_callback, ServerConnCallback},
+        {stream_acceptors, 2}
+        | default_conn_opts()
+    ],
+    StreamOpts = [
+        {stream_callback, ServerStreamCallback}
+        | default_stream_opts()
+    ],
+    Options = {ListenerOpts, ConnectionOpts, StreamOpts},
+
+    %% GIVEN: Two QUIC listeners
+    {ok, QuicApp} = quicer:spawn_listener(sample, Port0, Options),
+
+    {ok, QuicApp2} = quicer:spawn_listener(sample2, Port1, Options),
+
+    ClientConnOpts = default_conn_opts_verify(Config, ca),
+
+    %% WHEN: a client is connected to the first listener
+    {ok, ClientConnPid} = example_client_connection:start_link(
+        "localhost",
+        Port0,
+        {ClientConnOpts, default_stream_opts()}
+    ),
+    #{is_resumed := false} = snabbkaffe:retry(
+        50,
+        20,
+        fun() ->
+            #{is_resumed := false} = quicer_connection:get_cb_state(ClientConnPid)
+        end
+    ),
+
+    %% Then the first listener has one connection and other has none
+    ?assertEqual({1, 0}, {
+        quicer_listener:count_conns(QuicApp), quicer_listener:count_conns(QuicApp2)
+    }),
+
+    %% WHEN: client is stopped
+    gen_server:stop(ClientConnPid),
+
+    %% THEN: both listeners have no connections
+    {0, 0} = snabbkaffe:retry(
+        10,
+        100,
+        fun() ->
+            {0, 0} =
+                {quicer_listener:count_conns(QuicApp), quicer_listener:count_conns(QuicApp2)}
+        end
+    ),
+
+    quicer:terminate_listener(sample),
+    quicer:terminate_listener(sample2).
+
+%%% Helpers
+
 select_port() ->
     Port = select_free_port(quic),
     timer:sleep(100),