refactor: add mock-bus test harness, fix potential data race for gripper connection close

arm.go

@@ -87,11 +87,12 @@ func (cfg *SO101ArmConfig) Validate(path string) ([]string, []string, error) {
 type so101 struct {
 	resource.AlwaysRebuild
 
-	name       resource.Name
-	logger     logging.Logger
-	cfg        *SO101ArmConfig
-	opMgr      *operation.SingleOperationManager
-	controller *SafeSoArmController
+	name           resource.Name
+	logger         logging.Logger
+	cfg            *SO101ArmConfig
+	opMgr          *operation.SingleOperationManager
+	controller     *SafeSoArmController
+	controllerPort string // port path used to acquire the shared controller
 
 	mu       sync.RWMutex
 	moveLock sync.Mutex
@@ -105,10 +106,6 @@ type so101 struct {
 	defaultAcc   float32
 
 	motion motion.Service
-
-	cancelCtx  context.Context
-	cancelFunc func()
-	initCtx    context.Context // Context for initialization operations
 }
 
 func makeSO101ModelFrame() (referenceframe.Model, error) {
@@ -227,51 +224,50 @@ func NewSO101(ctx context.Context, deps resource.Dependencies, name resource.Nam
 
 	model, err := makeSO101ModelFrame()
 	if err != nil {
-		ReleaseSharedController() // Clean up on error
+		globalRegistry.ReleaseController(controllerConfig.Port)
 		return nil, fmt.Errorf("failed to create kinematic model: %w", err)
 	}
 
 	var ms motion.Service
 	if conf.Motion != "" {
 		if deps == nil {
+			globalRegistry.ReleaseController(controllerConfig.Port)
 			return nil, fmt.Errorf("no deps")
 		}
 		ms, err = motion.FromProvider(deps, conf.Motion)
 		if err != nil {
+			globalRegistry.ReleaseController(controllerConfig.Port)
 			return nil, err
 		}
 	} else {
 		ms, err = motion.FromProvider(deps, "builtin")
 		if err != nil {
+			globalRegistry.ReleaseController(controllerConfig.Port)
 			return nil, err
 		}
 	}
 
-	cancelCtx, cancelFunc := context.WithCancel(context.Background())
-
 	arm := &so101{
-		name:         name,
-		cfg:          conf,
-		opMgr:        operation.NewSingleOperationManager(),
-		logger:       logger,
-		controller:   controller,
-		model:        model,
-		armServoIDs:  conf.ServoIDs, // Store which servos this arm controls
-		defaultSpeed: speedDegsPerSec,
-		defaultAcc:   accelerationDegsPerSec,
-		motion:       ms,
-		cancelCtx:    cancelCtx,
-		cancelFunc:   cancelFunc,
-		initCtx:      ctx, // Store initialization context
+		name:           name,
+		cfg:            conf,
+		opMgr:          operation.NewSingleOperationManager(),
+		logger:         logger,
+		controller:     controller,
+		controllerPort: controllerConfig.Port,
+		model:          model,
+		armServoIDs:    conf.ServoIDs, // Store which servos this arm controls
+		defaultSpeed:   speedDegsPerSec,
+		defaultAcc:     accelerationDegsPerSec,
+		motion:         ms,
 	}
 
 	logger.Debugf("SO-101 configured with speed: %.1f deg/s, acceleration: %.1f deg/s²",
 		speedDegsPerSec, accelerationDegsPerSec)
 	logger.Debugf("Arm controlling servo IDs: %v", arm.armServoIDs)
 
 	// Initialize and verify servo connections
-	if err := arm.initializeServos(); err != nil {
-		ReleaseSharedController() // Clean up on error
+	if err := arm.initializeServos(ctx); err != nil {
+		globalRegistry.ReleaseController(controllerConfig.Port)
 		return nil, fmt.Errorf("failed to initialize servos: %w", err)
 	}
 
@@ -470,14 +466,14 @@ func (s *so101) DoCommand(ctx context.Context, cmd map[string]interface{}) (map[
 		}, nil
 
 	case "diagnose":
-		err := s.diagnoseConnection()
+		err := s.diagnoseConnection(ctx)
 		return map[string]interface{}{
 			"success": err == nil,
 			"error":   fmt.Sprintf("%v", err),
 		}, nil
 
 	case "verify_config":
-		err := s.verifyServoConfig()
+		err := s.verifyServoConfig(ctx)
 		return map[string]interface{}{
 			"success": err == nil,
 			"error":   fmt.Sprintf("%v", err),
@@ -488,7 +484,7 @@ func (s *so101) DoCommand(ctx context.Context, cmd map[string]interface{}) (map[
 		if r, ok := cmd["retries"].(float64); ok {
 			retries = int(r)
 		}
-		err := s.initializeServosWithRetry(retries)
+		err := s.initializeServosWithRetry(ctx, retries)
 		return map[string]interface{}{
 			"success": err == nil,
 			"error":   fmt.Sprintf("%v", err),
@@ -618,25 +614,24 @@ func (s *so101) Geometries(ctx context.Context, extra map[string]interface{}) ([
 }
 
 func (s *so101) Close(context.Context) error {
-	s.cancelFunc()
-	ReleaseSharedController()
+	globalRegistry.ReleaseController(s.controllerPort)
 	return nil
 }
 
 // initializeServos pings each servo and enables torque to ensure proper communication
-func (s *so101) initializeServos() error {
-	return s.initializeServosWithRetry(3)
+func (s *so101) initializeServos(ctx context.Context) error {
+	return s.initializeServosWithRetry(ctx, 3)
 }
 
 // initializeServosWithRetry attempts servo initialization with retries
-func (s *so101) initializeServosWithRetry(maxRetries int) error {
+func (s *so101) initializeServosWithRetry(ctx context.Context, maxRetries int) error {
 	s.logger.Debug("Initializing SO-101 arm servos...")
 
 	var lastErr error
 	for attempt := 1; attempt <= maxRetries; attempt++ {
 		s.logger.Debugf("Arm servo initialization attempt %d/%d", attempt, maxRetries)
 
-		if err := s.doServoInitialization(); err != nil {
+		if err := s.doServoInitialization(ctx); err != nil {
 			lastErr = err
 			s.logger.Warnf("Initialization attempt %d failed: %v", attempt, err)
 
@@ -656,10 +651,7 @@ func (s *so101) initializeServosWithRetry(maxRetries int) error {
 }
 
 // doServoInitialization performs the actual initialization steps
-func (s *so101) doServoInitialization() error {
-	// Use stored initialization context instead of creating new one
-	ctx := s.initCtx
-
+func (s *so101) doServoInitialization(ctx context.Context) error {
 	// Ping all servos to ensure they're responding
 	s.logger.Debug("Pinging all servos...")
 	if err := s.controller.Ping(ctx); err != nil {
@@ -690,10 +682,7 @@ func (s *so101) doServoInitialization() error {
 }
 
 // diagnoseConnection provides detailed diagnostics for troubleshooting
-func (s *so101) diagnoseConnection() error {
-	// Use stored initialization context instead of creating new one
-	ctx := s.initCtx
-
+func (s *so101) diagnoseConnection(ctx context.Context) error {
 	s.logger.Debug("Starting SO-101 arm connection diagnosis...")
 
 	// Test overall ping
@@ -718,10 +707,7 @@ func (s *so101) diagnoseConnection() error {
 }
 
 // verifyServoConfig checks servo configuration
-func (s *so101) verifyServoConfig() error {
-	// Use stored initialization context instead of creating new one
-	ctx := s.initCtx
-
+func (s *so101) verifyServoConfig(ctx context.Context) error {
 	s.logger.Debug("Verifying arm servo configuration...")
 
 	positions, err := s.controller.GetJointPositionsForServos(ctx, s.armServoIDs)

calibration.go

@@ -108,10 +108,11 @@ func (cfg *SO101CalibrationSensorConfig) Validate(path string) ([]string, []stri
 type so101CalibrationSensor struct {
 	resource.AlwaysRebuild
 
-	name       resource.Name
-	logger     logging.Logger
-	cfg        *SO101CalibrationSensorConfig
-	controller *SafeSoArmController
+	name           resource.Name
+	logger         logging.Logger
+	cfg            *SO101CalibrationSensorConfig
+	controller     *SafeSoArmController
+	controllerPort string // port path used to acquire the shared controller
 
 	// Calibration state
 	mu               sync.RWMutex
@@ -205,6 +206,7 @@ func NewSO101CalibrationSensor(
 		logger:          logger,
 		cfg:             conf,
 		controller:      controller,
+		controllerPort:  controllerConfig.Port,
 		state:           StateIdle,
 		joints:          joints,
 		servoNames:      servoNames,
@@ -1230,7 +1232,7 @@ func (cs *so101CalibrationSensor) Close(ctx context.Context) error {
 	cs.recordingActive = false
 
 	if cs.controller != nil {
-		ReleaseSharedController()
+		globalRegistry.ReleaseController(cs.controllerPort)
 	}
 
 	return nil

gripper.go

@@ -58,11 +58,12 @@ func (cfg *SO101GripperConfig) Validate(path string) ([]string, []string, error)
 type so101Gripper struct {
 	resource.AlwaysRebuild
 
-	name       resource.Name
-	logger     logging.Logger
-	controller *SafeSoArmController
-	geometries []spatialmath.Geometry
-	servoID    int
+	name           resource.Name
+	logger         logging.Logger
+	controller     *SafeSoArmController
+	controllerPort string // port path used to acquire the shared controller
+	geometries     []spatialmath.Geometry
+	servoID        int
 
 	mu       sync.Mutex
 	isMoving atomic.Bool
@@ -131,6 +132,7 @@ func newSO101Gripper(ctx context.Context, deps resource.Dependencies, conf resou
 		name:           conf.ResourceName(),
 		logger:         logger,
 		controller:     controller,
+		controllerPort: controllerConfig.Port,
 		geometries:     geometries,
 		servoID:        cfg.ServoID,
 		speed:          30,
@@ -339,7 +341,7 @@ func (g *so101Gripper) DoCommand(ctx context.Context, cmd map[string]interface{}
 }
 
 func (g *so101Gripper) Close(ctx context.Context) error {
-	ReleaseSharedController()
+	globalRegistry.ReleaseController(g.controllerPort)
 	return nil
 }