github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/client/executor_service_client_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 client
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"github.com/golang/mock/gomock"
    21  	"github.com/pingcap/tiflow/engine/enginepb"
    22  	pbMock "github.com/pingcap/tiflow/engine/enginepb/mock"
    23  	"github.com/pingcap/tiflow/pkg/errors"
    24  	"github.com/stretchr/testify/require"
    25  	"go.uber.org/atomic"
    26  	"google.golang.org/grpc"
    27  	"google.golang.org/grpc/codes"
    28  	"google.golang.org/grpc/status"
    29  )
    30  
    31  func TestDispatchTaskNormal(t *testing.T) {
    32  	t.Parallel()
    33  
    34  	ctrl := gomock.NewController(t)
    35  	client := pbMock.NewMockExecutorServiceClient(ctrl)
    36  	serviceCli := NewExecutorServiceClient(client)
    37  
    38  	var (
    39  		requestID           string
    40  		preDispatchComplete atomic.Bool
    41  		cbCalled            atomic.Bool
    42  	)
    43  
    44  	args := &DispatchTaskArgs{
    45  		WorkerID:     "worker-1",
    46  		MasterID:     "master-1",
    47  		WorkerType:   1,
    48  		WorkerConfig: []byte("testtest"),
    49  	}
    50  
    51  	client.EXPECT().PreDispatchTask(gomock.Any(), matchPreDispatchArgs(args)).
    52  		Do(func(_ context.Context, arg1 *enginepb.PreDispatchTaskRequest, _ ...grpc.CallOption) {
    53  			requestID = arg1.RequestId
    54  			preDispatchComplete.Store(true)
    55  		}).Return(&enginepb.PreDispatchTaskResponse{}, nil).Times(1)
    56  
    57  	client.EXPECT().ConfirmDispatchTask(gomock.Any(), matchConfirmDispatch(&requestID, "worker-1")).
    58  		Return(&enginepb.ConfirmDispatchTaskResponse{}, nil).Do(
    59  		func(_ context.Context, _ *enginepb.ConfirmDispatchTaskRequest, _ ...grpc.CallOption) {
    60  			require.True(t, preDispatchComplete.Load())
    61  			require.True(t, cbCalled.Load())
    62  		}).Times(1)
    63  
    64  	err := serviceCli.DispatchTask(context.Background(), args, func() {
    65  		require.True(t, preDispatchComplete.Load())
    66  		require.False(t, cbCalled.Swap(true))
    67  	})
    68  	require.NoError(t, err)
    69  }
    70  
    71  func TestPreDispatchAborted(t *testing.T) {
    72  	t.Parallel()
    73  
    74  	ctrl := gomock.NewController(t)
    75  	client := pbMock.NewMockExecutorServiceClient(ctrl)
    76  	serviceCli := NewExecutorServiceClient(client)
    77  
    78  	args := &DispatchTaskArgs{
    79  		WorkerID:     "worker-1",
    80  		MasterID:     "master-1",
    81  		WorkerType:   1,
    82  		WorkerConfig: []byte("testtest"),
    83  	}
    84  
    85  	unknownRPCError := status.Error(codes.Unknown, "fake error")
    86  	client.EXPECT().PreDispatchTask(gomock.Any(), matchPreDispatchArgs(args)).
    87  		Return((*enginepb.PreDispatchTaskResponse)(nil), unknownRPCError).Times(1)
    88  
    89  	err := serviceCli.DispatchTask(context.Background(), args, func() {
    90  		t.Fatalf("unexpected callback")
    91  	})
    92  	require.Error(t, err)
    93  	require.Regexp(t, "fake error", err)
    94  }
    95  
    96  func TestConfirmDispatchErrorFailFast(t *testing.T) {
    97  	t.Parallel()
    98  
    99  	// Only those errors that indicates a server-side failure can
   100  	// make DispatchTask fail fast. Otherwise, no error should be
   101  	// reported and at least a timeout should be waited for.
   102  	testCases := []struct {
   103  		err        error
   104  		isFailFast bool
   105  	}{
   106  		{
   107  			err:        errors.ErrRuntimeIncomingQueueFull.FastGenByArgs(),
   108  			isFailFast: true,
   109  		},
   110  		{
   111  			err:        errors.ErrDispatchTaskRequestIDNotFound.FastGenByArgs(),
   112  			isFailFast: true,
   113  		},
   114  		{
   115  			err:        errors.ErrInvalidArgument.FastGenByArgs("request-id"),
   116  			isFailFast: false,
   117  		},
   118  		{
   119  			err:        errors.ErrUnknown.FastGenByArgs(),
   120  			isFailFast: false,
   121  		},
   122  	}
   123  
   124  	ctrl := gomock.NewController(t)
   125  	client := pbMock.NewMockExecutorServiceClient(ctrl)
   126  	serviceCli := NewExecutorServiceClient(client)
   127  
   128  	for _, tc := range testCases {
   129  		var (
   130  			requestID           string
   131  			preDispatchComplete atomic.Bool
   132  			timerStarted        atomic.Bool
   133  		)
   134  
   135  		args := &DispatchTaskArgs{
   136  			WorkerID:     "worker-1",
   137  			MasterID:     "master-1",
   138  			WorkerType:   1,
   139  			WorkerConfig: []byte("testtest"),
   140  		}
   141  
   142  		client.EXPECT().PreDispatchTask(gomock.Any(), matchPreDispatchArgs(args)).
   143  			Do(func(_ context.Context, arg1 *enginepb.PreDispatchTaskRequest, _ ...grpc.CallOption) {
   144  				requestID = arg1.RequestId
   145  				preDispatchComplete.Store(true)
   146  			}).Return(&enginepb.PreDispatchTaskResponse{}, nil).Times(1)
   147  
   148  		client.EXPECT().ConfirmDispatchTask(gomock.Any(), matchConfirmDispatch(&requestID, "worker-1")).
   149  			Return((*enginepb.ConfirmDispatchTaskResponse)(nil), tc.err).Do(
   150  			func(_ context.Context, _ *enginepb.ConfirmDispatchTaskRequest, _ ...grpc.CallOption) {
   151  				require.True(t, preDispatchComplete.Load())
   152  				require.True(t, timerStarted.Load())
   153  			}).Times(1)
   154  
   155  		err := serviceCli.DispatchTask(context.Background(), args, func() {
   156  			require.True(t, preDispatchComplete.Load())
   157  			require.False(t, timerStarted.Swap(true))
   158  		})
   159  
   160  		if tc.isFailFast {
   161  			require.Error(t, err)
   162  		} else {
   163  			require.NoError(t, err)
   164  		}
   165  	}
   166  }