github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/executor/dm/unitholder_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 dm
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  	"sync"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/DATA-DOG/go-sqlmock"
    24  	"github.com/go-mysql-org/go-mysql/mysql"
    25  	"github.com/pingcap/tiflow/dm/config"
    26  	"github.com/pingcap/tiflow/dm/dumpling"
    27  	"github.com/pingcap/tiflow/dm/pb"
    28  	"github.com/pingcap/tiflow/dm/pkg/binlog"
    29  	"github.com/pingcap/tiflow/dm/pkg/conn"
    30  	"github.com/pingcap/tiflow/dm/pkg/log"
    31  	"github.com/pingcap/tiflow/dm/syncer"
    32  	"github.com/pingcap/tiflow/engine/jobmaster/dm/metadata"
    33  	dmpkg "github.com/pingcap/tiflow/engine/pkg/dm"
    34  	"github.com/pingcap/tiflow/pkg/errors"
    35  	"github.com/stretchr/testify/mock"
    36  	"github.com/stretchr/testify/require"
    37  	"go.uber.org/zap"
    38  )
    39  
    40  func TestUnitHolder(t *testing.T) {
    41  	unitHolder := &unitHolderImpl{
    42  		cfg: &config.SubTaskConfig{FrameworkLogger: zap.NewNop()},
    43  	}
    44  	u := &mockUnit{}
    45  	unitHolder.unit = u
    46  	_, _, err := conn.InitMockDBFull()
    47  	require.NoError(t, err)
    48  
    49  	u.On("Init").Return(errors.New("error")).Once()
    50  	require.Error(t, unitHolder.Init(context.Background()))
    51  	u.On("Init").Return(nil).Once()
    52  	require.NoError(t, unitHolder.Init(context.Background()))
    53  	stage, result := unitHolder.Stage()
    54  	require.Nil(t, result)
    55  	require.Equal(t, metadata.StageRunning, stage)
    56  
    57  	// mock error
    58  	time.Sleep(time.Second)
    59  	u.setResult(pb.ProcessResult{Errors: []*pb.ProcessError{{
    60  		ErrCode: 0,
    61  	}}})
    62  	unitHolder.processWg.Wait()
    63  	stage, result = unitHolder.Stage()
    64  	require.Equal(t, 0, int(result.Errors[0].ErrCode))
    65  	require.Equal(t, metadata.StageError, stage)
    66  
    67  	// mock auto resume
    68  	require.NoError(t, unitHolder.Resume(context.Background()))
    69  	stage, result = unitHolder.Stage()
    70  	require.Nil(t, result)
    71  	require.Equal(t, metadata.StageRunning, stage)
    72  
    73  	// mock paused
    74  	go func() {
    75  		time.Sleep(time.Second)
    76  		u.setResult(pb.ProcessResult{Errors: []*pb.ProcessError{{
    77  			Message: "context canceled",
    78  		}}})
    79  	}()
    80  	// mock pausing
    81  	go func() {
    82  		require.Eventually(t, func() bool {
    83  			stage, _ := unitHolder.Stage()
    84  			return stage == metadata.StagePausing
    85  		}, 5*time.Second, 100*time.Millisecond)
    86  	}()
    87  	require.NoError(t, unitHolder.Pause(context.Background()))
    88  	stage, result = unitHolder.Stage()
    89  	require.Len(t, result.Errors, 0)
    90  	require.Equal(t, metadata.StagePaused, stage)
    91  	// pause again
    92  	require.Error(t, unitHolder.Pause(context.Background()))
    93  	stage, result = unitHolder.Stage()
    94  	require.Len(t, result.Errors, 0)
    95  	require.Equal(t, metadata.StagePaused, stage)
    96  
    97  	// mock manually resume
    98  	require.NoError(t, unitHolder.Resume(context.Background()))
    99  	stage, result = unitHolder.Stage()
   100  	require.Nil(t, result)
   101  	require.Equal(t, metadata.StageRunning, stage)
   102  	// resume again
   103  	require.Error(t, unitHolder.Resume(context.Background()))
   104  
   105  	// mock finished
   106  	time.Sleep(time.Second)
   107  	u.setResult(pb.ProcessResult{Errors: []*pb.ProcessError{{
   108  		Message: "context canceled",
   109  	}}})
   110  	unitHolder.processWg.Wait()
   111  	stage, result = unitHolder.Stage()
   112  	require.Len(t, result.Errors, 0)
   113  	require.Equal(t, metadata.StageFinished, stage)
   114  
   115  	u.On("Status").Return(&pb.DumpStatus{})
   116  	status := unitHolder.Status(context.Background())
   117  	require.Equal(t, &pb.DumpStatus{}, status)
   118  
   119  	// mock close
   120  	require.NoError(t, unitHolder.Close(context.Background()))
   121  
   122  	// mock pause after close
   123  	require.EqualError(t, unitHolder.Pause(context.Background()), fmt.Sprintf("failed to pause unit with stage %s", metadata.StagePaused))
   124  	// mock resume after close
   125  	// depend on unit.Resume
   126  	require.NoError(t, unitHolder.Resume(context.Background()))
   127  }
   128  
   129  func TestUnitHolderBinlog(t *testing.T) {
   130  	unitHolder := &unitHolderImpl{}
   131  	unitHolder.unit = &dumpling.Dumpling{}
   132  
   133  	// wrong type
   134  	msg, err := unitHolder.Binlog(context.Background(), &dmpkg.BinlogTaskRequest{})
   135  	require.Error(t, err)
   136  	require.Equal(t, "", msg)
   137  	// no binlog error
   138  	unitHolder.unit = syncer.NewSyncer(&config.SubTaskConfig{Flavor: mysql.MySQLFlavor}, nil, nil)
   139  	unitHolder.runCtx = context.Background()
   140  	msg, err = unitHolder.Binlog(context.Background(), &dmpkg.BinlogTaskRequest{})
   141  	require.EqualError(t, err, "source '' has no error")
   142  	require.Equal(t, "", msg)
   143  }
   144  
   145  func TestUnitHolderBinlogSchema(t *testing.T) {
   146  	unitHolder := &unitHolderImpl{}
   147  	unitHolder.unit = &dumpling.Dumpling{}
   148  
   149  	// wrong type
   150  	msg, err := unitHolder.BinlogSchema(context.Background(), &dmpkg.BinlogSchemaTaskRequest{})
   151  	require.Error(t, err)
   152  	require.Equal(t, "", msg)
   153  	// wrong stage
   154  	unitHolder.unit = syncer.NewSyncer(&config.SubTaskConfig{Flavor: mysql.MySQLFlavor}, nil, nil)
   155  	unitHolder.runCtx = context.Background()
   156  	msg, err = unitHolder.BinlogSchema(context.Background(), &dmpkg.BinlogSchemaTaskRequest{})
   157  	require.EqualError(t, err, fmt.Sprintf("current stage is %s but not paused, invalid", metadata.StageRunning))
   158  	require.Equal(t, "", msg)
   159  	// binlog schema list
   160  	unitHolder.result = &pb.ProcessResult{Errors: []*pb.ProcessError{{ErrCode: 1}}}
   161  	msg, err = unitHolder.BinlogSchema(context.Background(), &dmpkg.BinlogSchemaTaskRequest{Op: pb.SchemaOp_RemoveSchema})
   162  	require.Nil(t, err)
   163  	require.Equal(t, "", msg)
   164  }
   165  
   166  func TestUnitHolderCheckAndUpdateStatus(t *testing.T) {
   167  	unitHolder := &unitHolderImpl{
   168  		cfg: &config.SubTaskConfig{
   169  			Flavor:          mysql.MySQLFlavor,
   170  			FrameworkLogger: zap.NewNop(),
   171  		},
   172  		logger: log.Logger{Logger: zap.NewNop()},
   173  	}
   174  	u := &mockUnit{}
   175  	unitHolder.unit = u
   176  	db, mock, err := conn.InitMockDBFull()
   177  	require.NoError(t, err)
   178  	unitHolder.upstreamDB = conn.NewBaseDBForTest(db)
   179  
   180  	u.On("Status").Return(&pb.DumpStatus{})
   181  	mock.ExpectQuery("SHOW MASTER STATUS").WillReturnRows(
   182  		sqlmock.NewRows([]string{"File", "Position", "Binlog_Do_DB", "Binlog_Ignore_DB", "Executed_Gtid_Set"}).
   183  			AddRow("mysql-bin.000001", "2345", "", "", ""),
   184  	)
   185  	mock.ExpectQuery("SHOW BINARY LOGS").WillReturnRows(
   186  		sqlmock.NewRows([]string{"File", "Position"}).AddRow("mysql-bin.000001", "2345"),
   187  	)
   188  	unitHolder.CheckAndUpdateStatus()
   189  	unitHolder.bgWg.Wait()
   190  	u.AssertExpectations(t)
   191  	require.NotNil(t, unitHolder.sourceStatus)
   192  	require.NoError(t, mock.ExpectationsWereMet())
   193  
   194  	// the second time CheckAndUpdateStatus, will not query upstreamDB
   195  	unitHolder.upstreamDB = nil
   196  	unitHolder.CheckAndUpdateStatus()
   197  	unitHolder.bgWg.Wait()
   198  	u.AssertExpectations(t)
   199  
   200  	// imitate pass refresh interval
   201  	backup := sourceStatusRefreshInterval
   202  	sourceStatusRefreshInterval = 0
   203  	defer func() {
   204  		sourceStatusRefreshInterval = backup
   205  	}()
   206  
   207  	unitHolder.upstreamDB = conn.NewBaseDBForTest(db)
   208  	lastUpdateTime := unitHolder.sourceStatus.UpdateTime
   209  	mock.ExpectQuery("SHOW MASTER STATUS").WillReturnRows(
   210  		sqlmock.NewRows([]string{"File", "Position", "Binlog_Do_DB", "Binlog_Ignore_DB", "Executed_Gtid_Set"}).
   211  			AddRow("mysql-bin.000001", "2345", "", "", ""),
   212  	)
   213  	mock.ExpectQuery("SHOW BINARY LOGS").WillReturnRows(
   214  		sqlmock.NewRows([]string{"File", "Position"}).AddRow("mysql-bin.000001", "2345"),
   215  	)
   216  	u.On("Status").Return(&pb.DumpStatus{})
   217  	unitHolder.CheckAndUpdateStatus()
   218  	unitHolder.bgWg.Wait()
   219  	u.AssertExpectations(t)
   220  	require.NoError(t, mock.ExpectationsWereMet())
   221  	require.NotEqual(t, lastUpdateTime, unitHolder.sourceStatus.UpdateTime)
   222  }
   223  
   224  type mockUnit struct {
   225  	sync.Mutex
   226  	mock.Mock
   227  	resultCh chan pb.ProcessResult
   228  }
   229  
   230  func (u *mockUnit) Init(ctx context.Context) error {
   231  	u.Lock()
   232  	defer u.Unlock()
   233  	return u.Called().Error(0)
   234  }
   235  
   236  func (u *mockUnit) Process(ctx context.Context, pr chan pb.ProcessResult) {
   237  	u.Lock()
   238  	defer u.Unlock()
   239  	u.resultCh = pr
   240  }
   241  
   242  func (u *mockUnit) setResult(r pb.ProcessResult) {
   243  	u.Lock()
   244  	defer u.Unlock()
   245  	u.resultCh <- r
   246  }
   247  
   248  func (u *mockUnit) Close() {}
   249  
   250  func (u *mockUnit) Kill() {}
   251  
   252  func (u *mockUnit) Pause() {}
   253  
   254  func (u *mockUnit) Resume(ctx context.Context, pr chan pb.ProcessResult) {
   255  	u.Lock()
   256  	defer u.Unlock()
   257  	u.resultCh = pr
   258  }
   259  
   260  func (u *mockUnit) Update(ctx context.Context, cfg *config.SubTaskConfig) error {
   261  	return nil
   262  }
   263  
   264  func (u *mockUnit) Status(sourceStatus *binlog.SourceStatus) interface{} {
   265  	u.Lock()
   266  	defer u.Unlock()
   267  	return u.Called().Get(0)
   268  }
   269  
   270  func (u *mockUnit) Type() pb.UnitType {
   271  	u.Lock()
   272  	defer u.Unlock()
   273  	return u.Called().Get(0).(pb.UnitType)
   274  }
   275  
   276  func (u *mockUnit) IsFreshTask(ctx context.Context) (bool, error) {
   277  	return false, nil
   278  }
   279  
   280  // mockUnitHolder implement Holder
   281  type mockUnitHolder struct {
   282  	sync.Mutex
   283  	mock.Mock
   284  }
   285  
   286  var _ unitHolder = &mockUnitHolder{}
   287  
   288  // Init implement Holder.Init
   289  func (m *mockUnitHolder) Init(ctx context.Context) error {
   290  	return nil
   291  }
   292  
   293  // Close implement Holder.Close
   294  func (m *mockUnitHolder) Close(ctx context.Context) error {
   295  	return nil
   296  }
   297  
   298  // Pause implement Holder.Pause
   299  func (m *mockUnitHolder) Pause(ctx context.Context) error {
   300  	m.Lock()
   301  	defer m.Unlock()
   302  	return m.Called().Error(0)
   303  }
   304  
   305  // Resume implement Holder.Resume
   306  func (m *mockUnitHolder) Resume(ctx context.Context) error {
   307  	m.Lock()
   308  	defer m.Unlock()
   309  	return m.Called().Error(0)
   310  }
   311  
   312  // Stage implement Holder.Stage
   313  func (m *mockUnitHolder) Stage() (metadata.TaskStage, *pb.ProcessResult) {
   314  	m.Lock()
   315  	defer m.Unlock()
   316  	args := m.Called()
   317  	if result := args.Get(1); result != nil {
   318  		return args.Get(0).(metadata.TaskStage), result.(*pb.ProcessResult)
   319  	}
   320  	return args.Get(0).(metadata.TaskStage), nil
   321  }
   322  
   323  // Status implement Holder.Status
   324  func (m *mockUnitHolder) Status(ctx context.Context) interface{} {
   325  	m.Lock()
   326  	defer m.Unlock()
   327  	args := m.Called()
   328  	return args.Get(0)
   329  }
   330  
   331  // CheckAndUpdateStatus implement Holder.CheckAndUpdateStatus
   332  func (m *mockUnitHolder) CheckAndUpdateStatus() {
   333  	m.Lock()
   334  	defer m.Unlock()
   335  	m.Called()
   336  }
   337  
   338  // Binlog implement Holder.Binlog
   339  func (m *mockUnitHolder) Binlog(ctx context.Context, req *dmpkg.BinlogTaskRequest) (string, error) {
   340  	m.Lock()
   341  	defer m.Unlock()
   342  	args := m.Called()
   343  	return args.Get(0).(string), args.Error(1)
   344  }
   345  
   346  // BinlogSchema implement Holder.BinlogSchema
   347  func (m *mockUnitHolder) BinlogSchema(ctx context.Context, req *dmpkg.BinlogSchemaTaskRequest) (string, error) {
   348  	m.Lock()
   349  	defer m.Unlock()
   350  	args := m.Called()
   351  	return args.Get(0).(string), args.Error(1)
   352  }