Additional semgents on /nodes endpoint to troubleshoot latency issues (…

…#148)

Co-authored-by: James Kwon <[email protected]>

go.sum

@@ -194,6 +194,8 @@ github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/
 github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0=
+github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
 github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y=
 github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
 github.com/mitchellh/go-wordwrap v1.0.1 h1:TLuKupo69TCn6TQSyGxwI1EblZZEsQ0vMlAFQflz0v0=
@@ -218,6 +220,8 @@ github.com/newrelic/go-agent/v3/integrations/nrecho-v4 v1.1.3 h1:+p/67kyG/ySqYgb
 github.com/newrelic/go-agent/v3/integrations/nrecho-v4 v1.1.3/go.mod h1:ANGVlWG7HgaLvW3npXhu074jiSB4dA6HAux/fmVTvJE=
 github.com/oapi-codegen/runtime v1.1.1 h1:EXLHh0DXIJnWhdRPN2w4MXAzFyE4CskzhNLUmtpMYro=
 github.com/oapi-codegen/runtime v1.1.1/go.mod h1:SK9X900oXmPWilYR5/WKPzt3Kqxn/uS/+lbpREv+eCg=
+github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
+github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.0-rc5 h1:Ygwkfw9bpDvs+c9E34SdgGOj41dX/cbdlwvlWt0pnFI=
@@ -244,6 +248,10 @@ github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU=
 github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
+github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spkg/bom v0.0.0-20160624110644-59b7046e48ad/go.mod h1:qLr4V1qq6nMqFKkMo8ZTx3f+BZEkzsRUY10Xsm2mwU0=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=

server/implementation/registry.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"registry-backend/drip"
 	"registry-backend/ent"
 	"registry-backend/ent/publisher"
@@ -237,6 +238,11 @@ func (s *DripStrictServerImplementation) ListNodesForPublisher(
 
 func (s *DripStrictServerImplementation) ListAllNodes(
 	ctx context.Context, request drip.ListAllNodesRequestObject) (drip.ListAllNodesResponseObject, error) {
+	if txn := newrelic.FromContext(ctx); txn != nil {
+		segment := txn.StartSegment("DripStrictServerImplementation.ListAllNodes")
+		defer segment.End()
+	}
+
 	err := s.MixpanelService.Track(ctx, []*mixpanel.Event{
 		s.MixpanelService.NewEvent("List All Nodes", "", map[string]any{
 			"page":  request.Params.Page,

server/middleware/request_logger.go

@@ -75,6 +75,11 @@ func RequestLoggerMiddleware() echo.MiddlewareFunc {
 
 	mw := func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) error {
+			if txn := newrelic.FromContext(c.Request().Context()); txn != nil {
+				segment := txn.StartSegment("RequestLoggerMiddleware")
+				defer segment.End()
+			}
+
 			req := c.Request()
 			reader := &teeReader{ReadCloser: req.Body}
 			req.Body = io.NopCloser(reader)

server/server.go

@@ -56,7 +56,6 @@ func NewServer(client *ent.Client, config *config.Config) (*Server, error) {
 		newrelic.ConfigDistributedTracerEnabled(true),
 		newrelic.ConfigEnabled(true),
 	)
-
 	if err != nil {
 		log.Error().Err(err).Msg("Failed to initialize NewRelic application")
 	}

services/registry/registry_svc.go

@@ -66,14 +66,18 @@ func NewRegistryService(storageSvc storage.StorageService, pubsubService pubsub.
 
 // ListNodes retrieves a paginated list of nodes with optional filtering.
 func (s *RegistryService) ListNodes(ctx context.Context, client *ent.Client, page, limit int, filter *entity.NodeFilter) (*entity.ListNodesResult, error) {
-	if txn := newrelic.FromContext(ctx); txn != nil {
+	// Start New Relic transaction segment
+	var txn *newrelic.Transaction
+	if txnCtx := newrelic.FromContext(ctx); txnCtx != nil {
+		txn = txnCtx
 		txn.Application().RecordCustomMetric(
 			"Custom/ListNodes/Limit",
 			float64(limit),
 		)
 		segment := txn.StartSegment("RegistryService.ListNodes")
 		defer segment.End()
 	}
+
 	// Ensure valid pagination parameters
 	if page < 1 {
 		page = 1
@@ -120,15 +124,40 @@ func (s *RegistryService) ListNodes(ctx context.Context, client *ent.Client, pag
 	// Calculate pagination offset
 	offset := (page - 1) * limit
 
-	// Count total nodes
-	total, err := query.Count(ctx)
+	// Count total nodes with New Relic datastore segment
+	var total int
+	var err error
+	if txn != nil {
+		segment := newrelic.DatastoreSegment{
+			Product:    newrelic.DatastorePostgres, // Change based on your DB
+			Collection: node.Table,                 // Table name
+			Operation:  "COUNT",
+			StartTime:  txn.StartSegmentNow(),
+		}
+		total, err = query.Count(ctx)
+		segment.End()
+	} else {
+		total, err = query.Count(ctx)
+	}
 	if err != nil {
 		return nil, fmt.Errorf("failed to count nodes: %w", err)
 	}
 
-	// Fetch nodes with pagination
+	// Fetch nodes with pagination and New Relic datastore segment
 	query = s.decorateNodeQueryWithLatestVersion(query).Offset(offset).Limit(limit)
-	nodes, err := query.All(ctx)
+	var nodes []*ent.Node
+	if txn != nil {
+		segment := newrelic.DatastoreSegment{
+			Product:    newrelic.DatastorePostgres,
+			Collection: node.Table,
+			Operation:  "SELECT",
+			StartTime:  txn.StartSegmentNow(),
+		}
+		nodes, err = query.All(ctx)
+		segment.End()
+	} else {
+		nodes, err = query.All(ctx)
+	}
 	if err != nil {
 		return nil, fmt.Errorf("failed to list nodes: %w", err)
 	}