pyroscope: trace_id propagation and loging (#4354)

Author: korniltsev-grafanista

CHANGELOG.md

@@ -70,6 +70,8 @@ Main (unreleased)
 
 - Fix `pyroscope.write` component's `AppendIngest` method to respect configured timeout and implement retry logic. The method now properly uses the configured `remote_timeout`, includes retry logic with exponential backoff, and tracks metrics for sent/dropped bytes and profiles consistently with the `Append` method. (@korniltsev)
 
+- `pyroscope.write`, `pyroscope.receive_http` components include `trace_id` in logs and propagate it downstream. (@korniltsev)
+
 - Improve logging in `pyroscope.write` component. (@korniltsev)
 
 - Add comprehensive latency metrics to `pyroscope.write` component with endpoint-specific tracking for both push and ingest operations. (@korniltsev, @claude)

internal/component/pyroscope/receive_http/receive_http.go

@@ -11,10 +11,12 @@ import (
 	"sync"
 
 	"connectrpc.com/connect"
-
+	"github.com/go-kit/log"
 	"github.com/gorilla/mux"
+	pyroutil "github.com/grafana/alloy/internal/component/pyroscope/util"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/model/labels"
+	"go.opentelemetry.io/otel/trace"
 
 	"github.com/grafana/alloy/internal/component"
 	fnet "github.com/grafana/alloy/internal/component/common/net"
@@ -35,7 +37,8 @@ func init() {
 		Stability: featuregate.StabilityGenerallyAvailable,
 		Args:      Arguments{},
 		Build: func(opts component.Options, args component.Arguments) (component.Component, error) {
-			return New(opts, args.(Arguments))
+			tracer := opts.Tracer.Tracer("pyroscope.receive_http")
+			return New(opts.Logger, tracer, opts.Registerer, args.(Arguments))
 		},
 	})
 }
@@ -54,20 +57,22 @@ func (a *Arguments) SetToDefault() {
 }
 
 type Component struct {
-	opts               component.Options
 	server             *fnet.TargetServer
 	serverConfig       *fnet.HTTPConfig
 	uncheckedCollector *util.UncheckedCollector
 	appendables        []pyroscope.Appendable
 	mut                sync.Mutex
+	logger             log.Logger
+	tracer             trace.Tracer
 }
 
