Add OTeL middleware to Exporter APIs

ssvlabs · Dec 4, 2024 · 0b7a31f · 0b7a31f
1 parent b64e1dd
commit 0b7a31f
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 31 deletions.
diff --git a/exporter/api/conn.go b/exporter/api/conn.go
@@ -139,7 +139,6 @@ func (c *conn) WriteLoop(logger *zap.Logger) {
 			c.writeLock.Lock()
 			n, err := c.sendMsg(message)
 			c.writeLock.Unlock()
-			reportStreamOutbound(c.ws.RemoteAddr().String(), err)
 			if err != nil {
 				logger.Warn("failed to send message", zap.Error(err))
 				return

diff --git a/exporter/api/metrics.go b/exporter/api/metrics.go
diff --git a/exporter/api/server.go b/exporter/api/server.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/gorilla/websocket"
 	"github.com/prysmaticlabs/prysm/v4/async/event"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.uber.org/zap"
 
 	"github.com/ssvlabs/ssv/logging"
@@ -60,8 +61,8 @@ func (ws *wsServer) UseQueryHandler(handler QueryMessageHandler) {
 func (ws *wsServer) Start(logger *zap.Logger, addr string) error {
 	logger = logger.Named(logging.NameWSServer)
 
-	ws.RegisterHandler(logger, "/query", ws.handleQuery)
-	ws.RegisterHandler(logger, "/stream", ws.handleStream)
+	ws.RegisterHandler(logger, "query", "/query", ws.handleQuery)
+	ws.RegisterHandler(logger, "stream", "/stream", ws.handleStream)
 
 	go func() {
 		if err := ws.broadcaster.FromFeed(logger, ws.out); err != nil {
@@ -93,8 +94,8 @@ func (ws *wsServer) BroadcastFeed() *event.Feed {
 }
 
 // RegisterHandler registers an end point
-func (ws *wsServer) RegisterHandler(logger *zap.Logger, endPoint string, handler func(logger *zap.Logger, conn *websocket.Conn)) {
-	ws.router.HandleFunc(endPoint, func(w http.ResponseWriter, r *http.Request) {
+func (ws *wsServer) RegisterHandler(logger *zap.Logger, name, endPoint string, handler func(logger *zap.Logger, conn *websocket.Conn)) {
+	wrappedHandler := func(w http.ResponseWriter, r *http.Request) {
 		conn, err := upgrader.Upgrade(w, r, w.Header())
 		logger := logger.With(zap.String("remote addr", conn.RemoteAddr().String()))
 		if err != nil {
@@ -110,7 +111,10 @@ func (ws *wsServer) RegisterHandler(logger *zap.Logger, endPoint string, handler
 			}
 		}()
 		handler(logger, conn)
-	})
+	}
+
+	otelHandler := otelhttp.NewHandler(http.HandlerFunc(wrappedHandler), name)
+	ws.router.Handle(endPoint, otelHandler)
 }
 
 // handleQuery receives query message and respond async

diff --git a/go.mod b/go.mod
@@ -57,6 +57,8 @@ require (
 	tailscale.com v1.72.0
 )
 
+require github.com/felixge/httpsnoop v1.0.4 // indirect
+
 require (
 	github.com/BurntSushi/toml v1.4.1-0.20240526193622-a339e1f7089c // indirect
 	github.com/DataDog/zstd v1.5.2 // indirect
@@ -222,6 +224,7 @@ require (
 	github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 // indirect
 	github.com/yusufpapurcu/wmi v1.2.2 // indirect
 	go.opencensus.io v0.24.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.57.0
 	go.opentelemetry.io/otel v1.32.0
 	go.opentelemetry.io/otel/exporters/prometheus v0.54.0
 	go.opentelemetry.io/otel/metric v1.32.0

diff --git a/go.sum b/go.sum
@@ -150,6 +150,8 @@ github.com/ethereum/go-verkle v0.1.1-0.20240306133620-7d920df305f0 h1:KrE8I4reeV
 github.com/ethereum/go-verkle v0.1.1-0.20240306133620-7d920df305f0/go.mod h1:D9AJLVXSyZQXJQVk8oh1EwjISE+sJTn2duYIZC0dy3w=
 github.com/fatih/color v1.17.0 h1:GlRw1BRJxkpqUCBKzKOw098ed57fEsKeNjpTe3cSjK4=
 github.com/fatih/color v1.17.0/go.mod h1:YZ7TlrGPkiz6ku9fK3TLD/pl3CpsiFyu8N92HLgmosI=
+github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
+github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/ferranbt/fastssz v0.0.0-20210120143747-11b9eff30ea9/go.mod h1:DyEu2iuLBnb/T51BlsiO3yLYdJC6UbGMrIkqK1KmQxM=
 github.com/ferranbt/fastssz v0.1.3 h1:ZI+z3JH05h4kgmFXdHuR1aWYsgrg7o+Fw7/NCzM16Mo=
 github.com/ferranbt/fastssz v0.1.3/go.mod h1:0Y9TEd/9XuFlh7mskMPfXiI2Dkw4Ddg9EyXt1W7MRvE=
@@ -816,6 +818,8 @@ github.com/yusufpapurcu/wmi v1.2.2/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQ
 go.opencensus.io v0.18.0/go.mod h1:vKdFvxhtzZ9onBp9VKHK8z/sRpBMnKAsufL7wlDrCOA=
 go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0=
 go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo=
+go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.57.0 h1:DheMAlT6POBP+gh8RUH19EOTnQIor5QE0uSRPtzCpSw=
+go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.57.0/go.mod h1:wZcGmeVO9nzP67aYSLDqXNWK87EZWhi7JWj1v7ZXf94=
 go.opentelemetry.io/otel v1.32.0 h1:WnBN+Xjcteh0zdk01SVqV55d/m62NJLJdIyb4y/WO5U=
 go.opentelemetry.io/otel v1.32.0/go.mod h1:00DCVSB0RQcnzlwyTfqtxSm+DRr9hpYrHjNGiBHVQIg=
 go.opentelemetry.io/otel/exporters/prometheus v0.54.0 h1:rFwzp68QMgtzu9PgP3jm9XaMICI6TsofWWPcBDKwlsU=