Atomic Transactions handling with Online DDL

go/vt/vttablet/onlineddl/executor.go

@@ -178,7 +178,7 @@ type Executor struct {
 	ts                    *topo.Server
 	lagThrottler          *throttle.Throttler
 	toggleBufferTableFunc func(cancelCtx context.Context, tableName string, timeout time.Duration, bufferQueries bool)
-	IsPreparedPoolEmpty   func(tableName string) bool
+	isPreparedPoolEmpty   func(tableName string) bool
 	requestGCChecksFunc   func()
 	tabletAlias           *topodatapb.TabletAlias
 
@@ -257,7 +257,7 @@ func NewExecutor(env tabletenv.Env, tabletAlias *topodatapb.TabletAlias, ts *top
 		ts:                    ts,
 		lagThrottler:          lagThrottler,
 		toggleBufferTableFunc: toggleBufferTableFunc,
-		IsPreparedPoolEmpty:   isPreparedPoolEmpty,
+		isPreparedPoolEmpty:   isPreparedPoolEmpty,
 		requestGCChecksFunc:   requestGCChecksFunc,
 		ticks:                 timer.NewTimer(migrationCheckInterval),
 		// Gracefully return an error if any caller tries to execute
@@ -1110,7 +1110,7 @@ func (e *Executor) cutOverVReplMigration(ctx context.Context, s *VReplStream, sh
 	if shouldForceCutOver {
 		// We should only proceed with forceful cut over if there is no pending atomic transaction for the table.
 		// This will help in keeping the atomicity guarantee of a prepared transaction.
-		if err := e.checkOnPreparedPool(onlineDDL.Table, 100*time.Millisecond, 1); err != nil {
+		if err := e.checkOnPreparedPool(ctx, onlineDDL.Table, 100*time.Millisecond); err != nil {
 			return err
 		}
 		if err := e.killTableLockHoldersAndAccessors(ctx, onlineDDL.Table); err != nil {
@@ -5356,14 +5356,19 @@ func (e *Executor) OnSchemaMigrationStatus(ctx context.Context,
 	return e.onSchemaMigrationStatus(ctx, uuidParam, status, dryRun, progressPct, etaSeconds, rowsCopied, hint)
 }
 
-func (e *Executor) checkOnPreparedPool(table string, waitTime time.Duration, retries int) error {
-	for i := 0; i <= retries; i++ {
-		if e.IsPreparedPoolEmpty(table) {
+func (e *Executor) checkOnPreparedPool(ctx context.Context, table string, waitTime time.Duration) error {
+	if e.isPreparedPoolEmpty(table) {
+		return nil
+	}
+
+	select {
+	case <-ctx.Done():
+		// Return context error if context is done
+		return ctx.Err()
+	case <-time.After(waitTime):
+		if e.isPreparedPoolEmpty(table) {
 			return nil
 		}
-		if i < retries {
-			time.Sleep(waitTime)
-		}
+		return vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "cannot force cut-over on non-empty prepared pool for table: %s", table)
 	}
-	return vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "cannot force cut-over on non-empty prepared pool for table: %s", table)
 }