go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add stats support for tunnel handler

Browse Source
This commit is contained in:
ginuerzh
2024-08-28 13:34:05 +08:00
parent bc0d6953bc
commit ff20711c25
7 changed files with 120 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
handler/tunnel/bind.go
+6 -1
View File
@@ -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
+1 -1
View File
@@ -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
+38
View File
@@ -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
+3
View File
@@ -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
+43 -22
View File
@@ -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{