Refactor handshake duration tracking in lb.Node

internal/cmgr/cmgr.go

@@ -2,7 +2,6 @@ package cmgr
 
 import (
 	"context"
-	"net/http"
 	"sort"
 	"sync"
 	"time"
@@ -175,32 +174,3 @@ func (cm *cmgrImpl) Start(ctx context.Context, errCH chan error) {
 		}
 	}
 }
-
-type syncReq struct {
-	RelayLabel string     `json:"relay_label"`
-	Stats      conn.Stats `json:"stats"`
-}
-
-func (cm *cmgrImpl) syncOnce() error {
-	cm.l.Infof("sync once total closed connections: %d", cm.countClosedConnection())
-	// todo: opt lock
-	cm.lock.Lock()
-
-	reqs := []syncReq{}
-	for label, conns := range cm.closedConnectionsMap {
-		for _, c := range conns {
-			reqs = append(reqs, syncReq{
-				RelayLabel: label,
-				Stats:      *c.GetStats(),
-			})
-		}
-	}
-	cm.closedConnectionsMap = make(map[string][]conn.RelayConn)
-	cm.lock.Unlock()
-	if cm.cfg.NeedSync() {
-		return myhttp.PostJson(http.DefaultClient, cm.cfg.SyncURL, &reqs)
-	} else {
-		cm.l.Debugf("remove %d closed connections", len(reqs))
-	}
-	return nil
-}

internal/cmgr/metric_sync.go

@@ -0,0 +1,53 @@
+package cmgr
+
+import (
+	"net/http"
+
+	"github.com/Ehco1996/ehco/internal/conn"
+	myhttp "github.com/Ehco1996/ehco/pkg/http"
+)
+
+type StatsPerRule struct {
+	RelayLabel string `json:"relay_label"`
+
+	Up               int64 `json:"up_bytes"`
+	Down             int64 `json:"down_bytes"`
+	ConnectionCnt    int   `json:"connection_count"`
+	HandShakeLatency int64 `json:"latency_in_ms"`
+}
+
+type syncReq struct {
+	Stats []StatsPerRule `json:"stats"`
+}
+
+func (cm *cmgrImpl) syncOnce() error {
+	cm.l.Infof("sync once total closed connections: %d", cm.countClosedConnection())
+	// todo: opt lock
+	cm.lock.Lock()
+
+	req := syncReq{Stats: []StatsPerRule{}}
+	for label, conns := range cm.closedConnectionsMap {
+		s := StatsPerRule{
+			RelayLabel: label,
+		}
+		var totalLatency int64
+		for _, c := range conns {
+			s.ConnectionCnt++
+			s.Up += c.GetStats().Up
+			s.Down += c.GetStats().Down
+			totalLatency += c.GetStats().HandShakeLatency.Milliseconds()
+		}
+		if s.ConnectionCnt > 0 {
+			s.HandShakeLatency = totalLatency / int64(s.ConnectionCnt)
+		}
+		req.Stats = append(req.Stats, s)
+	}
+	cm.closedConnectionsMap = make(map[string][]conn.RelayConn)
+	cm.lock.Unlock()
+	if cm.cfg.NeedSync() {
+		return myhttp.PostJson(http.DefaultClient, cm.cfg.SyncURL, &req)
+	} else {
+		cm.l.Debugf("remove %d closed connections", len(req.Stats))
+	}
+	return nil
+}

internal/conn/conn.go

