github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/msg/consumer/handlers_test.go (about)

     1  // Copyright (c) 2018 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package consumer
    22  
    23  import (
    24  	"net"
    25  	"sort"
    26  	"sync"
    27  	"testing"
    28  
    29  	"github.com/m3db/m3/src/msg/generated/proto/msgpb"
    30  	"github.com/m3db/m3/src/msg/protocol/proto"
    31  	"github.com/m3db/m3/src/x/server"
    32  	xtest "github.com/m3db/m3/src/x/test"
    33  
    34  	"github.com/fortytw2/leaktest"
    35  	"github.com/golang/mock/gomock"
    36  	"github.com/stretchr/testify/require"
    37  )
    38  
    39  func TestServerWithSingletonMessageProcessor(t *testing.T) {
    40  	defer leaktest.Check(t)()
    41  
    42  	var (
    43  		data []string
    44  		wg   sync.WaitGroup
    45  		mu   sync.Mutex
    46  	)
    47  
    48  	ctrl := gomock.NewController(t)
    49  	defer ctrl.Finish()
    50  
    51  	p := NewMockMessageProcessor(ctrl)
    52  	p.EXPECT().Process(gomock.Any()).Do(
    53  		func(m Message) {
    54  			mu.Lock()
    55  			data = append(data, string(m.Bytes()))
    56  			mu.Unlock()
    57  			m.Ack()
    58  			wg.Done()
    59  		},
    60  	).Times(3)
    61  	// Set a large ack buffer size to make sure the background go routine
    62  	// can flush it.
    63  	opts := testOptions().SetAckBufferSize(100)
    64  	l, err := net.Listen("tcp", "127.0.0.1:0")
    65  	require.NoError(t, err)
    66  
    67  	s := server.NewServer("a", NewMessageHandler(SingletonMessageProcessor(p), opts), server.NewOptions())
    68  	defer s.Close()
    69  	require.NoError(t, s.Serve(l))
    70  
    71  	conn1, err := net.Dial("tcp", l.Addr().String())
    72  	require.NoError(t, err)
    73  	conn2, err := net.Dial("tcp", l.Addr().String())
    74  	require.NoError(t, err)
    75  
    76  	wg.Add(3)
    77  	err = produce(conn1, &testMsg1)
    78  	require.NoError(t, err)
    79  	err = produce(conn1, &testMsg2)
    80  	require.NoError(t, err)
    81  	err = produce(conn2, &testMsg2)
    82  	require.NoError(t, err)
    83  
    84  	wg.Wait()
    85  	sort.Strings(data)
    86  	require.Equal(t, string(testMsg2.Value), data[0])
    87  	require.Equal(t, string(testMsg2.Value), data[1])
    88  	require.Equal(t, string(testMsg1.Value), data[2])
    89  
    90  	var ack msgpb.Ack
    91  	testDecoder := proto.NewDecoder(conn1, opts.DecoderOptions(), 10)
    92  	err = testDecoder.Decode(&ack)
    93  	require.NoError(t, err)
    94  	testDecoder = proto.NewDecoder(conn2, opts.DecoderOptions(), 10)
    95  	err = testDecoder.Decode(&ack)
    96  	require.NoError(t, err)
    97  	require.Equal(t, 3, len(ack.Metadata))
    98  	sort.Slice(ack.Metadata, func(i, j int) bool {
    99  		return ack.Metadata[i].Id < ack.Metadata[j].Id
   100  	})
   101  	require.Equal(t, testMsg1.Metadata, ack.Metadata[0])
   102  	require.Equal(t, testMsg2.Metadata, ack.Metadata[1])
   103  	require.Equal(t, testMsg2.Metadata, ack.Metadata[2])
   104  	p.EXPECT().Close()
   105  }
   106  
   107  func TestServerMessageDifferentConnections(t *testing.T) {
   108  	defer leaktest.Check(t)()
   109  
   110  	ctrl := xtest.NewController(t)
   111  	defer ctrl.Finish()
   112  
   113  	l, err := net.Listen("tcp", "127.0.0.1:0")
   114  	require.NoError(t, err)
   115  
   116  	var wg sync.WaitGroup
   117  	wg.Add(2)
   118  	handleMessage := func(m Message) {
   119  		wg.Done()
   120  	}
   121  
   122  	mp1 := NewMockMessageProcessor(ctrl)
   123  	mp2 := NewMockMessageProcessor(ctrl)
   124  	mp1.EXPECT().Process(gomock.Any()).Do(handleMessage)
   125  	mp1.EXPECT().Close()
   126  	mp2.EXPECT().Process(gomock.Any()).Do(handleMessage)
   127  	mp2.EXPECT().Close()
   128  
   129  	// Set a large ack buffer size to make sure the background go routine
   130  	// can flush it.
   131  	opts := testOptions().SetAckBufferSize(100)
   132  	first := true
   133  	var mu sync.Mutex
   134  	newMessageProcessor := func() MessageProcessor {
   135  		mu.Lock()
   136  		defer mu.Unlock()
   137  		if first {
   138  			first = false
   139  			return mp1
   140  		}
   141  		return mp2
   142  	}
   143  
   144  	s := server.NewServer("a",
   145  		NewMessageHandler(NewMessageProcessorFactory(newMessageProcessor), opts), server.NewOptions())
   146  	require.NoError(t, err)
   147  	require.NoError(t, s.Serve(l))
   148  
   149  	conn1, err := net.Dial("tcp", l.Addr().String())
   150  	require.NoError(t, err)
   151  	conn2, err := net.Dial("tcp", l.Addr().String())
   152  	require.NoError(t, err)
   153  
   154  	err = produce(conn1, &testMsg1)
   155  	require.NoError(t, err)
   156  	err = produce(conn2, &testMsg1)
   157  	require.NoError(t, err)
   158  
   159  	wg.Wait()
   160  	s.Close()
   161  }