Add nr segment for more middleware.

Comfy-Org · Jan 20, 2025 · 09da6fa · 09da6fa
1 parent aec1923
commit 09da6fa
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 0 deletions.
diff --git a/server/middleware/metric/metric_middleware.go b/server/middleware/metric/metric_middleware.go
@@ -14,6 +14,7 @@ import (
 
 	"cloud.google.com/go/monitoring/apiv3/v2/monitoringpb"
 	"github.com/labstack/echo/v4"
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/rs/zerolog/log"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -36,6 +37,10 @@ func init() {
 func MetricsMiddleware(client *monitoring.MetricClient, config *config.Config) echo.MiddlewareFunc {
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) error {
+			if txn := newrelic.FromContext(c.Request().Context()); txn != nil {
+				segment := txn.StartSegment("MetricsMiddleware")
+				defer segment.End()
+			}
 			startTime := time.Now()
 			err := next(c)
 			endTime := time.Now()

diff --git a/server/middleware/request_logger.go b/server/middleware/request_logger.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/labstack/echo/v4"
 	echo_middleware "github.com/labstack/echo/v4/middleware"
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/rs/zerolog/log"
 )
 
@@ -40,6 +41,10 @@ func RequestLoggerMiddleware() echo.MiddlewareFunc {
 		LogURI:    true,
 		LogStatus: true,
 		LogValuesFunc: func(c echo.Context, v echo_middleware.RequestLoggerValues) error {
+			if txn := newrelic.FromContext(c.Request().Context()); txn != nil {
+				segment := txn.StartSegment("RequestLoggerMiddleware")
+				defer segment.End()
+			}
 			// Get the teeReader
 			reader, ok := c.Request().Context().Value(ctxKey).(*teeReader)
 			if !ok {

diff --git a/server/middleware/response_logger.go b/server/middleware/response_logger.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 
 	"github.com/labstack/echo/v4"
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/rs/zerolog/log"
 )
 
@@ -29,6 +30,10 @@ func (rw *responseWriter) Write(p []byte) (n int, err error) {
 func ResponseLoggerMiddleware() echo.MiddlewareFunc {
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) error {
+			if txn := newrelic.FromContext(c.Request().Context()); txn != nil {
+				segment := txn.StartSegment("ResponseLoggerMiddleware")
+				defer segment.End()
+			}
 			// Create a custom response writer to capture the response body
 			rw := &responseWriter{
 				ResponseWriter: c.Response().Writer,

diff --git a/server/middleware/tracing_middleware.go b/server/middleware/tracing_middleware.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/labstack/echo/v4"
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/rs/zerolog"
 )
 
@@ -19,6 +20,11 @@ func generateFallbackCorrelationID() string {
 // TracingMiddleware is a middleware that adds a trace ID to the context
 func TracingMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
 	return func(c echo.Context) error {
+		if txn := newrelic.FromContext(c.Request().Context()); txn != nil {
+			segment := txn.StartSegment("TracingMiddleware")
+			defer segment.End()
+		}
+
 		traceID := c.Request().Header.Get("X-Cloud-Trace-Context")
 
 		if traceID == "" {