github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/errctx/center_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 errctx
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  	"testing"
    20  
    21  	"github.com/pingcap/tiflow/pkg/errors"
    22  	"github.com/stretchr/testify/require"
    23  )
    24  
    25  func TestErrCenterMultipleErrCtx(t *testing.T) {
    26  	center := NewErrCenter()
    27  
    28  	var wg sync.WaitGroup
    29  	for i := 0; i < 10; i++ {
    30  		wg.Add(1)
    31  		go func() {
    32  			defer wg.Done()
    33  
    34  			ctx, cancel := center.WithCancelOnFirstError(context.Background())
    35  			defer cancel()
    36  			<-ctx.Done()
    37  			require.Error(t, ctx.Err())
    38  			require.EqualError(t, ctx.Err(), "fake error")
    39  		}()
    40  	}
    41  
    42  	center.OnError(errors.New("fake error"))
    43  	wg.Wait()
    44  
    45  	require.Error(t, center.CheckError())
    46  	require.EqualError(t, center.CheckError(), "fake error")
    47  	require.Empty(t, center.children)
    48  }
    49  
    50  func TestErrCenterSingleErrCtx(t *testing.T) {
    51  	center := NewErrCenter()
    52  	ctx, cancel := center.WithCancelOnFirstError(context.Background())
    53  	defer cancel()
    54  
    55  	var wg sync.WaitGroup
    56  	for i := 0; i < 10; i++ {
    57  		wg.Add(1)
    58  		go func() {
    59  			defer wg.Done()
    60  
    61  			<-ctx.Done()
    62  			require.Error(t, ctx.Err())
    63  			require.EqualError(t, ctx.Err(), "fake error")
    64  		}()
    65  	}
    66  
    67  	center.OnError(errors.New("fake error"))
    68  	wg.Wait()
    69  
    70  	require.Error(t, center.CheckError())
    71  	require.EqualError(t, center.CheckError(), "fake error")
    72  	require.Empty(t, center.children)
    73  }
    74  
    75  func TestErrCenterRepeatedErrors(t *testing.T) {
    76  	center := NewErrCenter()
    77  	ctx, cancel := center.WithCancelOnFirstError(context.Background())
    78  	defer cancel()
    79  
    80  	center.OnError(nil) // no-op
    81  	center.OnError(errors.New("first error"))
    82  	center.OnError(errors.New("second error")) // no-op
    83  
    84  	<-ctx.Done()
    85  	require.Error(t, ctx.Err())
    86  	require.EqualError(t, ctx.Err(), "first error")
    87  
    88  	require.Error(t, center.CheckError())
    89  	require.EqualError(t, center.CheckError(), "first error")
    90  	require.Empty(t, center.children)
    91  }
    92  
    93  func TestErrCtxPropagate(t *testing.T) {
    94  	center := NewErrCenter()
    95  	parentCtx, cancelParent := context.WithCancel(context.Background())
    96  	ctx, cancel := center.WithCancelOnFirstError(parentCtx)
    97  	defer cancel()
    98  
    99  	cancelParent()
   100  	center.OnError(errors.New("fake error"))
   101  	<-ctx.Done()
   102  	require.Error(t, ctx.Err())
   103  	require.EqualError(t, ctx.Err(), "context canceled")
   104  	require.Empty(t, center.children)
   105  }
   106  
   107  func TestErrCtxDoneChSame(t *testing.T) {
   108  	center := NewErrCenter()
   109  	ctx, cancel := center.WithCancelOnFirstError(context.Background())
   110  	cancel()
   111  
   112  	doneCh1 := ctx.Done()
   113  	doneCh2 := ctx.Done()
   114  
   115  	require.Equal(t, doneCh1, doneCh2)
   116  }
   117  
   118  func TestErrCtxManualCancel(t *testing.T) {
   119  	center := NewErrCenter()
   120  	ctx, cancel := center.WithCancelOnFirstError(context.Background())
   121  	cancel()
   122  
   123  	<-ctx.Done()
   124  	require.Error(t, ctx.Err())
   125  	require.EqualError(t, ctx.Err(), "context canceled")
   126  	require.Empty(t, center.children)
   127  
   128  	// Multiple cancels should not panic.
   129  	cancel()
   130  	cancel()
   131  }