cleanup connection

src/connection/connection.jl

@@ -130,12 +130,16 @@ function stats(connection::Connection)
     connection.stats
 end
 
-function stats(connection::Connection, sub::Sub)
-    sub_data = @lock connection.lock get(connection.sub_data, sub.sid, nothing)
+function stats(connection::Connection, sid::String)
+    sub_data = @lock connection.lock get(connection.sub_data, sid, nothing)
     isnothing(sub_data) && return
     sub_data.stats
 end
 
+function stats(connection::Connection, sub::Sub)
+    stats(connection, sub.sid)
+end
+
 function status_change(f, nc::Connection)
     while true
         st = status(nc)

src/connection/drain.jl

@@ -23,7 +23,7 @@ function _do_drain(nc::Connection, is_connected; timeout = Timer(nc.drain_timeou
         send(nc, Unsub(sid, 0))
     end
     sleep(nc.drain_poll)
-    conn_stats = @lock nc.lock nc.stats
+    conn_stats = stats(nc)
     while !is_every_message_handled(conn_stats)
         if !isopen(timeout)
             @error "Timeout for drain exceeded, not all subs might be drained."

src/pubsub/drain.jl

@@ -17,13 +17,10 @@
 ### Code:
 
 function drain(connection::Connection, sid::String; timer = Timer(connection.drain_timeout))
-    sub_data = @lock connection.lock begin
-        get(connection.sub_data, sid, nothing)
-    end
-    if isnothing(sub_data)
+    sub_stats = stats(connection, sid)
+    if isnothing(sub_stats)
         return # Already drained.
     end
-    sub_stats = sub_data.stats
     send(connection, Unsub(sid, 0))
     sleep(connection.drain_poll)
     while !is_every_message_handled(sub_stats)