Replace log.Fatale with context-based error propagation

This patch modifies gh-ost to use a cancellable context instead of
log.Fatale() in listenOnPanicAbort. When using gh-ost as a library, this
allows the calling application to recover from aborts (e.g. log the
failure reason) instead of having the entire process terminate via
os.Exit(). Now we store the error and cancel a context to signal all
goroutines to stop gracefully.

Author: ggilder

go/base/context.go

@@ -6,6 +6,7 @@
 package base
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"os"
@@ -225,6 +226,16 @@ type MigrationContext struct {
 	InCutOverCriticalSectionFlag           int64
 	PanicAbort                             chan error
 
+	// Context for cancellation signaling across all goroutines
+	// Stored in struct as it spans the entire migration lifecycle, not per-function.
+	// context.Context is safe for concurrent use by multiple goroutines.
+	ctx        context.Context //nolint:containedctx
+	cancelFunc context.CancelFunc
+
+	// Stores the fatal error that triggered abort
+	AbortError error
+	abortMutex *sync.Mutex
+
 	OriginalTableColumnsOnApplier    *sql.ColumnList
 	OriginalTableColumns             *sql.ColumnList
 	OriginalTableVirtualColumns      *sql.ColumnList
@@ -293,6 +304,7 @@ type ContextConfig struct {
 }
 
 func NewMigrationContext() *MigrationContext {
+	ctx, cancelFunc := context.WithCancel(context.Background())
 	return &MigrationContext{
 		Uuid:                                uuid.NewString(),
 		defaultNumRetries:                   60,
@@ -313,6 +325,9 @@ func NewMigrationContext() *MigrationContext {
 		lastHeartbeatOnChangelogMutex:       &sync.Mutex{},
 		ColumnRenameMap:                     make(map[string]string),
 		PanicAbort:                          make(chan error),
+		ctx:                                 ctx,
+		cancelFunc:                          cancelFunc,
+		abortMutex:                          &sync.Mutex{},
 		Log:                                 NewDefaultLogger(),
 	}
 }
@@ -982,3 +997,33 @@ func (this *MigrationContext) GetGhostTriggerName(triggerName string) string {
 func (this *MigrationContext) ValidateGhostTriggerLengthBelowMaxLength(triggerName string) bool {
 	return utf8.RuneCountInString(triggerName) <= mysql.MaxTableNameLength
 }
+
+// GetContext returns the migration context for cancellation checking
+func (this *MigrationContext) GetContext() context.Context {
+	return this.ctx
+}
+
+// SetAbortError stores the fatal error that triggered abort
+// Only the first error is stored (subsequent errors are ignored)
+func (this *MigrationContext) SetAbortError(err error) {
+	this.abortMutex.Lock()
+	defer this.abortMutex.Unlock()
+	if this.AbortError == nil {
+		this.AbortError = err
+	}
+}
+
+// GetAbortError retrieves the stored abort error
+func (this *MigrationContext) GetAbortError() error {
+	this.abortMutex.Lock()
+	defer this.abortMutex.Unlock()
+	return this.AbortError
+}
+
+// CancelContext cancels the migration context to signal all goroutines to stop
+// The cancel function is safe to call multiple times and from multiple goroutines.
+func (this *MigrationContext) CancelContext() {
+	if this.cancelFunc != nil {
+		this.cancelFunc()
+	}
+}

go/base/context_test.go

@@ -6,8 +6,10 @@
 package base
 
 import (
+	"errors"
 	"os"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -213,3 +215,58 @@ func TestReadConfigFile(t *testing.T) {
 		}
 	}
 }
+
+func TestSetAbortError_StoresFirstError(t *testing.T) {
+	ctx := NewMigrationContext()
+
+	err1 := errors.New("first error")
+	err2 := errors.New("second error")
+
+	ctx.SetAbortError(err1)
+	ctx.SetAbortError(err2)
+
+	got := ctx.GetAbortError()
+	if got != err1 { //nolint:errorlint // Testing pointer equality for sentinel error
+		t.Errorf("Expected first error %v, got %v", err1, got)
+	}
+}
+
+func TestSetAbortError_ThreadSafe(t *testing.T) {
+	ctx := NewMigrationContext()
+
+	var wg sync.WaitGroup
+	errors := []error{
+		errors.New("error 1"),
+		errors.New("error 2"),
+		errors.New("error 3"),
+	}
+
+	// Launch 3 goroutines trying to set error concurrently
+	for _, err := range errors {
+		wg.Add(1)
+		go func(e error) {
+			defer wg.Done()
+			ctx.SetAbortError(e)
+		}(err)
+	}
+
+	wg.Wait()
+
+	// Should store exactly one of the errors
+	got := ctx.GetAbortError()
+	if got == nil {
+		t.Fatal("Expected error to be stored, got nil")
+	}
+
+	// Verify it's one of the errors we sent
+	found := false
+	for _, err := range errors {
+		if got == err { //nolint:errorlint // Testing pointer equality for sentinel error
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Errorf("Stored error %v not in list of sent errors", got)
+	}
+}

go/logic/applier.go

@@ -695,7 +695,17 @@ func (this *Applier) InitiateHeartbeat() {
 
 	ticker := time.NewTicker(time.Duration(this.migrationContext.HeartbeatIntervalMilliseconds) * time.Millisecond)
 	defer ticker.Stop()
-	for range ticker.C {
+	for {
+		// Check for context cancellation each iteration
+		ctx := this.migrationContext.GetContext()
+		select {
+		case <-ctx.Done():
+			this.migrationContext.Log.Debugf("Heartbeat injection cancelled")
+			return
+		case <-ticker.C:
+			// Process heartbeat
+		}
+
 		if atomic.LoadInt64(&this.finishedMigrating) > 0 {
 			return
 		}
@@ -706,7 +716,12 @@ func (this *Applier) InitiateHeartbeat() {
 			continue
 		}
 		if err := injectHeartbeat(); err != nil {
-			this.migrationContext.PanicAbort <- fmt.Errorf("injectHeartbeat writing failed %d times, last error: %w", numSuccessiveFailures, err)
+			select {
+			case this.migrationContext.PanicAbort <- fmt.Errorf("injectHeartbeat writing failed %d times, last error: %w", numSuccessiveFailures, err):
+				// Error sent successfully
+			case <-this.migrationContext.GetContext().Done():
+				// Context cancelled, someone else already reported an error
+			}
 			return
 		}
 	}