pyroscope.java: allow custom asprof distributions instead of embeded one.

.dockerignore

@@ -7,3 +7,4 @@ internal/web/ui/build/
 packaging/windows/LICENSE
 packaging/windows/agent-windows-amd64.exe
 cmd/grafana-agent/Dockerfile
+alloy

.github/workflows/test_pyroscope_pr.yml

@@ -26,4 +26,4 @@ jobs:
           go-version-file: go.mod
           cache: false
 
-      - run: make GO_TAGS="nodocker" test-pyroscope
+      - run: sudo make test-pyroscope

.gitignore

@@ -33,3 +33,4 @@ node_modules
 # file of the pair will detect a dirty work tree and detect the wrong tag name.
 .tag-only
 .image-tag
+alloy

internal/component/pyroscope/java/args.go

@@ -14,6 +14,9 @@ type Arguments struct {
 
 	TmpDir          string          `alloy:"tmp_dir,attr,optional"`
 	ProfilingConfig ProfilingConfig `alloy:"profiling_config,block,optional"`
+
+	// undocumented
+	Dist string `alloy:"dist,attr,optional"`
 }
 
 type ProfilingConfig struct {
@@ -29,7 +32,7 @@ type ProfilingConfig struct {
 }
 
 func (rc *Arguments) UnmarshalAlloy(f func(interface{}) error) error {
-	*rc = defaultArguments()
+	*rc = DefaultArguments()
 	type config Arguments
 	return f((*config)(rc))
 }
@@ -43,7 +46,7 @@ func (arg *Arguments) Validate() error {
 	}
 }
 
