RSDK-4116: Ensure access to wheeled base left/right motors are synchronized. (viamrobotics#2667)

Author: dgottlieb

operation/manager.go

@@ -10,22 +10,24 @@ import (
 	"go.viam.com/utils"
 )
 
-// SingleOperationManager ensures only 1 operation is happening a time
+// SingleOperationManager ensures only 1 operation is happening at a time.
 // An operation can be nested, so if there is already an operation in progress,
 // it can have sub-operations without an issue.
 type SingleOperationManager struct {
 	mu        sync.Mutex
 	currentOp *anOp
 }
 
-// CancelRunning cancel's a current operation unless it's mine.
+// CancelRunning cancels a current operation unless it's mine.
 func (sm *SingleOperationManager) CancelRunning(ctx context.Context) {
 	if ctx.Value(somCtxKeySingleOp) != nil {
 		return
 	}
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
-	sm.cancelInLock(ctx)
+	if sm.currentOp != nil {
+		sm.currentOp.cancelAndWaitFunc()
+	}
 }
 
 // OpRunning returns if there is a current operation.
@@ -41,32 +43,49 @@ const somCtxKeySingleOp = somCtxKey(iota)
 
 // New creates a new operation, cancels previous, returns a new context and function to call when done.
 func (sm *SingleOperationManager) New(ctx context.Context) (context.Context, func()) {
-	// handle nested ops
+	// Handle nested ops. Note an operation set on a context by one `SingleOperationManager` can be
+	// observed on a different instance of a `SingleOperationManager`.
 	if ctx.Value(somCtxKeySingleOp) != nil {
 		return ctx, func() {}
 	}
 
 	sm.mu.Lock()
 
-	// first cancel any old operation
-	sm.cancelInLock(ctx)
+	// Cancel any existing operation. This blocks until the operation is completed.
+	if sm.currentOp != nil {
+		sm.currentOp.cancelAndWaitFunc()
+	}
 
-	theOp := &anOp{}
+	theOp := &anOp{
+		closedCond: sync.NewCond(&sm.mu),
+	}
 
 	ctx = context.WithValue(ctx, somCtxKeySingleOp, theOp)
 
-	theOp.ctx, theOp.cancelFunc = context.WithCancel(ctx)
+	var newUserCtx context.Context
+	newUserCtx, theOp.interruptFunc = context.WithCancel(ctx)
+	theOp.cancelAndWaitFunc = func() {
+		// Precondition: Caller must be holding `sm.mu`.
+		//
+		// If there are two threads competing to win a race, it's not sufficient to return once the
+		// condition variable is signaled. We must re-check that a new operation didn't beat us to
+		// getting the next operation slot.
+		//
+		// Ironically, "winning the race" in this scenario just means the "loser" is going to
+		// immediately interrupt the winner. A future optimization could avoid this unnecessary
+		// starting/stopping.
+		for sm.currentOp != nil {
+			sm.currentOp.interruptFunc()
+			sm.currentOp.closedCond.Wait()
+		}
+	}
 	sm.currentOp = theOp
 	sm.mu.Unlock()
 
-	return theOp.ctx, func() {
-		if !theOp.closed {
-			theOp.closed = true
-		}
+	return newUserCtx, func() {
 		sm.mu.Lock()
-		if theOp == sm.currentOp {
-			sm.currentOp = nil
-		}
+		theOp.closedCond.Broadcast()
+		sm.currentOp = nil
 		sm.mu.Unlock()
 	}
 }
@@ -93,17 +112,11 @@ func (sm *SingleOperationManager) WaitTillNotPowered(ctx context.Context, pollTi
 ) (err error) {
 	// Defers a function that will stop and clean up if the context errors
 	defer func(ctx context.Context) {
-		var errStop error
 		if errors.Is(ctx.Err(), context.Canceled) {
-			sm.mu.Lock()
-			oldOp := sm.currentOp == ctx.Value(somCtxKeySingleOp)
-			sm.mu.Unlock()
-
-			if oldOp || sm.currentOp == nil {
-				errStop = stop(ctx, map[string]interface{}{})
-			}
+			err = multierr.Combine(ctx.Err(), stop(ctx, map[string]interface{}{}))
+		} else {
+			err = ctx.Err()
 		}
-		err = multierr.Combine(ctx.Err(), errStop)
 	}(ctx)
 	return sm.WaitForSuccess(
 		ctx,
@@ -139,21 +152,12 @@ func (sm *SingleOperationManager) WaitForSuccess(
 	}
 }
 
-func (sm *SingleOperationManager) cancelInLock(ctx context.Context) {
-	myOp := ctx.Value(somCtxKeySingleOp)
-	op := sm.currentOp
-
-	if op == nil || myOp == op {
-		return
-	}
-
-	op.cancelFunc()
-
-	sm.currentOp = nil
-}
-
 type anOp struct {
-	ctx        context.Context
-	cancelFunc context.CancelFunc
-	closed     bool
+	// cancelAndWaitFunc waits until the `SingleOperationManager.currentOp` is empty. This will
+	// interrupt any existing operations as necessary.
+	cancelAndWaitFunc func()
+	// Cancels the context of what's currently running an operation.
+	interruptFunc context.CancelFunc
+	// Used with `SingleOperationManager.mu`.
+	closedCond *sync.Cond
 }

operation/manager_test.go

@@ -97,22 +97,39 @@ func TestDontCancel(t *testing.T) {
 }
 
 func TestCancelRace(t *testing.T) {
+	// First set up the "worker" context and register an operation on
+	// the `SingleOperationManager` instance.
 	som := SingleOperationManager{}
-	ctx, done := som.New(context.Background())
-	defer done()
-
-	var wg sync.WaitGroup
-
-	wg.Add(1)
-	go func() {
-		_, done := som.New(context.Background())
-		wg.Done()
-		defer done()
-	}()
+	workerError := make(chan error, 1)
+	workerCtx, somCleanupFunc := som.New(context.Background())
+
+	// Spin up a separate go-routine for the worker to listen for cancelation. Canceling an
+	// operation blocks until the operation completes. This goroutine is responsible for running
+	// `somCleanupFunc` to signal that canceling has completed.
+	workerFunc := func() {
+		defer somCleanupFunc()
+
+		select {
+		case <-workerCtx.Done():
+			workerError <- nil
+		case <-time.After(5 * time.Second):
+			workerError <- errors.New("Failed to be signaled via a cancel")
+		}
 
-	som.CancelRunning(ctx)
-	wg.Wait()
-	test.That(t, ctx.Err(), test.ShouldNotBeNil)
+		close(workerError)
+	}
+	go workerFunc()
+
+	// Set up a "test" context to cancel the worker.
+	testCtx, testCleanupFunc := context.WithTimeout(context.Background(), time.Second)
+	defer testCleanupFunc()
+	som.CancelRunning(testCtx)
+	// If `workerCtx.Done` was observed to be closed, the worker thread will pass a `nil` error back.
+	test.That(t, <-workerError, test.ShouldBeNil)
+	// When `SingleOperationManager` cancels an operation, the operation's context should be in a
+	// "context canceled" error state.
+	test.That(t, workerCtx.Err(), test.ShouldNotBeNil)
+	test.That(t, workerCtx.Err(), test.ShouldEqual, context.Canceled)
 }
 
 func TestStopCalled(t *testing.T) {
@@ -155,24 +172,6 @@ func TestErrorContainsStopAndCancel(t *testing.T) {
 	test.That(t, errRet.Error(), test.ShouldEqual, "context canceled; Stop failed")
 }
 
-func TestStopNotCalledOnOldContext(t *testing.T) {
-	som := SingleOperationManager{}
-	ctx, done := som.New(context.Background())
-	defer done()
-	mock := &mock{stopCount: 0}
-	var wg sync.WaitGroup
-
-	wg.Add(1)
-	go func() {
-		som.WaitTillNotPowered(ctx, time.Second, mock, mock.stop)
-		wg.Done()
-	}()
-	som.New(context.Background())
-	wg.Wait()
-	test.That(t, ctx.Err(), test.ShouldNotBeNil)
-	test.That(t, mock.stopCount, test.ShouldEqual, 0)
-}
-
 type mock struct {
 	stopCount int
 }