feat: Add ActiveConnections function to track usage (#390)

kylecarbs · web-flow · commit f3378973e65e · 2021-07-20T17:12:10.000-06:00
* feat: Add ActiveConnections function tto track usage

* Fix linting errors

* Fix syntax err
diff --git a/wsnet/dial.go b/wsnet/dial.go
@@ -232,6 +232,17 @@ func (d *Dialer) Closed() <-chan struct{} {
 	return d.closedChan
 }
 
+// ActiveConnections returns the amount of active connections.
+// DialContext opens a connection, and close will end it.
+func (d *Dialer) ActiveConnections() int {
+	stats, ok := d.rtc.GetStats().GetConnectionStats(d.rtc)
+	if !ok {
+		return -1
+	}
+	// Subtract 1 for the control channel.
+	return int(stats.DataChannelsRequested-stats.DataChannelsClosed) - 1
+}
+
 // Close closes the RTC connection.
 // All data channels dialed will be closed.
 func (d *Dialer) Close() error {
diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
@@ -243,6 +243,38 @@ func TestDial(t *testing.T) {
 			t.Error("didn't close in time")
 		}
 	})
+
+	t.Run("Active Connections", func(t *testing.T) {
+		t.Parallel()
+
+		listener, err := net.Listen("tcp", "0.0.0.0:0")
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		go func() {
+			_, _ = listener.Accept()
+		}()
+		connectAddr, listenAddr := createDumbBroker(t)
+		_, err = Listen(context.Background(), slogtest.Make(t, nil), listenAddr, "")
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		dialer, err := DialWebsocket(context.Background(), connectAddr, nil)
+		if err != nil {
+			t.Error(err)
+		}
+		conn, _ := dialer.DialContext(context.Background(), listener.Addr().Network(), listener.Addr().String())
+		assert.Equal(t, 1, dialer.ActiveConnections())
+		_ = conn.Close()
+		assert.Equal(t, 0, dialer.ActiveConnections())
+		_, _ = dialer.DialContext(context.Background(), listener.Addr().Network(), listener.Addr().String())
+		conn, _ = dialer.DialContext(context.Background(), listener.Addr().Network(), listener.Addr().String())
+		assert.Equal(t, 2, dialer.ActiveConnections())
+		_ = conn.Close()
+		assert.Equal(t, 1, dialer.ActiveConnections())
+	})
 }
 
 func BenchmarkThroughput(b *testing.B) {
