github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/msg/consumer/consumer_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  	"errors"
    25  	"io"
    26  	"net"
    27  	"testing"
    28  	"time"
    29  
    30  	"github.com/m3db/m3/src/msg/generated/proto/msgpb"
    31  	"github.com/m3db/m3/src/msg/protocol/proto"
    32  	"github.com/m3db/m3/src/x/pool"
    33  
    34  	"github.com/fortytw2/leaktest"
    35  	"github.com/golang/mock/gomock"
    36  	"github.com/stretchr/testify/require"
    37  )
    38  
    39  var (
    40  	testMsg1 = msgpb.Message{
    41  		Metadata: msgpb.Metadata{
    42  			Shard: 100,
    43  			Id:    200,
    44  		},
    45  		Value: []byte("foooooooo"),
    46  	}
    47  
    48  	testMsg2 = msgpb.Message{
    49  		Metadata: msgpb.Metadata{
    50  			Shard: 0,
    51  			Id:    45678,
    52  		},
    53  		Value: []byte("barrrrrrr"),
    54  	}
    55  )
    56  
    57  func TestConsumerWithMessagePool(t *testing.T) {
    58  	defer leaktest.Check(t)()
    59  
    60  	opts := testOptions()
    61  	l, err := NewListener("127.0.0.1:0", opts)
    62  	require.NoError(t, err)
    63  	defer l.Close()
    64  
    65  	conn, err := net.Dial("tcp", l.Addr().String())
    66  	require.NoError(t, err)
    67  
    68  	c, err := l.Accept()
    69  	require.NoError(t, err)
    70  
    71  	err = produce(conn, &testMsg1)
    72  	require.NoError(t, err)
    73  
    74  	m1, err := c.Message()
    75  	require.NoError(t, err)
    76  	require.Equal(t, testMsg1.Value, m1.Bytes())
    77  
    78  	// Acking m1 making it available for reuse.
    79  	m1.Ack()
    80  
    81  	err = produce(conn, &testMsg2)
    82  	require.NoError(t, err)
    83  
    84  	m2, err := c.Message()
    85  	require.NoError(t, err)
    86  	require.Equal(t, testMsg2.Metadata, m2.(*message).Message.Metadata)
    87  	require.Equal(t, testMsg2.Value, m2.Bytes())
    88  
    89  	require.Equal(t, m1, m2)
    90  
    91  	err = produce(conn, &testMsg1)
    92  	require.NoError(t, err)
    93  
    94  	m3, err := c.Message()
    95  	require.NoError(t, err)
    96  	require.Equal(t, testMsg1.Metadata, m3.(*message).Message.Metadata)
    97  	require.Equal(t, testMsg1.Value, m3.Bytes())
    98  
    99  	// m2 was not acked so m3 will alloc a new message from pool.
   100  	require.NotEqual(t, m2, m3)
   101  }
   102  
   103  func TestConsumerAckReusedMessage(t *testing.T) {
   104  	defer leaktest.Check(t)()
   105  
   106  	opts := testOptions().SetAckBufferSize(100)
   107  	l, err := NewListener("127.0.0.1:0", opts)
   108  	require.NoError(t, err)
   109  	defer l.Close()
   110  
   111  	conn, err := net.Dial("tcp", l.Addr().String())
   112  	require.NoError(t, err)
   113  
   114  	c, err := l.Accept()
   115  	require.NoError(t, err)
   116  
   117  	err = produce(conn, &testMsg1)
   118  	require.NoError(t, err)
   119  
   120  	m1, err := c.Message()
   121  	require.NoError(t, err)
   122  	require.Equal(t, testMsg1.Value, m1.Bytes())
   123  
   124  	// Acking m1 making it available for reuse.
   125  	m1.Ack()
   126  
   127  	err = produce(conn, &testMsg2)
   128  	require.NoError(t, err)
   129  
   130  	m2, err := c.Message()
   131  	require.NoError(t, err)
   132  	require.Equal(t, testMsg2.Metadata, m2.(*message).Message.Metadata)
   133  	require.Equal(t, testMsg2.Value, m2.Bytes())
   134  
   135  	require.Equal(t, m1, m2)
   136  	m2.Ack()
   137  
   138  	cc := c.(*consumer)
   139  	require.Equal(t, 2, len(cc.ackPb.Metadata))
   140  	require.Equal(t, testMsg1.Metadata, cc.ackPb.Metadata[0])
   141  	require.Equal(t, testMsg2.Metadata, cc.ackPb.Metadata[1])
   142  }
   143  
   144  func TestConsumerWriteAckError(t *testing.T) {
   145  	defer leaktest.Check(t)()
   146  
   147  	opts := testOptions()
   148  	l, err := NewListener("127.0.0.1:0", opts)
   149  	require.NoError(t, err)
   150  	defer l.Close()
   151  
   152  	ctrl := gomock.NewController(t)
   153  	defer ctrl.Finish()
   154  
   155  	conn, err := net.Dial("tcp", l.Addr().String())
   156  	require.NoError(t, err)
   157  
   158  	c, err := l.Accept()
   159  	require.NoError(t, err)
   160  
   161  	mockEncoder := proto.NewMockEncoder(ctrl)
   162  	cc := c.(*consumer)
   163  	cc.encoder = mockEncoder
   164  
   165  	err = produce(conn, &testMsg1)
   166  	require.NoError(t, err)
   167  
   168  	m, err := cc.Message()
   169  	require.NoError(t, err)
   170  	require.Equal(t, testMsg1.Value, m.Bytes())
   171  
   172  	mockEncoder.EXPECT().Encode(gomock.Any())
   173  	mockEncoder.EXPECT().Bytes().Return([]byte("foo"))
   174  	// force a bad write
   175  	cc.w.Reset(&badWriter{})
   176  	m.Ack()
   177  
   178  	// connection can no longer be used.
   179  	_, err = cc.Message()
   180  	require.Error(t, err)
   181  }
   182  
   183  type badWriter struct{}
   184  
   185  func (w *badWriter) Write([]byte) (int, error) {
   186  	return 0, errors.New("fail")
   187  }
   188  
   189  func TestConsumerDecodeAckError(t *testing.T) {
   190  	defer leaktest.Check(t)()
   191  
   192  	opts := testOptions()
   193  	l, err := NewListener("127.0.0.1:0", opts)
   194  	require.NoError(t, err)
   195  	defer func() {
   196  		require.NoError(t, l.Close())
   197  	}()
   198  
   199  	ctrl := gomock.NewController(t)
   200  	defer ctrl.Finish()
   201  
   202  	conn, err := net.Dial("tcp", l.Addr().String())
   203  	require.NoError(t, err)
   204  
   205  	c, err := l.Accept()
   206  	require.NoError(t, err)
   207  
   208  	mockEncoder := proto.NewMockEncoder(ctrl)
   209  	cc := c.(*consumer)
   210  	cc.encoder = mockEncoder
   211  
   212  	err = produce(conn, &testMsg1)
   213  	require.NoError(t, err)
   214  
   215  	m, err := cc.Message()
   216  	require.NoError(t, err)
   217  	require.Equal(t, testMsg1.Value, m.Bytes())
   218  
   219  	mockEncoder.EXPECT().Encode(gomock.Any())
   220  	mockEncoder.EXPECT().Bytes()
   221  	m.Ack()
   222  
   223  	// can still use the connection after failing to decode an ack.
   224  	err = produce(conn, &testMsg1)
   225  	require.NoError(t, err)
   226  }
   227  
   228  func TestConsumerMessageError(t *testing.T) {
   229  	defer leaktest.Check(t)()
   230  
   231  	opts := testOptions()
   232  	l, err := NewListener("127.0.0.1:0", opts)
   233  	require.NoError(t, err)
   234  	defer l.Close()
   235  
   236  	ctrl := gomock.NewController(t)
   237  	defer ctrl.Finish()
   238  
   239  	conn, err := net.Dial("tcp", l.Addr().String())
   240  	require.NoError(t, err)
   241  
   242  	c, err := l.Accept()
   243  	require.NoError(t, err)
   244  
   245  	mockDecoder := proto.NewMockDecoder(ctrl)
   246  	cc := c.(*consumer)
   247  	cc.decoder = mockDecoder
   248  
   249  	err = produce(conn, &testMsg1)
   250  	require.NoError(t, err)
   251  
   252  	mockDecoder.EXPECT().Decode(gomock.Any()).Return(errors.New("mock encode err"))
   253  
   254  	_, err = cc.Message()
   255  	require.Error(t, err)
   256  }
   257  
   258  func TestConsumerAckBuffer(t *testing.T) {
   259  	defer leaktest.Check(t)()
   260  
   261  	opts := testOptions().SetAckBufferSize(2)
   262  	l, err := NewListener("127.0.0.1:0", opts)
   263  	require.NoError(t, err)
   264  	defer l.Close()
   265  
   266  	ctrl := gomock.NewController(t)
   267  	defer ctrl.Finish()
   268  
   269  	conn, err := net.Dial("tcp", l.Addr().String())
   270  	require.NoError(t, err)
   271  
   272  	c, err := l.Accept()
   273  	require.NoError(t, err)
   274  
   275  	mockEncoder := proto.NewMockEncoder(ctrl)
   276  	cc := c.(*consumer)
   277  	cc.encoder = mockEncoder
   278  
   279  	err = produce(conn, &testMsg1)
   280  	require.NoError(t, err)
   281  
   282  	err = produce(conn, &testMsg2)
   283  	require.NoError(t, err)
   284  
   285  	m1, err := cc.Message()
   286  	require.NoError(t, err)
   287  	require.Equal(t, testMsg1.Value, m1.Bytes())
   288  
   289  	m2, err := cc.Message()
   290  	require.NoError(t, err)
   291  	require.Equal(t, testMsg2.Value, m2.Bytes())
   292  
   293  	// First ack won't trigger encode because we buffer 2 acks before writing out.
   294  	m1.Ack()
   295  
   296  	// Second ack will trigger encode.
   297  	mockEncoder.EXPECT().Encode(gomock.Any())
   298  	mockEncoder.EXPECT().Bytes()
   299  	m2.Ack()
   300  }
   301  
   302  func TestConsumerAckAfterClosed(t *testing.T) {
   303  	defer leaktest.Check(t)()
   304  
   305  	opts := testOptions().SetAckBufferSize(1)
   306  	l, err := NewListener("127.0.0.1:0", opts)
   307  	require.NoError(t, err)
   308  	defer l.Close()
   309  
   310  	ctrl := gomock.NewController(t)
   311  	defer ctrl.Finish()
   312  
   313  	conn, err := net.Dial("tcp", l.Addr().String())
   314  	require.NoError(t, err)
   315  
   316  	c, err := l.Accept()
   317  	require.NoError(t, err)
   318  
   319  	mockEncoder := proto.NewMockEncoder(ctrl)
   320  	cc := c.(*consumer)
   321  	cc.encoder = mockEncoder
   322  
   323  	err = produce(conn, &testMsg1)
   324  	require.NoError(t, err)
   325  
   326  	err = produce(conn, &testMsg2)
   327  	require.NoError(t, err)
   328  
   329  	m1, err := cc.Message()
   330  	require.NoError(t, err)
   331  	require.Equal(t, testMsg1.Value, m1.Bytes())
   332  
   333  	m2, err := cc.Message()
   334  	require.NoError(t, err)
   335  	require.Equal(t, testMsg2.Value, m2.Bytes())
   336  
   337  	mockEncoder.EXPECT().Encode(gomock.Any())
   338  	mockEncoder.EXPECT().Bytes()
   339  	m1.Ack()
   340  
   341  	cc.Close()
   342  	// Second ack will not trigger encode since the consumer is closed.
   343  	m2.Ack()
   344  }
   345  
   346  func TestConsumerTimeBasedFlush(t *testing.T) {
   347  	defer leaktest.Check(t)()
   348  
   349  	opts := testOptions().SetAckBufferSize(2)
   350  	l, err := NewListener("127.0.0.1:0", opts)
   351  	require.NoError(t, err)
   352  	defer l.Close()
   353  
   354  	ctrl := gomock.NewController(t)
   355  	defer ctrl.Finish()
   356  
   357  	conn, err := net.Dial("tcp", l.Addr().String())
   358  	require.NoError(t, err)
   359  
   360  	c, err := l.Accept()
   361  	require.NoError(t, err)
   362  
   363  	mockEncoder := proto.NewMockEncoder(ctrl)
   364  	cc := c.(*consumer)
   365  	cc.encoder = mockEncoder
   366  
   367  	err = produce(conn, &testMsg1)
   368  	require.NoError(t, err)
   369  
   370  	err = produce(conn, &testMsg2)
   371  	require.NoError(t, err)
   372  
   373  	m1, err := cc.Message()
   374  	require.NoError(t, err)
   375  	require.Equal(t, testMsg1.Value, m1.Bytes())
   376  
   377  	m1.Ack()
   378  	require.Equal(t, 1, len(cc.ackPb.Metadata))
   379  
   380  	mockEncoder.EXPECT().Encode(gomock.Any())
   381  	mockEncoder.EXPECT().Bytes()
   382  	cc.Init()
   383  	cc.Close()
   384  }
   385  
   386  func TestConsumerFlushAcksOnClose(t *testing.T) {
   387  	defer leaktest.Check(t)()
   388  
   389  	opts := testOptions().SetAckBufferSize(2)
   390  	l, err := NewListener("127.0.0.1:0", opts)
   391  	require.NoError(t, err)
   392  	defer l.Close()
   393  
   394  	ctrl := gomock.NewController(t)
   395  	defer ctrl.Finish()
   396  
   397  	conn, err := net.Dial("tcp", l.Addr().String())
   398  	require.NoError(t, err)
   399  
   400  	c, err := l.Accept()
   401  	require.NoError(t, err)
   402  	c.Init()
   403  
   404  	mockEncoder := proto.NewMockEncoder(ctrl)
   405  	cc := c.(*consumer)
   406  	cc.encoder = mockEncoder
   407  
   408  	err = produce(conn, &testMsg1)
   409  	require.NoError(t, err)
   410  
   411  	err = produce(conn, &testMsg2)
   412  	require.NoError(t, err)
   413  
   414  	m1, err := cc.Message()
   415  	require.NoError(t, err)
   416  	require.Equal(t, testMsg1.Value, m1.Bytes())
   417  
   418  	m1.Ack()
   419  	require.Equal(t, 1, len(cc.ackPb.Metadata))
   420  
   421  	mockEncoder.EXPECT().Encode(gomock.Any()).Return(nil)
   422  	mockEncoder.EXPECT().Bytes()
   423  	cc.Close()
   424  }
   425  
   426  func TestListenerMultipleConnection(t *testing.T) {
   427  	defer leaktest.Check(t)()
   428  
   429  	opts := testOptions()
   430  	l, err := NewListener("127.0.0.1:0", opts)
   431  	require.NoError(t, err)
   432  	defer l.Close()
   433  
   434  	require.Equal(t, "tcp", l.Addr().Network())
   435  	testProduceAndReceiveAck(t, testMsg1, l, opts)
   436  	testProduceAndReceiveAck(t, testMsg2, l, opts)
   437  }
   438  
   439  func testProduceAndReceiveAck(t *testing.T, testMsg msgpb.Message, l Listener, opts Options) {
   440  	conn, err := net.Dial("tcp", l.Addr().String())
   441  	require.NoError(t, err)
   442  	produce(conn, &testMsg)
   443  	require.NoError(t, err)
   444  
   445  	c, err := l.Accept()
   446  	require.NoError(t, err)
   447  	defer c.Close()
   448  
   449  	m, err := c.Message()
   450  	require.NoError(t, err)
   451  	require.Equal(t, testMsg.Value, m.Bytes())
   452  
   453  	m.Ack()
   454  	var ack msgpb.Ack
   455  	err = proto.NewDecoder(conn, opts.DecoderOptions(), 10).Decode(&ack)
   456  	require.NoError(t, err)
   457  	require.Equal(t, 1, len(ack.Metadata))
   458  	require.Equal(t, testMsg.Metadata, ack.Metadata[0])
   459  }
   460  
   461  func testOptions() Options {
   462  	opts := NewOptions()
   463  	return opts.
   464  		SetAckBufferSize(1).
   465  		SetAckFlushInterval(50 * time.Millisecond).
   466  		SetMessagePoolOptions(MessagePoolOptions{
   467  			PoolOptions: pool.NewObjectPoolOptions().SetSize(1),
   468  		}).
   469  		SetConnectionWriteBufferSize(1)
   470  }
   471  
   472  func produce(w io.Writer, m proto.Marshaler) error {
   473  	encoder := proto.NewEncoder(nil)
   474  	err := encoder.Encode(m)
   475  	if err != nil {
   476  		return err
   477  	}
   478  	_, err = w.Write(encoder.Bytes())
   479  	return err
   480  }