@@ -18,8 +18,9 @@ var (
 )
 
 type Stats struct {
-	Up   int64 `json:"up"`
-	Down int64 `json:"down"`
+	Up               int64
+	Down             int64
+	HandShakeLatency time.Duration
 }
 
 func (s *Stats) Record(up, down int64) {
@@ -136,6 +137,8 @@ func copyConn(conn1, conn2 *innerConn) error {
 	return err2
 }
 
+type RelayConnOption func(*relayConnImpl)
+
 type RelayConn interface {
 	// Transport transports data between the client and the remote server.
 	// The remoteLabel is the label of the remote server.
@@ -149,14 +152,18 @@ type RelayConn interface {
 	Close() error
 }
 
-func NewRelayConn(relayName string, clientConn, remoteConn net.Conn) RelayConn {
-	return &relayConnImpl{
+func NewRelayConn(relayName string, clientConn, remoteConn net.Conn, opts ...RelayConnOption) RelayConn {
+	rci := &relayConnImpl{
 		RelayLabel: relayName,
-		Stats:      &Stats{Up: 0, Down: 0},
-
 		clientConn: clientConn,
 		remoteConn: remoteConn,
 	}
+	for _, opt := range opts {
+		opt(rci)
+	}
+	s := &Stats{Up: 0, Down: 0, HandShakeLatency: rci.HandshakeDuration}
+	rci.Stats = s
+	return rci
 }
 
 type relayConnImpl struct {
@@ -166,12 +173,19 @@ type relayConnImpl struct {
 	StartTime time.Time `json:"start_time"`
 	EndTime   time.Time `json:"end_time,omitempty"`
 
-	Stats *Stats `json:"stats"`
+	Stats             *Stats `json:"stats"`
+	HandshakeDuration time.Duration
 
 	clientConn net.Conn
 	remoteConn net.Conn
 }
 
+func WithHandshakeDuration(duration time.Duration) RelayConnOption {
+	return func(rci *relayConnImpl) {
+		rci.HandshakeDuration = duration
+	}
+}
+
 func (rc *relayConnImpl) Transport(remoteLabel string) error {
 	defer rc.Close() // nolint: errcheck
 	name := rc.Name()

internal/transporter/mtcp.go

@@ -26,17 +26,20 @@ func (s *MTCP) dialRemote(remote *lb.Node) (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(time.Since(t1).Milliseconds()))
+	latency := time.Since(t1)
+	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(latency.Milliseconds()))
+	remote.HandShakeDuration = latency
 	return mtcpc, nil
 }
 
 func (s *MTCP) HandleTCPConn(c net.Conn, remote *lb.Node) error {
-	mtcpc, err := s.dialRemote(remote)
+	clonedRemote := remote.Clone()
+	mtcpc, err := s.dialRemote(clonedRemote)
 	if err != nil {
 		return err
 	}
 	s.l.Infof("HandleTCPConn from:%s to:%s", c.LocalAddr(), remote.Address)
-	relayConn := conn.NewRelayConn(s.relayLabel, c, mtcpc)
+	relayConn := conn.NewRelayConn(s.relayLabel, c, mtcpc, conn.WithHandshakeDuration(clonedRemote.HandShakeDuration))
 	s.cmgr.AddConnection(relayConn)
 	defer s.cmgr.RemoveConnection(relayConn)
 	return relayConn.Transport(remote.Label)

internal/transporter/mwss.go

@@ -32,17 +32,21 @@ func (s *Mwss) dialRemote(remote *lb.Node) (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(time.Since(t1).Milliseconds()))
+
+	latency := time.Since(t1)
+	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(latency.Milliseconds()))
+	remote.HandShakeDuration = latency
 	return mwssc, nil
 }
 
 func (s *Mwss) HandleTCPConn(c net.Conn, remote *lb.Node) error {
-	mwsc, err := s.dialRemote(remote)
+	clonedRemote := remote.Clone()
+	mwsc, err := s.dialRemote(clonedRemote)
 	if err != nil {
 		return err
 	}
 	s.l.Infof("HandleTCPConn from:%s to:%s", c.LocalAddr(), remote.Address)
-	relayConn := conn.NewRelayConn(s.relayLabel, c, mwsc)
+	relayConn := conn.NewRelayConn(s.relayLabel, c, mwsc, conn.WithHandshakeDuration(clonedRemote.HandShakeDuration))
 	s.cmgr.AddConnection(relayConn)
 	defer s.cmgr.RemoveConnection(relayConn)
 	return relayConn.Transport(remote.Label)

internal/transporter/raw.go

@@ -150,7 +150,9 @@ func (raw *Raw) dialRemote(remote *lb.Node) (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(time.Since(t1).Milliseconds()))
+	latency := time.Since(t1)
+	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(latency.Milliseconds()))
+	remote.HandShakeDuration = latency
 	return rc, nil
 }
 
@@ -159,12 +161,13 @@ func (raw *Raw) HandleTCPConn(c net.Conn, remote *lb.Node) error {
 	metrics.CurConnectionCount.WithLabelValues(remote.Label, metrics.METRIC_CONN_TYPE_TCP).Inc()
 	defer metrics.CurConnectionCount.WithLabelValues(remote.Label, metrics.METRIC_CONN_TYPE_TCP).Dec()
 
-	rc, err := raw.dialRemote(remote)
+	clonedRemote := remote.Clone()
+	rc, err := raw.dialRemote(clonedRemote)
 	if err != nil {
 		return err
 	}
 	raw.l.Infof("HandleTCPConn from %s to %s", c.LocalAddr(), remote.Address)
-	relayConn := conn.NewRelayConn(raw.relayLabel, c, rc)
+	relayConn := conn.NewRelayConn(raw.relayLabel, c, rc, conn.WithHandshakeDuration(clonedRemote.HandShakeDuration))
 	raw.cmgr.AddConnection(relayConn)
 	defer raw.cmgr.RemoveConnection(relayConn)
 	return relayConn.Transport(remote.Label)

internal/transporter/ws.go

@@ -28,17 +28,22 @@ func (s *Ws) dialRemote(remote *lb.Node) (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(time.Since(t1).Milliseconds()))
+	latency := time.Since(t1)
+	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(latency.Milliseconds()))
+	remote.HandShakeDuration = latency
 	return wsc, nil
 }
 
 func (s *Ws) HandleTCPConn(c net.Conn, remote *lb.Node) error {
-	wsc, err := s.dialRemote(remote)
+	clonedRemote := remote.Clone()
+	wsc, err := s.dialRemote(clonedRemote)
 	if err != nil {
 		return err
 	}
 	s.l.Infof("HandleTCPConn from %s to %s", c.LocalAddr(), remote.Address)
-	relayConn := conn.NewRelayConn(s.relayLabel, c, wsc)
+	relayConn := conn.NewRelayConn(
+		s.relayLabel, c, wsc,
+		conn.WithHandshakeDuration(clonedRemote.HandShakeDuration))
 	s.cmgr.AddConnection(relayConn)
 	defer s.cmgr.RemoveConnection(relayConn)
 	return relayConn.Transport(remote.Label)

internal/transporter/wss.go

@@ -30,18 +30,21 @@ func (s *Wss) dialRemote(remote *lb.Node) (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(time.Since(t1).Milliseconds()))
+	latency := time.Since(t1)
+	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(latency.Milliseconds()))
+	remote.HandShakeDuration = latency
 	return wssc, nil
 }
 
 func (s *Wss) HandleTCPConn(c net.Conn, remote *lb.Node) error {
-	wssc, err := s.dialRemote(remote)
+	clonedRemote := remote.Clone()
+	wssc, err := s.dialRemote(clonedRemote)
 	if err != nil {
 		return err
 	}
 	s.l.Infof("HandleTCPConn from %s to %s", c.RemoteAddr(), remote.Address)
 
-	relayConn := conn.NewRelayConn(s.relayLabel, c, wssc)
+	relayConn := conn.NewRelayConn(s.relayLabel, c, wssc, conn.WithHandshakeDuration(clonedRemote.HandShakeDuration))
 	s.cmgr.AddConnection(relayConn)
 	defer s.cmgr.RemoveConnection(relayConn)
 	return relayConn.Transport(remote.Label)

pkg/lb/round_robin.go

@@ -1,12 +1,24 @@
 package lb
 
 import (
+	"time"
+
 	"go.uber.org/atomic"
 )
 
+// todo: move to internal/lb
 type Node struct {
-	Address string
-	Label   string
+	Address           string
+	Label             string
+	HandShakeDuration time.Duration
+}
+
+func (n *Node) Clone() *Node {
+	return &Node{
+		Address:           n.Address,
+		Label:             n.Label,
+		HandShakeDuration: n.HandShakeDuration,
+	}
 }
 
 // RoundRobin is an interface for representing round-robin balancing.