feat: Implement connection closing logic and enhance connection tracking

Author: M03ED

app/connectiontracker/tracker.go

@@ -29,6 +29,7 @@ type ConnEntry struct {
 	lastActivity int64 // atomic, Unix nanosecond timestamp
 	uplink       int64 // atomic, bytes received from client
 	downlink     int64 // atomic, bytes sent to client
+	closing      atomic.Bool
 }
 
 // ConnectionInfo is a read-only snapshot of an active connection's state.
@@ -295,6 +296,14 @@ func disconnectInfo(id uint32, entry *ConnEntry) ConnectionInfo {
 	}
 }
 
+func (e *ConnEntry) isClosing() bool {
+	return e != nil && e.closing.Load()
+}
+
+func (e *ConnEntry) beginClose() bool {
+	return e != nil && e.closing.CompareAndSwap(false, true)
+}
+
 // ListAllConnections returns a snapshot of every active connection across all
 // Tracker instances that were created by NewTracker.
 func (m *Manager) ListAllConnections() []ConnectionInfo {
@@ -315,7 +324,9 @@ func (m *Manager) GetUserStats(email string) (uplink, downlink int64, connCount
 		for _, e := range t.conns[email] {
 			uplink += atomic.LoadInt64(&e.uplink)
 			downlink += atomic.LoadInt64(&e.downlink)
-			connCount++
+			if !e.isClosing() {
+				connCount++
+			}
 		}
 		t.mu.Unlock()
 	}
@@ -393,19 +404,16 @@ func (t *Tracker) Unregister(email string, id uint32) {
 // CancelAll cancels every active connection belonging to email.
 func (t *Tracker) CancelAll(email string) {
 	t.mu.Lock()
-	entries := t.conns[email]
-	delete(t.conns, email)
-	for id := range entries {
-		delete(t.byID, id)
+	entries := make(map[uint32]*ConnEntry)
+	for id, entry := range t.conns[email] {
+		entries[id] = entry
 	}
 	t.mu.Unlock()
 
-	for id, entry := range entries {
-		t.manager.emit(WatchEvent{
-			Connected: false,
-			Info:      disconnectInfo(id, entry),
-		})
-		entry.Cancel()
+	for _, entry := range entries {
+		if entry.beginClose() && entry.Cancel != nil {
+			entry.Cancel()
+		}
 	}
 }
 
@@ -414,22 +422,9 @@ func (t *Tracker) CancelAll(email string) {
 func (t *Tracker) CloseConn(id uint32) bool {
 	t.mu.Lock()
 	entry, ok := t.byID[id]
-	if ok {
-		delete(t.byID, id)
-		if m := t.conns[entry.Email]; m != nil {
-			delete(m, id)
-			if len(m) == 0 {
-				delete(t.conns, entry.Email)
-			}
-		}
-	}
 	t.mu.Unlock()
 
-	if ok {
-		t.manager.emit(WatchEvent{
-			Connected: false,
-			Info:      disconnectInfo(id, entry),
-		})
+	if ok && entry.beginClose() && entry.Cancel != nil {
 		entry.Cancel()
 	}
 	return ok
@@ -438,7 +433,12 @@ func (t *Tracker) CloseConn(id uint32) bool {
 // GetConnCount returns the number of active connections for email.
 func (t *Tracker) GetConnCount(email string) int {
 	t.mu.Lock()
-	n := len(t.conns[email])
+	n := 0
+	for _, entry := range t.conns[email] {
+		if !entry.isClosing() {
+			n++
+		}
+	}
 	t.mu.Unlock()
 	return n
 }
@@ -448,6 +448,9 @@ func (t *Tracker) ListConnections() []ConnectionInfo {
 	t.mu.Lock()
 	result := make([]ConnectionInfo, 0, len(t.byID))
 	for id, entry := range t.byID {
+		if entry.isClosing() {
+			continue
+		}
 		info := disconnectInfo(id, entry)
 		info.Uplink = atomic.LoadInt64(&entry.uplink)
 		info.Downlink = atomic.LoadInt64(&entry.downlink)

app/connectiontracker/tracker_test.go

@@ -1,6 +1,7 @@
 package connectiontracker_test
 
 import (
+	"errors"
 	"net"
 	"sync"
 	"sync/atomic"
@@ -268,6 +269,10 @@ func TestCloseConnCancelsAndRemoves(t *testing.T) {
 	if atomic.LoadInt32(&cancelled) != 1 {
 		t.Error("CloseConn: cancel function was not called")
 	}
+	if len(tracker.ListConnections()) != 0 {
+		t.Error("closing connection should not appear in active list")
+	}
+	tracker.Unregister("user@example.com", id)
 	if len(tracker.ListConnections()) != 0 {
 		t.Error("connection still present after CloseConn")
 	}
@@ -327,6 +332,92 @@ func TestGetConnCountDecreasesAfterUnregister(t *testing.T) {
 	}
 }
 
+func TestCloseConnKeepsStatsUntilUnregister(t *testing.T) {
+	manager := connectiontracker.NewManager()
+	tracker := manager.NewTracker()
+
+	id, entry := tracker.RegisterWithMeta("user@example.com", func() {}, "", "")
+	conn := connectiontracker.WrapConn(&fakeConn{readData: make([]byte, 10)}, entry)
+	if _, err := conn.Read(make([]byte, 10)); err != nil {
+		t.Fatalf("read before close failed: %v", err)
+	}
+
+	if ok := tracker.CloseConn(id); !ok {
+		t.Fatal("CloseConn: expected true for existing connection")
+	}
+
+	if _, err := conn.Write(make([]byte, 5)); err != nil {
+		t.Fatalf("write after close failed: %v", err)
+	}
+
+	uplink, downlink, connCount := manager.GetUserStats("user@example.com")
+	if uplink != 10 || downlink != 5 {
+		t.Fatalf("GetUserStats after close: got uplink=%d downlink=%d, want 10/5", uplink, downlink)
+	}
+	if connCount != 0 {
+		t.Fatalf("GetUserStats connCount after close: got %d, want 0", connCount)
+	}
+
+	tracker.Unregister("user@example.com", id)
+}
+
+func TestDisconnectEventUsesFinalCountersAfterForcedClose(t *testing.T) {
+	manager := connectiontracker.NewManager()
+	tracker := manager.NewTracker()
+	ch := manager.Subscribe()
+	defer manager.Unsubscribe(ch)
+
+	id, entry := tracker.RegisterWithMeta("user@example.com", func() {}, "", "")
+	<-ch // connected event
+
+	conn := connectiontracker.WrapConn(&fakeConn{readData: make([]byte, 7)}, entry)
+	if _, err := conn.Read(make([]byte, 7)); err != nil {
+		t.Fatalf("read before close failed: %v", err)
+	}
+
+	if ok := tracker.CloseConn(id); !ok {
+		t.Fatal("CloseConn: expected true for existing connection")
+	}
+	if _, err := conn.Write(make([]byte, 9)); err != nil {
+		t.Fatalf("write after close failed: %v", err)
+	}
+
+	tracker.Unregister("user@example.com", id)
+
+	select {
+	case ev := <-ch:
+		if ev.Connected {
+			t.Fatal("expected disconnect event")
+		}
+		if ev.Info.Uplink != 7 || ev.Info.Downlink != 9 {
+			t.Fatalf("disconnect counters: got uplink=%d downlink=%d, want 7/9", ev.Info.Uplink, ev.Info.Downlink)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for disconnect event")
+	}
+}
+
+func TestRepeatedForcedCloseIsIdempotent(t *testing.T) {
+	tracker := connectiontracker.New()
+
+	var cancelled int32
+	id, _ := tracker.RegisterWithMeta("user@example.com", func() {
+		atomic.AddInt32(&cancelled, 1)
+	}, "", "")
+
+	if !tracker.CloseConn(id) {
+		t.Fatal("first CloseConn should find connection")
+	}
+	if !tracker.CloseConn(id) {
+		t.Fatal("second CloseConn should still find tracked connection before Unregister")
+	}
+	tracker.CancelAll("user@example.com")
+
+	if got := atomic.LoadInt32(&cancelled); got != 1 {
+		t.Fatalf("cancel count: got %d, want 1", got)
+	}
+}
+
 func TestListConnectionsMetadataFields(t *testing.T) {
 	tracker := connectiontracker.New()
 
@@ -558,3 +649,29 @@ func TestWrapPacketConnUpdatesLastActivity(t *testing.T) {
 		t.Errorf("LastActivity not updated: before=%v after=%v", before, after)
 	}
 }
+
+func TestWrapConnCountsBytesBeforeErrors(t *testing.T) {
+	tracker := connectiontracker.New()
+	_, entry := tracker.RegisterWithMeta("user@example.com", func() {}, "", "")
+
+	wrapped := connectiontracker.WrapConn(&fakeConn{
+		readData:  []byte("hello"),
+		readErr:   errors.New("read stopped"),
+		writeErr:  errors.New("write stopped"),
+	}, entry)
+
+	if _, err := wrapped.Read(make([]byte, 5)); err == nil {
+		t.Fatal("expected read error")
+	}
+	if _, err := wrapped.Write([]byte("bye")); err == nil {
+		t.Fatal("expected write error")
+	}
+
+	conns := tracker.ListConnections()
+	if len(conns) != 1 {
+		t.Fatalf("expected 1 connection")
+	}
+	if conns[0].Uplink != 5 || conns[0].Downlink != 3 {
+		t.Fatalf("counters with errors: got uplink=%d downlink=%d, want 5/3", conns[0].Uplink, conns[0].Downlink)
+	}
+}