github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/msg/producer/buffer/buffer_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 buffer
    22  
    23  import (
    24  	"sync"
    25  	"testing"
    26  	"time"
    27  
    28  	"github.com/m3db/m3/src/msg/producer"
    29  	"github.com/m3db/m3/src/x/retry"
    30  
    31  	"github.com/fortytw2/leaktest"
    32  	"github.com/golang/mock/gomock"
    33  	"github.com/stretchr/testify/require"
    34  )
    35  
    36  func TestOptionsValidation(t *testing.T) {
    37  	opts := NewOptions()
    38  	require.NoError(t, opts.Validate())
    39  
    40  	opts = opts.SetMaxMessageSize(100).SetMaxBufferSize(1)
    41  	require.Equal(t, errInvalidMaxMessageSize, opts.Validate())
    42  
    43  	opts = opts.SetMaxMessageSize(-1)
    44  	require.Equal(t, errNegativeMaxMessageSize, opts.Validate())
    45  
    46  	opts = opts.SetMaxBufferSize(-1)
    47  	require.Equal(t, errNegativeMaxBufferSize, opts.Validate())
    48  
    49  	opts = opts.SetScanBatchSize(0)
    50  	require.Equal(t, errInvalidScanBatchSize, opts.Validate())
    51  }
    52  
    53  func TestBuffer(t *testing.T) {
    54  	ctrl := gomock.NewController(t)
    55  	defer ctrl.Finish()
    56  
    57  	mm := producer.NewMockMessage(ctrl)
    58  	mm.EXPECT().Size().Return(100).AnyTimes()
    59  
    60  	b := mustNewBuffer(t, testOptions())
    61  	require.Equal(t, 0, int(b.size.Load()))
    62  	require.Equal(t, 0, b.bufferList.Len())
    63  
    64  	rm, err := b.Add(mm)
    65  	require.NoError(t, err)
    66  	require.Equal(t, mm.Size(), int(b.size.Load()))
    67  
    68  	mm.EXPECT().Finalize(producer.Consumed)
    69  	// Finalize the message will reduce the buffer size.
    70  	rm.IncRef()
    71  	rm.DecRef()
    72  	require.Equal(t, 0, int(b.size.Load()))
    73  }
    74  
    75  func TestBufferAddMessageTooLarge(t *testing.T) {
    76  	ctrl := gomock.NewController(t)
    77  	defer ctrl.Finish()
    78  
    79  	mm := producer.NewMockMessage(ctrl)
    80  	mm.EXPECT().Size().Return(100).AnyTimes()
    81  
    82  	b := mustNewBuffer(t, NewOptions().SetMaxMessageSize(1))
    83  	_, err := b.Add(mm)
    84  	require.Error(t, err)
    85  	require.Equal(t, errMessageTooLarge, err)
    86  }
    87  
    88  func TestBufferAddMessageLargerThanMaxBufferSize(t *testing.T) {
    89  	ctrl := gomock.NewController(t)
    90  	defer ctrl.Finish()
    91  
    92  	mm := producer.NewMockMessage(ctrl)
    93  	mm.EXPECT().Size().Return(100).AnyTimes()
    94  
    95  	b := mustNewBuffer(t, NewOptions().
    96  		SetMaxMessageSize(1).
    97  		SetMaxBufferSize(1),
    98  	)
    99  	_, err := b.Add(mm)
   100  	require.Error(t, err)
   101  	require.Equal(t, errMessageTooLarge, err)
   102  }
   103  
   104  func TestBufferCleanupOldest(t *testing.T) {
   105  	ctrl := gomock.NewController(t)
   106  	defer ctrl.Finish()
   107  
   108  	mm := producer.NewMockMessage(ctrl)
   109  	mm.EXPECT().Size().Return(100).AnyTimes()
   110  
   111  	b := mustNewBuffer(t, NewOptions())
   112  	rm, err := b.Add(mm)
   113  	require.NoError(t, err)
   114  	require.Equal(t, int(rm.Size()), mm.Size())
   115  	require.Equal(t, rm.Size(), b.size.Load())
   116  	require.Equal(t, 1, b.bufferList.Len())
   117  
   118  	mm.EXPECT().Finalize(producer.Dropped)
   119  	b.dropOldestUntilTarget(0)
   120  	require.Equal(t, uint64(0), b.size.Load())
   121  	require.Equal(t, 0, b.bufferList.Len())
   122  }
   123  
   124  func TestBufferCleanupOldestBatch(t *testing.T) {
   125  	ctrl := gomock.NewController(t)
   126  	defer ctrl.Finish()
   127  
   128  	mm := producer.NewMockMessage(ctrl)
   129  	mm.EXPECT().Size().Return(100).AnyTimes()
   130  
   131  	b := mustNewBuffer(t, NewOptions())
   132  	_, err := b.Add(mm)
   133  	require.NoError(t, err)
   134  	_, err = b.Add(mm)
   135  	require.NoError(t, err)
   136  	_, err = b.Add(mm)
   137  	require.NoError(t, err)
   138  	_, err = b.Add(mm)
   139  	require.NoError(t, err)
   140  	_, err = b.Add(mm)
   141  	require.NoError(t, err)
   142  	require.Equal(t, 5*mm.Size(), int(b.size.Load()))
   143  	require.Equal(t, 5, b.bufferList.Len())
   144  
   145  	mm.EXPECT().Finalize(producer.Dropped)
   146  	// Didn't need to iterate through the full batch size to get to target size.
   147  	shouldContinue := b.dropOldestBatchUntilTargetWithListLock(uint64(450), 2)
   148  	require.False(t, shouldContinue)
   149  	require.Equal(t, uint64(4*mm.Size()), b.size.Load())
   150  	require.Equal(t, 4, b.bufferList.Len())
   151  
   152  	mm.EXPECT().Finalize(producer.Dropped).Times(2)
   153  	shouldContinue = b.dropOldestBatchUntilTargetWithListLock(uint64(50), 2)
   154  	require.True(t, shouldContinue)
   155  	require.Equal(t, uint64(200), b.size.Load())
   156  	require.Equal(t, 2, b.bufferList.Len())
   157  
   158  	mm.EXPECT().Finalize(producer.Dropped).Times(2)
   159  	b.dropOldestUntilTarget(0)
   160  	require.Equal(t, uint64(0), b.size.Load())
   161  	require.Equal(t, 0, b.bufferList.Len())
   162  }
   163  
   164  func TestCleanupBatch(t *testing.T) {
   165  	defer leaktest.Check(t)()
   166  
   167  	ctrl := gomock.NewController(t)
   168  	defer ctrl.Finish()
   169  
   170  	mm1 := producer.NewMockMessage(ctrl)
   171  	mm1.EXPECT().Size().Return(1).AnyTimes()
   172  
   173  	mm2 := producer.NewMockMessage(ctrl)
   174  	mm2.EXPECT().Size().Return(2).AnyTimes()
   175  
   176  	mm3 := producer.NewMockMessage(ctrl)
   177  	mm3.EXPECT().Size().Return(3).AnyTimes()
   178  
   179  	b := mustNewBuffer(t, NewOptions().SetScanBatchSize(2))
   180  	_, err := b.Add(mm1)
   181  	require.NoError(t, err)
   182  	_, err = b.Add(mm2)
   183  	require.NoError(t, err)
   184  	_, err = b.Add(mm3)
   185  	require.NoError(t, err)
   186  
   187  	mm1.EXPECT().Finalize(gomock.Eq(producer.Dropped))
   188  	front := b.bufferList.Front()
   189  	front.Value.(*producer.RefCountedMessage).Drop()
   190  
   191  	require.Equal(t, 3, b.bufferLen())
   192  	e, removed := b.cleanupBatchWithListLock(front, 2, false)
   193  	require.Equal(t, 3, int(e.Value.(*producer.RefCountedMessage).Size()))
   194  	require.Equal(t, 2, b.bufferLen())
   195  	require.Equal(t, 1, removed)
   196  
   197  	e, removed = b.cleanupBatchWithListLock(e, 2, false)
   198  	require.Nil(t, e)
   199  	require.Equal(t, 2, b.bufferLen())
   200  	require.Equal(t, 0, removed)
   201  }
   202  
   203  func TestCleanupBatchWithElementBeingRemovedByOtherThread(t *testing.T) {
   204  	defer leaktest.Check(t)()
   205  
   206  	ctrl := gomock.NewController(t)
   207  	defer ctrl.Finish()
   208  
   209  	mm1 := producer.NewMockMessage(ctrl)
   210  	mm1.EXPECT().Size().Return(1).AnyTimes()
   211  
   212  	mm2 := producer.NewMockMessage(ctrl)
   213  	mm2.EXPECT().Size().Return(2).AnyTimes()
   214  
   215  	mm3 := producer.NewMockMessage(ctrl)
   216  	mm3.EXPECT().Size().Return(3).AnyTimes()
   217  
   218  	b := mustNewBuffer(t, NewOptions())
   219  	_, err := b.Add(mm1)
   220  	require.NoError(t, err)
   221  	_, err = b.Add(mm2)
   222  	require.NoError(t, err)
   223  	_, err = b.Add(mm3)
   224  	require.NoError(t, err)
   225  	require.Error(t, b.cleanup())
   226  
   227  	mm1.EXPECT().Finalize(gomock.Eq(producer.Dropped))
   228  	b.bufferList.Front().Value.(*producer.RefCountedMessage).Drop()
   229  
   230  	require.Equal(t, 3, b.bufferLen())
   231  	e, removed := b.cleanupBatchWithListLock(b.bufferList.Front(), 1, false)
   232  	// e stopped at message 2.
   233  	require.Equal(t, 2, int(e.Value.(*producer.RefCountedMessage).Size()))
   234  	require.Equal(t, 2, b.bufferLen())
   235  	require.Equal(t, 1, removed)
   236  	require.NotNil(t, e)
   237  
   238  	require.NotNil(t, e.Next())
   239  	// Mimic A new write triggered DropOldest and removed message 2.
   240  	b.bufferList.Remove(e)
   241  	require.Nil(t, e.Next())
   242  	require.Equal(t, 1, b.bufferLen())
   243  
   244  	// Mark message 3 as dropped, so it's ready to be removed.
   245  	mm3.EXPECT().Finalize(gomock.Eq(producer.Dropped))
   246  	b.bufferList.Front().Value.(*producer.RefCountedMessage).Drop()
   247  
   248  	// But next clean batch from the removed element is going to do nothing
   249  	// because the starting element is already removed.
   250  	e, removed = b.cleanupBatchWithListLock(e, 3, false)
   251  	require.Equal(t, 1, b.bufferLen())
   252  	require.Equal(t, 0, removed)
   253  	require.Nil(t, e)
   254  
   255  	// Next tick will start from the beginning again and will remove
   256  	// the dropped message.
   257  	require.NoError(t, b.cleanup())
   258  	require.Equal(t, 0, b.bufferLen())
   259  }
   260  
   261  func TestBufferCleanupBackground(t *testing.T) {
   262  	defer leaktest.Check(t)()
   263  
   264  	ctrl := gomock.NewController(t)
   265  	defer ctrl.Finish()
   266  
   267  	mm := producer.NewMockMessage(ctrl)
   268  	mm.EXPECT().Size().Return(100).AnyTimes()
   269  
   270  	b := mustNewBuffer(t, testOptions())
   271  	rm, err := b.Add(mm)
   272  	require.NoError(t, err)
   273  	require.Equal(t, rm.Size(), uint64(mm.Size()))
   274  	require.Equal(t, rm.Size(), b.size.Load())
   275  
   276  	b.Init()
   277  	mm.EXPECT().Finalize(producer.Consumed)
   278  	rm.IncRef()
   279  	rm.DecRef()
   280  
   281  	b.Close(producer.WaitForConsumption)
   282  	require.Equal(t, 0, int(b.size.Load()))
   283  	_, err = b.Add(mm)
   284  	require.Error(t, err)
   285  	// Safe to close again.
   286  	b.Close(producer.WaitForConsumption)
   287  }
   288  
   289  func TestListRemoveCleanupNextAndPrev(t *testing.T) {
   290  	defer leaktest.Check(t)()
   291  
   292  	ctrl := gomock.NewController(t)
   293  	defer ctrl.Finish()
   294  
   295  	mm := producer.NewMockMessage(ctrl)
   296  	mm.EXPECT().Size().Return(100).AnyTimes()
   297  
   298  	b := mustNewBuffer(t, NewOptions())
   299  	_, err := b.Add(mm)
   300  	require.NoError(t, err)
   301  	_, err = b.Add(mm)
   302  	require.NoError(t, err)
   303  	_, err = b.Add(mm)
   304  	require.NoError(t, err)
   305  
   306  	e := b.bufferList.Front().Next()
   307  	require.NotNil(t, e.Next())
   308  	require.NotNil(t, e.Prev())
   309  
   310  	b.bufferList.Remove(e)
   311  	require.Nil(t, e.Next())
   312  	require.Nil(t, e.Prev())
   313  }
   314  
   315  func TestBufferCloseDropEverything(t *testing.T) {
   316  	defer leaktest.Check(t)()
   317  
   318  	ctrl := gomock.NewController(t)
   319  	defer ctrl.Finish()
   320  
   321  	mm := producer.NewMockMessage(ctrl)
   322  	mm.EXPECT().Size().Return(100).AnyTimes()
   323  
   324  	b := mustNewBuffer(t, testOptions())
   325  	rm, err := b.Add(mm)
   326  	require.NoError(t, err)
   327  	require.Equal(t, rm.Size(), uint64(mm.Size()))
   328  	require.Equal(t, rm.Size(), b.size.Load())
   329  
   330  	b.Init()
   331  	mm.EXPECT().Finalize(producer.Dropped)
   332  	b.Close(producer.DropEverything)
   333  	for {
   334  		l := b.size.Load()
   335  		if l == 0 {
   336  			break
   337  		}
   338  		time.Sleep(100 * time.Millisecond)
   339  	}
   340  }
   341  
   342  func TestBufferDropOldestAsyncOnFull(t *testing.T) {
   343  	defer leaktest.Check(t)()
   344  
   345  	ctrl := gomock.NewController(t)
   346  	defer ctrl.Finish()
   347  
   348  	mm := producer.NewMockMessage(ctrl)
   349  	mm.EXPECT().Size().Return(100).AnyTimes()
   350  
   351  	b := mustNewBuffer(t,
   352  		testOptions().
   353  			SetAllowedSpilloverRatio(0.8).
   354  			SetMaxMessageSize(3*int(mm.Size())).
   355  			SetMaxBufferSize(3*int(mm.Size())),
   356  	)
   357  	require.Equal(t, 540, int(b.maxSpilloverSize))
   358  	rd1, err := b.Add(mm)
   359  	require.NoError(t, err)
   360  	require.Equal(t, 100, int(b.size.Load()))
   361  	rd2, err := b.Add(mm)
   362  	require.NoError(t, err)
   363  	require.Equal(t, 200, int(b.size.Load()))
   364  	rd3, err := b.Add(mm)
   365  	require.NoError(t, err)
   366  	require.Equal(t, 300, int(b.size.Load()))
   367  	require.False(t, rd1.IsDroppedOrConsumed())
   368  	require.False(t, rd2.IsDroppedOrConsumed())
   369  	require.False(t, rd3.IsDroppedOrConsumed())
   370  	require.Equal(t, b.maxBufferSize, b.size.Load())
   371  
   372  	mm2 := producer.NewMockMessage(ctrl)
   373  	mm2.EXPECT().Size().Return(2 * mm.Size()).AnyTimes()
   374  
   375  	require.Equal(t, 0, len(b.dropOldestCh))
   376  	_, err = b.Add(mm2)
   377  	require.NoError(t, err)
   378  	require.Equal(t, 500, int(b.size.Load()))
   379  	require.Equal(t, 1, len(b.dropOldestCh))
   380  
   381  	var wg sync.WaitGroup
   382  	wg.Add(2)
   383  	mm.EXPECT().Finalize(producer.Dropped).Do(func(interface{}) {
   384  		wg.Done()
   385  	}).Times(2)
   386  	b.Init()
   387  	wg.Wait()
   388  	require.True(t, rd1.IsDroppedOrConsumed())
   389  	require.True(t, rd2.IsDroppedOrConsumed())
   390  	require.False(t, rd3.IsDroppedOrConsumed())
   391  
   392  	mm.EXPECT().Finalize(producer.Dropped)
   393  	mm2.EXPECT().Finalize(producer.Dropped)
   394  	b.Close(producer.DropEverything)
   395  	require.Equal(t, 0, int(b.size.Load()))
   396  }
   397  
   398  func TestBufferDropOldestsyncOnFull(t *testing.T) {
   399  	defer leaktest.Check(t)()
   400  
   401  	ctrl := gomock.NewController(t)
   402  	defer ctrl.Finish()
   403  
   404  	mm := producer.NewMockMessage(ctrl)
   405  	mm.EXPECT().Size().Return(100).AnyTimes()
   406  
   407  	b := mustNewBuffer(t,
   408  		testOptions().
   409  			SetAllowedSpilloverRatio(0.2).
   410  			SetMaxMessageSize(3*int(mm.Size())).
   411  			SetMaxBufferSize(3*int(mm.Size())),
   412  	)
   413  	require.Equal(t, 360, int(b.maxSpilloverSize))
   414  	rd1, err := b.Add(mm)
   415  	require.NoError(t, err)
   416  	require.Equal(t, 100, int(b.size.Load()))
   417  	rd2, err := b.Add(mm)
   418  	require.NoError(t, err)
   419  	require.Equal(t, 200, int(b.size.Load()))
   420  	rd3, err := b.Add(mm)
   421  	require.NoError(t, err)
   422  	require.Equal(t, 300, int(b.size.Load()))
   423  	require.False(t, rd1.IsDroppedOrConsumed())
   424  	require.False(t, rd2.IsDroppedOrConsumed())
   425  	require.False(t, rd3.IsDroppedOrConsumed())
   426  	require.Equal(t, b.maxBufferSize, b.size.Load())
   427  
   428  	mm2 := producer.NewMockMessage(ctrl)
   429  	mm2.EXPECT().Size().Return(2 * mm.Size()).AnyTimes()
   430  
   431  	require.Equal(t, 0, len(b.dropOldestCh))
   432  	mm.EXPECT().Finalize(producer.Dropped).Times(2)
   433  	// This will trigger dropEarlist synchronizely as it reached
   434  	// max allowed spill over ratio.
   435  	_, err = b.Add(mm2)
   436  	require.NoError(t, err)
   437  	require.Equal(t, 300, int(b.size.Load()))
   438  	require.Equal(t, 0, len(b.dropOldestCh))
   439  }
   440  
   441  func TestBufferReturnErrorOnFull(t *testing.T) {
   442  	defer leaktest.Check(t)()
   443  
   444  	ctrl := gomock.NewController(t)
   445  	defer ctrl.Finish()
   446  
   447  	mm := producer.NewMockMessage(ctrl)
   448  	mm.EXPECT().Size().Return(100).AnyTimes()
   449  
   450  	b := mustNewBuffer(t, testOptions().
   451  		SetMaxMessageSize(int(mm.Size())).
   452  		SetMaxBufferSize(3*int(mm.Size())).
   453  		SetOnFullStrategy(ReturnError),
   454  	)
   455  
   456  	rd1, err := b.Add(mm)
   457  	require.NoError(t, err)
   458  	require.Equal(t, 100, int(b.size.Load()))
   459  	rd2, err := b.Add(mm)
   460  	require.NoError(t, err)
   461  	require.Equal(t, 200, int(b.size.Load()))
   462  	rd3, err := b.Add(mm)
   463  	require.NoError(t, err)
   464  	require.Equal(t, 300, int(b.size.Load()))
   465  	require.False(t, rd1.IsDroppedOrConsumed())
   466  	require.False(t, rd2.IsDroppedOrConsumed())
   467  	require.False(t, rd3.IsDroppedOrConsumed())
   468  
   469  	_, err = b.Add(mm)
   470  	require.Error(t, err)
   471  	require.Equal(t, 300, int(b.size.Load()))
   472  }
   473  
   474  func mustNewBuffer(t testing.TB, opts Options) *buffer {
   475  	b, err := NewBuffer(opts)
   476  	require.NoError(t, err)
   477  	return b.(*buffer)
   478  }
   479  
   480  func testOptions() Options {
   481  	return NewOptions().
   482  		SetCloseCheckInterval(100 * time.Millisecond).
   483  		SetDropOldestInterval(100 * time.Millisecond).
   484  		SetCleanupRetryOptions(retry.NewOptions().SetInitialBackoff(100 * time.Millisecond).SetMaxBackoff(200 * time.Millisecond))
   485  }
   486  
   487  func BenchmarkProduce(b *testing.B) {
   488  	ctrl := gomock.NewController(b)
   489  	defer ctrl.Finish()
   490  
   491  	mm := producer.NewMockMessage(ctrl)
   492  	mm.EXPECT().Size().Return(100).AnyTimes()
   493  	mm.EXPECT().Finalize(producer.Dropped).AnyTimes()
   494  
   495  	buffer := mustNewBuffer(b, NewOptions().
   496  		SetMaxBufferSize(1000*1000*200).
   497  		SetDropOldestInterval(100*time.Millisecond).
   498  		SetOnFullStrategy(DropOldest),
   499  	)
   500  
   501  	for n := 0; n < b.N; n++ {
   502  		_, err := buffer.Add(mm)
   503  		if err != nil {
   504  			b.FailNow()
   505  		}
   506  	}
   507  }