fix(pyroscope.java): Ensure the jfr is deleted on shutdown (#3630)

* fix(pyroscope.java): Ensure the jfr is deleted on shutdown

* Address review comments

* Improve log messages

Author: simonswine

internal/component/pyroscope/java/loop.go

@@ -38,7 +38,7 @@ type profilingLoop struct {
 	dist       *asprof.Distribution
 	jfrFile    string
 	startTime  time.Time
-	profiler   *asprof.Profiler
+	profiler   Profiler
 	sampleRate int
 
 	error            error
@@ -49,7 +49,13 @@ type profilingLoop struct {
 	totalSamples     int64
 }
 
-func newProfilingLoop(pid int, target discovery.Target, logger log.Logger, profiler *asprof.Profiler, output *pyroscope.Fanout, cfg ProfilingConfig) *profilingLoop {
+type Profiler interface {
+	CopyLib(dist *asprof.Distribution, pid int) error
+	Execute(dist *asprof.Distribution, argv []string) (string, string, error)
+	Distribution() *asprof.Distribution
+}
+
+func newProfilingLoop(pid int, target discovery.Target, logger log.Logger, profiler Profiler, output *pyroscope.Fanout, cfg ProfilingConfig) *profilingLoop {
 	ctx, cancel := context.WithCancel(context.Background())
 	dist := profiler.Distribution()
 	p := &profilingLoop{
@@ -114,15 +120,37 @@ func (p *profilingLoop) loop(ctx context.Context) {
 	}
 }
 
+func (p *profilingLoop) cleanupJFR() {
+	// first try to find through process path
+	jfrFile := asprof.ProcessPath(p.jfrFile, p.pid)
+	if err := os.Remove(jfrFile); os.IsNotExist(err) {
+		// the process path was not found, this is possible when the target process stopped in the meantime.
+
+		if jfrFile == p.jfrFile {
+			// nothing we can do, the process path was not actually a /proc path
+			return
+		}
+
+		jfrFile = p.jfrFile
+		if err := os.Remove(jfrFile); os.IsNotExist(err) {
+			_ = level.Debug(p.logger).Log("msg", "unable to delete jfr file, likely because target process is stopped and was containerized", "path", jfrFile, "err", err)
+			// file not found on the host system, process was likely containerized and we can't delete this file anymore
+			return
+		} else if err != nil {
+			_ = level.Warn(p.logger).Log("msg", "failed to delete jfr file at host path", "path", jfrFile, "err", err)
+		}
+	} else if err != nil {
+		_ = level.Warn(p.logger).Log("msg", "failed to delete jfr file at process path", "path", jfrFile, "err", err)
+	}
+}
+
 func (p *profilingLoop) reset() error {
 	jfrFile := asprof.ProcessPath(p.jfrFile, p.pid)
 	startTime := p.startTime
 	endTime := time.Now()
 	sampleRate := p.sampleRate
 	p.startTime = endTime
-	defer func() {
-		os.Remove(jfrFile)
-	}()
+	defer p.cleanupJFR()
 
 	err := p.stop()
 	if err != nil {
@@ -262,6 +290,7 @@ func (p *profilingLoop) update(target discovery.Target, config ProfilingConfig)
 func (p *profilingLoop) Close() error {
 	p.cancel()
 	p.wg.Wait()
+	p.cleanupJFR()
 	return nil
 }
 

internal/component/pyroscope/java/loop_test.go

@@ -0,0 +1,124 @@
+//go:build (linux || darwin) && (amd64 || arm64)
+
+package java
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/alloy/internal/component/discovery"
+	"github.com/grafana/alloy/internal/component/pyroscope"
+	"github.com/grafana/alloy/internal/component/pyroscope/java/asprof"
+)
+
+type mockProfiler struct {
+	mock.Mock
+	dist *asprof.Distribution
+}
+
+func (m *mockProfiler) CopyLib(dist *asprof.Distribution, pid int) error {
+	args := m.Called(dist, pid)
+	return args.Error(0)
+}
+
+func (m *mockProfiler) Execute(dist *asprof.Distribution, argv []string) (string, string, error) {
+	args := m.Called(dist, argv)
+	return args.String(0), args.String(1), args.Error(2)
+}
+
+func (m *mockProfiler) Distribution() *asprof.Distribution {
+	return m.dist
+}
+
+type mockAppendable struct {
+	mock.Mock
+}
+
+func (m *mockAppendable) Appender() pyroscope.Appender {
+	return m
+}
+
+func (m *mockAppendable) Append(ctx context.Context, labels labels.Labels, samples []*pyroscope.RawSample) error {
+	args := m.Called(ctx, labels, samples)
+	return args.Error(0)
+}
+
+func (m *mockAppendable) AppendIngest(ctx context.Context, profile *pyroscope.IncomingProfile) error {
+	args := m.Called(ctx, profile)
+	return args.Error(0)
+}
+
+func newTestProfilingLoop(_ *testing.T, profiler *mockProfiler, appendable pyroscope.Appendable) *profilingLoop {
+	reg := prometheus.NewRegistry()
+	output := pyroscope.NewFanout([]pyroscope.Appendable{appendable}, "test-appendable", reg)
+	logger := log.NewNopLogger()
+	cfg := ProfilingConfig{
+		Interval:   10 * time.Millisecond,
+		SampleRate: 1000,
+		CPU:        true,
+		Event:      "cpu",
+	}
+	target := discovery.NewTargetFromMap(map[string]string{"foo": "bar"})
+	return newProfilingLoop(os.Getpid(), target, logger, profiler, output, cfg)
+}
+
+func TestProfilingLoop_StartStop(t *testing.T) {
+	profiler := &mockProfiler{dist: &asprof.Distribution{}}
+	appendable := &mockAppendable{}
+	pid := os.Getpid()
+	jfrPath := fmt.Sprintf("/tmp/asprof-%d-%d.jfr", pid, pid)
+
+	pCh := make(chan *profilingLoop)
+
+	profiler.On("CopyLib", profiler.dist, pid).Return(nil).Once()
+
+	// expect the profiler to be executed with the correct arguments to start it
+	profiler.On("Execute", profiler.dist, []string{
+		"-f",
+		jfrPath,
+		"-o", "jfr",
+		"-e", "cpu",
+		"-i", "1000000",
+		"start",
+		"--timeout", "0",
+		strconv.Itoa(pid),
+	}).Run(func(args mock.Arguments) {
+		// wait for the profiling loop to be created
+		p := <-pCh
+
+		// create the jfr file
+		f, err := os.Create(p.jfrFile)
+		require.NoError(t, err)
+		defer f.Close()
+	}).Return("", "", nil).Once()
+
+	// expect the profiler to be executed with the correct arguments to stop it
+	profiler.On("Execute", profiler.dist, []string{
+		"stop",
+		"-o", "jfr",
+		strconv.Itoa(pid),
+	}).Return("", "", nil).Once()
+
+	p := newTestProfilingLoop(t, profiler, appendable)
+	pCh <- p
+
+	// wait for the profiling loop to finish
+	require.NoError(t, p.Close())
+
+	// expect the profiler to clean up the jfr file
+	_, err := os.Stat(p.jfrFile)
+	require.True(t, os.IsNotExist(err))
+
+	profiler.AssertExpectations(t)
+	appendable.AssertExpectations(t)
+}