move the helper to pyroscope package

Author: korniltsev-grafanista

internal/cmd/integration-tests/common/profiles_assert.go

@@ -1,28 +1,57 @@
 package main
 
 import (
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/grafana/alloy/internal/cmd/integration-tests/common"
+	pyroutil "github.com/grafana/alloy/internal/component/pyroscope/util/test"
+	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
+	"github.com/stretchr/testify/require"
 )
 
 func TestPyroscopeJavaKafka(t *testing.T) {
-	q := `{service_name="integration-test/java/kafka"}`
-	common.AssertPyroscopeProfile(
-		t,
-		`process_cpu:cpu:nanoseconds:cpu:nanoseconds`,
-		q,
-		[]string{
-			`kafka/server/KafkaRequestHandler.run`,
-			`libjvm.so.JavaThread::thread_main_inner`,
-		},
-	)
-	common.AssertPyroscopeProfile(
-		t,
-		"memory:alloc_in_new_tlab_bytes:bytes:space:bytes",
-		q,
-		[]string{
-			`kafka/server/KafkaRequestHandler.run`,
-		},
-	)
+
+	require.Eventually(t, func() bool {
+		req := &querierv1.SelectMergeProfileRequest{
+			ProfileTypeID: `process_cpu:cpu:nanoseconds:cpu:nanoseconds`,
+			LabelSelector: `{service_name="integration-test/java/kafka"}`,
+			Start:         time.Now().Add(-time.Hour).UnixMilli(),
+			End:           time.Now().UnixMilli(),
+		}
+		res, err := pyroutil.Query("http://localhost:4040", req)
+		if err != nil {
+			return false
+		}
+		ss := res.String()
+		if !strings.Contains(ss, `kafka/server/KafkaRequestHandler.run`) {
+			return false
+		}
+		if !strings.Contains(ss, `libjvm.so.JavaThread::thread_main_inner`) {
+			return false
+		}
+		return true
+	}, common.DefaultTimeout, common.DefaultRetryInterval)
+
+	require.Eventually(t, func() bool {
+		req := &querierv1.SelectMergeProfileRequest{
+			ProfileTypeID: `memory:alloc_in_new_tlab_bytes:bytes:space:bytes`,
+			LabelSelector: `{service_name="integration-test/java/kafka"}`,
+			Start:         time.Now().Add(-time.Hour).UnixMilli(),
+			End:           time.Now().UnixMilli(),
+		}
+		res, err := pyroutil.Query("http://localhost:4040", req)
+		if err != nil {
+			return false
+		}
+		ss := res.String()
+		if !strings.Contains(ss, `kafka/server/KafkaRequestHandler.run`) {
+			return false
+		}
+		if strings.Contains(ss, `libjvm.so.JavaThread::thread_main_inner`) {
+			return false
+		}
+		return true
+	}, common.DefaultTimeout, common.DefaultRetryInterval)
 }

internal/cmd/integration-tests/tests/pyroscope-java/pyroscope_java_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/go-kit/log"
 	"github.com/gorilla/mux"
 	pyroutil "github.com/grafana/alloy/internal/component/pyroscope/util"
+	"github.com/grafana/alloy/internal/component/pyroscope/util/tracelog"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/model/labels"
 	"go.opentelemetry.io/otel/trace"
@@ -170,7 +171,7 @@ func (c *Component) Push(ctx context.Context, req *connect.Request[pushv1.PushRe
 
 	ctx, sp := c.tracer.Start(ctx, "/push.v1.PusherService/Push")
 	defer sp.End()
-	l := pyroutil.TraceLog(c.logger, sp)
+	l := tracelog.TraceLog(c.logger, sp)
 
 	var wg sync.WaitGroup
 	var errs error
@@ -224,7 +225,7 @@ func (c *Component) handleIngest(w http.ResponseWriter, r *http.Request) {
 	ctx, sp := c.tracer.Start(ctx, "/ingest")
 	defer sp.End()
 
-	l := pyroutil.TraceLog(c.logger, sp)
+	l := tracelog.TraceLog(c.logger, sp)
 
 	// Parse labels early
 	var lbls labels.Labels

internal/component/pyroscope/receive_http/receive_http.go

@@ -0,0 +1,24 @@
+package test
+
+import (
+	"context"
+	"net/http"
+
+	"connectrpc.com/connect"
+	"github.com/google/pprof/profile"
+	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
+	"github.com/grafana/pyroscope/api/gen/proto/go/querier/v1/querierv1connect"
+)
+
+func Query(url string, q *querierv1.SelectMergeProfileRequest) (*profile.Profile, error) {
+	client := querierv1connect.NewQuerierServiceClient(http.DefaultClient, url)
+	res, err := client.SelectMergeProfile(context.Background(), connect.NewRequest(q))
+	if err != nil {
+		return nil, err
+	}
+	bs, err := res.Msg.MarshalVT()
+	if err != nil {
+		return nil, err
+	}
+	return profile.ParseData(bs)
+}

internal/component/pyroscope/util/test/query.go

@@ -1,4 +1,4 @@
-package util
+package tracelog
 
 import (
 	"github.com/go-kit/log"

internal/component/pyroscope/util/log.go renamed to internal/component/pyroscope/util/tracelog/log.go

@@ -22,6 +22,7 @@ import (
 	"github.com/grafana/alloy/internal/component/common/config"
 	"github.com/grafana/alloy/internal/component/pyroscope"
 	"github.com/grafana/alloy/internal/component/pyroscope/util"
+	"github.com/grafana/alloy/internal/component/pyroscope/util/tracelog"
 	"github.com/grafana/alloy/internal/featuregate"
 	"github.com/grafana/alloy/internal/useragent"
 	"github.com/grafana/dskit/backoff"
@@ -260,7 +261,7 @@ func (f *fanOutClient) Push(
 		dl                    any
 		ok                    bool
 		reqSize, profileCount = requestSize(req)
-		l                     = util.TraceLog(f.logger, sp)
+		l                     = tracelog.TraceLog(f.logger, sp)
 		st                    = time.Now()
 	)
 	if dl, ok = ctx.Deadline(); !ok {
@@ -466,7 +467,7 @@ func (f *fanOutClient) AppendIngest(ctx context.Context, profile *pyroscope.Inco
 		dl                    any
 		ok                    bool
 		reqSize, profileCount = int64(len(profile.RawBody)), int64(1)
-		l                     = util.TraceLog(f.logger, sp)
+		l                     = tracelog.TraceLog(f.logger, sp)
 		st                    = time.Now()
 	)
 	if dl, ok = ctx.Deadline(); !ok {