github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/cdc/processor/sinkmanager/table_sink_worker_test.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package sinkmanager
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  	"testing"
    20  	"time"
    21  
    22  	"github.com/pingcap/tiflow/cdc/entry"
    23  	"github.com/pingcap/tiflow/cdc/model"
    24  	"github.com/pingcap/tiflow/cdc/processor/memquota"
    25  	"github.com/pingcap/tiflow/cdc/processor/sourcemanager"
    26  	"github.com/pingcap/tiflow/cdc/processor/sourcemanager/sorter"
    27  	"github.com/pingcap/tiflow/cdc/processor/sourcemanager/sorter/memory"
    28  	"github.com/pingcap/tiflow/cdc/processor/tablepb"
    29  	"github.com/pingcap/tiflow/pkg/spanz"
    30  	"github.com/pingcap/tiflow/pkg/upstream"
    31  	"github.com/stretchr/testify/require"
    32  	"github.com/stretchr/testify/suite"
    33  )
    34  
    35  // testEventSize is the size of a test event.
    36  // It is used to calculate the memory quota.
    37  var (
    38  	emptyEvent    = model.RowChangedEvent{}
    39  	testEventSize = emptyEvent.ApproximateBytes()
    40  )
    41  
    42  //nolint:unparam
    43  func genPolymorphicEventWithNilRow(startTs,
    44  	commitTs uint64,
    45  ) *model.PolymorphicEvent {
    46  	return &model.PolymorphicEvent{
    47  		StartTs: startTs,
    48  		CRTs:    commitTs,
    49  		RawKV: &model.RawKVEntry{
    50  			OpType:  model.OpTypePut,
    51  			StartTs: startTs,
    52  			CRTs:    commitTs,
    53  		},
    54  		Row: nil,
    55  	}
    56  }
    57  
    58  func genPolymorphicResolvedEvent(resolvedTs uint64) *model.PolymorphicEvent {
    59  	return &model.PolymorphicEvent{
    60  		CRTs: resolvedTs,
    61  		RawKV: &model.RawKVEntry{
    62  			OpType: model.OpTypeResolved,
    63  			CRTs:   resolvedTs,
    64  		},
    65  	}
    66  }
    67  
    68  func genPolymorphicEvent(startTs, commitTs uint64, span tablepb.Span) *model.PolymorphicEvent {
    69  	return &model.PolymorphicEvent{
    70  		StartTs: startTs,
    71  		CRTs:    commitTs,
    72  		RawKV: &model.RawKVEntry{
    73  			OpType:  model.OpTypePut,
    74  			StartTs: startTs,
    75  			CRTs:    commitTs,
    76  		},
    77  		Row: genRowChangedEvent(startTs, commitTs, span),
    78  	}
    79  }
    80  
    81  func genRowChangedEvent(startTs, commitTs uint64, span tablepb.Span) *model.RowChangedEvent {
    82  	columns := []*model.Column{
    83  		{Name: "a", Value: 2},
    84  	}
    85  	preColumns := []*model.Column{
    86  		{Name: "a", Value: 1},
    87  	}
    88  	tableInfo := model.BuildTableInfo("table", "table", columns, nil)
    89  	return &model.RowChangedEvent{
    90  		StartTs:         startTs,
    91  		CommitTs:        commitTs,
    92  		PhysicalTableID: span.TableID,
    93  		TableInfo:       tableInfo,
    94  		Columns:         model.Columns2ColumnDatas(columns, tableInfo),
    95  		PreColumns:      model.Columns2ColumnDatas(preColumns, tableInfo),
    96  	}
    97  }
    98  
    99  type tableSinkWorkerSuite struct {
   100  	suite.Suite
   101  	testChangefeedID model.ChangeFeedID
   102  	testSpan         tablepb.Span
   103  }
   104  
   105  func TestTableSinkWorkerSuite(t *testing.T) {
   106  	suite.Run(t, new(tableSinkWorkerSuite))
   107  }
   108  
   109  func (suite *tableSinkWorkerSuite) SetupSuite() {
   110  	requestMemSize = uint64(testEventSize)
   111  	// For one batch size.
   112  	// Advance table sink per 2 events.
   113  	maxUpdateIntervalSize = uint64(testEventSize * 2)
   114  	suite.testChangefeedID = model.DefaultChangeFeedID("1")
   115  	suite.testSpan = spanz.TableIDToComparableSpan(1)
   116  }
   117  
   118  func (suite *tableSinkWorkerSuite) SetupTest() {
   119  	// reset batchID
   120  	batchID.Store(0)
   121  }
   122  
   123  func (suite *tableSinkWorkerSuite) TearDownSuite() {
   124  	requestMemSize = defaultRequestMemSize
   125  	maxUpdateIntervalSize = defaultMaxUpdateIntervalSize
   126  }
   127  
   128  func (suite *tableSinkWorkerSuite) createWorker(
   129  	ctx context.Context, memQuota uint64, splitTxn bool,
   130  ) (*sinkWorker, sorter.SortEngine) {
   131  	sortEngine := memory.New(context.Background())
   132  	// Only sourcemanager.FetcyByTable is used, so NewForTest is fine.
   133  	sm := sourcemanager.NewForTest(suite.testChangefeedID, upstream.NewUpstream4Test(&MockPD{}),
   134  		&entry.MockMountGroup{}, sortEngine, false)
   135  	go func() { sm.Run(ctx) }()
   136  
   137  	// To avoid refund or release panics.
   138  	quota := memquota.NewMemQuota(suite.testChangefeedID, memQuota, "sink")
   139  	// NOTICE: Do not forget the initial memory quota in the worker first time running.
   140  	quota.ForceAcquire(uint64(testEventSize))
   141  	quota.AddTable(suite.testSpan)
   142  
   143  	return newSinkWorker(suite.testChangefeedID, sm, quota, splitTxn), sortEngine
   144  }
   145  
   146  func (suite *tableSinkWorkerSuite) addEventsToSortEngine(
   147  	events []*model.PolymorphicEvent,
   148  	sortEngine sorter.SortEngine,
   149  ) {
   150  	sortEngine.AddTable(suite.testSpan, 0)
   151  	for _, event := range events {
   152  		sortEngine.Add(suite.testSpan, event)
   153  	}
   154  }
   155  
   156  func genUpperBoundGetter(commitTs model.Ts) func(_ model.Ts) sorter.Position {
   157  	return func(_ model.Ts) sorter.Position {
   158  		return sorter.Position{
   159  			StartTs:  commitTs - 1,
   160  			CommitTs: commitTs,
   161  		}
   162  	}
   163  }
   164  
   165  func genLowerBound() sorter.Position {
   166  	return sorter.Position{
   167  		StartTs:  0,
   168  		CommitTs: 1,
   169  	}
   170  }
   171  
   172  // Test Scenario:
   173  // Worker should ignore the filtered events(row is nil).
   174  func (suite *tableSinkWorkerSuite) TestHandleTaskWithSplitTxnAndGotSomeFilteredEvents() {
   175  	ctx, cancel := context.WithCancel(context.Background())
   176  	events := []*model.PolymorphicEvent{
   177  		genPolymorphicEvent(1, 2, suite.testSpan),
   178  		// This event will be filtered, so its Row will be nil.
   179  		genPolymorphicEventWithNilRow(1, 2),
   180  		genPolymorphicEventWithNilRow(1, 2),
   181  		genPolymorphicEvent(1, 3, suite.testSpan),
   182  		genPolymorphicEvent(1, 4, suite.testSpan),
   183  		genPolymorphicResolvedEvent(4),
   184  	}
   185  
   186  	// Only for three events.
   187  	eventSize := uint64(testEventSize * 3)
   188  	w, e := suite.createWorker(ctx, eventSize, true)
   189  	defer w.sinkMemQuota.Close()
   190  	suite.addEventsToSortEngine(events, e)
   191  
   192  	taskChan := make(chan *sinkTask)
   193  	var wg sync.WaitGroup
   194  	wg.Add(1)
   195  	go func() {
   196  		defer wg.Done()
   197  		err := w.handleTasks(ctx, taskChan)
   198  		require.Equal(suite.T(), context.Canceled, err)
   199  	}()
   200  
   201  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   202  	callback := func(lastWritePos sorter.Position) {
   203  		require.Equal(suite.T(), sorter.Position{
   204  			StartTs:  1,
   205  			CommitTs: 4,
   206  		}, lastWritePos)
   207  		require.Equal(suite.T(), sorter.Position{
   208  			StartTs:  2,
   209  			CommitTs: 4,
   210  		}, lastWritePos.Next())
   211  		cancel()
   212  	}
   213  	taskChan <- &sinkTask{
   214  		span:          suite.testSpan,
   215  		lowerBound:    genLowerBound(),
   216  		getUpperBound: genUpperBoundGetter(4),
   217  		tableSink:     wrapper,
   218  		callback:      callback,
   219  		isCanceled:    func() bool { return false },
   220  	}
   221  	wg.Wait()
   222  	require.Len(suite.T(), sink.GetEvents(), 3)
   223  }
   224  
   225  // Test Scenario:
   226  // worker will stop when no memory quota and meet the txn boundary.
   227  func (suite *tableSinkWorkerSuite) TestHandleTaskWithSplitTxnAndAbortWhenNoMemAndOneTxnFinished() {
   228  	ctx, cancel := context.WithCancel(context.Background())
   229  	events := []*model.PolymorphicEvent{
   230  		genPolymorphicEvent(1, 2, suite.testSpan),
   231  		genPolymorphicEvent(1, 2, suite.testSpan),
   232  		genPolymorphicEvent(1, 3, suite.testSpan),
   233  		genPolymorphicEvent(2, 4, suite.testSpan),
   234  		genPolymorphicResolvedEvent(4),
   235  	}
   236  
   237  	// Only for three events.
   238  	eventSize := uint64(testEventSize * 3)
   239  	w, e := suite.createWorker(ctx, eventSize, true)
   240  	defer w.sinkMemQuota.Close()
   241  	suite.addEventsToSortEngine(events, e)
   242  
   243  	taskChan := make(chan *sinkTask)
   244  	var wg sync.WaitGroup
   245  	wg.Add(1)
   246  	go func() {
   247  		defer wg.Done()
   248  		err := w.handleTasks(ctx, taskChan)
   249  		require.Equal(suite.T(), context.Canceled, err)
   250  	}()
   251  
   252  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   253  	callback := func(lastWritePos sorter.Position) {
   254  		require.Equal(suite.T(), sorter.Position{
   255  			StartTs:  1,
   256  			CommitTs: 3,
   257  		}, lastWritePos, "we only write 3 events because of the memory quota")
   258  		require.Equal(suite.T(), sorter.Position{
   259  			StartTs:  2,
   260  			CommitTs: 3,
   261  		}, lastWritePos.Next())
   262  		cancel()
   263  	}
   264  	taskChan <- &sinkTask{
   265  		span:          suite.testSpan,
   266  		lowerBound:    genLowerBound(),
   267  		getUpperBound: genUpperBoundGetter(4),
   268  		tableSink:     wrapper,
   269  		callback:      callback,
   270  		isCanceled:    func() bool { return false },
   271  	}
   272  	wg.Wait()
   273  	require.Len(suite.T(), sink.GetEvents(), 3)
   274  }
   275  
   276  // Test Scenario:
   277  // worker will block when no memory quota until the mem quota is aborted.
   278  func (suite *tableSinkWorkerSuite) TestHandleTaskWithSplitTxnAndAbortWhenNoMemAndBlocked() {
   279  	ctx, cancel := context.WithCancel(context.Background())
   280  	events := []*model.PolymorphicEvent{
   281  		genPolymorphicEvent(1, 10, suite.testSpan),
   282  		genPolymorphicEvent(1, 10, suite.testSpan),
   283  		genPolymorphicEvent(1, 10, suite.testSpan),
   284  		genPolymorphicEvent(1, 10, suite.testSpan),
   285  		genPolymorphicResolvedEvent(14),
   286  	}
   287  	// Only for three events.
   288  	eventSize := uint64(testEventSize * 3)
   289  	w, e := suite.createWorker(ctx, eventSize, true)
   290  	defer w.sinkMemQuota.Close()
   291  	suite.addEventsToSortEngine(events, e)
   292  
   293  	taskChan := make(chan *sinkTask)
   294  	var wg sync.WaitGroup
   295  	wg.Add(1)
   296  	go func() {
   297  		defer wg.Done()
   298  		err := w.handleTasks(ctx, taskChan)
   299  		require.ErrorIs(suite.T(), err, context.Canceled)
   300  	}()
   301  
   302  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   303  	callback := func(lastWritePos sorter.Position) {
   304  		require.Equal(suite.T(), sorter.Position{
   305  			StartTs:  0,
   306  			CommitTs: 0,
   307  		}, lastWritePos)
   308  	}
   309  	taskChan <- &sinkTask{
   310  		span:          suite.testSpan,
   311  		lowerBound:    genLowerBound(),
   312  		getUpperBound: genUpperBoundGetter(14),
   313  		tableSink:     wrapper,
   314  		callback:      callback,
   315  		isCanceled:    func() bool { return false },
   316  	}
   317  	require.Eventually(suite.T(), func() bool {
   318  		return len(sink.GetEvents()) == 2
   319  	}, 5*time.Second, 10*time.Millisecond)
   320  	// Abort the task when no memory quota and blocked.
   321  	w.sinkMemQuota.Close()
   322  	cancel()
   323  	wg.Wait()
   324  	require.Len(suite.T(), sink.GetEvents(), 2, "Only two events should be sent to sink")
   325  }
   326  
   327  // Test Scenario:
   328  // worker will advance the table sink only when it reaches the batch size.
   329  func (suite *tableSinkWorkerSuite) TestHandleTaskWithSplitTxnAndOnlyAdvanceWhenReachOneBatchSize() {
   330  	ctx, cancel := context.WithCancel(context.Background())
   331  	events := []*model.PolymorphicEvent{
   332  		genPolymorphicEvent(1, 2, suite.testSpan),
   333  		genPolymorphicEvent(1, 2, suite.testSpan),
   334  		genPolymorphicEvent(1, 2, suite.testSpan),
   335  		genPolymorphicEvent(1, 2, suite.testSpan),
   336  		genPolymorphicEvent(1, 2, suite.testSpan),
   337  		genPolymorphicEvent(1, 3, suite.testSpan),
   338  		genPolymorphicResolvedEvent(4),
   339  	}
   340  	// For five events.
   341  	eventSize := uint64(testEventSize * 5)
   342  	w, e := suite.createWorker(ctx, eventSize, true)
   343  	defer w.sinkMemQuota.Close()
   344  	suite.addEventsToSortEngine(events, e)
   345  
   346  	taskChan := make(chan *sinkTask)
   347  	var wg sync.WaitGroup
   348  	wg.Add(1)
   349  	go func() {
   350  		defer wg.Done()
   351  		err := w.handleTasks(ctx, taskChan)
   352  		require.ErrorIs(suite.T(), err, context.Canceled)
   353  	}()
   354  
   355  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   356  	callback := func(lastWritePos sorter.Position) {
   357  		require.Equal(suite.T(), sorter.Position{
   358  			StartTs:  1,
   359  			CommitTs: 2,
   360  		}, lastWritePos)
   361  		require.Equal(suite.T(), sorter.Position{
   362  			StartTs:  2,
   363  			CommitTs: 2,
   364  		}, lastWritePos.Next())
   365  		cancel()
   366  	}
   367  	taskChan <- &sinkTask{
   368  		span:          suite.testSpan,
   369  		lowerBound:    genLowerBound(),
   370  		getUpperBound: genUpperBoundGetter(2),
   371  		tableSink:     wrapper,
   372  		callback:      callback,
   373  		isCanceled:    func() bool { return false },
   374  	}
   375  	wg.Wait()
   376  	require.Len(suite.T(), sink.GetEvents(), 5, "All events should be sent to sink")
   377  	require.Equal(suite.T(), 3, sink.GetWriteTimes(), "Three txn batch should be sent to sink")
   378  }
   379  
   380  // Test Scenario:
   381  // worker will force consume only one Txn when the memory quota is not enough.
   382  func (suite *tableSinkWorkerSuite) TestHandleTaskWithoutSplitTxnAndAbortWhenNoMemAndForceConsume() {
   383  	ctx, cancel := context.WithCancel(context.Background())
   384  	events := []*model.PolymorphicEvent{
   385  		genPolymorphicEvent(1, 2, suite.testSpan),
   386  		genPolymorphicEvent(1, 2, suite.testSpan),
   387  		genPolymorphicEvent(1, 2, suite.testSpan),
   388  		genPolymorphicEvent(1, 2, suite.testSpan),
   389  		genPolymorphicEvent(1, 2, suite.testSpan),
   390  		genPolymorphicEvent(1, 4, suite.testSpan),
   391  		genPolymorphicResolvedEvent(5),
   392  	}
   393  
   394  	// Only for three events.
   395  	eventSize := uint64(testEventSize * 3)
   396  	// Disable split txn.
   397  	w, e := suite.createWorker(ctx, eventSize, false)
   398  	defer w.sinkMemQuota.Close()
   399  	suite.addEventsToSortEngine(events, e)
   400  
   401  	taskChan := make(chan *sinkTask)
   402  	var wg sync.WaitGroup
   403  	wg.Add(1)
   404  	go func() {
   405  		defer wg.Done()
   406  		err := w.handleTasks(ctx, taskChan)
   407  		require.Equal(suite.T(), context.Canceled, err)
   408  	}()
   409  
   410  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   411  	callback := func(lastWritePos sorter.Position) {
   412  		require.Equal(suite.T(), sorter.Position{
   413  			StartTs:  1,
   414  			CommitTs: 2,
   415  		}, lastWritePos)
   416  		require.Equal(suite.T(), sorter.Position{
   417  			StartTs:  2,
   418  			CommitTs: 2,
   419  		}, lastWritePos.Next())
   420  		cancel()
   421  	}
   422  	taskChan <- &sinkTask{
   423  		span:          suite.testSpan,
   424  		lowerBound:    genLowerBound(),
   425  		getUpperBound: genUpperBoundGetter(4),
   426  		tableSink:     wrapper,
   427  		callback:      callback,
   428  		isCanceled:    func() bool { return false },
   429  	}
   430  	wg.Wait()
   431  	require.Len(suite.T(), sink.GetEvents(), 5,
   432  		"All events from the first txn should be sent to sink")
   433  }
   434  
   435  // Test Scenario:
   436  // worker will advance the table sink only when it reaches the max update interval size.
   437  func (suite *tableSinkWorkerSuite) TestTaskWithoutSplitTxnOnlyAdvanceWhenReachMaxUpdateIntSize() {
   438  	ctx, cancel := context.WithCancel(context.Background())
   439  	events := []*model.PolymorphicEvent{
   440  		genPolymorphicEvent(1, 2, suite.testSpan),
   441  		genPolymorphicEvent(1, 3, suite.testSpan),
   442  		genPolymorphicEvent(1, 4, suite.testSpan),
   443  		genPolymorphicEvent(1, 4, suite.testSpan),
   444  		genPolymorphicEvent(1, 5, suite.testSpan),
   445  		genPolymorphicEvent(1, 6, suite.testSpan),
   446  		genPolymorphicResolvedEvent(6),
   447  	}
   448  	// Only for three events.
   449  	eventSize := uint64(testEventSize * 3)
   450  	w, e := suite.createWorker(ctx, eventSize, false)
   451  	defer w.sinkMemQuota.Close()
   452  	suite.addEventsToSortEngine(events, e)
   453  
   454  	taskChan := make(chan *sinkTask)
   455  	var wg sync.WaitGroup
   456  	wg.Add(1)
   457  	go func() {
   458  		defer wg.Done()
   459  		err := w.handleTasks(ctx, taskChan)
   460  		require.Equal(suite.T(), context.Canceled, err)
   461  	}()
   462  
   463  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   464  	callback := func(lastWritePos sorter.Position) {
   465  		require.Equal(suite.T(), sorter.Position{
   466  			StartTs:  1,
   467  			CommitTs: 4,
   468  		}, lastWritePos)
   469  		require.Equal(suite.T(), sorter.Position{
   470  			StartTs:  2,
   471  			CommitTs: 4,
   472  		}, lastWritePos.Next())
   473  		cancel()
   474  	}
   475  	taskChan <- &sinkTask{
   476  		span:          suite.testSpan,
   477  		lowerBound:    genLowerBound(),
   478  		getUpperBound: genUpperBoundGetter(6),
   479  		tableSink:     wrapper,
   480  		callback:      callback,
   481  		isCanceled:    func() bool { return false },
   482  	}
   483  	wg.Wait()
   484  	require.Len(suite.T(), sink.GetEvents(), 4, "All events should be sent to sink")
   485  	require.Equal(suite.T(), 2, sink.GetWriteTimes(), "Only two times write to sink, "+
   486  		"because the max update interval size is 2 * event size")
   487  }
   488  
   489  // Test Scenario:
   490  // worker will advance the table sink only when task is finished.
   491  func (suite *tableSinkWorkerSuite) TestHandleTaskWithSplitTxnAndAdvanceTableWhenTaskIsFinished() {
   492  	ctx, cancel := context.WithCancel(context.Background())
   493  	events := []*model.PolymorphicEvent{
   494  		genPolymorphicEvent(0, 1, suite.testSpan),
   495  		genPolymorphicResolvedEvent(4),
   496  	}
   497  	// Only for three events.
   498  	eventSize := uint64(testEventSize * 3)
   499  	w, e := suite.createWorker(ctx, eventSize, true)
   500  	defer w.sinkMemQuota.Close()
   501  	suite.addEventsToSortEngine(events, e)
   502  
   503  	taskChan := make(chan *sinkTask)
   504  	var wg sync.WaitGroup
   505  	wg.Add(1)
   506  	go func() {
   507  		defer wg.Done()
   508  		err := w.handleTasks(ctx, taskChan)
   509  		require.ErrorIs(suite.T(), err, context.Canceled)
   510  	}()
   511  
   512  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   513  	callback := func(lastWritePos sorter.Position) {
   514  		require.Equal(suite.T(), sorter.Position{
   515  			StartTs:  3,
   516  			CommitTs: 4,
   517  		}, lastWritePos)
   518  		require.Equal(suite.T(), sorter.Position{
   519  			StartTs:  4,
   520  			CommitTs: 4,
   521  		}, lastWritePos.Next())
   522  	}
   523  	taskChan <- &sinkTask{
   524  		span:          suite.testSpan,
   525  		lowerBound:    genLowerBound(),
   526  		getUpperBound: genUpperBoundGetter(4),
   527  		tableSink:     wrapper,
   528  		callback:      callback,
   529  		isCanceled:    func() bool { return false },
   530  	}
   531  	require.Eventually(suite.T(), func() bool {
   532  		return len(sink.GetEvents()) == 1
   533  	}, 5*time.Second, 10*time.Millisecond)
   534  	cancel()
   535  	wg.Wait()
   536  	receivedEvents := sink.GetEvents()
   537  	receivedEvents[0].Callback()
   538  	require.Len(suite.T(), sink.GetEvents(), 1, "No more events should be sent to sink")
   539  	checkpointTs := wrapper.getCheckpointTs()
   540  	require.Equal(suite.T(), uint64(4), checkpointTs.ResolvedMark())
   541  }
   542  
   543  // Test Scenario:
   544  // worker will advance the table sink directly when there are no events.
   545  func (suite *tableSinkWorkerSuite) TestHandleTaskWithSplitTxnAndAdvanceTableIfNoWorkload() {
   546  	ctx, cancel := context.WithCancel(context.Background())
   547  	events := []*model.PolymorphicEvent{
   548  		genPolymorphicResolvedEvent(4),
   549  	}
   550  	// Only for three events.
   551  	eventSize := uint64(testEventSize * 3)
   552  	w, e := suite.createWorker(ctx, eventSize, true)
   553  	defer w.sinkMemQuota.Close()
   554  	suite.addEventsToSortEngine(events, e)
   555  
   556  	taskChan := make(chan *sinkTask)
   557  	var wg sync.WaitGroup
   558  	wg.Add(1)
   559  	go func() {
   560  		defer wg.Done()
   561  		err := w.handleTasks(ctx, taskChan)
   562  		require.ErrorIs(suite.T(), err, context.Canceled)
   563  	}()
   564  
   565  	wrapper, _ := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   566  	callback := func(lastWritePos sorter.Position) {
   567  		require.Equal(suite.T(), sorter.Position{
   568  			StartTs:  3,
   569  			CommitTs: 4,
   570  		}, lastWritePos)
   571  		require.Equal(suite.T(), sorter.Position{
   572  			StartTs:  4,
   573  			CommitTs: 4,
   574  		}, lastWritePos.Next())
   575  	}
   576  	taskChan <- &sinkTask{
   577  		span:          suite.testSpan,
   578  		lowerBound:    genLowerBound(),
   579  		getUpperBound: genUpperBoundGetter(4),
   580  		tableSink:     wrapper,
   581  		callback:      callback,
   582  		isCanceled:    func() bool { return false },
   583  	}
   584  	require.Eventually(suite.T(), func() bool {
   585  		checkpointTs := wrapper.getCheckpointTs()
   586  		return checkpointTs.ResolvedMark() == 4
   587  	}, 5*time.Second, 10*time.Millisecond, "Directly advance resolved mark to 4")
   588  	cancel()
   589  	wg.Wait()
   590  }
   591  
   592  func (suite *tableSinkWorkerSuite) TestHandleTaskUseDifferentBatchIDEveryTime() {
   593  	ctx, cancel := context.WithCancel(context.Background())
   594  	events := []*model.PolymorphicEvent{
   595  		genPolymorphicEvent(1, 3, suite.testSpan),
   596  		genPolymorphicEvent(1, 3, suite.testSpan),
   597  		genPolymorphicEvent(1, 3, suite.testSpan),
   598  		genPolymorphicResolvedEvent(4),
   599  	}
   600  	// Only for three events.
   601  	eventSize := uint64(testEventSize * 3)
   602  	w, e := suite.createWorker(ctx, eventSize, true)
   603  	defer w.sinkMemQuota.Close()
   604  	suite.addEventsToSortEngine(events, e)
   605  
   606  	taskChan := make(chan *sinkTask)
   607  	var wg sync.WaitGroup
   608  	wg.Add(1)
   609  	go func() {
   610  		defer wg.Done()
   611  		err := w.handleTasks(ctx, taskChan)
   612  		require.Equal(suite.T(), context.Canceled, err)
   613  	}()
   614  
   615  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   616  	callback := func(lastWritePos sorter.Position) {
   617  		require.Equal(suite.T(), sorter.Position{
   618  			StartTs:  1,
   619  			CommitTs: 3,
   620  		}, lastWritePos)
   621  		require.Equal(suite.T(), sorter.Position{
   622  			StartTs:  2,
   623  			CommitTs: 3,
   624  		}, lastWritePos.Next())
   625  	}
   626  	taskChan <- &sinkTask{
   627  		span:          suite.testSpan,
   628  		lowerBound:    genLowerBound(),
   629  		getUpperBound: genUpperBoundGetter(4),
   630  		tableSink:     wrapper,
   631  		callback:      callback,
   632  		isCanceled:    func() bool { return false },
   633  	}
   634  	require.Eventually(suite.T(), func() bool {
   635  		// Only three events should be sent to sink
   636  		return len(sink.GetEvents()) == 3
   637  	}, 5*time.Second, 10*time.Millisecond)
   638  	require.Equal(suite.T(), 2, sink.GetWriteTimes(), "Only two times write to sink, "+
   639  		"because the max update interval size is 2 * event size")
   640  	require.Equal(suite.T(), uint64(3), batchID.Load())
   641  	sink.AckAllEvents()
   642  	require.Eventually(suite.T(), func() bool {
   643  		checkpointTs := wrapper.getCheckpointTs()
   644  		return checkpointTs.ResolvedMark() == 2
   645  	}, 5*time.Second, 10*time.Millisecond)
   646  
   647  	events = []*model.PolymorphicEvent{
   648  		genPolymorphicEvent(2, 5, suite.testSpan),
   649  		genPolymorphicResolvedEvent(6),
   650  	}
   651  	e.Add(suite.testSpan, events...)
   652  	// Send another task to make sure the batchID is started from 2.
   653  	callback = func(_ sorter.Position) {
   654  		cancel()
   655  	}
   656  	taskChan <- &sinkTask{
   657  		span: suite.testSpan,
   658  		lowerBound: sorter.Position{
   659  			StartTs:  2,
   660  			CommitTs: 3,
   661  		},
   662  		getUpperBound: genUpperBoundGetter(6),
   663  		tableSink:     wrapper,
   664  		callback:      callback,
   665  		isCanceled:    func() bool { return false },
   666  	}
   667  	wg.Wait()
   668  	require.Equal(suite.T(), uint64(5), batchID.Load(), "The batchID should be 5, "+
   669  		"because the first task has 3 events, the second task has 1 event")
   670  }
   671  
   672  // When starts to handle a task, advancer.lastPos should be set to a correct position.
   673  // Otherwise if advancer.lastPos isn't updated during scanning, callback will get an
   674  // invalid `advancer.lastPos`.
   675  func (suite *tableSinkWorkerSuite) TestHandleTaskWithoutMemory() {
   676  	ctx, cancel := context.WithCancel(context.Background())
   677  	events := []*model.PolymorphicEvent{
   678  		genPolymorphicEvent(1, 3, suite.testSpan),
   679  		genPolymorphicResolvedEvent(4),
   680  	}
   681  	w, e := suite.createWorker(ctx, 0, true)
   682  	defer w.sinkMemQuota.Close()
   683  	suite.addEventsToSortEngine(events, e)
   684  
   685  	taskChan := make(chan *sinkTask)
   686  	var wg sync.WaitGroup
   687  	wg.Add(1)
   688  	go func() {
   689  		defer wg.Done()
   690  		err := w.handleTasks(ctx, taskChan)
   691  		require.Equal(suite.T(), context.Canceled, err)
   692  	}()
   693  
   694  	wrapper, sink := createTableSinkWrapper(suite.testChangefeedID, suite.testSpan)
   695  	defer sink.Close()
   696  
   697  	chShouldBeClosed := make(chan struct{}, 1)
   698  	callback := func(lastWritePos sorter.Position) {
   699  		require.Equal(suite.T(), genLowerBound().Prev(), lastWritePos)
   700  		close(chShouldBeClosed)
   701  	}
   702  	taskChan <- &sinkTask{
   703  		span:          suite.testSpan,
   704  		lowerBound:    genLowerBound(),
   705  		getUpperBound: genUpperBoundGetter(4),
   706  		tableSink:     wrapper,
   707  		callback:      callback,
   708  		isCanceled:    func() bool { return true },
   709  	}
   710  
   711  	<-chShouldBeClosed
   712  	cancel()
   713  	wg.Wait()
   714  }