add observer

handler/http2/handler.go

@@ -26,8 +26,11 @@ import (
 	ctxvalue "github.com/go-gost/x/internal/ctx"
 	xio "github.com/go-gost/x/internal/io"
 	netpkg "github.com/go-gost/x/internal/net"
+	stats_util "github.com/go-gost/x/internal/util/stats"
 	"github.com/go-gost/x/limiter/traffic/wrapper"
 	"github.com/go-gost/x/registry"
+	"github.com/go-gost/x/stats"
+	stats_wrapper "github.com/go-gost/x/stats/wrapper"
 )
 
 func init() {
@@ -38,6 +41,8 @@ type http2Handler struct {
 	router  *chain.Router
 	md      metadata
 	options handler.Options
+	stats   *stats_util.HandlerStats
+	cancel  context.CancelFunc
 }
 
 func NewHandler(opts ...handler.Option) handler.Handler {
@@ -48,6 +53,7 @@ func NewHandler(opts ...handler.Option) handler.Handler {
 
 	return &http2Handler{
 		options: options,
+		stats:   stats_util.NewHandlerStats(options.Service),
 	}
 }
 
@@ -61,6 +67,12 @@ func (h *http2Handler) Init(md md.Metadata) error {
 		h.router = chain.NewRouter(chain.LoggerRouterOption(h.options.Logger))
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
+	h.cancel = cancel
+
+	if h.options.Observer != nil {
+		go h.observeStats(ctx)
+	}
 	return nil
 }
 
@@ -98,6 +110,13 @@ func (h *http2Handler) Handle(ctx context.Context, conn net.Conn, opts ...handle
 	)
 }
 
+func (h *http2Handler) Close() error {
+	if h.cancel != nil {
+		h.cancel()
+	}
+	return nil
+}
+
 // NOTE: there is an issue (golang/go#43989) will cause the client hangs
 // when server returns an non-200 status code,
 // May be fixed in go1.18.
@@ -204,12 +223,20 @@ func (h *http2Handler) roundTrip(ctx context.Context, w http.ResponseWriter, req
 			return nil
 		}
 
-		rw := wrapper.WrapReadWriter(h.options.Limiter, xio.NewReadWriter(req.Body, flushWriter{w}), req.RemoteAddr,
+		rw := wrapper.WrapReadWriter(h.options.Limiter, xio.NewReadWriter(req.Body, flushWriter{w}),
 			traffic.NetworkOption("tcp"),
 			traffic.AddrOption(addr),
 			traffic.ClientOption(clientID),
 			traffic.SrcOption(req.RemoteAddr),
 		)
+		if h.options.Observer != nil {
+			pstats := h.stats.Stats(clientID)
+			pstats.Add(stats.KindTotalConns, 1)
+			pstats.Add(stats.KindCurrentConns, 1)
+			defer pstats.Add(stats.KindCurrentConns, -1)
+			rw = stats_wrapper.WrapReadWriter(rw, pstats)
+		}
+
 		start := time.Now()
 		log.Infof("%s <-> %s", req.RemoteAddr, addr)
 		netpkg.Transport(rw, cc)
@@ -386,3 +413,21 @@ func (h *http2Handler) checkRateLimit(addr net.Addr) bool {
 
 	return true
 }
+
+func (h *http2Handler) observeStats(ctx context.Context) {
+	if h.options.Observer == nil {
+		return
+	}
+
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			h.options.Observer.Observe(ctx, h.stats.Events())
+		case <-ctx.Done():
+			return
+		}
+	}
+}