-func New(opts component.Options, args Arguments) (*Component, error) {
+func New(logger log.Logger, tracer trace.Tracer, reg prometheus.Registerer, args Arguments) (*Component, error) {
 	uncheckedCollector := util.NewUncheckedCollector(nil)
-	opts.Registerer.MustRegister(uncheckedCollector)
+	reg.MustRegister(uncheckedCollector)
 
 	c := &Component{
-		opts:               opts,
+		logger:             logger,
+		tracer:             tracer,
 		uncheckedCollector: uncheckedCollector,
 		appendables:        args.ForwardTo,
 	}
@@ -87,7 +92,7 @@ func (c *Component) Run(ctx context.Context) error {
 	}()
 
 	<-ctx.Done()
-	level.Info(c.opts.Logger).Log("msg", "terminating due to context done")
+	level.Info(c.logger).Log("msg", "terminating due to context done")
 	return nil
 }
 
@@ -122,7 +127,7 @@ func (c *Component) update(args component.Arguments) (bool, error) {
 	serverRegistry := prometheus.NewRegistry()
 	c.uncheckedCollector.SetCollector(serverRegistry)
 
-	srv, err := fnet.NewTargetServer(c.opts.Logger, "pyroscope_receive_http", serverRegistry, newArgs.Server)
+	srv, err := fnet.NewTargetServer(c.logger, "pyroscope_receive_http", serverRegistry, newArgs.Server)
 	if err != nil {
 		return shutdown, fmt.Errorf("failed to create server: %w", err)
 	}
@@ -163,6 +168,10 @@ func (c *Component) Push(ctx context.Context, req *connect.Request[pushv1.PushRe
 
 	appendables := c.getAppendables()
 
+	ctx, sp := c.tracer.Start(ctx, "/push.v1.PusherService/Push")
+	defer sp.End()
+	l := pyroutil.TraceLog(c.logger, sp)
+
 	var wg sync.WaitGroup
 	var errs error
 	var errorMut sync.Mutex
@@ -182,7 +191,7 @@ func (c *Component) Push(ctx context.Context, req *connect.Request[pushv1.PushRe
 				lbls := ensureServiceName(lb.Labels())
 				err := appendable.Append(ctx, lbls, apiToAlloySamples(req.Msg.Series[idx].Samples))
 				if err != nil {
-					util.ErrorsJoinConcurrent(
+					pyroutil.ErrorsJoinConcurrent(
 						&errs,
 						fmt.Errorf("unable to append series %s to appendable %d: %w", lb.Labels().String(), i, err),
 						&errorMut,
@@ -193,11 +202,10 @@ func (c *Component) Push(ctx context.Context, req *connect.Request[pushv1.PushRe
 	}
 	wg.Wait()
 	if errs != nil {
-		level.Error(c.opts.Logger).Log("msg", "Failed to forward profiles requests", "err", errs)
+		level.Warn(l).Log("msg", "Failed to forward profiles requests", "err", errs)
 		return nil, connect.NewError(connect.CodeInternal, errs)
 	}
 
-	level.Debug(c.opts.Logger).Log("msg", "Profiles successfully forwarded")
 	return connect.NewResponse(&pushv1.PushResponse{}), nil
 }
 
@@ -211,12 +219,19 @@ func (c *Component) getAppendables() []pyroscope.Appendable {
 func (c *Component) handleIngest(w http.ResponseWriter, r *http.Request) {
 	appendables := c.getAppendables()
 
+	ctx := r.Context()
+
+	ctx, sp := c.tracer.Start(ctx, "/ingest")
+	defer sp.End()
+
+	l := pyroutil.TraceLog(c.logger, sp)
+
 	// Parse labels early
 	var lbls labels.Labels
 	if nameParam := r.URL.Query().Get("name"); nameParam != "" {
 		ls, err := labelset.Parse(nameParam)
 		if err != nil {
-			level.Warn(c.opts.Logger).Log(
+			level.Warn(l).Log(
 				"msg", "Failed to parse labels from name parameter",
 				"name", nameParam,
 				"err", err,
@@ -229,7 +244,7 @@ func (c *Component) handleIngest(w http.ResponseWriter, r *http.Request) {
 			}
 			lbls = labels.New(labelPairs...)
 		}
-	}
+	} // todo this is a required parameter, treat absence as error
 
 	// Ensure service_name label is set
 	lbls = ensureServiceName(lbls)
@@ -239,7 +254,7 @@ func (c *Component) handleIngest(w http.ResponseWriter, r *http.Request) {
 	// but means the entire profile will be held in memory
 	var buf bytes.Buffer
 	if _, err := io.Copy(&buf, r.Body); err != nil {
-		level.Error(c.opts.Logger).Log("msg", "Failed to read request body", "err", err)
+		level.Warn(l).Log("msg", "Failed to read request body", "err", err)
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
@@ -260,23 +275,21 @@ func (c *Component) handleIngest(w http.ResponseWriter, r *http.Request) {
 				Labels:      lbls,
 			}
 
-			if err := appendable.Appender().AppendIngest(r.Context(), profile); err != nil {
-				level.Error(c.opts.Logger).Log("msg", "Failed to append profile", "appendable", i, "err", err)
-				util.ErrorsJoinConcurrent(&errs, err, &errorMut)
+			if err := appendable.Appender().AppendIngest(ctx, profile); err != nil {
+				err = fmt.Errorf("failed to ingest profile to appendable %d: %w", i, err)
+				pyroutil.ErrorsJoinConcurrent(&errs, err, &errorMut)
 			}
-
-			level.Debug(c.opts.Logger).Log("msg", "Profile appended successfully", "appendable", i)
 		}()
 	}
 
 	wg.Wait()
 
 	if errs != nil {
+		level.Warn(l).Log("msg", "Failed to ingest profiles", "err", errs)
 		var writeErr *write.PyroscopeWriteError
 		if errors.As(errs, &writeErr) {
 			http.Error(w, http.StatusText(writeErr.StatusCode), writeErr.StatusCode)
 		} else {
-			level.Error(c.opts.Logger).Log("msg", "Failed to process request", "err", errs)
 			http.Error(w, "Failed to process request", http.StatusInternalServerError)
 		}
 		return

internal/component/pyroscope/receive_http/receive_http_test.go

@@ -13,13 +13,13 @@ import (
 	"time"
 
 	"connectrpc.com/connect"
+	"go.opentelemetry.io/otel/trace/noop"
 
 	"github.com/phayes/freeport"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/stretchr/testify/require"
 
-	"github.com/grafana/alloy/internal/component"
 	fnet "github.com/grafana/alloy/internal/component/common/net"
 	"github.com/grafana/alloy/internal/component/pyroscope"
 	"github.com/grafana/alloy/internal/util"
@@ -425,7 +425,12 @@ func startComponent(t *testing.T, appendables []pyroscope.Appendable) int {
 		ForwardTo: appendables,
 	}
 
-	comp, err := New(testOptions(t), args)
+	comp, err := New(
+		util.TestAlloyLogger(t),
+		noop.Tracer{},
+		prometheus.NewRegistry(),
+		args,
+	)
 	require.NoError(t, err)
 
 	ctx, cancel := context.WithCancel(t.Context())
@@ -540,14 +545,6 @@ func (a *testAppender) AppendIngest(_ context.Context, profile *pyroscope.Incomi
 	return a.appendErr
 }
 
-func testOptions(t *testing.T) component.Options {
-	return component.Options{
-		ID:         "pyroscope.receive_http.test",
-		Logger:     util.TestAlloyLogger(t),
-		Registerer: prometheus.NewRegistry(),
-	}
-}
-
 // TestUpdateArgs verifies that the component can be updated with new arguments. This explicitly also makes sure that the server is restarted when the server configuration changes. And there are no metric registration conflicts.
 func TestUpdateArgs(t *testing.T) {
 	ports, err := freeport.GetFreePorts(2)
@@ -565,7 +562,12 @@ func TestUpdateArgs(t *testing.T) {
 		ForwardTo: forwardTo,
 	}
 
-	comp, err := New(testOptions(t), args)
+	comp, err := New(
+		util.TestAlloyLogger(t),
+		noop.Tracer{},
+		prometheus.NewRegistry(),
+		args,
+	)
 	require.NoError(t, err)
 
 	ctx, cancel := context.WithCancel(t.Context())

internal/util/errors.go renamed to internal/component/pyroscope/util/errors.go

@@ -0,0 +1,12 @@
+package util
+
+import (
+	"github.com/go-kit/log"
+	"go.opentelemetry.io/otel/trace"
+)
+
+func TraceLog(l log.Logger, sp trace.Span) log.Logger {
+	return log.With(l,
+		"trace_id", sp.SpanContext().TraceID().String(),
+	)
+}