Merge branch 'master' into 4.22.07.x

# Conflicts:
#	driver/driver.go
#	driver/mysql/applier_incr.go
#	g/g.go

Author: ffffwh

api/handler/v2/database.go

@@ -1,6 +1,7 @@
 package v2
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"strings"
@@ -123,7 +124,7 @@ func listOracleSchema(logger hclog.Logger, reqParam *models.ListDatabaseSchemasR
 		}
 		reqParam.Password = realPwd
 	}
-	oracleDb, err := config.NewDB(&config.OracleConfig{
+	oracleDb, err := config.NewDB(context.TODO(), &config.OracleConfig{
 		User:        reqParam.User,
 		Password:    reqParam.Password,
 		Host:        reqParam.Host,
@@ -206,7 +207,7 @@ func ListDatabaseColumnsV2(c echo.Context) error {
 			}
 			reqParam.Password = realPwd
 		}
-		oracleDb, err := config.NewDB(&config.OracleConfig{
+		oracleDb, err := config.NewDB(context.TODO(), &config.OracleConfig{
 			User:        reqParam.User,
 			Password:    reqParam.Password,
 			Host:        reqParam.Host,

driver/common/common.go

@@ -28,9 +28,7 @@ const (
 )
 
 const (
-	TaskStateComplete int = iota
-	TaskStateRestart
-	TaskStateDead
+	TaskStateDead = 2
 )
 
 const (

driver/common/dumper.go

@@ -7,6 +7,7 @@
 package common
 
 import (
+	"context"
 	"sync"
 
 	"github.com/actiontech/dtle/g"
@@ -16,6 +17,7 @@ type GetChunkDataFn func() (nRows int64, err error)
 type PrepareFn func() (err error)
 
 type Dumper struct {
+	Ctx                context.Context
 	Logger             g.LoggerType
 	ChunkSize          int64
 	TableSchema        string
@@ -36,9 +38,10 @@ type Dumper struct {
 	PrepareForDumping PrepareFn
 }
 
-func NewDumper(table *Table, chunkSize int64, logger g.LoggerType, memory *int64) *Dumper {
+func NewDumper(ctx context.Context, table *Table, chunkSize int64, logger g.LoggerType, memory *int64) *Dumper {
 
 	dumper := &Dumper{
+		Ctx:            ctx,
 		Logger:         logger,
 		TableSchema:    table.TableSchema,
 		TableName:      table.TableName,

driver/common/utils.go

@@ -8,6 +8,7 @@ package common
 
 import (
 	"fmt"
+	"github.com/hashicorp/nomad/plugins/drivers"
 	"github.com/pingcap/tidb/parser/format"
 	"regexp"
 	"strconv"
@@ -37,3 +38,10 @@ func MysqlVersionInDigit(v string) (int, error) {
 
 	return m0*10000 + m1*100 + m2, nil
 }
+
+func WriteWaitCh(ch chan<- *drivers.ExitResult, r *drivers.ExitResult) {
+	select {
+	case ch<-r:
+	default:
+	}
+}

driver/driver.go

@@ -655,22 +655,12 @@ func (d *Driver) handleWait(ctx context.Context, handle *taskHandle, ch chan *dr
 		return
 	case <-d.ctx.Done():
 		return
-	case <-handle.ctx.Done():
-		result := &drivers.ExitResult{
-			ExitCode:  0,
-			Signal:    0,
-			OOMKilled: false,
-			Err:       nil,
-		}
-		ch <- result
-	case result := <-handle.waitCh: // Do not refer to handle.runner.waitCh. It might be nil.
-		handle.stateLock.Lock()
-		handle.procState = drivers.TaskStateExited
-		handle.stateLock.Unlock()
-		ch <- result
+	case <-handle.doneCh:
+		ch <- handle.exitResult.Copy()
 	}
 }
 
+// StopTask will not be called if the task has already exited (e.g. onError).
 func (d *Driver) StopTask(taskID string, timeout time.Duration, signal string) error {
 	d.logger.Info("StopTask", "id", taskID, "signal", signal)
 	handle, ok := d.tasks.Get(taskID)
@@ -793,10 +783,15 @@ func (d *Driver) SignalTask(taskID string, signal string) error {
 			return errors.New(string(bs))
 		}
 	case "finish":
+		if h.runner == nil {
+			return fmt.Errorf("h.runner is nil")
+		}
 		return h.runner.Finish1()
 	case "pause":
 		d.logger.Info("pause a task", "taskID", taskID)
-		h := d.tasks.store[taskID]
+		if h.runner == nil {
+			return fmt.Errorf("h.runner is nil")
+		}
 		err := h.runner.Shutdown()
 		if err != nil {
 			d.logger.Error("error when pausing a task", "taskID", taskID, "err", err)
@@ -806,7 +801,6 @@ func (d *Driver) SignalTask(taskID string, signal string) error {
 		return nil
 	case "resume":
 		d.logger.Info("resume a task", "taskID", taskID)
-		h := d.tasks.store[taskID]
 		err := h.resumeTask(d)
 		if err != nil {
 			d.logger.Error("error when resuming a task", "taskID", taskID, "err", err)
@@ -817,7 +811,7 @@ func (d *Driver) SignalTask(taskID string, signal string) error {
 		return nil
 	}
 
-	return nil
+	return fmt.Errorf("DTLE_BUG SignalTask. should not reach here")
 }
 
 func (d *Driver) ExecTask(taskID string, cmd []string, timeout time.Duration) (*drivers.ExecTaskResult, error) {

driver/handle.go

@@ -1,7 +1,6 @@
 package mysql
 
 import (
-	"context"
 	"fmt"
 	"sync"
 	"time"
@@ -33,12 +32,12 @@ type taskHandle struct {
 
 	runner DriverHandle
 
-	ctx        context.Context
-	cancelFunc context.CancelFunc
-	waitCh     chan *drivers.ExitResult
-	stats      *common.TaskStatistics
+	waitCh chan *drivers.ExitResult
+	doneCh chan struct{}
+	stats  *common.TaskStatistics
 
 	driverConfig *common.MySQLDriverConfig
+	shutdown     bool
 }
 
 func newDtleTaskHandle(logger g.LoggerType, cfg *drivers.TaskConfig, state drivers.TaskState, started time.Time) *taskHandle {
@@ -50,12 +49,24 @@ func newDtleTaskHandle(logger g.LoggerType, cfg *drivers.TaskConfig, state drive
 		startedAt:   started,
 		completedAt: time.Time{},
 		exitResult:  nil,
-		waitCh:      make(chan *drivers.ExitResult, 1),
+		waitCh:      make(chan *drivers.ExitResult),
+		doneCh:      make(chan struct{}),
 	}
-	h.ctx, h.cancelFunc = context.WithCancel(context.TODO())
+	go h.watchWaitCh()
 	return h
 }
 
+func (h *taskHandle) watchWaitCh() {
+	select {
+	case r := <-h.waitCh:
+		h.stateLock.Lock()
+		h.exitResult = r
+		h.stateLock.Unlock()
+		close(h.doneCh)
+	case <-h.doneCh:
+	}
+}
+
 func (h *taskHandle) TaskStatus() (*drivers.TaskStatus, error) {
 	h.stateLock.RLock()
 	defer h.stateLock.RUnlock()
@@ -83,13 +94,14 @@ func (h *taskHandle) TaskStatus() (*drivers.TaskStatus, error) {
 	}, nil
 }
 
+// used when h.runner has not been setup
 func (h *taskHandle) onError(err error) {
-	h.waitCh <- &drivers.ExitResult{
+	common.WriteWaitCh(h.waitCh, &drivers.ExitResult{
 		ExitCode:  common.TaskStateDead,
 		Signal:    0,
 		OOMKilled: false,
 		Err:       err,
-	}
+	})
 }
 
 func (h *taskHandle) IsRunning() bool {
@@ -137,7 +149,8 @@ func (h *taskHandle) run(d *Driver) {
 		t := time.NewTimer(0)
 		for {
 			select {
-			case <-h.ctx.Done():
+			case <-h.doneCh:
+				if !t.Stop() { <-t.C }
 				return
 			case <-t.C:
 				if h.runner != nil {
@@ -168,12 +181,12 @@ func (h *taskHandle) NewRunner(d *Driver) (runner DriverHandle, err error) {
 	case common.TaskTypeSrc:
 		if h.driverConfig.OracleConfig != nil {
 			h.logger.Debug("found oracle src", "OracleConfig", h.driverConfig.OracleConfig)
-			runner, err = extractor.NewExtractorOracle(ctx, h.driverConfig, h.logger, d.storeManager, h.waitCh)
+			runner, err = extractor.NewExtractorOracle(ctx, h.driverConfig, h.logger, d.storeManager, h.waitCh, d.ctx)
 			if err != nil {
 				return nil, errors.Wrap(err, "NewExtractor")
 			}
 		} else {
-			runner, err = mysql.NewExtractor(ctx, h.driverConfig, h.logger, d.storeManager, h.waitCh, h.ctx)
+			runner, err = mysql.NewExtractor(ctx, h.driverConfig, h.logger, d.storeManager, h.waitCh, d.ctx)
 			if err != nil {
 				return nil, errors.Wrap(err, "NewOracleExtractor")
 			}
@@ -183,13 +196,13 @@ func (h *taskHandle) NewRunner(d *Driver) (runner DriverHandle, err error) {
 		if h.driverConfig.KafkaConfig != nil {
 			h.logger.Debug("found kafka", "KafkaConfig", h.driverConfig.KafkaConfig)
 			runner, err = kafka.NewKafkaRunner(ctx, h.driverConfig.KafkaConfig, h.logger,
-				d.storeManager, d.config.NatsAdvertise, h.waitCh, h.ctx)
+				d.storeManager, d.config.NatsAdvertise, h.waitCh, d.ctx)
 			if err != nil {
 				return nil, errors.Wrap(err, "NewKafkaRunner")
 			}
 		} else {
 			runner, err = mysql.NewApplier(ctx, h.driverConfig, h.logger, d.storeManager,
-				d.config.NatsAdvertise, h.waitCh, d.eventer, h.taskConfig, h.ctx)
+				d.config.NatsAdvertise, h.waitCh, d.eventer, h.taskConfig, d.ctx)
 			if err != nil {
 				return nil, errors.Wrap(err, "NewApplier")
 			}
@@ -265,23 +278,35 @@ func (h *taskHandle) emitStats(ru *common.TaskStatistics) {
 	}
 }
 
-func (h *taskHandle) Destroy() bool {
-	h.stateLock.RLock()
-	//driver.des
-	h.cancelFunc()
+func (h *taskHandle) Destroy() {
+	if h.shutdown {
+		return
+	}
+	h.stateLock.Lock()
+	h.shutdown = true
+	h.stateLock.Unlock()
+
+	common.WriteWaitCh(h.waitCh, &drivers.ExitResult{
+		ExitCode:  0,
+		Signal:    0,
+		OOMKilled: false,
+		Err:       nil,
+	})
+
 	if h.runner != nil {
 		err := h.runner.Shutdown()
 		if err != nil {
 			h.logger.Error("error in h.runner.Shutdown", "err", err)
 		}
 	}
-	return h.procState == drivers.TaskStateExited
 }
 
 type DriverHandle interface {
 	Run()
 
-	// Shutdown is used to stop the task
+	// Shutdown is used to stop the task.
+	// Do not send ExitResult in Shutdown().
+	// pause API will call Shutdown and the task should not exit.
 	Shutdown() error
 
 	// Stats returns aggregated stats of the driver

driver/kafka/kafka3.go

@@ -46,9 +46,10 @@ type KafkaRunner struct {
 	natsConn    *gonats.Conn
 	waitCh      chan *drivers.ExitResult
 
-	ctx        context.Context
-	shutdown   bool
-	shutdownCh chan struct{}
+	ctx          context.Context
+	shutdown     bool
+	shutdownCh   chan struct{}
+	shutdownLock sync.Mutex
 
 	kafkaConfig *common.KafkaConfig
 	kafkaMgr    *KafkaManager
@@ -157,9 +158,14 @@ func (kr *KafkaRunner) updateGtidLoop() {
 }
 
 func (kr *KafkaRunner) Shutdown() error {
+	kr.logger.Debug("Shutting down")
+	kr.shutdownLock.Lock()
+	defer kr.shutdownLock.Unlock()
+
 	if kr.shutdown {
 		return nil
 	}
+
 	if kr.natsConn != nil {
 		kr.natsConn.Close()
 	}
@@ -634,14 +640,14 @@ func (kr *KafkaRunner) initiateStreaming() error {
 }
 
 func (kr *KafkaRunner) onError(state int, err error) {
+	kr.logger.Info("onError", "err", err, "hasShutdown", kr.shutdown)
+
 	if kr.shutdown {
 		return
 	}
 
 	switch state {
-	case common.TaskStateComplete:
-		kr.logger.Info("Done migrating")
-	case common.TaskStateRestart, common.TaskStateDead:
+	case common.TaskStateDead:
 		msg := &common.ControlMsg{
 			Msg:  err.Error(),
 			Type: common.ControlMsgError,
@@ -660,13 +666,13 @@ func (kr *KafkaRunner) onError(state int, err error) {
 		}
 	}
 
-	kr.waitCh <- &drivers.ExitResult{
+	common.WriteWaitCh(kr.waitCh, &drivers.ExitResult{
 		ExitCode:  state,
 		Signal:    0,
 		OOMKilled: false,
 		Err:       err,
-	}
-	kr.Shutdown()
+	})
+	_ = kr.Shutdown()
 }
 
 func (kr *KafkaRunner) kafkaTransformSnapshotData(