overlord: start turning restart into a full state manager

overlord/devicestate/devicestate_test.go

@@ -165,13 +165,14 @@ func (s *deviceMgrBaseSuite) setupBaseTest(c *C, classic bool) {
 	s.o = overlord.Mock()
 	s.state = s.o.State()
 	s.state.Lock()
-	restart.Init(s.state, "boot-id-0", snapstatetest.MockRestartHandler(func(req restart.RestartType) {
+	_, err = restart.Manager(s.state, "boot-id-0", snapstatetest.MockRestartHandler(func(req restart.RestartType) {
 		s.restartRequests = append(s.restartRequests, req)
 		if s.restartObserve != nil {
 			s.restartObserve()
 		}
 	}))
 	s.state.Unlock()
+	c.Assert(err, IsNil)
 	s.se = s.o.StateEngine()
 
 	s.AddCleanup(sysdb.MockGenericClassicModel(s.storeSigning.GenericClassicModel))

overlord/hookstate/hookstate_test.go

@@ -76,8 +76,9 @@ func (s *baseHookManagerSuite) commonSetUpTest(c *C) {
 	s.o = overlord.Mock()
 	s.state = s.o.State()
 	s.state.Lock()
-	restart.Init(s.state, "boot-id-0", nil)
+	_, err := restart.Manager(s.state, "boot-id-0", nil)
 	s.state.Unlock()
+	c.Assert(err, IsNil)
 	manager, err := hookstate.Manager(s.state, s.o.TaskRunner())
 	c.Assert(err, IsNil)
 	s.manager = manager

overlord/overlord.go

@@ -1,7 +1,7 @@
 // -*- Mode: Go; indent-tabs-mode: t -*-
 
 /*
- * Copyright (C) 2016-2021 Canonical Ltd
+ * Copyright (C) 2016-2022 Canonical Ltd
  *
  * This program is free software: you can redistribute it and/or modify
  * it under the terms of the GNU General Public License version 3 as
@@ -99,6 +99,7 @@ type Overlord struct {
 	inited     bool
 	startedUp  bool
 	runner     *state.TaskRunner
+	restartMgr *restart.RestartManager
 	snapMgr    *snapstate.SnapManager
 	serviceMgr *servicestate.ServiceManager
 	assertMgr  *assertstate.AssertManager
@@ -124,7 +125,7 @@ func New(restartHandler restart.Handler) (*Overlord, error) {
 		path:         dirs.SnapStateFile,
 		ensureBefore: o.ensureBefore,
 	}
-	s, err := o.loadState(backend, restartHandler)
+	s, restartMgr, err := o.loadState(backend, restartHandler)
 	if err != nil {
 		return nil, err
 	}
@@ -138,6 +139,8 @@ func New(restartHandler restart.Handler) (*Overlord, error) {
 	}
 	o.runner.AddOptionalHandler(matchAnyUnknownTask, handleUnknownTask, nil)
 
+	o.addManager(restartMgr)
+
 	hookMgr, err := hookstate.Manager(s, o.runner)
 	if err != nil {
 		return nil, err
@@ -212,6 +215,8 @@ func (o *Overlord) addManager(mgr StateManager) {
 		o.cmdMgr = x
 	case *snapshotstate.SnapshotManager:
 		o.shotMgr = x
+	case *restart.RestartManager:
+		o.restartMgr = x
 	}
 	o.stateEng.AddManager(mgr)
 }
@@ -253,23 +258,23 @@ func lockWithTimeout(l *osutil.FileLock, timeout time.Duration) error {
 	}
 }
 
-func (o *Overlord) loadState(backend state.Backend, restartHandler restart.Handler) (*state.State, error) {
+func (o *Overlord) loadState(backend state.Backend, restartHandler restart.Handler) (*state.State, *restart.RestartManager, error) {
 	flock, err := initStateFileLock()
 	if err != nil {
-		return nil, fmt.Errorf("fatal: error opening lock file: %v", err)
+		return nil, nil, fmt.Errorf("fatal: error opening lock file: %v", err)
 	}
 	o.stateFLock = flock
 
 	logger.Noticef("Acquiring state lock file")
 	if err := lockWithTimeout(o.stateFLock, stateLockTimeout); err != nil {
 		logger.Noticef("Failed to lock state file")
-		return nil, fmt.Errorf("fatal: could not lock state file: %v", err)
+		return nil, nil, fmt.Errorf("fatal: could not lock state file: %v", err)
 	}
 	logger.Noticef("Acquired state lock file")
 
 	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)
 	}
 
 	perfTimings := timings.New(map[string]string{"startup": "load-state"})
@@ -279,17 +284,20 @@ func (o *Overlord) loadState(backend state.Backend, restartHandler restart.Handl
 		// by the snapd package
 		stateDir := filepath.Dir(dirs.SnapStateFile)
 		if !osutil.IsDirectory(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(dirs.SnapStateFile)
 	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()
 
@@ -298,29 +306,29 @@ func (o *Overlord) loadState(backend state.Backend, restartHandler restart.Handl
 		s, err = state.ReadState(backend, r)
 	})
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	s.Lock()
 	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) newStoreWithContext(storeCtx store.DeviceAndAuthContext) snapstate.StoreService {
@@ -613,6 +621,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
+}
+
 // SnapManager returns the snap manager responsible for snaps under
 // the overlord.
 func (o *Overlord) SnapManager() *snapstate.SnapManager {

overlord/overlord_test.go

@@ -116,6 +116,7 @@ func (ovs *overlordSuite) TestNew(c *C) {
 
 	c.Check(o.StateEngine(), NotNil)
 	c.Check(o.TaskRunner(), NotNil)
+	c.Check(o.RestartManager(), NotNil)
 	c.Check(o.SnapManager(), NotNil)
 	c.Check(o.ServiceManager(), NotNil)
 	c.Check(o.AssertManager(), NotNil)
@@ -175,6 +176,7 @@ func (ovs *overlordSuite) TestNewWithGoodState(c *C) {
 
 	o, err := overlord.New(nil)
 	c.Assert(err, IsNil)
+	c.Check(o.RestartManager(), NotNil)
 
 	state := o.State()
 	c.Assert(err, IsNil)

overlord/restart/restart.go

@@ -1,7 +1,7 @@
 // -*- Mode: Go; indent-tabs-mode: t -*-
 
 /*
- * Copyright (C) 2016-2021 Canonical Ltd
+ * Copyright (C) 2016-2022 Canonical Ltd
  *
  * This program is free software: you can redistribute it and/or modify
  * it under the terms of the GNU General Public License version 3 as
@@ -17,7 +17,9 @@
  *
  */
 
-// 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 (
@@ -57,108 +59,122 @@ type Handler interface {
 	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(restartStateKey{}, rs)
+	st.Cache(restartManagerKey{}, rm)
+	if err := rm.init(fromBootID, curBootID); err != nil {
+		return nil, err
+	}
+	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 (m *RestartManager) Ensure() error {
+	return nil
 }
 
-func (rs *restartState) handleRestart(t RestartType, rebootInfo *boot.RebootInfo) {
-	if rs.h != nil {
-		rs.h.HandleRestart(t, rebootInfo)
+func (rm *RestartManager) handleRestart(t RestartType, rebootInfo *boot.RebootInfo) {
+	if rm.h != nil {
+		rm.h.HandleRestart(t, rebootInfo)
 	}
 }
 
-func (rs *restartState) rebootAsExpected(st *state.State) error {
-	if rs.h != nil {
-		return rs.h.RebootAsExpected(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.RebootDidNotHappen(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, rebootInfo *boot.RebootInfo) {
-	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, rebootInfo)
+	rm.restarting = t
+	rm.handleRestart(t, rebootInfo)
 }
 
 // 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 MockPending(st *state.State, restarting RestartType) RestartType {
-	cached := st.Cached(restartStateKey{})
-	if cached == nil {
-		panic("internal error: cannot mock 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 mock 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
 }