improves rbspy + pyspy integrations, resolves #5

Author: petethepig

.gitignore

@@ -9,3 +9,4 @@ pkged.go
 /Icon*
 /scripts/packages/
 .eslintcache
+*.pyc

pkg/agent/pyspy/pyspy.go

@@ -8,8 +8,9 @@ package pyspy
 // #include "../../../third_party/rustdeps/pyspy.h"
 import "C"
 import (
-	"fmt"
+	"errors"
 	"unsafe"
+	"time"
 
 	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
 )
@@ -19,48 +20,55 @@ import (
 var bufferLength = 1024 * 64
 
 type PySpy struct {
-	cPointer unsafe.Pointer
-	data     []byte
+	dataPtr unsafe.Pointer
+	dataBuf     []byte
+
+	errorBuf []byte
+	errorPtr unsafe.Pointer
+
 	pid      int
 }
 
 func Start(pid int) (spy.Spy, error) {
-	data := make([]byte, bufferLength)
-	cPointer := unsafe.Pointer(&data[0])
+	dataBuf := make([]byte, bufferLength)
+	dataPtr := unsafe.Pointer(&dataBuf[0])
+
+	errorBuf := make([]byte, bufferLength)
+	errorPtr := unsafe.Pointer(&errorBuf[0])
+
+	// TODO: handle this better
+	time.Sleep(1 * time.Second)
 
-	r := C.pyspy_init(C.int(pid))
+	r := C.pyspy_init(C.int(pid), errorPtr, C.int(bufferLength))
 
-	if r == -1 {
-		return nil, fmt.Errorf("buffer too small, current size %d", bufferLength)
+	if r < 0 {
+		return nil, errors.New(string(errorBuf[:-r]))
 	}
 
 	return &PySpy{
-		cPointer: cPointer,
-		data:     data,
+		dataPtr: dataPtr,
+		dataBuf:     dataBuf,
+		errorBuf: errorBuf,
+		errorPtr: errorPtr,
 		pid:      pid,
 	}, nil
 }
 
 func (s *PySpy) Stop() error {
-	r := C.pyspy_cleanup(C.int(s.pid))
-	if r == -1 {
-		return fmt.Errorf("failed to close spy")
+	r := C.pyspy_cleanup(C.int(s.pid), s.errorPtr, C.int(bufferLength))
+	if r < 0 {
+		return errors.New(string(s.errorBuf[:-r]))
 	}
 	return nil
 }
 
 // Snapshot calls callback function with stack-trace or error.
 func (s *PySpy) Snapshot(cb func([]byte, error)) {
-	newL := C.pyspy_snapshot(C.int(s.pid), s.cPointer, C.int(bufferLength))
-	switch newL {
-	case -1:
-		cb(nil, fmt.Errorf("buffer too small, current size %d", bufferLength))
-	case -2:
-		cb(nil, fmt.Errorf("spy is not initialized yet"))
-	case -3:
-		cb(nil, fmt.Errorf("failed to get a trace"))
-	default:
-		cb(s.data[:newL], nil)
+	r := C.pyspy_snapshot(C.int(s.pid), s.dataPtr, C.int(bufferLength), s.errorPtr, C.int(bufferLength))
+	if r < 0 {
+		cb(nil, errors.New(string(s.errorBuf[:-r])))
+	} else {
+		cb(s.dataBuf[:r], nil)
 	}
 }
 

pkg/agent/rbspy/rbspy.go

@@ -8,7 +8,7 @@ package rbspy
 // #include "../../../third_party/rustdeps/rbspy.h"
 import "C"
 import (
-	"fmt"
+	"errors"
 	"unsafe"
 
 	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
@@ -19,48 +19,52 @@ import (
 var bufferLength = 1024 * 64
 
 type RbSpy struct {
-	cPointer unsafe.Pointer
-	data     []byte
-	pid      int
+	dataBuf []byte
+	dataPtr unsafe.Pointer
+
+	errorBuf []byte
+	errorPtr unsafe.Pointer
+
+	pid int
 }
 
 func Start(pid int) (spy.Spy, error) {
-	data := make([]byte, bufferLength)
-	cPointer := unsafe.Pointer(&data[0])
+	dataBuf := make([]byte, bufferLength)
+	dataPtr := unsafe.Pointer(&dataBuf[0])
+
+	errorBuf := make([]byte, bufferLength)
+	errorPtr := unsafe.Pointer(&errorBuf[0])
 
-	r := C.rbspy_init(C.int(pid))
+	r := C.rbspy_init(C.int(pid), errorPtr, C.int(bufferLength))
 
-	if r == -1 {
-		return nil, fmt.Errorf("buffer too small, current size %d", bufferLength)
+	if r < 0 {
+		return nil, errors.New(string(errorBuf[:-r]))
 	}
 
 	return &RbSpy{
-		cPointer: cPointer,
-		data:     data,
+		dataPtr:  dataPtr,
+		dataBuf:  dataBuf,
+		errorBuf: errorBuf,
+		errorPtr: errorPtr,
 		pid:      pid,
 	}, nil
 }
 
 func (s *RbSpy) Stop() error {
-	r := C.rbspy_cleanup(C.int(s.pid))
-	if r == -1 {
-		return fmt.Errorf("failed to close spy")
+	r := C.rbspy_cleanup(C.int(s.pid), s.errorPtr, C.int(bufferLength))
+	if r < 0 {
+		return errors.New(string(s.errorBuf[:-r]))
 	}
 	return nil
 }
 
 // Snapshot calls callback function with stack-trace or error.
 func (s *RbSpy) Snapshot(cb func([]byte, error)) {
-	newL := C.rbspy_snapshot(C.int(s.pid), s.cPointer, C.int(bufferLength))
-	switch newL {
-	case -1:
-		cb(nil, fmt.Errorf("buffer too small, current size %d", bufferLength))
-	case -2:
-		cb(nil, fmt.Errorf("spy is not initialized yet"))
-	case -3:
-		cb(nil, fmt.Errorf("failed to get a trace"))
-	default:
-		cb(s.data[:newL], nil)
+	r := C.rbspy_snapshot(C.int(s.pid), s.dataPtr, C.int(bufferLength), s.errorPtr, C.int(bufferLength))
+	if r < 0 {
+		cb(nil, errors.New(string(s.errorBuf[:-r])))
+	} else {
+		cb(s.dataBuf[:r], nil)
 	}
 }
 

pkg/agent/session.go

@@ -134,8 +134,9 @@ func (ps *ProfileSession) Stop() {
 	}
 	close(ps.stopCh)
 
-	now := time.Now()
-	ps.upstream.Upload(ps.appName, ps.startTime, now, ps.spyName, ps.sampleRate, ps.trie)
+	if ps.trie != nil {
+		ps.upstream.Upload(ps.appName, ps.startTime, time.Now(), ps.spyName, ps.sampleRate, ps.trie)
+	}
 }
 
 func (ps *ProfileSession) addSubprocesses() {

pkg/exec/cli.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
-	"os/signal"
 	"os/user"
 	"path"
 	"runtime"
@@ -20,6 +19,7 @@ import (
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
 	"github.com/pyroscope-io/pyroscope/pkg/config"
 	"github.com/pyroscope-io/pyroscope/pkg/util/atexit"
+	"github.com/pyroscope-io/pyroscope/pkg/util/names"
 	"github.com/sirupsen/logrus"
 )
 
@@ -53,8 +53,15 @@ func Cli(cfg *config.Config, args []string) error {
 		return err
 	}
 
-	signal.Ignore(syscall.SIGCHLD)
+	if cfg.Exec.ApplicationName == "" {
+		logrus.Infof("we recommend specifying application name via %s flag or env variable %s", color.YellowString("-application-name"), color.YellowString("PYROSCOPE_APPLICATION_NAME"))
+		cfg.Exec.ApplicationName = spyName + "." + names.GetRandomName(generateSeed(args))
+		logrus.Infof("for now we chose the name for you and it's \"%s\"", color.BlueString(cfg.Exec.ApplicationName))
+	}
 
+	logrus.WithFields(logrus.Fields{
+		"args": fmt.Sprintf("%q", args),
+	}).Debug("starting command")
 	cmd := exec.Command(args[0], args[1:]...)
 	cmd.Stderr = os.Stderr
 	cmd.Stdout = os.Stdout
@@ -72,12 +79,19 @@ func Cli(cfg *config.Config, args []string) error {
 	})
 	defer u.Stop()
 
-	// TODO: improve this logic, basically we need a smart way of detecting that an app successfully loaded.
-	//   Maybe do this on some ticker (every 100 ms) with a timeout (20 s). Make this configurable too
-	time.Sleep(5 * time.Second)
+	logrus.WithFields(logrus.Fields{
+		"app-name":            cfg.Exec.ApplicationName,
+		"spy-name":            spyName,
+		"pid":                 cmd.Process.Pid,
+		"detect-subprocesses": cfg.Exec.DetectSubprocesses,
+	}).Debug("starting agent session")
+
 	// TODO: add sample rate, make it configurable
 	sess := agent.NewSession(u, cfg.Exec.ApplicationName, spyName, 100, cmd.Process.Pid, cfg.Exec.DetectSubprocesses)
-	sess.Start()
+	err = sess.Start()
+	if err != nil {
+		logrus.Errorf("error when starting session: %q", err)
+	}
 	defer sess.Stop()
 
 	waitForProcessToExit(cmd)
@@ -89,6 +103,10 @@ func Cli(cfg *config.Config, args []string) error {
 func waitForProcessToExit(cmd *exec.Cmd) {
 	sigc := make(chan struct{})
 
+	go func(){
+		cmd.Wait()
+	}()
+
 	atexit.Register(func() {
 		sigc <- struct{}{}
 	})
@@ -97,11 +115,13 @@ func waitForProcessToExit(cmd *exec.Cmd) {
 	for {
 		select {
 		case <-sigc:
+			logrus.Debug("received a signal, killing subprocess")
 			cmd.Process.Kill()
 			return
 		case <-t.C:
 			p, err := ps.FindProcess(cmd.Process.Pid)
 			if p == nil || err != nil {
+				logrus.WithField("err", err).Debug("could not find subprocess, it might be dead")
 				return
 			}
 		}
@@ -152,3 +172,11 @@ func armMessage() string {
 	}
 	return ""
 }
+
+func generateSeed(args []string) string{
+	path, err := os.Getwd()
+	if err != nil {
+		path = "<unknown>"
+	}
+	return path + "|" + strings.Join(args, "&")
+}