improves exec child process handling

Author: petethepig

pkg/agent/spy/spy.go

@@ -25,6 +25,7 @@ var autoDetectionMapping = map[string]string{
 	"ruby":   "rbspy",
 	"bundle": "rbspy",
 	"rails":  "rbspy",
+	"rake":   "rbspy",
 }
 
 func init() {
@@ -46,3 +47,14 @@ func SpyFromName(name string, pid int) (Spy, error) {
 func ResolveAutoName(s string) string {
 	return autoDetectionMapping[s]
 }
+
+func SupportedExecSpies() []string {
+	supportedSpies := []string{}
+	for _, s := range SupportedSpies {
+		if s != "gospy" {
+			supportedSpies = append(supportedSpies, s)
+		}
+	}
+
+	return supportedSpies
+}

pkg/agent/upstream/remote/remote.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -25,7 +26,7 @@ type uploadJob struct {
 type Remote struct {
 	cfg    RemoteConfig
 	todo   chan *uploadJob
-	done   chan struct{}
+	done   chan *sync.WaitGroup
 	client *http.Client
 }
 
@@ -39,7 +40,7 @@ func New(cfg RemoteConfig) *Remote {
 	r := &Remote{
 		cfg:  cfg,
 		todo: make(chan *uploadJob, 100),
-		done: make(chan struct{}, cfg.UpstreamThreads),
+		done: make(chan *sync.WaitGroup, cfg.UpstreamThreads),
 		client: &http.Client{
 			Transport: &http.Transport{
 				MaxConnsPerHost: cfg.UpstreamThreads,
@@ -58,9 +59,12 @@ func (u *Remote) start() {
 }
 
 func (u *Remote) Stop() {
+	wg := sync.WaitGroup{}
+	wg.Add(u.cfg.UpstreamThreads)
 	for i := 0; i < u.cfg.UpstreamThreads; i++ {
-		u.done <- struct{}{}
+		u.done <- &wg
 	}
+	wg.Wait()
 }
 
 // TODO: this metadata class should be unified
@@ -112,7 +116,8 @@ func (u *Remote) uploadLoop() {
 		select {
 		case j := <-u.todo:
 			u.uploadProfile(j)
-		case <-u.done:
+		case wg := <-u.done:
+			wg.Done()
 			return
 		}
 	}

pkg/exec/cli.go

@@ -19,6 +19,7 @@ import (
 	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
 	"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/sirupsen/logrus"
 )
 
@@ -32,7 +33,7 @@ func Cli(cfg *config.Config, args []string) error {
 		baseName := path.Base(args[0])
 		spyName = spy.ResolveAutoName(baseName)
 		if spyName == "" {
-			supportedSpies := supportedSpiesWithoutGospy()
+			supportedSpies := spy.SupportedExecSpies()
 			suggestedCommand := fmt.Sprintf("pyroscope exec -spy-name %s %s", supportedSpies[0], strings.Join(args, " "))
 			return fmt.Errorf(
 				"could not automatically find a spy for program \"%s\". Pass spy name via %s argument, for example: \n  %s\n\nAvailable spies are: %s\n%s\nIf you believe this is a mistake, please submit an issue at %s",
@@ -69,36 +70,42 @@ func Cli(cfg *config.Config, args []string) error {
 		UpstreamThreads:        cfg.Exec.UpstreamThreads,
 		UpstreamRequestTimeout: cfg.Exec.UpstreamRequestTimeout,
 	})
+	defer u.Stop()
 
-	// TODO: make configurable?
+	// 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)
-	// TODO: add sample rate
+	// TODO: add sample rate, make it configurable
 	sess := agent.NewSession(u, cfg.Exec.ApplicationName, spyName, 100, cmd.Process.Pid, cfg.Exec.DetectSubprocesses)
 	sess.Start()
 	defer sess.Stop()
 
-	// TODO: very hacky, at some point we'll need to make wait work
-	// cmd.Wait()
-
-	for {
-		time.Sleep(time.Second)
-		p, err := ps.FindProcess(cmd.Process.Pid)
-		if p == nil || err != nil {
-			break
-		}
-	}
+	waitForProcessToExit(cmd)
 	return nil
 }
 
-func supportedSpiesWithoutGospy() []string {
-	supportedSpies := []string{}
-	for _, s := range spy.SupportedSpies {
-		if s != "gospy" {
-			supportedSpies = append(supportedSpies, s)
+// TODO: very hacky, at some point we'll need to make `cmd.Wait()` work
+//   Currently the issue is that on Linux it often thinks the process exited when it did not.
+func waitForProcessToExit(cmd *exec.Cmd) {
+	sigc := make(chan struct{})
+
+	atexit.Register(func() {
+		sigc <- struct{}{}
+	})
+
+	t := time.NewTicker(time.Second)
+	for {
+		select {
+		case <-sigc:
+			cmd.Process.Kill()
+			return
+		case <-t.C:
+			p, err := ps.FindProcess(cmd.Process.Pid)
+			if p == nil || err != nil {
+				return
+			}
 		}
 	}
-
-	return supportedSpies
 }
 
 func performChecks(spyName string) error {
@@ -113,7 +120,7 @@ func performChecks(spyName string) error {
 	}
 
 	if !stringsContains(spy.SupportedSpies, spyName) {
-		supportedSpies := supportedSpiesWithoutGospy()
+		supportedSpies := spy.SupportedExecSpies()
 		return fmt.Errorf(
 			"Spy \"%s\" is not supported. Available spies are: %s\n%s",
 			color.BlueString(spyName),

pkg/util/atexit/atexit.go

@@ -16,7 +16,6 @@ func init() {
 		for _, cb := range callbacks {
 			cb()
 		}
-		os.Exit(0)
 	}()
 }