Fix tests and issues

spn/instance.go

@@ -34,10 +34,17 @@ import (
 
 // Instance is an instance of a Portmaster service.
 type Instance struct {
-	ctx          context.Context
-	cancelCtx    context.CancelFunc
+	ctx       context.Context
+	cancelCtx context.CancelFunc
+
+	shutdownCtx       context.Context
+	cancelShutdownCtx context.CancelFunc
+
 	serviceGroup *mgr.Group
 
+	binDir  string
+	dataDir string
+
 	exitCode atomic.Int32
 
 	base     *base.Base
@@ -67,6 +74,7 @@ type Instance struct {
 	terminal  *terminal.TerminalModule
 
 	CommandLineOperation func() error
+	ShouldRestart        bool
 }
 
 // New returns a new Portmaster service instance.
@@ -74,6 +82,7 @@ func New() (*Instance, error) {
 	// Create instance to pass it to modules.
 	instance := &Instance{}
 	instance.ctx, instance.cancelCtx = context.WithCancel(context.Background())
+	instance.shutdownCtx, instance.cancelShutdownCtx = context.WithCancel(context.Background())
 
 	binaryUpdateIndex := updates.Config{
 		// FIXME: fill
@@ -234,6 +243,18 @@ func (i *Instance) SetSleep(enabled bool) {
 	}
 }
 
+// BinDir returns the directory for binaries.
+// This directory may be read-only.
+func (i *Instance) BinDir() string {
+	return i.binDir
+}
+
+// DataDir returns the directory for variable data.
+// This directory is expected to be read/writeable.
+func (i *Instance) DataDir() string {
+	return i.dataDir
+}
+
 // Database returns the database module.
 func (i *Instance) Database() *dbmodule.DBModule {
 	return i.database
@@ -379,12 +400,6 @@ func (i *Instance) Ready() bool {
 	return i.serviceGroup.Ready()
 }
 
-// Ctx returns the instance context.
-// It is only canceled on shutdown.
-func (i *Instance) Ctx() context.Context {
-	return i.ctx
-}
-
 // Start starts the instance.
 func (i *Instance) Start() error {
 	return i.serviceGroup.Start()
@@ -392,7 +407,6 @@ func (i *Instance) Start() error {
 
 // Stop stops the instance and cancels the instance context when done.
 func (i *Instance) Stop() error {
-	defer i.cancelCtx()
 	return i.serviceGroup.Stop()
 }
 
@@ -406,6 +420,8 @@ func (i *Instance) Restart() {
 	i.core.EventRestart.Submit(struct{}{})
 	time.Sleep(10 * time.Millisecond)
 
+	// Set the restart flag and shutdown.
+	i.ShouldRestart = true
 	i.shutdown(RestartExitCode)
 }
 
@@ -419,32 +435,63 @@ func (i *Instance) Shutdown() {
 }
 
 func (i *Instance) shutdown(exitCode int) {
+	// Only shutdown once.
+	if i.IsShuttingDown() {
+		return
+	}
+
+	// Cancel main  context.
+	i.cancelCtx()
+
 	// Set given exit code.
 	i.exitCode.Store(int32(exitCode))
 
+	// Start shutdown asynchronously in a separate manager.
 	m := mgr.New("instance")
 	m.Go("shutdown", func(w *mgr.WorkerCtx) error {
-		for {
-			if err := i.Stop(); err != nil {
-				w.Error("failed to shutdown", "err", err, "retry", "1s")
-				time.Sleep(1 * time.Second)
-			} else {
-				return nil
-			}
+		// Stop all modules.
+		if err := i.Stop(); err != nil {
+			w.Error("failed to shutdown", "err", err)
 		}
+
+		// Cancel shutdown process context.
+		i.cancelShutdownCtx()
+		return nil
 	})
 }
 
-// Stopping returns whether the instance is shutting down.
-func (i *Instance) Stopping() bool {
-	return i.ctx.Err() == nil
+// Ctx returns the instance context.
+// It is canceled when shutdown is started.
+func (i *Instance) Ctx() context.Context {
+	return i.ctx
 }
 
-// Stopped returns a channel that is triggered when the instance has shut down.
-func (i *Instance) Stopped() <-chan struct{} {
+// IsShuttingDown returns whether the instance is shutting down.
+func (i *Instance) IsShuttingDown() bool {
+	return i.ctx.Err() != nil
+}
+
+// ShuttingDown returns a channel that is triggered when the instance starts shutting down.
+func (i *Instance) ShuttingDown() <-chan struct{} {
 	return i.ctx.Done()
 }
 
+// ShutdownCtx returns the instance shutdown context.
+// It is canceled when shutdown is complete.
+func (i *Instance) ShutdownCtx() context.Context {
+	return i.shutdownCtx
+}
+
+// IsShutDown returns whether the instance has stopped.
+func (i *Instance) IsShutDown() bool {
+	return i.shutdownCtx.Err() != nil
+}
+
+// ShutDownComplete returns a channel that is triggered when the instance has shut down.
+func (i *Instance) ShutdownComplete() <-chan struct{} {
+	return i.shutdownCtx.Done()
+}
+
 // ExitCode returns the set exit code of the instance.
 func (i *Instance) ExitCode() int {
 	return int(i.exitCode.Load())