Simplify supervisor, support agent-triggered broker restart

Author: shawnburke

agent/cmd/relay.go

@@ -66,7 +66,7 @@ var RelayCommand = &cobra.Command{
 		}
 
 		fmt.Println("Starting agent")
-		startAgent(cmd, buildRelayStack(cmd, config, info))
+		startAgent(buildRelayStack(cmd, config, info))
 	},
 }
 

agent/cmd/serve.go

@@ -35,7 +35,7 @@ var serveCmd = &cobra.Command{
 		}
 
 		config.Print()
-		startAgent(cmd, buildServeStack(cmd, config))
+		startAgent(buildServeStack(cmd, config))
 		fmt.Println("Server stopped")
 	},
 }

agent/cmd/stack.go

@@ -15,28 +15,11 @@ import (
 	"github.com/spf13/cobra"
 	"go.uber.org/fx"
 	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
 )
 
-func getLoggingLevel(cmd *cobra.Command) zapcore.Level {
-
-	if ok, _ := cmd.Flags().GetBool("verbose"); ok {
-		return zap.DebugLevel
-	}
-	return zap.InfoLevel
-}
-
-func startAgent(cmd *cobra.Command, opts fx.Option) {
-	options := []fx.Option{
-		opts,
-	}
-
-	if getLoggingLevel(cmd) != zap.DebugLevel {
-		options = append(options, fx.NopLogger)
-	}
-
+func startAgent(opts fx.Option) {
 	app := fx.New(
-		options...,
+		opts,
 	)
 
 	noBanner := os.Getenv("NO_BANNER")
@@ -47,17 +30,34 @@ func startAgent(cmd *cobra.Command, opts fx.Option) {
 }
 
 func buildCoreAgentStack(cmd *cobra.Command, cfg config.AgentConfig) fx.Option {
-	return fx.Options(
-		fx.Provide(func() *zap.Logger {
+
+	if ok, _ := cmd.Flags().GetBool("verbose"); ok {
+		cfg.VerboseOutput = true
+	}
+
+	options := []fx.Option{}
+
+	if !cfg.VerboseOutput {
+		options = append(options, fx.NopLogger)
+	}
+
+	stackOptions := fx.Options(
+		fx.Supply(cfg),
+		fx.Provide(func(config config.AgentConfig) *zap.Logger {
 			cfg := zap.NewDevelopmentConfig()
-			cfg.Level = zap.NewAtomicLevelAt(getLoggingLevel(cmd))
+
+			loggingLevel := zap.InfoLevel
+			if config.VerboseOutput {
+				loggingLevel = zap.DebugLevel
+			}
+
+			cfg.Level = zap.NewAtomicLevelAt(loggingLevel)
 			logger, err := cfg.Build()
 			if err != nil {
 				panic(err)
 			}
 			return logger
 		}),
-		fx.Supply(cfg),
 		fx.Invoke(func(config config.AgentConfig, logger *zap.Logger) {
 			if config.CortexApiToken == "" && !config.DryRun {
 				logger.Fatal("Cannot start agent: either CORTEX_API_TOKEN or DRYRUN is required")
@@ -67,6 +67,13 @@ func buildCoreAgentStack(cmd *cobra.Command, cfg config.AgentConfig) fx.Option {
 		fx.Provide(handler.NewHandlerManager),
 		fx.Invoke(createAxonAgent),
 	)
+
+	options = append(options, stackOptions)
+
+	return fx.Options(
+		options...,
+	)
+
 }
 
 func createAxonAgent(

agent/config/config.go

@@ -28,6 +28,7 @@ type AgentConfig struct {
 	WebhookServerPort int
 	EnableApiProxy    bool
 	FailWaitTime      time.Duration
+	VerboseOutput     bool
 }
 
 func (ac AgentConfig) Print() {

agent/server/snykbroker/relay_instance_manager.go

@@ -21,6 +21,7 @@ const defaultSnykBroker = "snyk-broker"
 
 type RelayInstanceManager interface {
 	Start() error
+	Restart() error
 	Close() error
 }
 
@@ -63,7 +64,7 @@ func NewRelayInstanceManager(
 }
 
 func (r *relayInstanceManager) RegisterRoutes(mux *http.ServeMux) error {
-	mux.Handle(fmt.Sprintf("%s/reregister", cortexHttp.AxonPathRoot), r)
+	mux.Handle(fmt.Sprintf("%s/broker/restart", cortexHttp.AxonPathRoot), r)
 	return nil
 }
 
@@ -73,7 +74,7 @@ func (r *relayInstanceManager) ServeHTTP(w http.ResponseWriter, req *http.Reques
 		return
 	}
 
-	err := r.restart()
+	err := r.Restart()
 	if err != nil {
 		r.logger.Error("Unable to reregister", zap.Error(err))
 		w.WriteHeader(http.StatusInternalServerError)
@@ -85,16 +86,19 @@ func (r *relayInstanceManager) ServeHTTP(w http.ResponseWriter, req *http.Reques
 
 var errSkipBroker = errors.New("NoBrokerToken")
 
-func (r *relayInstanceManager) restart() error {
+func (r *relayInstanceManager) Restart() error {
+
+	r.logger.Info("Restarting broker, shutting down existing broker")
 	// re-register and restart supervisor
 	err := r.Close()
 	if err != nil {
-		r.logger.Error("unable to close supervisor on /reregister", zap.Error(err))
+		r.logger.Error("unable to close supervisor on Restart", zap.Error(err))
 	}
 
+	r.logger.Info("Restarting broker")
 	err = r.Start()
 	if err != nil {
-		return fmt.Errorf("unable to start supervisor on /reregister: %w", err)
+		return fmt.Errorf("unable to start supervisor on Restart: %w", err)
 	}
 	return nil
 }
@@ -196,15 +200,21 @@ func (r *relayInstanceManager) Start() error {
 			zap.String("acceptFile", acceptFile),
 		)
 
+		brokerEnv := map[string]string{
+			"ACCEPT":            acceptFile,
+			"BROKER_SERVER_URL": uri,
+			"BROKER_TOKEN":      token,
+			"PORT":              "7343",
+		}
+
+		if r.config.VerboseOutput {
+			brokerEnv["LOG_LEVEL"] = "debug"
+		}
+
 		r.supervisor = NewSupervisor(
 			executable,
 			args,
-			map[string]string{
-				"ACCEPT":            acceptFile,
-				"BROKER_SERVER_URL": uri,
-				"BROKER_TOKEN":      token,
-				"PORT":              "7343",
-			},
+			brokerEnv,
 			r.config.FailWaitTime,
 		)
 		r.startCount.Add(1)

agent/server/snykbroker/supervisor.go

@@ -27,6 +27,10 @@ type Supervisor struct {
 	cmd        *exec.Cmd
 	env        map[string]string
 	done       chan struct{}
+	stopFunc   func()
+	pid        int
+	killed     bool
+	runCount   int
 	lastError  error
 }
 
@@ -43,32 +47,19 @@ func NewSupervisor(
 		args:              args,
 		env:               env,
 		fastFailTime:      fastFailTime * 2,
+		done:              make(chan struct{}),
+		stopFunc:          func() {},
 	}
 }
 
 var errKilled = errors.New("killed")
 var errMaxRetries = errors.New("max retries reached")
 
-func (b *Supervisor) trigger() (func(error), error) {
-	b.Lock()
-	if b.done != nil {
-		b.Unlock()
-		return nil, fmt.Errorf("can't call start when already running")
-	}
-	b.done = make(chan struct{})
-	b.Unlock()
+func (b *Supervisor) Start(maxRetries int, window time.Duration) error {
 
-	finish := func(err error) {
-		b.Lock()
-		defer b.Unlock()
-		b.lastError = err
-		close(b.done)
-		b.done = nil
+	if b.runCount > 0 {
+		return errors.New("already started, cannot start again")
 	}
-	return finish, nil
-}
-
-func (b *Supervisor) Start(maxRetries int, window time.Duration) error {
 
 	if err := b.runExecutionLoop(maxRetries, window); err != nil {
 		return err
@@ -80,32 +71,34 @@ func (b *Supervisor) runExecutionLoop(maxRetries int, window time.Duration) erro
 
 	tracker := newEventTracker()
 	startTime := time.Now()
-	runCount := 0
 
-	finish, err := b.trigger()
-	if err != nil {
-		return err
+	finish := func(err error) {
+		b.Lock()
+		defer b.Unlock()
+		b.lastError = err
+		if b.done != nil {
+			close(b.done)
+		}
 	}
 
 	fastfail := make(chan struct{})
 	// we run this off thread, looping to restart
 	// the process if it crashes, but exiting
-	// if too many happen in the restart windo
+	// if too many happen in the restart window
+	b.runCount = 1
 	go func() {
 		defer close(fastfail)
 		for maxRetries > 0 {
 			tracker.AddEvent()
-			runCount++
 			err := b.runCommand()
 			runTime := time.Since(startTime)
 
 			if errors.Is(err, errKilled) {
-				fmt.Println("Process killed")
 				finish(nil)
 				return
 			}
 
-			if err != nil && runCount == 1 && runTime < b.fastFailTime {
+			if err != nil && b.runCount == 1 && runTime < b.fastFailTime {
 				finish(fmt.Errorf("run failed immediately: %v", err))
 				return
 			}
@@ -123,7 +116,7 @@ func (b *Supervisor) runExecutionLoop(maxRetries int, window time.Duration) erro
 				}
 				return
 			}
-
+			b.runCount++
 		}
 	}()
 
@@ -170,12 +163,10 @@ func (b *Supervisor) runWatchdog(pid int) func() {
 	}
 }
 
+// runCommand runs a single command and returns the result of the command after it exits,
+// or errKilled if it was intentionally shut down.
 func (b *Supervisor) runCommand() error {
 
-	if b.cmd != nil {
-		panic("Command already running")
-	}
-
 	cmd := exec.Command(b.executable)
 	cmd.Args = append(cmd.Args, b.args...)
 
@@ -199,38 +190,67 @@ func (b *Supervisor) runCommand() error {
 		}
 	}()
 
+	// sigChan triggers shutdown
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-	killed := false
+
+	running := make(chan struct{})
+
 	go func() {
 		<-sigChan
-		err := cmd.Process.Kill()
-		if err != nil {
-			fmt.Printf("Error killing process: %v\n", err)
-		}
-		killed = true
+		b.stopFunc()
 	}()
 
-	b.cmd = cmd
+	killed := false
+
+	// our stopfunc allows anyone to close the process
+	// and wait for it to finish
+	b.stopFunc = func() {
+		killed = true
+
+		// if process is running trigger a kill
+		if cmd.Process != nil {
+			err := cmd.Process.Kill()
+			if err != nil {
+				fmt.Printf("Error killing process: %v\n", err)
+			}
+		}
+
+		// wait for it to actually finish
+		<-running
+	}
+
 	err := cmd.Start()
 
 	if err != nil {
 		return err
 	}
+	pid := cmd.Process.Pid
+	b.pid = pid
 
 	// We want to make sure the broker is killed if the agent dies or is killed.  This is
 	// mostly useful in the debugger but prevents port from being held open.
 	cancelWatchdog := b.runWatchdog(cmd.Process.Pid)
 	defer cancelWatchdog()
 
-	cmd.Wait()
-	b.cmd = nil
+	// here we wait for it to start then fully finish
+	go func() {
+		cmd.Wait()
+		b.pid = 0
+		b.stopFunc = func() {}
+		close(running)
+		running = nil
+	}()
+
+	// block until the process is done
+	<-running
 	stopStdOut()
 	stopStdErr()
 	wg.Wait()
 
 	if killed {
 		err = errKilled
+		fmt.Printf("Process %v (pid=%v) killed\n", cmd.Path, pid)
 	}
 
 	if err == nil && cmd.ProcessState.ExitCode() != 0 {
@@ -239,11 +259,12 @@ func (b *Supervisor) runCommand() error {
 	return err
 }
 
+func (b *Supervisor) Pid() int {
+	return b.pid
+}
+
 func (b *Supervisor) Close() error {
-	if b.cmd != nil && b.cmd.Process != nil {
-		b.cmd.Process.Kill()
-	}
-	b.cmd = nil
+	b.stopFunc()
 	return nil
 }