github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/framework/message_router_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 framework
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  	"testing"
    20  	"time"
    21  
    22  	"github.com/pingcap/tiflow/engine/pkg/p2p"
    23  	"github.com/pingcap/tiflow/pkg/workerpool"
    24  	"github.com/stretchr/testify/require"
    25  	"go.uber.org/atomic"
    26  )
    27  
    28  type testMessage struct {
    29  	id int
    30  }
    31  
    32  type messageSuite struct {
    33  	sendCount int
    34  	expected  []testMessage
    35  }
    36  
    37  func testMessageRouter(t *testing.T, suite *messageSuite) {
    38  	received := &struct {
    39  		sync.Mutex
    40  		msgs []testMessage
    41  	}{
    42  		msgs: make([]testMessage, 0),
    43  	}
    44  	msgCounter := atomic.NewInt32(0)
    45  	routeFn := func(topic p2p.Topic, msg p2p.MessageValue) error {
    46  		msgCounter.Add(1)
    47  		received.Lock()
    48  		defer received.Unlock()
    49  		tmsg, ok := msg.(testMessage)
    50  		require.True(t, ok)
    51  		received.msgs = append(received.msgs, tmsg)
    52  		return nil
    53  	}
    54  
    55  	// send suite.sendCount messages to mesasge router
    56  	pool := workerpool.NewDefaultAsyncPool(1)
    57  	router := NewMessageRouter("test-worker", pool, defaultMessageRouterBufferSize, routeFn)
    58  	for i := 0; i < suite.sendCount; i++ {
    59  		router.AppendMessage(p2p.Topic("test-topic"), testMessage{id: i})
    60  	}
    61  
    62  	ctx, cancel := context.WithCancel(context.Background())
    63  	var wg sync.WaitGroup
    64  	wg.Add(1)
    65  	go func() {
    66  		defer wg.Done()
    67  		_ = pool.Run(ctx)
    68  	}()
    69  
    70  	wg.Add(1)
    71  	go func() {
    72  		defer wg.Done()
    73  		for {
    74  			if err := router.Tick(ctx); err != nil {
    75  				return
    76  			}
    77  		}
    78  	}()
    79  
    80  	require.Eventually(t, func() bool {
    81  		if int(msgCounter.Load()) != len(suite.expected) {
    82  			return false
    83  		}
    84  		received.Lock()
    85  		defer received.Unlock()
    86  		require.Equal(t, suite.expected, received.msgs)
    87  		return true
    88  	}, time.Second, time.Millisecond*10)
    89  
    90  	cancel()
    91  	wg.Wait()
    92  }
    93  
    94  func TestMessageRouter(t *testing.T) {
    95  	t.Parallel()
    96  
    97  	suite := &messageSuite{
    98  		sendCount: defaultMessageRouterBufferSize,
    99  		expected: []testMessage{
   100  			{0}, {1}, {2}, {3},
   101  		},
   102  	}
   103  	testMessageRouter(t, suite)
   104  }
   105  
   106  func TestMessageRouterOverflow(t *testing.T) {
   107  	t.Parallel()
   108  
   109  	// old messages that are not processed in time will be dropped
   110  	suite := &messageSuite{
   111  		sendCount: 4 * defaultMessageRouterBufferSize,
   112  		expected: []testMessage{
   113  			{12}, {13}, {14}, {15},
   114  		},
   115  	}
   116  	testMessageRouter(t, suite)
   117  }