diff --git a/go/logic/migrator.go b/go/logic/migrator.go index aaacbdd2e..6aca663fb 100644 --- a/go/logic/migrator.go +++ b/go/logic/migrator.go @@ -234,10 +234,16 @@ func (mgtr *Migrator) retryOperationWithExponentialBackoff(operation func() erro // consumeRowCopyComplete blocks on the rowCopyComplete channel once, and then // consumes and drops any further incoming events that may be left hanging. func (mgtr *Migrator) consumeRowCopyComplete() { - if err := <-mgtr.rowCopyComplete; err != nil { - // Abort synchronously to ensure checkAbort() sees the error immediately - mgtr.abort(err) - // Don't mark row copy as complete if there was an error + select { + case err := <-mgtr.rowCopyComplete: + if err != nil { + // Abort synchronously to ensure checkAbort() sees the error immediately + mgtr.abort(err) + // Don't mark row copy as complete if there was an error + return + } + case <-mgtr.migrationContext.GetContext().Done(): + // Abort cancelled the context return } atomic.StoreInt64(&mgtr.rowCopyCompleteFlag, 1)