go.temporal.io/server@v1.23.0/common/quotas/multi_rate_limiter_impl_test.go (about)

     1  // The MIT License
     2  //
     3  // Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
     4  //
     5  // Copyright (c) 2020 Uber Technologies, Inc.
     6  //
     7  // Permission is hereby granted, free of charge, to any person obtaining a copy
     8  // of this software and associated documentation files (the "Software"), to deal
     9  // in the Software without restriction, including without limitation the rights
    10  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    11  // copies of the Software, and to permit persons to whom the Software is
    12  // furnished to do so, subject to the following conditions:
    13  //
    14  // The above copyright notice and this permission notice shall be included in
    15  // all copies or substantial portions of the Software.
    16  //
    17  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    18  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    19  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    20  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    21  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    22  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    23  // THE SOFTWARE.
    24  
    25  package quotas
    26  
    27  import (
    28  	"context"
    29  	"testing"
    30  	"time"
    31  
    32  	"github.com/golang/mock/gomock"
    33  	"github.com/stretchr/testify/require"
    34  	"github.com/stretchr/testify/suite"
    35  )
    36  
    37  type (
    38  	multiStageRateLimiterSuite struct {
    39  		suite.Suite
    40  		*require.Assertions
    41  
    42  		controller        *gomock.Controller
    43  		firstRateLimiter  *MockRateLimiter
    44  		secondRateLimiter *MockRateLimiter
    45  		firstReservation  *MockReservation
    46  		secondReservation *MockReservation
    47  
    48  		rateLimiter *MultiRateLimiterImpl
    49  	}
    50  )
    51  
    52  func TestMultiStageRateLimiterSuite(t *testing.T) {
    53  	s := new(multiStageRateLimiterSuite)
    54  	suite.Run(t, s)
    55  }
    56  
    57  func (s *multiStageRateLimiterSuite) SetupSuite() {
    58  
    59  }
    60  
    61  func (s *multiStageRateLimiterSuite) TearDownSuite() {
    62  
    63  }
    64  
    65  func (s *multiStageRateLimiterSuite) SetupTest() {
    66  	s.Assertions = require.New(s.T())
    67  
    68  	s.controller = gomock.NewController(s.T())
    69  	s.firstRateLimiter = NewMockRateLimiter(s.controller)
    70  	s.secondRateLimiter = NewMockRateLimiter(s.controller)
    71  	s.firstReservation = NewMockReservation(s.controller)
    72  	s.secondReservation = NewMockReservation(s.controller)
    73  
    74  	s.rateLimiter = NewMultiRateLimiter([]RateLimiter{s.firstRateLimiter, s.secondRateLimiter})
    75  }
    76  
    77  func (s *multiStageRateLimiterSuite) TearDownTest() {
    78  	s.controller.Finish()
    79  }
    80  
    81  func (s *multiStageRateLimiterSuite) TestAllowN_NonSuccess() {
    82  	now := time.Now()
    83  	numToken := 2
    84  
    85  	s.firstReservation.EXPECT().OK().Return(false).AnyTimes()
    86  	s.firstRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.firstReservation)
    87  
    88  	result := s.rateLimiter.AllowN(now, numToken)
    89  	s.False(result)
    90  }
    91  
    92  func (s *multiStageRateLimiterSuite) TestAllowN_SomeSuccess_Case1() {
    93  	now := time.Now()
    94  	numToken := 2
    95  
    96  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
    97  	s.firstReservation.EXPECT().DelayFrom(now).Return(time.Duration(0)).AnyTimes()
    98  	s.firstReservation.EXPECT().CancelAt(now)
    99  	s.firstRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.firstReservation)
   100  
   101  	s.secondReservation.EXPECT().OK().Return(false).AnyTimes()
   102  	s.secondReservation.EXPECT().DelayFrom(now).Return(time.Duration(0)).AnyTimes()
   103  	s.secondRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.secondReservation)
   104  
   105  	result := s.rateLimiter.AllowN(now, numToken)
   106  	s.False(result)
   107  }
   108  
   109  func (s *multiStageRateLimiterSuite) TestAllowN_SomeSuccess_Case2() {
   110  	now := time.Now()
   111  	numToken := 2
   112  
   113  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   114  	s.firstReservation.EXPECT().DelayFrom(now).Return(time.Duration(0)).AnyTimes()
   115  	s.firstReservation.EXPECT().CancelAt(now)
   116  	s.firstRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.firstReservation)
   117  
   118  	s.secondReservation.EXPECT().OK().Return(true).AnyTimes()
   119  	s.secondReservation.EXPECT().DelayFrom(now).Return(time.Duration(1)).AnyTimes()
   120  	s.secondReservation.EXPECT().CancelAt(now)
   121  	s.secondRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.secondReservation)
   122  
   123  	result := s.rateLimiter.AllowN(now, numToken)
   124  	s.False(result)
   125  }
   126  
   127  func (s *multiStageRateLimiterSuite) TestAllowN_AllSuccess() {
   128  	now := time.Now()
   129  	numToken := 2
   130  
   131  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   132  	s.firstReservation.EXPECT().DelayFrom(now).Return(time.Duration(0)).AnyTimes()
   133  	s.firstRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.firstReservation)
   134  
   135  	s.secondReservation.EXPECT().OK().Return(true).AnyTimes()
   136  	s.secondReservation.EXPECT().DelayFrom(now).Return(time.Duration(0)).AnyTimes()
   137  	s.secondRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.secondReservation)
   138  
   139  	result := s.rateLimiter.AllowN(now, numToken)
   140  	s.True(result)
   141  }
   142  
   143  func (s *multiStageRateLimiterSuite) TestReserveN_NonSuccess() {
   144  	now := time.Now()
   145  	numToken := 4
   146  
   147  	s.firstReservation.EXPECT().OK().Return(false).AnyTimes()
   148  	s.firstRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.firstReservation)
   149  
   150  	result := s.rateLimiter.ReserveN(now, numToken)
   151  	s.Equal(&MultiReservationImpl{
   152  		ok:           false,
   153  		reservations: nil,
   154  	}, result)
   155  }
   156  
   157  func (s *multiStageRateLimiterSuite) TestReserveN_SomeSuccess() {
   158  	now := time.Now()
   159  	numToken := 4
   160  
   161  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   162  	s.firstReservation.EXPECT().CancelAt(now)
   163  	s.firstRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.firstReservation)
   164  
   165  	s.secondReservation.EXPECT().OK().Return(false).AnyTimes()
   166  	s.secondRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.secondReservation)
   167  
   168  	result := s.rateLimiter.ReserveN(now, numToken)
   169  	s.Equal(&MultiReservationImpl{
   170  		ok:           false,
   171  		reservations: nil,
   172  	}, result)
   173  }
   174  
   175  func (s *multiStageRateLimiterSuite) TestReserveN_AllSuccess() {
   176  	now := time.Now()
   177  	numToken := 4
   178  
   179  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   180  	s.firstRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.firstReservation)
   181  
   182  	s.secondReservation.EXPECT().OK().Return(true).AnyTimes()
   183  	s.secondRateLimiter.EXPECT().ReserveN(now, numToken).Return(s.secondReservation)
   184  
   185  	result := s.rateLimiter.ReserveN(now, numToken)
   186  	s.Equal(&MultiReservationImpl{
   187  		ok:           true,
   188  		reservations: []Reservation{s.firstReservation, s.secondReservation},
   189  	}, result)
   190  }
   191  
   192  func (s *multiStageRateLimiterSuite) TestWaitN_AlreadyExpired() {
   193  	ctx, cancel := context.WithCancel(context.Background())
   194  	cancel()
   195  	numToken := 4
   196  
   197  	result := s.rateLimiter.WaitN(ctx, numToken)
   198  	s.Error(result)
   199  }
   200  
   201  func (s *multiStageRateLimiterSuite) TestWaitN_NotExpired_WithExpiration_Error() {
   202  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
   203  	defer cancel()
   204  	numToken := 4
   205  
   206  	firstReservationDelay := 2 * time.Second
   207  	secondReservationDelay := 3 * time.Second
   208  	s.firstReservation.EXPECT().DelayFrom(gomock.Any()).Return(firstReservationDelay).AnyTimes()
   209  	s.secondReservation.EXPECT().DelayFrom(gomock.Any()).Return(secondReservationDelay).AnyTimes()
   210  	s.firstReservation.EXPECT().CancelAt(gomock.Any())
   211  	s.secondReservation.EXPECT().CancelAt(gomock.Any())
   212  
   213  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   214  	s.firstRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.firstReservation)
   215  	s.secondReservation.EXPECT().OK().Return(true).AnyTimes()
   216  	s.secondRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.secondReservation)
   217  
   218  	result := s.rateLimiter.WaitN(ctx, numToken)
   219  	s.Error(result)
   220  }
   221  
   222  func (s *multiStageRateLimiterSuite) TestWaitN_NotExpired_WithExpiration_Cancelled() {
   223  	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
   224  	numToken := 4
   225  
   226  	go func() {
   227  		time.Sleep(4 * time.Second)
   228  		cancel()
   229  	}()
   230  
   231  	firstReservationDelay := 20 * time.Second
   232  	secondReservationDelay := 30 * time.Second
   233  	s.firstReservation.EXPECT().DelayFrom(gomock.Any()).Return(firstReservationDelay).AnyTimes()
   234  	s.secondReservation.EXPECT().DelayFrom(gomock.Any()).Return(secondReservationDelay).AnyTimes()
   235  	s.firstReservation.EXPECT().CancelAt(gomock.Any())
   236  	s.secondReservation.EXPECT().CancelAt(gomock.Any())
   237  
   238  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   239  	s.firstRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.firstReservation)
   240  	s.secondReservation.EXPECT().OK().Return(true).AnyTimes()
   241  	s.secondRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.secondReservation)
   242  
   243  	result := s.rateLimiter.WaitN(ctx, numToken)
   244  	s.Error(result)
   245  }
   246  
   247  func (s *multiStageRateLimiterSuite) TestWaitN_NotExpired_WithExpiration_NoError() {
   248  	ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
   249  	defer cancel()
   250  	numToken := 4
   251  
   252  	firstReservationDelay := 2 * time.Second
   253  	secondReservationDelay := 3 * time.Second
   254  	s.firstReservation.EXPECT().DelayFrom(gomock.Any()).Return(firstReservationDelay).AnyTimes()
   255  	s.secondReservation.EXPECT().DelayFrom(gomock.Any()).Return(secondReservationDelay).AnyTimes()
   256  
   257  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   258  	s.firstRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.firstReservation)
   259  	s.secondReservation.EXPECT().OK().Return(true).AnyTimes()
   260  	s.secondRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.secondReservation)
   261  
   262  	result := s.rateLimiter.WaitN(ctx, numToken)
   263  	s.NoError(result)
   264  }
   265  
   266  func (s *multiStageRateLimiterSuite) TestWaitN_NotExpired_WithoutExpiration() {
   267  	ctx := context.Background()
   268  	numToken := 4
   269  
   270  	firstReservationDelay := 2 * time.Second
   271  	secondReservationDelay := 3 * time.Second
   272  	s.firstReservation.EXPECT().DelayFrom(gomock.Any()).Return(firstReservationDelay).AnyTimes()
   273  	s.secondReservation.EXPECT().DelayFrom(gomock.Any()).Return(secondReservationDelay).AnyTimes()
   274  
   275  	s.firstReservation.EXPECT().OK().Return(true).AnyTimes()
   276  	s.firstRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.firstReservation)
   277  	s.secondReservation.EXPECT().OK().Return(true).AnyTimes()
   278  	s.secondRateLimiter.EXPECT().ReserveN(gomock.Any(), numToken).Return(s.secondReservation)
   279  
   280  	result := s.rateLimiter.WaitN(ctx, numToken)
   281  	s.NoError(result)
   282  }
   283  
   284  func (s *multiStageRateLimiterSuite) TestRate() {
   285  	firstRateLimiterRate := float64(10)
   286  	secondRateLimiterRate := float64(5)
   287  
   288  	s.firstRateLimiter.EXPECT().Rate().Return(firstRateLimiterRate).AnyTimes()
   289  	s.secondRateLimiter.EXPECT().Rate().Return(secondRateLimiterRate).AnyTimes()
   290  
   291  	result := s.rateLimiter.Rate()
   292  	s.Equal(secondRateLimiterRate, result)
   293  }
   294  
   295  func (s *multiStageRateLimiterSuite) TestBurst() {
   296  	firstRateLimiterBurst := 5
   297  	secondRateLimiterBurst := 10
   298  
   299  	s.firstRateLimiter.EXPECT().Burst().Return(firstRateLimiterBurst).AnyTimes()
   300  	s.secondRateLimiter.EXPECT().Burst().Return(secondRateLimiterBurst).AnyTimes()
   301  
   302  	result := s.rateLimiter.Burst()
   303  	s.Equal(firstRateLimiterBurst, result)
   304  }