more cleanups

gmegidish · gmegidish · commit ffeb2baab700 · 2026-01-26T22:51:35.000+01:00
diff --git a/cli/device.go b/cli/device.go
@@ -49,7 +49,7 @@ var deviceInfoCmd = &cobra.Command{
 
 		// Start agent
 		err = targetDevice.StartAgent(devices.StartAgentConfig{
-			Registry: commands.GetRegistry(),
+			Hook: commands.GetShutdownHook(),
 		})
 		if err != nil {
 			response := commands.NewErrorResponse(fmt.Errorf("error starting agent: %v", err))
diff --git a/cli/screenshot.go b/cli/screenshot.go
@@ -80,7 +80,7 @@ var screencaptureCmd = &cobra.Command{
 			OnProgress: func(message string) {
 				utils.Verbose(message)
 			},
-			Registry: commands.GetRegistry(),
+			Hook: commands.GetShutdownHook(),
 		})
 		if err != nil {
 			response := commands.NewErrorResponse(fmt.Errorf("error starting agent: %v", err))
diff --git a/commands/commands.go b/commands/commands.go
@@ -33,24 +33,24 @@ func NewErrorResponse(err error) *CommandResponse {
 // DeviceCache provides a simple cache for devices to avoid repeated lookups
 var deviceCache = make(map[string]devices.ControllableDevice)
 
-// deviceRegistry holds the registry for device cleanup tracking.
-// It is set once at application startup via SetRegistry and used by commands
-// to register devices for graceful shutdown cleanup.
-var deviceRegistry *devices.DeviceRegistry
+// shutdownHook holds the shutdown hook for resource cleanup tracking.
+// It is set once at application startup via SetShutdownHook and used by commands
+// to register cleanup functions for graceful shutdown.
+var shutdownHook *devices.ShutdownHook
 
-// SetRegistry sets the global device registry for cleanup tracking.
+// SetShutdownHook sets the global shutdown hook for resource cleanup.
 // This should be called once at application startup (main.go or server.go).
-// The registry is used to track active devices and clean up their resources
+// The hook is used to register cleanup functions that will be called
 // during graceful shutdown (SIGINT/SIGTERM).
-func SetRegistry(registry *devices.DeviceRegistry) {
-	deviceRegistry = registry
+func SetShutdownHook(hook *devices.ShutdownHook) {
+	shutdownHook = hook
 }
 
-// GetRegistry returns the current device registry.
-// Returns nil if SetRegistry has not been called yet.
-// Commands use this to register devices when StartAgent is called.
-func GetRegistry() *devices.DeviceRegistry {
-	return deviceRegistry
+// GetShutdownHook returns the current shutdown hook.
+// Returns nil if SetShutdownHook has not been called yet.
+// Commands use this to register cleanup functions.
+func GetShutdownHook() *devices.ShutdownHook {
+	return shutdownHook
 }
 
 // FindDevice finds a device by ID, using cache when possible
diff --git a/commands/dump.go b/commands/dump.go
@@ -27,7 +27,7 @@ func DumpUICommand(req DumpUIRequest) *CommandResponse {
 
 	// Start agent if needed
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %w", targetDevice.ID(), err))
diff --git a/commands/info.go b/commands/info.go
@@ -18,7 +18,7 @@ func InfoCommand(deviceID string) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("error starting agent: %v", err))
diff --git a/commands/input.go b/commands/input.go
@@ -62,7 +62,7 @@ func TapCommand(req TapRequest) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
@@ -90,7 +90,7 @@ func LongPressCommand(req LongPressRequest) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
@@ -118,7 +118,7 @@ func TextCommand(req TextRequest) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
@@ -146,7 +146,7 @@ func ButtonCommand(req ButtonRequest) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
@@ -174,7 +174,7 @@ func GestureCommand(req GestureRequest) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
@@ -213,7 +213,7 @@ func SwipeCommand(req SwipeRequest) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
diff --git a/commands/orientation.go b/commands/orientation.go
@@ -31,7 +31,7 @@ func OrientationGetCommand(req OrientationGetRequest) *CommandResponse {
 
 	// start agent if needed
 	err = device.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", device.ID(), err))
@@ -63,7 +63,7 @@ func OrientationSetCommand(req OrientationSetRequest) *CommandResponse {
 
 	// start agent if needed
 	err = device.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", device.ID(), err))
diff --git a/commands/screenshot.go b/commands/screenshot.go
@@ -55,7 +55,7 @@ func ScreenshotCommand(req ScreenshotRequest) *CommandResponse {
 
 	// Start agent if needed
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
diff --git a/commands/url.go b/commands/url.go
@@ -24,7 +24,7 @@ func URLCommand(req URLRequest) *CommandResponse {
 	}
 
 	err = targetDevice.StartAgent(devices.StartAgentConfig{
-		Registry: GetRegistry(),
+		Hook: GetShutdownHook(),
 	})
 	if err != nil {
 		return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))
diff --git a/devices/common.go b/devices/common.go
@@ -31,7 +31,7 @@ type ScreenCaptureConfig struct {
 // StartAgentConfig contains configuration for agent startup operations
 type StartAgentConfig struct {
 	OnProgress func(message string) // optional progress callback
-	Registry   *DeviceRegistry       // optional registry for cleanup tracking
+	Hook       *ShutdownHook         // optional shutdown hook for cleanup tracking
 }
 
 // ScreenElementRect represents the rectangle coordinates and dimensions
diff --git a/devices/ios.go b/devices/ios.go
@@ -382,9 +382,10 @@ func (d *IOSDevice) startTunnel() error {
 }
 
 func (d *IOSDevice) StartAgent(config StartAgentConfig) error {
-	// register device for cleanup tracking
-	if config.Registry != nil {
-		config.Registry.Register(d)
+	// register cleanup hook for this device
+	if config.Hook != nil {
+		hookName := fmt.Sprintf("ios-device-%s", d.Udid)
+		config.Hook.Register(hookName, d.Cleanup)
 	}
 
 	// starting an agent on a real device requires quite a few things to happen in the right order:
@@ -1152,10 +1153,11 @@ func filterButtons(elements []ScreenElement) []ScreenElement {
 // StartDeviceKit starts the devicekit-ios XCUITest which provides:
 // - An HTTP server for tap/dumpUI commands (port 12004)
 // - A broadcast extension for H.264 screen streaming (port 12005)
-func (d *IOSDevice) StartDeviceKit(registry *DeviceRegistry) (*DeviceKitInfo, error) {
-	// register device for cleanup tracking
-	if registry != nil {
-		registry.Register(d)
+func (d *IOSDevice) StartDeviceKit(hook *ShutdownHook) (*DeviceKitInfo, error) {
+	// register cleanup hook for this device
+	if hook != nil {
+		hookName := fmt.Sprintf("ios-devicekit-%s", d.Udid)
+		hook.Register(hookName, d.Cleanup)
 	}
 
 	// Start tunnel if needed (iOS 17+)
diff --git a/devices/registry.go b/devices/registry.go
diff --git a/devices/shutdown.go b/devices/shutdown.go
@@ -0,0 +1,78 @@
+package devices
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/mobile-next/mobilecli/utils"
+)
+
+// ShutdownHook manages graceful shutdown of application resources.
+// It allows registering cleanup functions that will be called during
+// application shutdown (SIGINT/SIGTERM).
+type ShutdownHook struct {
+	mu    sync.RWMutex
+	hooks []namedHook
+}
+
+type namedHook struct {
+	name string
+	fn   func() error
+}
+
+// NewShutdownHook creates a new shutdown hook registry
+func NewShutdownHook() *ShutdownHook {
+	return &ShutdownHook{
+		hooks: make([]namedHook, 0),
+	}
+}
+
+// Register adds a cleanup function to be called during shutdown.
+// The name parameter is used for logging and error reporting.
+// Cleanup functions are called in the order they were registered.
+func (s *ShutdownHook) Register(name string, cleanupFn func() error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.hooks = append(s.hooks, namedHook{name: name, fn: cleanupFn})
+	utils.Verbose("Registered shutdown hook: %s", name)
+}
+
+// Shutdown executes all registered cleanup functions.
+// Returns an error if any cleanup function fails, but continues
+// executing remaining hooks to ensure best-effort cleanup.
+func (s *ShutdownHook) Shutdown() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if len(s.hooks) == 0 {
+		return nil
+	}
+
+	utils.Verbose("Executing %d shutdown hook(s)", len(s.hooks))
+	var errs []error
+
+	for _, hook := range s.hooks {
+		utils.Verbose("Running shutdown hook: %s", hook.name)
+		if err := hook.fn(); err != nil {
+			errs = append(errs, fmt.Errorf("%s: %w", hook.name, err))
+			utils.Verbose("Shutdown hook %s failed: %v", hook.name, err)
+		}
+	}
+
+	// clear hooks after shutdown
+	s.hooks = make([]namedHook, 0)
+
+	if len(errs) > 0 {
+		return fmt.Errorf("shutdown failed with %d error(s): %v", len(errs), errs)
+	}
+
+	utils.Verbose("All shutdown hooks completed successfully")
+	return nil
+}
+
+// Count returns the number of registered hooks (useful for testing)
+func (s *ShutdownHook) Count() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return len(s.hooks)
+}
diff --git a/devices/shutdown_test.go b/devices/shutdown_test.go
diff --git a/main.go b/main.go
diff --git a/server/server.go b/server/server.go

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ func InfoCommand(deviceID string) *CommandResponse {`
`18`	`18`	`}`
`19`	`19`
`20`	`20`	`err = targetDevice.StartAgent(devices.StartAgentConfig{`
`21`		`- Registry: GetRegistry(),`
	`21`	`+ Hook: GetShutdownHook(),`
`22`	`22`	`})`
`23`	`23`	`if err != nil {`
`24`	`24`	`return NewErrorResponse(fmt.Errorf("error starting agent: %v", err))`
Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ func URLCommand(req URLRequest) *CommandResponse {`
`24`	`24`	`}`
`25`	`25`
`26`	`26`	`err = targetDevice.StartAgent(devices.StartAgentConfig{`
`27`		`- Registry: GetRegistry(),`
	`27`	`+ Hook: GetShutdownHook(),`
`28`	`28`	`})`
`29`	`29`	`if err != nil {`
`30`	`30`	`return NewErrorResponse(fmt.Errorf("failed to start agent on device %s: %v", targetDevice.ID(), err))`
Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ type ScreenCaptureConfig struct {`
`31`	`31`	`// StartAgentConfig contains configuration for agent startup operations`
`32`	`32`	`type StartAgentConfig struct {`
`33`	`33`	`OnProgress func(message string) // optional progress callback`
`34`		`- Registry *DeviceRegistry // optional registry for cleanup tracking`
	`34`	`+ Hook *ShutdownHook // optional shutdown hook for cleanup tracking`
`35`	`35`	`}`
`36`	`36`
`37`	`37`	`// ScreenElementRect represents the rectangle coordinates and dimensions`