From 2c66795b865309766d78a7bd87cdefa13f1d3efc Mon Sep 17 00:00:00 2001
From: Tiexin Guo <guotiexin@gmail.com>
Date: Tue, 24 Sep 2024 19:48:04 +0800
Subject: [PATCH] poc: allow stopping services that are starting and restart
 services that quickly exist within the 1s okay delay

---
 internals/overlord/servstate/handlers.go | 16 ++++++++++++++++
 internals/overlord/servstate/manager.go  |  2 +-
 tests/run_test.go                        |  6 +-----
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/internals/overlord/servstate/handlers.go b/internals/overlord/servstate/handlers.go
index 2bcdd6945..937c4cbbc 100644
--- a/internals/overlord/servstate/handlers.go
+++ b/internals/overlord/servstate/handlers.go
@@ -512,6 +512,12 @@ func (s *serviceData) exited(exitCode int) error {
 	case stateStarting:
 		s.started <- fmt.Errorf("exited quickly with code %d", exitCode)
 		s.transition(stateExited) // not strictly necessary as doStart will return, but doesn't hurt
+		action, onType := getAction(s.config, exitCode == 0)
+		switch action {
+		case plan.ActionRestart:
+			s.doBackoff(action, onType)
+		default:
+		}
 
 	case stateRunning:
 		logger.Noticef("Service %q stopped unexpectedly with code %d", s.config.Name, exitCode)
@@ -682,6 +688,16 @@ func (s *serviceData) stop() error {
 	defer s.manager.servicesLock.Unlock()
 
 	switch s.state {
+	case stateStarting:
+		s.started <- fmt.Errorf("stopped before the 1 second okay delay")
+		logger.Debugf("Attempting to stop service %q by sending SIGTERM", s.config.Name)
+		err := syscall.Kill(-s.cmd.Process.Pid, syscall.SIGTERM)
+		if err != nil {
+			logger.Noticef("Cannot send SIGTERM to process: %v", err)
+		}
+		s.transition(stateTerminating)
+		time.AfterFunc(s.killDelay(), func() { logError(s.terminateTimeElapsed()) })
+
 	case stateRunning:
 		logger.Debugf("Attempting to stop service %q by sending SIGTERM", s.config.Name)
 		// First send SIGTERM to try to terminate it gracefully.
diff --git a/internals/overlord/servstate/manager.go b/internals/overlord/servstate/manager.go
index 43a379d27..7b319aab1 100644
--- a/internals/overlord/servstate/manager.go
+++ b/internals/overlord/servstate/manager.go
@@ -375,7 +375,7 @@ func servicesToStop(m *ServiceManager) ([][]string, error) {
 		var notStopped []string
 		for _, name := range services {
 			s := m.services[name]
-			if s != nil && (s.state == stateRunning || s.state == stateBackoff) {
+			if s != nil && (s.state == stateStarting || s.state == stateRunning || s.state == stateBackoff) {
 				notStopped = append(notStopped, name)
 			}
 		}
diff --git a/tests/run_test.go b/tests/run_test.go
index c98c6937e..d9b0073e4 100644
--- a/tests/run_test.go
+++ b/tests/run_test.go
@@ -49,9 +49,7 @@ services:
 
 	createLayer(t, pebbleDir, "001-simple-layer.yaml", layerYAML)
 
-	_, stderrCh := pebbleRun(t, pebbleDir)
-	waitForLog(t, stderrCh, "pebble", "Started default services", 3*time.Second)
-
+	_, _ = pebbleRun(t, pebbleDir)
 	waitForFile(t, filepath.Join(pebbleDir, "svc1"), 3*time.Second)
 	waitForFile(t, filepath.Join(pebbleDir, "svc2"), 3*time.Second)
 }
@@ -142,7 +140,6 @@ services:
 	stdoutCh, stderrCh := pebbleRun(t, pebbleDir, "--verbose")
 	waitForLog(t, stderrCh, "pebble", "Started daemon", 3*time.Second)
 	waitForLog(t, stdoutCh, "svc1", "hello world", 3*time.Second)
-	waitForLog(t, stderrCh, "pebble", "Started default services", 3*time.Second)
 }
 
 // TestArgs tests that Pebble provides additional arguments to a service
@@ -168,7 +165,6 @@ services:
 	)
 	waitForLog(t, stderrCh, "pebble", "Started daemon", 3*time.Second)
 	waitForLog(t, stdoutCh, "svc1", "hello world", 3*time.Second)
-	waitForLog(t, stderrCh, "pebble", "Started default services", 3*time.Second)
 }
 
 // TestIdentities tests that Pebble seeds identities from a file