github.com/lrita/ratelimit@v0.0.0-20190723030019-81504bd89bc5/ratelimit_test.go (about)

     1  // Copyright 2014 Canonical Ltd.
     2  // Licensed under the LGPLv3 with static-linking exception.
     3  // See LICENCE file for details.
     4  
     5  package ratelimit
     6  
     7  import (
     8  	"math"
     9  	"testing"
    10  	"time"
    11  
    12  	gc "gopkg.in/check.v1"
    13  )
    14  
    15  func TestPackage(t *testing.T) {
    16  	gc.TestingT(t)
    17  }
    18  
    19  type rateLimitSuite struct{}
    20  
    21  var _ = gc.Suite(rateLimitSuite{})
    22  
    23  type takeReq struct {
    24  	time       time.Duration
    25  	count      int64
    26  	expectWait time.Duration
    27  }
    28  
    29  var takeTests = []struct {
    30  	about        string
    31  	fillInterval time.Duration
    32  	capacity     int64
    33  	reqs         []takeReq
    34  }{{
    35  	about:        "serial requests",
    36  	fillInterval: 250 * time.Millisecond,
    37  	capacity:     10,
    38  	reqs: []takeReq{{
    39  		time:       0,
    40  		count:      0,
    41  		expectWait: 0,
    42  	}, {
    43  		time:       0,
    44  		count:      10,
    45  		expectWait: 0,
    46  	}, {
    47  		time:       0,
    48  		count:      1,
    49  		expectWait: 250 * time.Millisecond,
    50  	}, {
    51  		time:       250 * time.Millisecond,
    52  		count:      1,
    53  		expectWait: 250 * time.Millisecond,
    54  	}},
    55  }, {
    56  	about:        "concurrent requests",
    57  	fillInterval: 250 * time.Millisecond,
    58  	capacity:     10,
    59  	reqs: []takeReq{{
    60  		time:       0,
    61  		count:      10,
    62  		expectWait: 0,
    63  	}, {
    64  		time:       0,
    65  		count:      2,
    66  		expectWait: 500 * time.Millisecond,
    67  	}, {
    68  		time:       0,
    69  		count:      2,
    70  		expectWait: 1000 * time.Millisecond,
    71  	}, {
    72  		time:       0,
    73  		count:      1,
    74  		expectWait: 1250 * time.Millisecond,
    75  	}},
    76  }, {
    77  	about:        "more than capacity",
    78  	fillInterval: 1 * time.Millisecond,
    79  	capacity:     10,
    80  	reqs: []takeReq{{
    81  		time:       0,
    82  		count:      10,
    83  		expectWait: 0,
    84  	}, {
    85  		time:       20 * time.Millisecond,
    86  		count:      15,
    87  		expectWait: 5 * time.Millisecond,
    88  	}},
    89  }, {
    90  	about:        "sub-quantum time",
    91  	fillInterval: 10 * time.Millisecond,
    92  	capacity:     10,
    93  	reqs: []takeReq{{
    94  		time:       0,
    95  		count:      10,
    96  		expectWait: 0,
    97  	}, {
    98  		time:       7 * time.Millisecond,
    99  		count:      1,
   100  		expectWait: 3 * time.Millisecond,
   101  	}, {
   102  		time:       8 * time.Millisecond,
   103  		count:      1,
   104  		expectWait: 12 * time.Millisecond,
   105  	}},
   106  }, {
   107  	about:        "within capacity",
   108  	fillInterval: 10 * time.Millisecond,
   109  	capacity:     5,
   110  	reqs: []takeReq{{
   111  		time:       0,
   112  		count:      5,
   113  		expectWait: 0,
   114  	}, {
   115  		time:       60 * time.Millisecond,
   116  		count:      5,
   117  		expectWait: 0,
   118  	}, {
   119  		time:       60 * time.Millisecond,
   120  		count:      1,
   121  		expectWait: 10 * time.Millisecond,
   122  	}, {
   123  		time:       80 * time.Millisecond,
   124  		count:      2,
   125  		expectWait: 10 * time.Millisecond,
   126  	}},
   127  }}
   128  
   129  var availTests = []struct {
   130  	about        string
   131  	capacity     int64
   132  	fillInterval time.Duration
   133  	take         int64
   134  	sleep        time.Duration
   135  
   136  	expectCountAfterTake  int64
   137  	expectCountAfterSleep int64
   138  }{{
   139  	about:                 "should fill tokens after interval",
   140  	capacity:              5,
   141  	fillInterval:          time.Second,
   142  	take:                  5,
   143  	sleep:                 time.Second,
   144  	expectCountAfterTake:  0,
   145  	expectCountAfterSleep: 1,
   146  }, {
   147  	about:                 "should fill tokens plus existing count",
   148  	capacity:              2,
   149  	fillInterval:          time.Second,
   150  	take:                  1,
   151  	sleep:                 time.Second,
   152  	expectCountAfterTake:  1,
   153  	expectCountAfterSleep: 2,
   154  }, {
   155  	about:                 "shouldn't fill before interval",
   156  	capacity:              2,
   157  	fillInterval:          2 * time.Second,
   158  	take:                  1,
   159  	sleep:                 time.Second,
   160  	expectCountAfterTake:  1,
   161  	expectCountAfterSleep: 1,
   162  }, {
   163  	about:                 "should fill only once after 1*interval before 2*interval",
   164  	capacity:              2,
   165  	fillInterval:          2 * time.Second,
   166  	take:                  1,
   167  	sleep:                 3 * time.Second,
   168  	expectCountAfterTake:  1,
   169  	expectCountAfterSleep: 2,
   170  }}
   171  
   172  func (rateLimitSuite) TestTake(c *gc.C) {
   173  	for i, test := range takeTests {
   174  		tb := NewBucket(test.fillInterval, test.capacity)
   175  		for j, req := range test.reqs {
   176  			d, ok := tb.take(tb.startTime+int64(req.time), req.count, infinityDuration)
   177  			c.Assert(ok, gc.Equals, true)
   178  			if d != req.expectWait {
   179  				c.Fatalf("test %d.%d, %s, got %v want %v", i, j, test.about, d, req.expectWait)
   180  			}
   181  		}
   182  	}
   183  }
   184  
   185  func (rateLimitSuite) TestTakeMaxDuration(c *gc.C) {
   186  	for i, test := range takeTests {
   187  		tb := NewBucket(test.fillInterval, test.capacity)
   188  		for j, req := range test.reqs {
   189  			if req.expectWait > 0 {
   190  				d, ok := tb.take(tb.startTime+int64(req.time), req.count, req.expectWait-1)
   191  				c.Assert(ok, gc.Equals, false)
   192  				c.Assert(d, gc.Equals, time.Duration(0))
   193  			}
   194  			d, ok := tb.take(tb.startTime+int64(req.time), req.count, req.expectWait)
   195  			c.Assert(ok, gc.Equals, true)
   196  			if d != req.expectWait {
   197  				c.Fatalf("test %d.%d, %s, got %v want %v", i, j, test.about, d, req.expectWait)
   198  			}
   199  		}
   200  	}
   201  }
   202  
   203  type takeAvailableReq struct {
   204  	time   time.Duration
   205  	count  int64
   206  	expect int64
   207  }
   208  
   209  var takeAvailableTests = []struct {
   210  	about        string
   211  	fillInterval time.Duration
   212  	capacity     int64
   213  	reqs         []takeAvailableReq
   214  }{{
   215  	about:        "serial requests",
   216  	fillInterval: 250 * time.Millisecond,
   217  	capacity:     10,
   218  	reqs: []takeAvailableReq{{
   219  		time:   0,
   220  		count:  0,
   221  		expect: 0,
   222  	}, {
   223  		time:   0,
   224  		count:  10,
   225  		expect: 10,
   226  	}, {
   227  		time:   0,
   228  		count:  1,
   229  		expect: 0,
   230  	}, {
   231  		time:   250 * time.Millisecond,
   232  		count:  1,
   233  		expect: 1,
   234  	}},
   235  }, {
   236  	about:        "concurrent requests",
   237  	fillInterval: 250 * time.Millisecond,
   238  	capacity:     10,
   239  	reqs: []takeAvailableReq{{
   240  		time:   0,
   241  		count:  5,
   242  		expect: 5,
   243  	}, {
   244  		time:   0,
   245  		count:  2,
   246  		expect: 2,
   247  	}, {
   248  		time:   0,
   249  		count:  5,
   250  		expect: 3,
   251  	}, {
   252  		time:   0,
   253  		count:  1,
   254  		expect: 0,
   255  	}},
   256  }, {
   257  	about:        "more than capacity",
   258  	fillInterval: 1 * time.Millisecond,
   259  	capacity:     10,
   260  	reqs: []takeAvailableReq{{
   261  		time:   0,
   262  		count:  10,
   263  		expect: 10,
   264  	}, {
   265  		time:   20 * time.Millisecond,
   266  		count:  15,
   267  		expect: 10,
   268  	}},
   269  }, {
   270  	about:        "within capacity",
   271  	fillInterval: 10 * time.Millisecond,
   272  	capacity:     5,
   273  	reqs: []takeAvailableReq{{
   274  		time:   0,
   275  		count:  5,
   276  		expect: 5,
   277  	}, {
   278  		time:   60 * time.Millisecond,
   279  		count:  5,
   280  		expect: 5,
   281  	}, {
   282  		time:   70 * time.Millisecond,
   283  		count:  1,
   284  		expect: 1,
   285  	}},
   286  }}
   287  
   288  func (rateLimitSuite) TestTakeAvailable(c *gc.C) {
   289  	for i, test := range takeAvailableTests {
   290  		tb := NewBucket(test.fillInterval, test.capacity)
   291  		for j, req := range test.reqs {
   292  			d := tb.takeAvailable(tb.startTime+int64(req.time), req.count)
   293  			if d != req.expect {
   294  				c.Fatalf("test %d.%d, %s, got %v want %v", i, j, test.about, d, req.expect)
   295  			}
   296  		}
   297  	}
   298  }
   299  
   300  func (rateLimitSuite) TestPanics(c *gc.C) {
   301  	c.Assert(func() { NewBucket(0, 1) }, gc.PanicMatches, "token bucket fill interval is not > 0")
   302  	c.Assert(func() { NewBucket(-2, 1) }, gc.PanicMatches, "token bucket fill interval is not > 0")
   303  	c.Assert(func() { NewBucket(1, 0) }, gc.PanicMatches, "token bucket capacity is not > 0")
   304  	c.Assert(func() { NewBucket(1, -2) }, gc.PanicMatches, "token bucket capacity is not > 0")
   305  	c.Assert(func() { NewBucketWithQuantum(1, 1, 2) }, gc.PanicMatches, "token capacity need large than quantum")
   306  }
   307  
   308  func (rateLimitSuite) TestAvailable(c *gc.C) {
   309  	tb := NewBucketWithQuantum(time.Hour, 2, 2)
   310  	c.Assert(tb.Available(), gc.Equals, int64(2))
   311  	tb.Take(1)
   312  	c.Assert(tb.Available(), gc.Equals, int64(1))
   313  	c.Assert(tb.TakeAvailable(1), gc.Equals, int64(1))
   314  	c.Assert(tb.Available(), gc.Equals, int64(0))
   315  	c.Assert(tb.WaitMaxDuration(1, time.Millisecond), gc.Equals, false)
   316  }
   317  
   318  func isCloseTo(x, y, tolerance float64) bool {
   319  	return math.Abs(x-y)/y < tolerance
   320  }
   321  
   322  func (rateLimitSuite) TestRate(c *gc.C) {
   323  	tb := NewBucket(1, 1)
   324  	if !isCloseTo(tb.Rate(), 1e9, 0.00001) {
   325  		c.Fatalf("got %v want 1e9", tb.Rate())
   326  	}
   327  	tb = NewBucket(2*time.Second, 1)
   328  	if !isCloseTo(tb.Rate(), 0.5, 0.00001) {
   329  		c.Fatalf("got %v want 0.5", tb.Rate())
   330  	}
   331  	tb = NewBucketWithQuantum(100*time.Millisecond, 5, 5)
   332  	if !isCloseTo(tb.Rate(), 50, 0.00001) {
   333  		c.Fatalf("got %v want 50", tb.Rate())
   334  	}
   335  	tb.ResetRate(500.0, 50)
   336  	if !isCloseTo(tb.Rate(), 500.0, 0.00001) {
   337  		c.Fatalf("got %v want 500", tb.Rate())
   338  	}
   339  }
   340  
   341  func (rateLimitSuite) TestResetRate(c *gc.C) {
   342  	tb := NewBucketWithRate(500.0, 50)
   343  	go func() {
   344  		time.Sleep(500 * time.Millisecond)
   345  		tb.ResetRate(1000.0, 100)
   346  	}()
   347  	count := 0
   348  	begin := time.Now()
   349  	for time.Since(begin) < time.Second {
   350  		if tb.WaitMaxDuration(1, time.Millisecond) {
   351  			count++
   352  		}
   353  	}
   354  	if count < 600 {
   355  		c.Fatalf("got %v want > 550", count)
   356  	}
   357  }
   358  
   359  func checkRate(c *gc.C, rate float64) {
   360  	tb := NewBucketWithRate(rate, 1<<62)
   361  	if !isCloseTo(tb.Rate(), rate, rateMargin) {
   362  		c.Fatalf("got %g want %v", tb.Rate(), rate)
   363  	}
   364  	d, ok := tb.take(tb.startTime, 1<<62, infinityDuration)
   365  	c.Assert(ok, gc.Equals, true)
   366  	c.Assert(d, gc.Equals, time.Duration(0))
   367  
   368  	// Check that the actual rate is as expected by
   369  	// asking for a not-quite multiple of the bucket's
   370  	// quantum and checking that the wait time
   371  	// correct.
   372  	d, ok = tb.take(tb.startTime, tb.quantum*2-tb.quantum/2, infinityDuration)
   373  	c.Assert(ok, gc.Equals, true)
   374  	expectTime := 1e9 * float64(tb.quantum) * 2 / rate
   375  	if !isCloseTo(float64(d), expectTime, rateMargin) {
   376  		c.Fatalf("rate %g: got %g want %v", rate, float64(d), expectTime)
   377  	}
   378  }
   379  
   380  func (rateLimitSuite) TestNewBucketWithRate(c *gc.C) {
   381  	for rate := float64(1); rate < 1e6; rate += 7 {
   382  		checkRate(c, rate)
   383  	}
   384  	for _, rate := range []float64{
   385  		1024 * 1024 * 1024,
   386  		1e-5,
   387  		0.9e-5,
   388  		0.5,
   389  		0.9,
   390  		0.9e8,
   391  		3e12,
   392  		4e18,
   393  		float64(1<<63 - 1),
   394  	} {
   395  		checkRate(c, rate)
   396  		checkRate(c, rate/3)
   397  		checkRate(c, rate*1.3)
   398  	}
   399  }
   400  
   401  func TestAvailable(t *testing.T) {
   402  	for i, tt := range availTests {
   403  		tb := NewBucket(tt.fillInterval, tt.capacity)
   404  		if tb.Capacity() != tt.capacity {
   405  			t.Fatalf("#%d: %s, take = %d, want = %d",
   406  				i, tt.about, tb.Capacity(), tt.capacity)
   407  		}
   408  		if c := tb.takeAvailable(tb.startTime, tt.take); c != tt.take {
   409  			t.Fatalf("#%d: %s, take = %d, want = %d", i, tt.about, c, tt.take)
   410  		}
   411  		if c := tb.available(tb.startTime); c != tt.expectCountAfterTake {
   412  			t.Fatalf("#%d: %s, after take, available = %d, want = %d", i, tt.about, c, tt.expectCountAfterTake)
   413  		}
   414  		if c := tb.available(tb.startTime + int64(tt.sleep)); c != tt.expectCountAfterSleep {
   415  			t.Fatalf("#%d: %s, after some time it should fill in new tokens, available = %d, want = %d",
   416  				i, tt.about, c, tt.expectCountAfterSleep)
   417  		}
   418  	}
   419  
   420  }
   421  
   422  func BenchmarkWait(b *testing.B) {
   423  	tb := NewBucket(1, 16*1024)
   424  	for i := b.N - 1; i >= 0; i-- {
   425  		tb.Wait(1)
   426  	}
   427  }
   428  
   429  func BenchmarkNewBucket(b *testing.B) {
   430  	for i := b.N - 1; i >= 0; i-- {
   431  		NewBucketWithRate(4e18, 1<<62)
   432  	}
   433  }
   434  
   435  func BenchmarkTakeParallel(b *testing.B) {
   436  	tb := NewBucket(1, 16*1024)
   437  	b.RunParallel(func(pb *testing.PB) {
   438  		for pb.Next() {
   439  			tb.Take(1)
   440  		}
   441  	})
   442  }
   443  
   444  func BenchmarkTakeAvailableParallel(b *testing.B) {
   445  	tb := NewBucket(1, 16*1024)
   446  	b.RunParallel(func(pb *testing.PB) {
   447  		for pb.Next() {
   448  			tb.TakeAvailable(1)
   449  		}
   450  	})
   451  }
   452  
   453  func BenchmarkResetRate(b *testing.B) {
   454  	tb := NewBucketWithRate(500.0, 500)
   455  	for i := 0; i < b.N; i++ {
   456  		tb.ResetRate(float64(i+1), int64(i+1))
   457  	}
   458  }