Add new stats metric

wrappers/common.go

@@ -0,0 +1,21 @@
+package wrappers
+
+import (
+	"net"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+func NewClientConn(parent net.Conn, connID conntypes.ConnID) conntypes.StreamReadWriteCloser {
+	conn := newConn(parent, connID, connPurposeClient)
+	conn = NewTrafficStats(conn)
+
+	return conn
+}
+
+func NewTelegramConn(dc conntypes.DC, parent net.Conn) conntypes.StreamReadWriteCloser {
+	conn := newConn(parent, conntypes.ConnID{}, connPurposeTelegram)
+	conn = NewTelegramStats(dc, conn)
+
+	return conn
+}

wrappers/conn.go

@@ -91,11 +91,11 @@ func newConn(parent net.Conn,
 	localAddr := *parent.LocalAddr().(*net.TCPAddr)
 
 	if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil {
-		if config.C.PublicIPv4Addr.IP != nil {
-			localAddr.IP = config.C.PublicIPv4Addr.IP
+		if config.C.PublicIPv4.IP != nil {
+			localAddr.IP = config.C.PublicIPv4.IP
 		}
-	} else if config.C.PublicIPv6Addr.IP != nil {
-		localAddr.IP = config.C.PublicIPv6Addr.IP
+	} else if config.C.PublicIPv6.IP != nil {
+		localAddr.IP = config.C.PublicIPv6.IP
 	}
 
 	logger := zap.S().With(
@@ -115,12 +115,3 @@ func newConn(parent net.Conn,
 		localAddr:  &localAddr,
 	}
 }
-
-func NewClientConn(parent net.Conn,
-	connID conntypes.ConnID) conntypes.StreamReadWriteCloser {
-	return newConn(parent, connID, connPurposeClient)
-}
-
-func NewTelegramConn(parent net.Conn) conntypes.StreamReadWriteCloser {
-	return newConn(parent, conntypes.ConnID{}, connPurposeTelegram)
-}

wrappers/stats.go

@@ -1,67 +0,0 @@
-package wrappers
-
-import (
-	"net"
-	"time"
-
-	"go.uber.org/zap"
-
-	"github.com/9seconds/mtg/conntypes"
-	"github.com/9seconds/mtg/stats"
-)
-
-type wrapperStats struct {
-	parent conntypes.StreamReadWriteCloser
-}
-
-func (w *wrapperStats) Write(p []byte) (int, error) {
-	n, err := w.parent.Write(p)
-	stats.S.EgressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
-	n, err := w.parent.WriteTimeout(p, timeout)
-	stats.S.EgressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) Read(p []byte) (int, error) {
-	n, err := w.parent.Read(p)
-	stats.S.IngressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
-	n, err := w.parent.ReadTimeout(p, timeout)
-	stats.S.IngressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) Conn() net.Conn {
-	return w.parent.Conn()
-}
-
-func (w *wrapperStats) Logger() *zap.SugaredLogger {
-	return w.parent.Logger().Named("traffic")
-}
-
-func (w *wrapperStats) LocalAddr() *net.TCPAddr {
-	return w.parent.LocalAddr()
-}
-
-func (w *wrapperStats) RemoteAddr() *net.TCPAddr {
-	return w.parent.RemoteAddr()
-}
-
-func (w *wrapperStats) Close() error {
-	return w.parent.Close()
-}
-
-func NewTraffic(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
-	return &wrapperStats{parent}
-}

wrappers/stats_telegram.go

@@ -0,0 +1,70 @@
+package wrappers
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/stats"
+	"go.uber.org/zap"
+)
+
+type wrapperTelegramStats struct {
+	parent conntypes.StreamReadWriteCloser
+	dc     conntypes.DC
+	once   sync.Once
+}
+
+func (w *wrapperTelegramStats) Write(p []byte) (int, error) {
+	return w.parent.Write(p)
+}
+
+func (w *wrapperTelegramStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	return w.parent.WriteTimeout(p, timeout)
+}
+
+func (w *wrapperTelegramStats) Read(p []byte) (int, error) {
+	return w.parent.Read(p)
+}
+
+func (w *wrapperTelegramStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	return w.parent.ReadTimeout(p, timeout)
+}
+
+func (w *wrapperTelegramStats) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperTelegramStats) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("stats-telegram")
+}
+
+func (w *wrapperTelegramStats) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperTelegramStats) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperTelegramStats) Close() error {
+	var err error
+
+	w.once.Do(func() {
+		err = w.parent.Close()
+		stats.Stats.TelegramDisconnected(w.dc, w.RemoteAddr())
+	})
+
+	return err
+}
+
+func NewTelegramStats(dc conntypes.DC, parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
+	conn := &wrapperTelegramStats{
+		parent: parent,
+		dc:     dc,
+	}
+	stats.Stats.TelegramConnected(dc, parent.RemoteAddr())
+
+	return conn
+}

wrappers/stats_traffic.go

@@ -0,0 +1,67 @@
+package wrappers
+
+import (
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/stats"
+)
+
+type wrapperTrafficStats struct {
+	parent conntypes.StreamReadWriteCloser
+}
+
+func (w *wrapperTrafficStats) Write(p []byte) (int, error) {
+	n, err := w.parent.Write(p)
+	stats.Stats.EgressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.WriteTimeout(p, timeout)
+	stats.Stats.EgressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) Read(p []byte) (int, error) {
+	n, err := w.parent.Read(p)
+	stats.Stats.IngressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.ReadTimeout(p, timeout)
+	stats.Stats.IngressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperTrafficStats) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("stats-traffic")
+}
+
+func (w *wrapperTrafficStats) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperTrafficStats) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperTrafficStats) Close() error {
+	return w.parent.Close()
+}
+
+func NewTrafficStats(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
+	return &wrapperTrafficStats{parent}
+}