sdk/go/arvados/limiter_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package arvados
   6
   7 import (
   8         "context"
   9         "errors"
  10         "net/http"
  11         "sync"
  12         "time"
  13
  14         . "gopkg.in/check.v1"
  15 )
  16
  17 var _ = Suite(&limiterSuite{})
  18
  19 type limiterSuite struct{}
  20
  21 func (*limiterSuite) TestInitialLimit(c *C) {
  22         ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Minute))
  23         defer cancel()
  24         rl := requestLimiter{}
  25
  26         var wg sync.WaitGroup
  27         wg.Add(int(requestLimiterInitialLimit))
  28         for i := int64(0); i < requestLimiterInitialLimit; i++ {
  29                 go func() {
  30                         rl.Acquire(ctx)
  31                         wg.Done()
  32                 }()
  33         }
  34         wg.Wait()
  35         c.Check(rl.current, Equals, requestLimiterInitialLimit)
  36         wg.Add(int(requestLimiterInitialLimit))
  37         for i := int64(0); i < requestLimiterInitialLimit; i++ {
  38                 go func() {
  39                         rl.Release()
  40                         wg.Done()
  41                 }()
  42         }
  43         wg.Wait()
  44         c.Check(rl.current, Equals, int64(0))
  45 }
  46
  47 func (*limiterSuite) TestCancelWhileWaitingForAcquire(c *C) {
  48         ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Minute))
  49         defer cancel()
  50         rl := requestLimiter{}
  51
  52         rl.Acquire(ctx)
  53         rl.limit = 1
  54         ctxShort, cancel := context.WithDeadline(ctx, time.Now().Add(time.Millisecond))
  55         defer cancel()
  56         rl.Acquire(ctxShort)
  57         c.Check(rl.current, Equals, int64(2))
  58         c.Check(ctxShort.Err(), NotNil)
  59         rl.Release()
  60         rl.Release()
  61         c.Check(rl.current, Equals, int64(0))
  62 }
  63
  64 func (*limiterSuite) TestReducedLimitAndQuietPeriod(c *C) {
  65         ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Minute))
  66         defer cancel()
  67         rl := requestLimiter{}
  68
  69         // Use a short quiet period to make tests faster
  70         defer func(orig time.Duration) { requestLimiterQuietPeriod = orig }(requestLimiterQuietPeriod)
  71         requestLimiterQuietPeriod = time.Second / 10
  72
  73         for i := 0; i < 5; i++ {
  74                 rl.Acquire(ctx)
  75         }
  76         rl.Report(&http.Response{StatusCode: http.StatusServiceUnavailable}, nil)
  77         c.Check(rl.limit, Equals, requestLimiterInitialLimit/2)
  78         for i := 0; i < 5; i++ {
  79                 rl.Release()
  80         }
  81
  82         // Even with all slots released, we can't Acquire in the quiet
  83         // period.
  84
  85         // (a) If our context expires before the end of the quiet
  86         // period, we get back DeadlineExceeded -- without waiting for
  87         // the end of the quiet period.
  88         acquire := time.Now()
  89         ctxShort, cancel := context.WithDeadline(ctx, time.Now().Add(requestLimiterQuietPeriod/10))
  90         defer cancel()
  91         rl.Acquire(ctxShort)
  92         c.Check(ctxShort.Err(), Equals, context.DeadlineExceeded)
  93         c.Check(time.Since(acquire) < requestLimiterQuietPeriod/2, Equals, true)
  94         c.Check(rl.quietUntil.Sub(time.Now()) > requestLimiterQuietPeriod/2, Equals, true)
  95         rl.Release()
  96
  97         // (b) If our context does not expire first, Acquire waits for
  98         // the end of the quiet period.
  99         ctxLong, cancel := context.WithDeadline(ctx, time.Now().Add(requestLimiterQuietPeriod*2))
 100         defer cancel()
 101         acquire = time.Now()
 102         rl.Acquire(ctxLong)
 103         c.Check(time.Since(acquire) > requestLimiterQuietPeriod/10, Equals, true)
 104         c.Check(time.Since(acquire) < requestLimiterQuietPeriod, Equals, true)
 105         c.Check(ctxLong.Err(), IsNil)
 106         rl.Release()
 107
 108         // OK to call Report() with nil Response and non-nil error.
 109         rl.Report(nil, errors.New("network error"))
 110 }