Fix ISCP index consumer cancellation during DROP INDEX

pkg/cnservice/server_query.go

@@ -26,6 +26,7 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/container/types"
 	"github.com/matrixorigin/matrixone/pkg/defines"
 	"github.com/matrixorigin/matrixone/pkg/fileservice"
+	"github.com/matrixorigin/matrixone/pkg/iscp"
 	"github.com/matrixorigin/matrixone/pkg/lockservice"
 	"github.com/matrixorigin/matrixone/pkg/logutil"
 	"github.com/matrixorigin/matrixone/pkg/objectio"
@@ -102,6 +103,7 @@ func (s *service) initQueryCommandHandler() {
 	s.queryService.AddHandleFunc(query.CmdMethod_WorkspaceThreshold, s.handleWorkspaceThresholdRequest, false)
 	s.queryService.AddHandleFunc(query.CmdMethod_MinTimestamp, s.handleGetMinTimestamp, false)
 	s.queryService.AddHandleFunc(query.CmdMethod_CtlPrefetchOnSubscribed, s.handleCtlPrefetchOnSubscribed, false)
+	s.queryService.AddHandleFunc(query.CmdMethod_ISCPDrainConsumer, s.handleISCPDrainConsumer, false)
 }
 
 func (s *service) handleKillConn(ctx context.Context, req *query.Request, resp *query.Response, _ *morpc.Buffer) error {
@@ -192,6 +194,47 @@ func (s *service) handleCtlPrefetchOnSubscribed(ctx context.Context, req *query.
 	return nil
 }
 
+func (s *service) handleISCPDrainConsumer(ctx context.Context, req *query.Request, resp *query.Response, _ *morpc.Buffer) error {
+	if req == nil || req.ISCPDrainConsumerRequest == nil {
+		return moerr.NewInternalError(ctx, "bad request")
+	}
+	r := req.ISCPDrainConsumerRequest
+	exec, ok := iscp.GetExecutorRuntime(s.cfg.UUID)
+	if !ok || exec == nil {
+		return moerr.NewInternalErrorf(
+			ctx,
+			"cannot confirm ISCP consumer quiescence on CN %s for tableID=%d jobName=%s jobID=%d",
+			s.cfg.UUID,
+			r.TableID,
+			r.JobName,
+			r.JobID,
+		)
+	}
+	key := iscp.NewJobRuntimeKey(r.AccountID, r.TableID, r.JobName, r.JobID)
+	if r.RemoveFenceOnly {
+		if _, msg, injected := fault.TriggerFault(objectio.FJ_ISCPCancelRemoveFenceError); injected {
+			if msg == "" {
+				msg = objectio.FJ_ISCPCancelRemoveFenceError
+			}
+			return moerr.NewInternalErrorNoCtxf("injected ISCP remove fence error: %s", msg)
+		}
+		exec.RemoveJobFence(key)
+		resp.ISCPDrainConsumerResponse = &query.ISCPDrainConsumerResponse{Success: true}
+		return nil
+	}
+	if r.RenewFenceOnly {
+		exec.RenewJobFence(key, iscp.RollbackFenceTTL())
+		resp.ISCPDrainConsumerResponse = &query.ISCPDrainConsumerResponse{Success: true}
+		return nil
+	}
+	if err := exec.CancelAndDrainJobConsumer(ctx, r.AccountID, r.TableID, r.JobName, r.JobID); err != nil {
+		exec.RemoveJobFence(key)
+		return err
+	}
+	resp.ISCPDrainConsumerResponse = &query.ISCPDrainConsumerResponse{Success: true}
+	return nil
+}
+
 // handleGetLockInfo sends the lock info on current cn to another cn that needs.
 func (s *service) handleGetLockInfo(ctx context.Context, req *query.Request, resp *query.Response, _ *morpc.Buffer) error {
 	resp.GetLockInfoResponse = new(query.GetLockInfoResponse)

pkg/cnservice/server_query_test.go

@@ -21,6 +21,7 @@ import (
 	goruntime "runtime"
 	"runtime/debug"
 	"testing"
+	"time"
 	"unsafe"
 
 	"github.com/golang/mock/gomock"
@@ -40,7 +41,9 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/frontend/test/mock_shard"
 	"github.com/matrixorigin/matrixone/pkg/frontend/test/mock_task"
 	"github.com/matrixorigin/matrixone/pkg/incrservice"
+	"github.com/matrixorigin/matrixone/pkg/iscp"
 	"github.com/matrixorigin/matrixone/pkg/lockservice"
+	"github.com/matrixorigin/matrixone/pkg/objectio"
 	"github.com/matrixorigin/matrixone/pkg/pb/lock"
 	"github.com/matrixorigin/matrixone/pkg/pb/query"
 	"github.com/matrixorigin/matrixone/pkg/pb/statsinfo"
@@ -50,13 +53,59 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/taskservice"
 	"github.com/matrixorigin/matrixone/pkg/testutil"
 	"github.com/matrixorigin/matrixone/pkg/txn/client"
+	"github.com/matrixorigin/matrixone/pkg/util/fault"
 	"github.com/matrixorigin/matrixone/pkg/util/trace"
 	"github.com/matrixorigin/matrixone/pkg/vm/engine"
 )
 
 var dummyBadRequestErr = moerr.NewInternalError(context.TODO(), "bad request")
 var dummyErr = moerr.NewInternalError(context.TODO(), "dummy error")
 
+func Test_service_handleISCPDrainConsumerRenewFenceOnly(t *testing.T) {
+	require.True(t, fault.Enable())
+	defer fault.Disable()
+	require.NoError(t, fault.AddFaultPoint(context.Background(), objectio.FJ_ISCPCancelRollbackFenceTTL, ":::", "echo", 1, "", false))
+	defer func() {
+		_, _ = fault.RemoveFaultPoint(context.Background(), objectio.FJ_ISCPCancelRollbackFenceTTL)
+	}()
+
+	exec := &iscp.ISCPTaskExecutor{}
+	iscp.RegisterExecutorRuntime("runner-cn", exec)
+	defer iscp.UnregisterExecutorRuntime("runner-cn", exec)
+
+	s := &service{cfg: &Config{UUID: "runner-cn"}}
+	key := iscp.NewJobRuntimeKey(1, 42, "index_idx1", 7)
+
+	resp := &query.Response{}
+	require.NoError(t, s.handleISCPDrainConsumer(context.Background(), &query.Request{
+		ISCPDrainConsumerRequest: &query.ISCPDrainConsumerRequest{
+			AccountID:      key.AccountID,
+			TableID:        key.TableID,
+			JobName:        key.JobName,
+			JobID:          key.JobID,
+			RenewFenceOnly: true,
+		},
+	}, resp, nil))
+	require.NotNil(t, resp.ISCPDrainConsumerResponse)
+	require.False(t, exec.IsJobFenced(key))
+
+	require.NoError(t, exec.CancelAndDrainJobConsumer(context.Background(), key.AccountID, key.TableID, key.JobName, key.JobID))
+	time.Sleep(700 * time.Millisecond)
+
+	resp = &query.Response{}
+	require.NoError(t, s.handleISCPDrainConsumer(context.Background(), &query.Request{
+		ISCPDrainConsumerRequest: &query.ISCPDrainConsumerRequest{
+			AccountID:      key.AccountID,
+			TableID:        key.TableID,
+			JobName:        key.JobName,
+			JobID:          key.JobID,
+			RenewFenceOnly: true,
+		},
+	}, resp, nil))
+	time.Sleep(700 * time.Millisecond)
+	require.True(t, exec.IsJobFenced(key))
+}
+
 func Test_service_handleGoMaxProcs(t *testing.T) {
 	ctx := context.Background()
 	type fields struct{}

pkg/iscp/data_retriever.go

@@ -81,6 +81,20 @@ func (d *ISCPData) Done() {
 	}
 }
 
+func (d *ISCPData) Close() {
+	if d == nil {
+		return
+	}
+	if d.insertBatch != nil {
+		d.insertBatch.Close()
+		d.insertBatch = nil
+	}
+	if d.deleteBatch != nil {
+		d.deleteBatch.Close()
+		d.deleteBatch = nil
+	}
+}
+
 const (
 	ISCPDataType_Snapshot int8 = iota
 	ISCPDataType_Tail
@@ -145,6 +159,9 @@ func (r *DataRetrieverImpl) UpdateWatermark(ctx context.Context,
 	cnUUID string,
 	txn client.TxnOperator) error {
 
+	if r.IsCanceled() {
+		return r.terminalError()
+	}
 	ctxWithSysAccount := context.WithValue(ctx, defines.TenantIDKey{}, catalog.System_Account)
 	ctxWithSysAccount, cancel := context.WithTimeout(ctxWithSysAccount, time.Minute*5)
 	defer cancel()
@@ -187,17 +204,24 @@ func (r *DataRetrieverImpl) GetTableID() uint64 {
 	return r.tableID
 }
 
-func (r *DataRetrieverImpl) SetNextBatch(data *ISCPData) {
+func (r *DataRetrieverImpl) SetNextBatch(data *ISCPData) bool {
 	if r.hasError() {
-		data.Done()
-		return
+		return false
 	}
 	select {
 	case r.insertDataCh <- data:
-		return
+		if r.IsCanceled() {
+			select {
+			case queued := <-r.insertDataCh:
+				if queued == data {
+					return false
+				}
+			default:
+			}
+		}
+		return true
 	case <-r.ctx.Done():
-		data.Done()
-		return
+		return false
 	}
 }
 
@@ -218,6 +242,10 @@ func (r *DataRetrieverImpl) terminalError() error {
 
 // after error occurs, the data retriever won't consume any more data
 func (r *DataRetrieverImpl) SetError(err error) {
+	r.Cancel(err)
+}
+
+func (r *DataRetrieverImpl) Cancel(err error) {
 	if r.hasError() {
 		return
 	}
@@ -228,6 +256,26 @@ func (r *DataRetrieverImpl) SetError(err error) {
 	}
 	r.cancel()
 	r.err = err
+	for {
+		select {
+		case data := <-r.insertDataCh:
+			data.Done()
+		default:
+			return
+		}
+	}
+}
+
+func (r *DataRetrieverImpl) IsCanceled() bool {
+	if r.hasError() {
+		return true
+	}
+	select {
+	case <-r.ctx.Done():
+		return true
+	default:
+		return false
+	}
 }
 
 func (r *DataRetrieverImpl) Close() {

pkg/iscp/executor.go

@@ -187,6 +187,10 @@ func NewISCPTaskExecutor(
 		tableMu:     sync.RWMutex{},
 		option:      option,
 		mp:          mp,
+		fencedJobs:  make(map[JobRuntimeKey]JobFence),
+		runningConsumers: make(
+			map[JobRuntimeKey]map[uint64]*RunningJobConsumer,
+		),
 	}
 	return exec, nil
 }
@@ -341,9 +345,10 @@ func (exec *ISCPTaskExecutor) initStateLocked(parent context.Context) (err error
 
 	// Publish a fully initialized generation only after replay and state repair
 	// have both succeeded. NewWorker seals all worker goroutines before return.
-	exec.worker = NewWorker(exec.cnUUID, exec.txnEngine, exec.cnTxnClient, exec.mp)
+	exec.worker = NewWorker(exec, exec.cnUUID, exec.txnEngine, exec.cnTxnClient, exec.mp)
 	exec.wg.Add(1)
 	exec.running = true
+	RegisterExecutorRuntime(exec.cnUUID, exec)
 	logutil.Info(
 		"ISCP-Task Start",
 	)
@@ -367,6 +372,7 @@ func (exec *ISCPTaskExecutor) stopLocked() {
 	exec.cancel()
 	exec.worker.Stop()
 	exec.wg.Wait()
+	UnregisterExecutorRuntime(exec.cnUUID, exec)
 	exec.ctx, exec.cancel = nil, nil
 	exec.worker = nil
 }
@@ -530,6 +536,12 @@ func (exec *ISCPTaskExecutor) run(ctx context.Context, worker Worker) {
 						go exec.Cancel()
 						return
 					}
+					if objectio.WaitInjected(objectio.FJ_ISCPCancelAfterSubmit) {
+						logutil.Infof("ISCP-Task cancel fault wait %s", objectio.FJ_ISCPCancelAfterSubmit)
+					}
+					if msg, injected := objectio.ISCPExecutorInjected(); injected && msg == "iscp:after-submit" {
+						logutil.Infof("ISCP-Task injected hook %s", msg)
+					}
 				} else {
 					table.UpdateWatermark(iter)
 				}
@@ -960,6 +972,7 @@ func (exec *ISCPTaskExecutor) addOrUpdateJob(
 	notPrint bool,
 ) error {
 	var newCreate bool
+	fenceKey := NewJobRuntimeKey(accountID, tableID, jobName, jobID)
 
 	var watermark types.TS
 	defer func() {
@@ -990,6 +1003,7 @@ func (exec *ISCPTaskExecutor) addOrUpdateJob(
 	table, ok := exec.getTable(accountID, tableID)
 	if !ok {
 		if dropAt != 0 {
+			exec.RemoveJobFence(fenceKey)
 			return nil
 		}
 		table = NewTableEntry(
@@ -1003,6 +1017,9 @@ func (exec *ISCPTaskExecutor) addOrUpdateJob(
 		exec.setTable(table)
 	}
 	newCreate = table.AddOrUpdateSinker(exec.ctx, jobName, jobSpec, jobStatus, jobID, watermark, state, dropAt)
+	if dropAt != 0 {
+		exec.RemoveJobFence(fenceKey)
+	}
 	return nil
 }
 
@@ -1018,6 +1035,7 @@ func (exec *ISCPTaskExecutor) GCInMemoryJob(threshold time.Duration) {
 	tids := make([]uint64, 0, len(tablesToDelete))
 	for _, table := range tablesToDelete {
 		exec.deleteTableEntry(table)
+		exec.RemoveTableJobFences(table.accountID, table.tableID)
 		tids = append(tids, table.tableID)
 	}
 	logutil.Infof("ISCP-Task delete table %v", tids)