github.com/larrabee/ratelimit@v1.0.6-0.20191102113931-712217ec4fdc/real_bucket_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  	_, err := NewBucket(0, 1)
   302  	switch err.(type) {
   303  	case *ValueError:
   304  		if err.(*ValueError).Field != FieldFillInterval {
   305  			c.Fatalf("got error for incorrect field")
   306  		}
   307  		break
   308  	default:
   309  		c.Fatalf("got invalid error")
   310  	}
   311  
   312  	_, err = NewBucket(-2, 1)
   313  	switch err.(type) {
   314  	case *ValueError:
   315  		if err.(*ValueError).Field != FieldFillInterval {
   316  			c.Fatalf("got error for incorrect field")
   317  		}
   318  		break
   319  	default:
   320  		c.Fatalf("got invalid error")
   321  	}
   322  
   323  	_, err = NewBucket(1, 0)
   324  	switch err.(type) {
   325  	case *ValueError:
   326  		if err.(*ValueError).Field != FieldCapacity {
   327  			c.Fatalf("got error for incorrect field")
   328  		}
   329  		break
   330  	default:
   331  		c.Fatalf("got invalid error")
   332  	}
   333  
   334  	_, err = NewBucket(1, -2)
   335  	switch err.(type) {
   336  	case *ValueError:
   337  		if err.(*ValueError).Field != FieldCapacity {
   338  			c.Fatalf("got error for incorrect field")
   339  		}
   340  		break
   341  	default:
   342  		c.Fatalf("got invalid error")
   343  	}
   344  }
   345  
   346  func isCloseTo(x, y, tolerance float64) bool {
   347  	return math.Abs(x-y)/y < tolerance
   348  }
   349  
   350  func (rateLimitSuite) TestRate(c *gc.C) {
   351  	tb, _ := NewBucket(1, 1)
   352  	if !isCloseTo(tb.Rate(), 1e9, 0.00001) {
   353  		c.Fatalf("got %v want 1e9", tb.Rate())
   354  	}
   355  	tb, _ = NewBucket(2*time.Second, 1)
   356  	if !isCloseTo(tb.Rate(), 0.5, 0.00001) {
   357  		c.Fatalf("got %v want 0.5", tb.Rate())
   358  	}
   359  	tb, _ = NewBucketWithQuantum(100*time.Millisecond, 1, 5)
   360  	if !isCloseTo(tb.Rate(), 50, 0.00001) {
   361  		c.Fatalf("got %v want 50", tb.Rate())
   362  	}
   363  }
   364  
   365  func checkRate(c *gc.C, rate float64) {
   366  	tb, _ := NewBucketWithRate(rate, 1<<62)
   367  	if !isCloseTo(tb.Rate(), rate, rateMargin) {
   368  		c.Fatalf("got %g want %v", tb.Rate(), rate)
   369  	}
   370  	d, ok := tb.take(tb.startTime, 1<<62, infinityDuration)
   371  	c.Assert(ok, gc.Equals, true)
   372  	c.Assert(d, gc.Equals, time.Duration(0))
   373  
   374  	// Check that the actual rate is as expected by
   375  	// asking for a not-quite multiple of the bucket's
   376  	// quantum and checking that the wait time
   377  	// correct.
   378  	d, ok = tb.take(tb.startTime, tb.quantum*2-tb.quantum/2, infinityDuration)
   379  	c.Assert(ok, gc.Equals, true)
   380  	expectTime := 1e9 * float64(tb.quantum) * 2 / rate
   381  	if !isCloseTo(float64(d), expectTime, rateMargin) {
   382  		c.Fatalf("rate %g: got %g want %v", rate, float64(d), expectTime)
   383  	}
   384  }
   385  
   386  func (rateLimitSuite) TestNewBucketWithRate(c *gc.C) {
   387  	for rate := float64(1); rate < 1e6; rate += 7 {
   388  		checkRate(c, rate)
   389  	}
   390  	for _, rate := range []float64{
   391  		1024 * 1024 * 1024,
   392  		1e-5,
   393  		0.9e-5,
   394  		0.5,
   395  		0.9,
   396  		0.9e8,
   397  		3e12,
   398  		4e18,
   399  		float64(1<<63 - 1),
   400  	} {
   401  		checkRate(c, rate)
   402  		checkRate(c, rate/3)
   403  		checkRate(c, rate*1.3)
   404  	}
   405  }
   406  
   407  func TestAvailable(t *testing.T) {
   408  	for i, tt := range availTests {
   409  		tb, _ := NewBucket(tt.fillInterval, tt.capacity)
   410  		if tb.Capacity() != tt.capacity {
   411  			t.Fatalf("#%d: %s, take = %d, want = %d",
   412  				i, tt.about, tb.Capacity(), tt.capacity)
   413  		}
   414  		if c := tb.takeAvailable(tb.startTime, tt.take); c != tt.take {
   415  			t.Fatalf("#%d: %s, take = %d, want = %d", i, tt.about, c, tt.take)
   416  		}
   417  		if c := tb.available(tb.startTime); c != tt.expectCountAfterTake {
   418  			t.Fatalf("#%d: %s, after take, available = %d, want = %d", i, tt.about, c, tt.expectCountAfterTake)
   419  		}
   420  		if c := tb.available(tb.startTime + int64(tt.sleep)); c != tt.expectCountAfterSleep {
   421  			t.Fatalf("#%d: %s, after some time it should fill in new tokens, available = %d, want = %d",
   422  				i, tt.about, c, tt.expectCountAfterSleep)
   423  		}
   424  	}
   425  
   426  }
   427  
   428  func BenchmarkWait(b *testing.B) {
   429  	tb, _ := NewBucket(1, 16*1024)
   430  	for i := b.N - 1; i >= 0; i-- {
   431  		tb.Wait(1)
   432  	}
   433  }
   434  
   435  func BenchmarkNewBucket(b *testing.B) {
   436  	for i := b.N - 1; i >= 0; i-- {
   437  		NewBucketWithRate(4e18, 1<<62)
   438  	}
   439  }
   440  
   441  func BenchmarkTakeParallel(b *testing.B) {
   442  	tb, _ := NewBucket(1, 16*1024)
   443  	b.RunParallel(func(pb *testing.PB) {
   444  		for pb.Next() {
   445  			tb.Take(1)
   446  		}
   447  	})
   448  }
   449  
   450  func BenchmarkTakeAvailableParallel(b *testing.B) {
   451  	tb, _ := NewBucket(1, 16*1024)
   452  	b.RunParallel(func(pb *testing.PB) {
   453  		for pb.Next() {
   454  			tb.TakeAvailable(1)
   455  		}
   456  	})
   457  }