refactor: pull in the snapd RestartManager refactoring (#446)

This is based on canonical/snapd#12166 (though
the diff was applied more or less manually, as the patch couldn't be
applied directly).

It includes the name changes to the restart.Handler interface methods.

internals/daemon/api_test.go

@@ -108,8 +108,9 @@ func (s *apiSuite) TestSysInfo(c *check.C) {
 	d.Version = "42b1"
 	state := d.overlord.State()
 	state.Lock()
-	restart.Init(state, "ffffffff-ffff-ffff-ffff-ffffffffffff", nil)
+	_, err := restart.Manager(state, "ffffffff-ffff-ffff-ffff-ffffffffffff", nil)
 	state.Unlock()
+	c.Assert(err, check.IsNil)
 
 	sysInfoCmd.GET(sysInfoCmd, nil, nil).ServeHTTP(rec, nil)
 	c.Check(rec.Code, check.Equals, 200)

internals/daemon/daemon.go

@@ -787,16 +787,16 @@ func clearReboot(st *state.State) {
 	st.Set("daemon-system-restart-tentative", nil)
 }
 
-// RebootIsFine implements part of overlord.RestartBehavior.
-func (d *Daemon) RebootIsFine(st *state.State) error {
+// RebootAsExpected implements part of overlord.RestartBehavior.
+func (d *Daemon) RebootAsExpected(st *state.State) error {
 	clearReboot(st)
 	return nil
 }
 
 var errExpectedReboot = errors.New("expected reboot did not happen")
 
-// RebootIsMissing implements part of overlord.RestartBehavior.
-func (d *Daemon) RebootIsMissing(st *state.State) error {
+// RebootDidNotHappen implements part of overlord.RestartBehavior.
+func (d *Daemon) RebootDidNotHappen(st *state.State) error {
 	var nTentative int
 	err := st.Get("daemon-system-restart-tentative", &nTentative)
 	if err != nil && !errors.Is(err, state.ErrNoState) {

internals/overlord/overlord.go

@@ -96,6 +96,7 @@ type Overlord struct {
 	inited     bool
 	startedUp  bool
 	runner     *state.TaskRunner
+	restartMgr *restart.RestartManager
 	planMgr    *planstate.PlanManager
 	serviceMgr *servstate.ServiceManager
 	commandMgr *cmdstate.CommandManager
@@ -127,7 +128,7 @@ func New(opts *Options) (*Overlord, error) {
 		path:         statePath,
 		ensureBefore: o.ensureBefore,
 	}
-	s, err := loadState(statePath, opts.RestartHandler, backend)
+	s, restartMgr, err := loadState(statePath, opts.RestartHandler, backend)
 	if err != nil {
 		return nil, err
 	}
@@ -141,6 +142,9 @@ func New(opts *Options) (*Overlord, error) {
 	}
 	o.runner.AddOptionalHandler(matchAnyUnknownTask, handleUnknownTask, nil)
 
+	o.restartMgr = restartMgr
+	o.stateEng.AddManager(restartMgr)
+
 	o.planMgr, err = planstate.NewManager(o.pebbleDir)
 	if err != nil {
 		return nil, fmt.Errorf("cannot create plan manager: %w", err)
@@ -213,12 +217,12 @@ func (o *Overlord) Extension() Extension {
 	return o.extension
 }
 
-func loadState(statePath string, restartHandler restart.Handler, backend state.Backend) (*state.State, error) {
+func loadState(statePath string, restartHandler restart.Handler, backend state.Backend) (*state.State, *restart.RestartManager, error) {
 	timings := timing.Start("", "", map[string]string{"startup": "load-state"})
 
 	curBootID, err := osutil.BootID()
 	if err != nil {
-		return nil, fmt.Errorf("fatal: cannot find current boot ID: %v", err)
+		return nil, nil, fmt.Errorf("fatal: cannot find current boot ID: %v", err)
 	}
 	// If pebble is PID 1 we don't care about /proc/sys/kernel/random/boot_id
 	// as we are most likely running in a container. LXD mounts it's own boot_id
@@ -228,25 +232,28 @@ func loadState(statePath string, restartHandler restart.Handler, backend state.B
 	if os.Getpid() == 1 {
 		curBootID, err = randutil.RandomKernelUUID()
 		if err != nil {
-			return nil, fmt.Errorf("fatal: cannot generate psuedo boot-id: %v", err)
+			return nil, nil, fmt.Errorf("fatal: cannot generate psuedo boot-id: %v", err)
 		}
 	}
 
 	if !osutil.CanStat(statePath) {
 		// fail fast, mostly interesting for tests, this dir is set up by pebble
 		stateDir := filepath.Dir(statePath)
 		if !osutil.IsDir(stateDir) {
-			return nil, fmt.Errorf("fatal: directory %q must be present", stateDir)
+			return nil, nil, fmt.Errorf("fatal: directory %q must be present", stateDir)
 		}
 		s := state.New(backend)
-		initRestart(s, curBootID, restartHandler)
+		restartMgr, err := initRestart(s, curBootID, restartHandler)
+		if err != nil {
+			return nil, nil, err
+		}
 		patch.Init(s)
-		return s, nil
+		return s, restartMgr, nil
 	}
 
 	r, err := os.Open(statePath)
 	if err != nil {
-		return nil, fmt.Errorf("cannot read the state file: %s", err)
+		return nil, nil, fmt.Errorf("cannot read the state file: %s", err)
 	}
 	defer r.Close()
 
@@ -255,7 +262,7 @@ func loadState(statePath string, restartHandler restart.Handler, backend state.B
 	s, err = state.ReadState(backend, r)
 	span.Stop()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	timings.Stop()
@@ -264,23 +271,23 @@ func loadState(statePath string, restartHandler restart.Handler, backend state.B
 	//perfTimings.Save(s)
 	//s.Unlock()
 
-	err = initRestart(s, curBootID, restartHandler)
+	restartMgr, err := initRestart(s, curBootID, restartHandler)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	// one-shot migrations
 	err = patch.Apply(s)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return s, nil
+	return s, restartMgr, nil
 }
 
-func initRestart(s *state.State, curBootID string, restartHandler restart.Handler) error {
+func initRestart(s *state.State, curBootID string, restartHandler restart.Handler) (*restart.RestartManager, error) {
 	s.Lock()
 	defer s.Unlock()
-	return restart.Init(s, curBootID, restartHandler)
+	return restart.Manager(s, curBootID, restartHandler)
 }
 
 func (o *Overlord) StartUp() error {
@@ -492,6 +499,11 @@ func (o *Overlord) TaskRunner() *state.TaskRunner {
 	return o.runner
 }
 
+// RestartManager returns the manager responsible for restart state.
+func (o *Overlord) RestartManager() *restart.RestartManager {
+	return o.restartMgr
+}
+
 // ServiceManager returns the service manager responsible for services
 // under the overlord.
 func (o *Overlord) ServiceManager() *servstate.ServiceManager {

internals/overlord/overlord_test.go

@@ -83,6 +83,7 @@ func (ovs *overlordSuite) TestNew(c *C) {
 
 	c.Check(o.StateEngine(), NotNil)
 	c.Check(o.TaskRunner(), NotNil)
+	c.Check(o.RestartManager(), NotNil)
 
 	s := o.State()
 	c.Check(s, NotNil)
@@ -112,6 +113,7 @@ func (ovs *overlordSuite) TestNewWithGoodState(c *C) {
 
 	o, err := overlord.New(&overlord.Options{PebbleDir: ovs.dir})
 	c.Assert(err, IsNil)
+	c.Check(o.RestartManager(), NotNil)
 
 	state := o.State()
 	c.Assert(err, IsNil)
@@ -932,12 +934,12 @@ func (rb *testRestartHandler) HandleRestart(t restart.RestartType) {
 	rb.restartRequested = t
 }
 
-func (rb *testRestartHandler) RebootIsFine(_ *state.State) error {
+func (rb *testRestartHandler) RebootAsExpected(_ *state.State) error {
 	rb.rebootState = "as-expected"
 	return rb.rebootVerifiedErr
 }
 
-func (rb *testRestartHandler) RebootIsMissing(_ *state.State) error {
+func (rb *testRestartHandler) RebootDidNotHappen(_ *state.State) error {
 	rb.rebootState = "did-not-happen"
 	return rb.rebootVerifiedErr
 }

internals/overlord/restart/restart.go

@@ -12,7 +12,9 @@
 // You should have received a copy of the GNU General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-// Package restart implements requesting restarts from any part of the code that has access to state.
+// Package restart implements requesting restarts from any part of the
+// code that has access to state. It also implements a mimimal manager
+// to take care of restart state.
 package restart
 
 import (
@@ -45,125 +47,135 @@ const (
 // Handler can handle restart requests and whether expected reboots happen.
 type Handler interface {
 	HandleRestart(t RestartType)
-	// RebootIsFine is called early when either a reboot was
+	// RebootAsExpected is called early when either a reboot was
 	// requested by snapd and happened or no reboot was expected at all.
-	RebootIsFine(st *state.State) error
-	// RebootIsMissing is called early instead when a reboot was
+	RebootAsExpected(st *state.State) error
+	// RebootDidNotHappen is called early instead when a reboot was
 	// requested by snapd but did not happen.
-	RebootIsMissing(st *state.State) error
+	RebootDidNotHappen(st *state.State) error
 }
 
-// Init initializes the support for restarts requests.
-// It takes the current boot id to track and verify reboots and a
-// Handler that handles the actual requests and reacts to reboot
-// happening.
-// It must be called with the state lock held.
-func Init(st *state.State, curBootID string, h Handler) error {
-	rs := &restartState{
+type restartManagerKey struct{}
+
+// RestartManager takes care of restart-related state.
+type RestartManager struct {
+	state      *state.State
+	restarting RestartType
+	h          Handler
+	bootID     string
+}
+
+// Manager returns a new restart manager and initializes the support
+// for restarts requests. It takes the current boot id to track and
+// verify reboots and a Handler that handles the actual requests and
+// reacts to reboot happening. It must be called with the state lock
+// held.
+func Manager(st *state.State, curBootID string, h Handler) (*RestartManager, error) {
+	rm := &RestartManager{
+		state:  st,
 		h:      h,
 		bootID: curBootID,
 	}
 	var fromBootID string
 	err := st.Get("system-restart-from-boot-id", &fromBootID)
 	if err != nil && !errors.Is(err, state.ErrNoState) {
-		return err
+		return nil, err
+	}
+	st.Cache(restartManagerKey{}, rm)
+	if err := rm.init(fromBootID, curBootID); err != nil {
+		return nil, err
 	}
-	st.Cache(restartStateKey{}, rs)
+	return rm, nil
+}
+
+func (rm *RestartManager) init(fromBootID, curBootID string) error {
 	if fromBootID == "" {
-		return rs.rebootAsExpected(st)
+		// We didn't need a reboot, it might have happened or
+		// not but things are fine in either case.
+		return rm.rebootAsExpected()
 	}
 	if fromBootID == curBootID {
-		return rs.rebootDidNotHappen(st)
+		return rm.rebootDidNotHappen()
 	}
 	// we rebooted alright
-	ClearReboot(st)
-	return rs.rebootAsExpected(st)
+	ClearReboot(rm.state)
+	return rm.rebootAsExpected()
 }
 
 // ClearReboot clears state information about tracking requested reboots.
 func ClearReboot(st *state.State) {
 	st.Set("system-restart-from-boot-id", nil)
 }
 
-type restartStateKey struct{}
-
-type restartState struct {
-	restarting RestartType
-	h          Handler
-	bootID     string
+// Ensure implements StateManager.Ensure. Required by StateEngine, we
+// actually do nothing here.
+func (rm *RestartManager) Ensure() error {
+	return nil
 }
 
-func (rs *restartState) handleRestart(t RestartType) {
-	if rs.h != nil {
-		rs.h.HandleRestart(t)
+func (rm *RestartManager) handleRestart(t RestartType) {
+	if rm.h != nil {
+		rm.h.HandleRestart(t)
 	}
 }
 
-func (rs *restartState) rebootAsExpected(st *state.State) error {
-	if rs.h != nil {
-		return rs.h.RebootIsFine(st)
+func (rm *RestartManager) rebootAsExpected() error {
+	if rm.h != nil {
+		return rm.h.RebootAsExpected(rm.state)
 	}
 	return nil
 }
 
-func (rs *restartState) rebootDidNotHappen(st *state.State) error {
-	if rs.h != nil {
-		return rs.h.RebootIsMissing(st)
+func (rm *RestartManager) rebootDidNotHappen() error {
+	if rm.h != nil {
+		return rm.h.RebootDidNotHappen(rm.state)
 	}
 	return nil
 }
 
+func restartManager(st *state.State, errMsg string) *RestartManager {
+	cached := st.Cached(restartManagerKey{})
+	if cached == nil {
+		panic(errMsg)
+	}
+	return cached.(*RestartManager)
+}
+
 // Request asks for a restart of the managing process.
 // The state needs to be locked to request a restart.
 func Request(st *state.State, t RestartType) {
-	cached := st.Cached(restartStateKey{})
-	if cached == nil {
-		panic("internal error: cannot request a restart before restart.Init")
-	}
-	rs := cached.(*restartState)
+	rm := restartManager(st, "internal error: cannot request a restart before RestartManager initialization")
 	switch t {
 	case RestartSystem, RestartSystemNow, RestartSystemHaltNow, RestartSystemPoweroffNow:
-		st.Set("system-restart-from-boot-id", rs.bootID)
+		st.Set("system-restart-from-boot-id", rm.bootID)
 	}
-	rs.restarting = t
-	rs.handleRestart(t)
+	rm.restarting = t
+	rm.handleRestart(t)
 }
 
 // Pending returns whether a restart was requested with Request and of which type.
 func Pending(st *state.State) (bool, RestartType) {
-	cached := st.Cached(restartStateKey{})
+	cached := st.Cached(restartManagerKey{})
 	if cached == nil {
 		return false, RestartUnset
 	}
-	rs := cached.(*restartState)
-	return rs.restarting != RestartUnset, rs.restarting
+	rm := cached.(*RestartManager)
+	return rm.restarting != RestartUnset, rm.restarting
 }
 
 func FakePending(st *state.State, restarting RestartType) RestartType {
-	cached := st.Cached(restartStateKey{})
-	if cached == nil {
-		panic("internal error: cannot fake a restart request before restart.Init")
-	}
-	rs := cached.(*restartState)
-	old := rs.restarting
-	rs.restarting = restarting
+	rm := restartManager(st, "internal error: cannot mock a restart request before RestartManager initialization")
+	old := rm.restarting
+	rm.restarting = restarting
 	return old
 }
 
 func ReplaceBootID(st *state.State, bootID string) {
-	cached := st.Cached(restartStateKey{})
-	if cached == nil {
-		panic("internal error: cannot fake a restart request before restart.Init")
-	}
-	rs := cached.(*restartState)
-	rs.bootID = bootID
+	rm := restartManager(st, "internal error: cannot mock a restart request before RestartManager initialization")
+	rm.bootID = bootID
 }
 
 func BootID(st *state.State) string {
-	cached := st.Cached(restartStateKey{})
-	if cached == nil {
-		panic("internal error: cannot fake a restart request before restart.Init")
-	}
-	rs := cached.(*restartState)
-	return rs.bootID
+	rm := restartManager(st, "internal error: cannot get boot ID before RestartManager initialization")
+	return rm.bootID
 }