add stats support for tunnel handler

handler/tunnel/bind.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/go-gost/core/ingress"
 	"github.com/go-gost/core/logger"
+	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/core/sd"
 	"github.com/go-gost/relay"
 	"github.com/go-gost/x/internal/util/mux"
@@ -60,7 +61,11 @@ func (h *tunnelHandler) handleBind(ctx context.Context, conn net.Conn, network,
 		return
 	}
 
-	h.pool.Add(tunnelID, NewConnector(connectorID, tunnelID, h.id, session, h.md.sd), h.md.tunnelTTL)
+	var stats *stats.Stats
+	if h.stats != nil {
+		stats = h.stats.Stats(tunnelID.String())
+	}
+	h.pool.Add(tunnelID, NewConnector(connectorID, tunnelID, h.id, session, h.md.sd, stats), h.md.tunnelTTL)
 	if h.md.ingress != nil {
 		h.md.ingress.SetRule(ctx, &ingress.Rule{
 			Hostname: endpoint,

handler/tunnel/dialer.go

@@ -30,7 +30,7 @@ func (d *Dialer) Dial(ctx context.Context, network string, tid string) (conn net
 			break
 		}
 
-		conn, err = c.Session().GetConn()
+		conn, err = c.GetConn()
 		if err != nil {
 			d.log.Error(err)
 			continue

handler/tunnel/handler.go

@@ -17,6 +17,7 @@ import (
 	"github.com/go-gost/relay"
 	ctxvalue "github.com/go-gost/x/ctx"
 	xnet "github.com/go-gost/x/internal/net"
+	stats_util "github.com/go-gost/x/internal/util/stats"
 	xrecorder "github.com/go-gost/x/recorder"
 	"github.com/go-gost/x/registry"
 	xservice "github.com/go-gost/x/service"
@@ -45,6 +46,8 @@ type tunnelHandler struct {
 	ep       *entrypoint
 	md       metadata
 	log      logger.Logger
+	stats    *stats_util.HandlerStats
+	cancel   context.CancelFunc
 }
 
 func NewHandler(opts ...handler.Option) handler.Handler {
@@ -55,6 +58,7 @@ func NewHandler(opts ...handler.Option) handler.Handler {
 
 	return &tunnelHandler{
 		options: options,
+		stats:   stats_util.NewHandlerStats(options.Service),
 	}
 }
 
@@ -97,6 +101,13 @@ func (h *tunnelHandler) Init(md md.Metadata) (err error) {
 		return
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
+	h.cancel = cancel
+
+	if h.options.Observer != nil {
+		go h.observeStats(ctx)
+	}
+
 	return nil
 }
 
@@ -268,6 +279,11 @@ func (h *tunnelHandler) Close() error {
 		h.epSvc.Close()
 	}
 	h.pool.Close()
+
+	if h.cancel != nil {
+		h.cancel()
+	}
+
 	return nil
 }
 
@@ -282,3 +298,25 @@ func (h *tunnelHandler) checkRateLimit(addr net.Addr) bool {
 
 	return true
 }
+
+func (h *tunnelHandler) observeStats(ctx context.Context) {
+	if h.options.Observer == nil {
+		return
+	}
+
+	d := h.md.observePeriod
+	if d < time.Millisecond {
+		d = 5 * time.Second
+	}
+	ticker := time.NewTicker(d)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			h.options.Observer.Observe(ctx, h.stats.Events())
+		case <-ctx.Done():
+			return
+		}
+	}
+}

handler/tunnel/metadata.go

@@ -29,6 +29,7 @@ type metadata struct {
 	ingress                 ingress.Ingress
 	sd                      sd.SD
 	muxCfg                  *mux.Config
+	observePeriod           time.Duration
 }
 
 func (h *tunnelHandler) parseMetadata(md mdata.Metadata) (err error) {
@@ -81,5 +82,7 @@ func (h *tunnelHandler) parseMetadata(md mdata.Metadata) (err error) {
 		h.md.muxCfg.Version = 2
 	}
 
+	h.md.observePeriod = mdutil.GetDuration(md, "observePeriod")
+
 	return
 }

handler/tunnel/tunnel.go

@@ -2,6 +2,7 @@ package tunnel
 
 import (
 	"context"
+	"net"
 	"sync"
 	"time"
 
@@ -9,6 +10,9 @@ import (
 	"github.com/go-gost/core/sd"
 	"github.com/go-gost/relay"
 	"github.com/go-gost/x/internal/util/mux"
+
+	"github.com/go-gost/core/observer/stats"
+	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
 	"github.com/go-gost/x/selector"
 	"github.com/google/uuid"
 )
@@ -24,14 +28,16 @@ type Connector struct {
 	sd    sd.SD
 	t     time.Time
 	s     *mux.Session
+	stats *stats.Stats
 }
 
-func NewConnector(id relay.ConnectorID, tid relay.TunnelID, node string, s *mux.Session, sd sd.SD) *Connector {
+func NewConnector(id relay.ConnectorID, tid relay.TunnelID, node string, s *mux.Session, sd sd.SD, stats *stats.Stats) *Connector {
 	c := &Connector{
 		id:    id,
 		tid:   tid,
 		node:  node,
 		sd:    sd,
+		stats: stats,
 		t:     time.Now(),
 		s:     s,
 	}
@@ -62,8 +68,20 @@ func (c *Connector) ID() relay.ConnectorID {
 	return c.id
 }
 
-func (c *Connector) Session() *mux.Session {
+func (c *Connector) GetConn() (net.Conn, error) {
-	return c.s
+	if c == nil || c.s == nil {
+		return nil, nil
+	}
+
+	conn, err := c.s.GetConn()
+	if err != nil {
+		return nil, err
+	}
+
+	if c.stats != nil {
+		conn = stats_wrapper.WrapConn(conn, c.stats)
+	}
+	return conn, nil
 }
 
 func (c *Connector) Close() error {
@@ -74,6 +92,14 @@ func (c *Connector) Close() error {
 	return c.s.Close()
 }
 
+func (c *Connector) IsClosed() bool {
+	if c == nil || c.s == nil {
+		return true
+	}
+
+	return c.s.IsClosed()
+}
+
 type Tunnel struct {
 	node       string
 	id         relay.TunnelID
@@ -83,7 +109,6 @@ type Tunnel struct {
 	mu         sync.RWMutex
 	sd         sd.SD
 	ttl        time.Duration
-	// rw         *selector.RandomWeighted[*Connector]
 }
 
 func NewTunnel(node string, tid relay.TunnelID, ttl time.Duration) *Tunnel {
@@ -93,7 +118,6 @@ func NewTunnel(node string, tid relay.TunnelID, ttl time.Duration) *Tunnel {
 		t:     time.Now(),
 		close: make(chan struct{}),
 		ttl:   ttl,
-		// rw:    selector.NewRandomWeighted[*Connector](),
 	}
 	if t.ttl <= 0 {
 		t.ttl = defaultTTL
@@ -125,9 +149,6 @@ func (t *Tunnel) GetConnector(network string) *Connector {
 	t.mu.RLock()
 	defer t.mu.RUnlock()
 
-	// rw := t.rw
-	// rw.Reset()
-
 	if len(t.connectors) == 1 {
 		return t.connectors[0]
 	}
@@ -136,7 +157,7 @@ func (t *Tunnel) GetConnector(network string) *Connector {
 
 	found := false
 	for _, c := range t.connectors {
-		if c.Session().IsClosed() {
+		if c.IsClosed() {
 			continue
 		}
 
@@ -206,7 +227,7 @@ func (t *Tunnel) clean() {
 			}
 			var connectors []*Connector
 			for _, c := range t.connectors {
-				if c.Session().IsClosed() {
+				if c.IsClosed() {
 					logger.Default().Debugf("remove tunnel: %s, connector: %s", t.id, c.id)
 					if t.sd != nil {
 						t.sd.Deregister(context.Background(), &sd.Service{

observer/stats/wrapper/conn.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync"
 	"syscall"
 
 	"github.com/go-gost/core/metadata"
@@ -19,16 +20,22 @@ var (
 type conn struct {
 	net.Conn
 	stats  *stats.Stats
+	closed chan struct{}
+	mu     sync.Mutex
 }
 
-func WrapConn(c net.Conn, stats *stats.Stats) net.Conn {
+func WrapConn(c net.Conn, pStats *stats.Stats) net.Conn {
-	if stats == nil {
+	if pStats == nil {
 		return c
 	}
 
+	pStats.Add(stats.KindTotalConns, 1)
+	pStats.Add(stats.KindCurrentConns, 1)
+
 	return &conn{
 		Conn:   c,
-		stats: stats,
+		stats:  pStats,
+		closed: make(chan struct{}),
 	}
 }
 
@@ -44,6 +51,21 @@ func (c *conn) Write(b []byte) (n int, err error) {
 	return
 }
 
+func (c *conn) Close() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	select {
+	case <-c.closed:
+		return nil
+	default:
+		close(c.closed)
+	}
+
+	c.stats.Add(stats.KindCurrentConns, -1)
+	return c.Conn.Close()
+}
+
 func (c *conn) SyscallConn() (rc syscall.RawConn, err error) {
 	if sc, ok := c.Conn.(syscall.Conn); ok {
 		rc, err = sc.SyscallConn()

service/service.go

@@ -187,8 +187,6 @@ func (s *defaultService) Serve() error {
 			s.setState(StateReady)
 		}
 
-		s.status.stats.Add(stats.KindTotalConns, 1)
-
 		clientAddr := conn.RemoteAddr().String()
 		clientIP := clientAddr
 		if h, _, _ := net.SplitHostPort(clientAddr); h != "" {
@@ -208,16 +206,13 @@ func (s *defaultService) Serve() error {
 			}
 		}
 		if s.options.admission != nil &&
-			!s.options.admission.Admit(ctx, conn.RemoteAddr().String()) {
+			!s.options.admission.Admit(ctx, clientAddr) {
 			conn.Close()
-			s.options.logger.Debugf("admission: %s is denied", conn.RemoteAddr())
+			s.options.logger.Debugf("admission: %s is denied", clientAddr)
 			continue
 		}
 
 		go func() {
-			s.status.stats.Add(stats.KindCurrentConns, 1)
-			defer s.status.stats.Add(stats.KindCurrentConns, -1)
-
 			if v := xmetrics.GetCounter(xmetrics.MetricServiceRequestsCounter,
 				metrics.Labels{"service": s.name, "client": clientIP}); v != nil {
 				v.Inc()