-func defaultArguments() Arguments {
+func DefaultArguments() Arguments {
 	return Arguments{
 		TmpDir: "/tmp",
 		ProfilingConfig: ProfilingConfig{

internal/component/pyroscope/java/asprof/asprof.go

@@ -20,49 +20,47 @@ var fsMutex sync.Mutex
 
 type Distribution struct {
 	extractedDir string
-	version      int
 }
 
-func (d *Distribution) LauncherPath() string {
-	return filepath.Join(d.extractedDir, "bin/asprof")
+func NewExtractedDistribution(extractedDir string) (Distribution, error) {
+	d := Distribution{extractedDir: extractedDir}
+	if _, err := os.Stat(d.LauncherPath()); err != nil {
+		return d, fmt.Errorf("asprof launcher not found: %w", err)
+	}
+	if _, err := os.Stat(d.LibPath()); err != nil {
+		return d, fmt.Errorf("asprof lib not found: %w", err)
+	}
+	return d, nil
 }
 
-type Profiler struct {
-	tmpDir       string
-	extractOnce  sync.Once
-	dist         *Distribution
-	extractError error
-	tmpDirMarker any
-	archiveHash  string
-	archive      Archive
+func (d Distribution) LauncherPath() string {
+	return filepath.Join(d.extractedDir, "bin/asprof")
 }
 
 type Archive struct {
-	data    []byte
-	version int
-	format  int
+	data   []byte
+	format int
+}
+
+func (a *Archive) sha1() string {
+	sum := sha1.Sum(a.data)
+	return hex.EncodeToString(sum[:])
+}
+
+func (a *Archive) DistName() string {
+	return fmt.Sprintf("alloy-asprof-%s", a.sha1())
 }
 
 const (
 	ArchiveFormatTarGz = iota
 	ArchiveFormatZip
 )
 
-func NewProfiler(tmpDir string, archive Archive) *Profiler {
-	res := &Profiler{tmpDir: tmpDir, dist: new(Distribution), tmpDirMarker: "alloy-asprof"}
-	sum := sha1.Sum(archive.data)
-	hexSum := hex.EncodeToString(sum[:])
-	res.archiveHash = hexSum
-	res.dist.version = archive.version
-	res.archive = archive
-	return res
-}
-
-func (p *Profiler) Execute(dist *Distribution, argv []string) (string, string, error) {
+func (d Distribution) Execute(argv []string) (string, string, error) {
 	stdout := bytes.NewBuffer(nil)
 	stderr := bytes.NewBuffer(nil)
 
-	exe := dist.LauncherPath()
+	exe := d.LauncherPath()
 	cmd := exec.Command(exe, argv...)
 
 	cmd.Stdout = stdout
@@ -78,61 +76,46 @@ func (p *Profiler) Execute(dist *Distribution, argv []string) (string, string, e
 	return stdout.String(), stderr.String(), nil
 }
 
-func (p *Profiler) Distribution() *Distribution {
-	return p.dist
-}
-
-func (p *Profiler) ExtractDistributions() error {
-	p.extractOnce.Do(func() {
-		p.extractError = p.extractDistributions()
-	})
-	return p.extractError
-}
-
-func (p *Profiler) extractDistributions() error {
+func ExtractDistribution(a Archive, tmpDir, distName string) (Distribution, error) {
+	d := Distribution{}
 	fsMutex.Lock()
 	defer fsMutex.Unlock()
-	distName := p.getDistName()
 
-	var launcher, dist []byte
-	err := readArchive(p.archive.data, p.archive.format, func(name string, fi fs.FileInfo, data []byte) error {
+	var launcher, lib []byte
+	err := readArchive(a.data, a.format, func(name string, fi fs.FileInfo, data []byte) error {
 		if strings.Contains(name, "asprof") {
 			launcher = data
 		}
 		if strings.Contains(name, "libasyncProfiler") {
-			dist = data
+			lib = data
 		}
 		return nil
 	})
 	if err != nil {
-		return err
+		return d, err
 	}
-	if launcher == nil || dist == nil {
-		return fmt.Errorf("failed to find libasyncProfiler in archive %s", distName)
+	if launcher == nil || lib == nil {
+		return d, fmt.Errorf("failed to find libasyncProfiler in archive %s", distName)
 	}
 
 	fileMap := map[string][]byte{}
-	fileMap[filepath.Join(distName, p.dist.LauncherPath())] = launcher
-	fileMap[filepath.Join(distName, p.dist.LibPath())] = dist
-	tmpDirFile, err := os.Open(p.tmpDir)
+	fileMap[filepath.Join(distName, d.LauncherPath())] = launcher
+	fileMap[filepath.Join(distName, d.LibPath())] = lib
+	tmpDirFile, err := os.Open(tmpDir)
 	if err != nil {
-		return fmt.Errorf("failed to open tmp dir %s: %w", p.tmpDir, err)
+		return d, fmt.Errorf("failed to open tmp dir %s: %w", tmpDir, err)
 	}
 	defer tmpDirFile.Close()
 
 	if err = checkTempDirPermissions(tmpDirFile); err != nil {
-		return err
+		return d, err
 	}
 
 	for path, data := range fileMap {
 		if err = writeFile(tmpDirFile, path, data, true); err != nil {
-			return err
+			return d, err
 		}
 	}
-	p.dist.extractedDir = filepath.Join(p.tmpDir, distName)
-	return nil
-}
-
-func (p *Profiler) getDistName() string {
-	return fmt.Sprintf("%s-%s", p.tmpDirMarker, p.archiveHash)
+	d.extractedDir = filepath.Join(tmpDir, distName)
+	return d, nil
 }

internal/component/pyroscope/java/asprof/asprof_darwin.go

@@ -13,15 +13,13 @@ var embeddedArchiveData []byte
 // bin/asprof
 // lib/libasyncProfiler.dylib
 
-var embeddedArchiveVersion = 300
+var EmbeddedArchive = Archive{data: embeddedArchiveData, format: ArchiveFormatZip}
 
-var EmbeddedArchive = Archive{data: embeddedArchiveData, version: embeddedArchiveVersion, format: ArchiveFormatZip}
-
-func (d *Distribution) LibPath() string {
+func (d Distribution) LibPath() string {
 	return filepath.Join(d.extractedDir, "lib/libasyncProfiler.dylib")
 }
 
-func (p *Profiler) CopyLib(dist *Distribution, pid int) error {
+func (d Distribution) CopyLib(pid int) error {
 	return nil
 }
 

internal/component/pyroscope/java/asprof/asprof_linux.go

@@ -10,22 +10,20 @@ import (
 	"strings"
 )
 
-var embeddedArchiveVersion = 300
+var EmbeddedArchive = Archive{data: embeddedArchiveData, format: ArchiveFormatTarGz}
 
-var EmbeddedArchive = Archive{data: embeddedArchiveData, version: embeddedArchiveVersion, format: ArchiveFormatTarGz}
-
-func (d *Distribution) LibPath() string {
+func (d Distribution) LibPath() string {
 	return filepath.Join(d.extractedDir, "lib/libasyncProfiler.so")
 }
 
-func (p *Profiler) CopyLib(dist *Distribution, pid int) error {
+func (d Distribution) CopyLib(pid int) error {
 	fsMutex.Lock()
 	defer fsMutex.Unlock()
-	libData, err := os.ReadFile(dist.LibPath())
+	libData, err := os.ReadFile(d.LibPath())
 	if err != nil {
 		return err
 	}
-	launcherData, err := os.ReadFile(dist.LauncherPath())
+	launcherData, err := os.ReadFile(d.LauncherPath())
 	if err != nil {
 		return err
 	}
@@ -35,8 +33,8 @@ func (p *Profiler) CopyLib(dist *Distribution, pid int) error {
 		return fmt.Errorf("failed to open proc root %s: %w", procRoot, err)
 	}
 	defer procRootFile.Close()
-	dstLibPath := strings.TrimPrefix(dist.LibPath(), "/")
-	dstLauncherPath := strings.TrimPrefix(dist.LauncherPath(), "/")
+	dstLibPath := strings.TrimPrefix(d.LibPath(), "/")
+	dstLauncherPath := strings.TrimPrefix(d.LauncherPath(), "/")
 	if err = writeFile(procRootFile, dstLibPath, libData, false); err != nil {
 		return err
 	}
@@ -48,15 +46,5 @@ func (p *Profiler) CopyLib(dist *Distribution, pid int) error {
 }
 
 func ProcessPath(path string, pid int) string {
-	f := procFile{path, pid}
-	return f.procRootPath()
-}
-
-type procFile struct {
-	path string
-	pid  int
-}
-
-func (f *procFile) procRootPath() string {
-	return filepath.Join("/proc", strconv.Itoa(f.pid), "root", f.path)
+	return filepath.Join("/proc", strconv.Itoa(pid), "root", path)
 }

internal/component/pyroscope/java/asprof/asprof_test.go

@@ -9,7 +9,7 @@ import (
 	"testing"
 
 	"github.com/google/uuid"
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // extracting to /tmp
@@ -23,45 +23,45 @@ import (
 // write skippable tests with uid=0
 func TestStickyDir(t *testing.T) {
 	dir := "/tmp"
-	p := NewProfiler(dir, EmbeddedArchive)
-	p.tmpDirMarker = fmt.Sprintf("alloy-asprof-%s", uuid.NewString())
-	t.Logf("tmpDirMarker: %s", p.tmpDirMarker)
-	err := p.ExtractDistributions()
-	assert.NoError(t, err)
+	tmpDirMarker := fmt.Sprintf("alloy-asprof-%s", uuid.NewString())
+	t.Logf("tmpDirMarker: %s", tmpDirMarker)
+	dist, err := ExtractDistribution(EmbeddedArchive, dir, tmpDirMarker)
+	require.NoError(t, err)
+	require.NotNil(t, dist)
 }
 
 func TestOwnedDir(t *testing.T) {
 	dir := t.TempDir()
 	err := os.Chmod(dir, 0755)
-	assert.NoError(t, err)
-	p := NewProfiler(dir, EmbeddedArchive)
-	err = p.ExtractDistributions()
-	assert.NoError(t, err)
+	require.NoError(t, err)
+	dist, err := ExtractDistribution(EmbeddedArchive, dir, "alloy-asprof")
+	require.NoError(t, err)
+	require.NotNil(t, dist)
 }
 
 func TestOwnedDirWrongPermission(t *testing.T) {
 	dir := t.TempDir()
 	err := os.Chmod(dir, 0777)
-	assert.NoError(t, err)
-	p := NewProfiler(dir, EmbeddedArchive)
-	err = p.ExtractDistributions()
-	assert.Error(t, err)
+	require.NoError(t, err)
+	dist, err := ExtractDistribution(EmbeddedArchive, dir, "alloy-asprof-")
+	require.Error(t, err)
+	require.Empty(t, dist.extractedDir)
 }
 
 func TestDistSymlink(t *testing.T) {
 	root := t.TempDir()
 	err := os.Chmod(root, 0755)
-	assert.NoError(t, err)
+	require.NoError(t, err)
 	manipulated := t.TempDir()
 	err = os.Chmod(manipulated, 0755)
-	assert.NoError(t, err)
-	p := NewProfiler(root, EmbeddedArchive)
-	distName := p.getDistName()
+	require.NoError(t, err)
+	distName := "dist"
 
 	err = os.Symlink(manipulated, filepath.Join(root, distName))
-	assert.NoError(t, err)
+	require.NoError(t, err)
 
-	err = p.ExtractDistributions()
+	dist, err := ExtractDistribution(EmbeddedArchive, root, distName)
 	t.Logf("expected %s", err)
-	assert.Error(t, err)
+	require.Error(t, err)
+	require.Empty(t, dist.extractedDir)
 }