From ff20711c25cd4286b4b4bd1955b46b65571fc598 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Wed, 28 Aug 2024 13:34:05 +0800 Subject: [PATCH] add stats support for tunnel handler --- handler/tunnel/bind.go | 7 +++- handler/tunnel/dialer.go | 2 +- handler/tunnel/handler.go | 38 ++++++++++++++++++++ handler/tunnel/metadata.go | 3 ++ handler/tunnel/tunnel.go | 65 ++++++++++++++++++++++------------ observer/stats/wrapper/conn.go | 32 ++++++++++++++--- service/service.go | 9 ++--- 7 files changed, 120 insertions(+), 36 deletions(-) diff --git a/handler/tunnel/bind.go b/handler/tunnel/bind.go index aa76849b..a5bb319f 100644 --- a/handler/tunnel/bind.go +++ b/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, diff --git a/handler/tunnel/dialer.go b/handler/tunnel/dialer.go index 34d9e767..f86f387a 100644 --- a/handler/tunnel/dialer.go +++ b/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 diff --git a/handler/tunnel/handler.go b/handler/tunnel/handler.go index b177520e..aebc3dc5 100644 --- a/handler/tunnel/handler.go +++ b/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 + } + } +} diff --git a/handler/tunnel/metadata.go b/handler/tunnel/metadata.go index 14ee9bec..0ada237f 100644 --- a/handler/tunnel/metadata.go +++ b/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 } diff --git a/handler/tunnel/tunnel.go b/handler/tunnel/tunnel.go index ebc9041e..fe6bc01d 100644 --- a/handler/tunnel/tunnel.go +++ b/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" ) @@ -18,22 +22,24 @@ const ( ) type Connector struct { - id relay.ConnectorID - tid relay.TunnelID - node string - sd sd.SD - t time.Time - s *mux.Session + id relay.ConnectorID + tid relay.TunnelID + node string + 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, - t: time.Now(), - s: s, + id: id, + tid: tid, + node: node, + sd: sd, + stats: stats, + t: time.Now(), + s: s, } go c.accept() return c @@ -62,8 +68,20 @@ func (c *Connector) ID() relay.ConnectorID { return c.id } -func (c *Connector) Session() *mux.Session { - return c.s +func (c *Connector) GetConn() (net.Conn, error) { + 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{ diff --git a/observer/stats/wrapper/conn.go b/observer/stats/wrapper/conn.go index 55e072a4..1724cbd4 100644 --- a/observer/stats/wrapper/conn.go +++ b/observer/stats/wrapper/conn.go @@ -4,6 +4,7 @@ import ( "errors" "io" "net" + "sync" "syscall" "github.com/go-gost/core/metadata" @@ -18,17 +19,23 @@ var ( type conn struct { net.Conn - stats *stats.Stats + stats *stats.Stats + closed chan struct{} + mu sync.Mutex } -func WrapConn(c net.Conn, stats *stats.Stats) net.Conn { - if stats == nil { +func WrapConn(c net.Conn, pStats *stats.Stats) net.Conn { + if pStats == nil { return c } + pStats.Add(stats.KindTotalConns, 1) + pStats.Add(stats.KindCurrentConns, 1) + return &conn{ - Conn: c, - stats: stats, + Conn: c, + 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() diff --git a/service/service.go b/service/service.go index d3c6491c..2c93b787 100644 --- a/service/service.go +++ b/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()