Fix supervisor exiting after restart-all

goreman_test.go

@@ -189,6 +189,42 @@ func TestGoremanReleasesRPCPort(t *testing.T) {
 	t.Fatalf("RPC port was not released after goreman stopped: %v", err)
 }
 
+func TestGoremanRestartAllKeepsSupervisorAlive(t *testing.T) {
+	f, err := os.CreateTemp("", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := f.Write([]byte("web1: sleep 10\nweb2: sleep 10\n")); err != nil {
+		t.Fatal(err)
+	}
+	cfg := &config{
+		Procfile: f.Name(),
+		Port:     18556,
+	}
+	sc := make(chan os.Signal, 1)
+	done := make(chan struct{}, 1)
+	go func() {
+		start(context.TODO(), sc, cfg)
+		done <- struct{}{}
+	}()
+	for i := 0; ; i++ {
+		if err = run("restart-all", nil, cfg.Port); err == nil {
+			break
+		}
+		if i > 100 {
+			t.Fatalf("could not reach RPC server: %v", err)
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	select {
+	case <-done:
+		t.Fatal("restart-all should not have stopped the supervisor")
+	case <-time.After(500 * time.Millisecond):
+	}
+	sc <- os.Interrupt
+	<-done
+}
+
 func TestGoremanRestartDoesntLeakGoroutines(t *testing.T) {
 	var file = []byte(`
 web1: sleep 10
@@ -216,7 +252,7 @@ web1: sleep 10
 	waitRunning()
 	before := runtime.NumGoroutine()
 	for i := 0; i < 20; i++ {
-		if err := restartProc("web1"); err != nil {
+		if err := restartProc("web1", nil, nil); err != nil {
 			t.Fatal(err)
 		}
 		waitRunning()

proc.go

@@ -117,12 +117,18 @@ func startProc(name string, wg *sync.WaitGroup, errCh chan<- error) error {
 }
 
 // restart specified proc.
-func restartProc(name string) error {
+func restartProc(name string, wg *sync.WaitGroup, errCh chan<- error) error {
+	if wg != nil {
+		// keep the proc counted between stop and start so the supervisor
+		// does not see "all procs done" in the middle of a restart.
+		wg.Add(1)
+		defer wg.Done()
+	}
 	err := stopProc(name, nil)
 	if err != nil {
 		return err
 	}
-	return startProc(name, nil, nil)
+	return startProc(name, wg, errCh)
 }
 
 // stopProcs attempts to stop every running process and returns any non-nil
@@ -160,6 +166,30 @@ func startProcs(sc <-chan os.Signal, rpcCh <-chan *rpcMessage, exitOnError bool)
 		case rpcMsg := <-rpcCh:
 			switch rpcMsg.Msg {
 			// TODO: add more events here.
+			case "start":
+				for _, proc := range rpcMsg.Args {
+					if err := startProc(proc, &wg, errCh); err != nil {
+						rpcMsg.ErrCh <- err
+						break
+					}
+				}
+				close(rpcMsg.ErrCh)
+			case "restart":
+				names := rpcMsg.Args
+				if len(names) == 0 {
+					mu.Lock()
+					for _, proc := range procs {
+						names = append(names, proc.name)
+					}
+					mu.Unlock()
+				}
+				for _, proc := range names {
+					if err := restartProc(proc, &wg, errCh); err != nil {
+						rpcMsg.ErrCh <- err
+						break
+					}
+				}
+				close(rpcMsg.ErrCh)
 			case "stop":
 				for _, proc := range rpcMsg.Args {
 					if err := stopProc(proc, nil); err != nil {

rpc.go

@@ -22,19 +22,25 @@ type rpcMessage struct {
 	ErrCh chan error
 }
 
+// rpcExec sends the message to the supervisor loop and waits for the result.
+func (r *Goreman) rpcExec(msg string, args []string) error {
+	errChan := make(chan error, 1)
+	r.rpcChan <- &rpcMessage{
+		Msg:   msg,
+		Args:  args,
+		ErrCh: errChan,
+	}
+	return <-errChan
+}
+
 // Start do start
 func (r *Goreman) Start(args []string, ret *string) (err error) {
 	defer func() {
 		if r := recover(); r != nil {
 			err = fmt.Errorf("%v", r)
 		}
 	}()
-	for _, arg := range args {
-		if err = startProc(arg, nil, nil); err != nil {
-			break
-		}
-	}
-	return err
+	return r.rpcExec("start", args)
 }
 
 // Stop do stop
@@ -44,14 +50,7 @@ func (r *Goreman) Stop(args []string, ret *string) (err error) {
 			err = fmt.Errorf("%v", r)
 		}
 	}()
-	errChan := make(chan error, 1)
-	r.rpcChan <- &rpcMessage{
-		Msg:   "stop",
-		Args:  args,
-		ErrCh: errChan,
-	}
-	err = <-errChan
-	return
+	return r.rpcExec("stop", args)
 }
 
 // StopAll do stop all
@@ -76,12 +75,7 @@ func (r *Goreman) Restart(args []string, ret *string) (err error) {
 			err = fmt.Errorf("%v", r)
 		}
 	}()
-	for _, arg := range args {
-		if err = restartProc(arg); err != nil {
-			break
-		}
-	}
-	return err
+	return r.rpcExec("restart", args)
 }
 
 // RestartAll do restart all
@@ -91,12 +85,7 @@ func (r *Goreman) RestartAll(args []string, ret *string) (err error) {
 			err = fmt.Errorf("%v", r)
 		}
 	}()
-	for _, proc := range procs {
-		if err = restartProc(proc.name); err != nil {
-			break
-		}
-	}
-	return err
+	return r.rpcExec("restart", nil)
 }
 
 // List do list
@@ -174,7 +163,12 @@ func startServer(ctx context.Context, rpcChan chan<- *rpcMessage, listenPort uin
 	gm := &Goreman{
 		rpcChan: rpcChan,
 	}
-	rpc.Register(gm)
+	// use a dedicated rpc.Server so repeated startServer calls in one
+	// process do not keep serving a previously registered instance.
+	rpcServer := rpc.NewServer()
+	if err := rpcServer.Register(gm); err != nil {
+		return err
+	}
 	server, err := net.Listen("tcp", fmt.Sprintf("%s:%d", defaultAddr(), listenPort))
 	if err != nil {
 		return err
@@ -202,7 +196,7 @@ func startServer(ctx context.Context, rpcChan chan<- *rpcMessage, listenPort uin
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
-				rpc.ServeConn(client)
+				rpcServer.ServeConn(client)
 			}()
 		}
 	}