sdk/go/httpserver/request_limiter_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package httpserver
   6
   7 import (
   8         "net/http"
   9         "net/http/httptest"
  10         "sync"
  11         "testing"
  12         "time"
  13 )
  14
  15 type testHandler struct {
  16         inHandler   chan struct{}
  17         okToProceed chan struct{}
  18 }
  19
  20 func (h *testHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  21         h.inHandler <- struct{}{}
  22         <-h.okToProceed
  23 }
  24
  25 func newTestHandler(maxReqs int) *testHandler {
  26         return &testHandler{
  27                 inHandler:   make(chan struct{}),
  28                 okToProceed: make(chan struct{}),
  29         }
  30 }
  31
  32 func TestRequestLimiter1(t *testing.T) {
  33         h := newTestHandler(10)
  34         l := NewRequestLimiter(1, h)
  35         var wg sync.WaitGroup
  36         resps := make([]*httptest.ResponseRecorder, 10)
  37         for i := 0; i < 10; i++ {
  38                 wg.Add(1)
  39                 resps[i] = httptest.NewRecorder()
  40                 go func(i int) {
  41                         l.ServeHTTP(resps[i], &http.Request{})
  42                         wg.Done()
  43                 }(i)
  44         }
  45         done := make(chan struct{})
  46         go func() {
  47                 // Make sure one request has entered the handler
  48                 <-h.inHandler
  49                 // Make sure all unsuccessful requests finish (but don't wait
  50                 // for the one that's still waiting for okToProceed)
  51                 wg.Add(-1)
  52                 wg.Wait()
  53                 // Wait for the last goroutine
  54                 wg.Add(1)
  55                 h.okToProceed <- struct{}{}
  56                 wg.Wait()
  57                 done <- struct{}{}
  58         }()
  59         select {
  60         case <-done:
  61         case <-time.After(10 * time.Second):
  62                 t.Fatal("test timed out, probably deadlocked")
  63         }
  64         n200 := 0
  65         n503 := 0
  66         for i := 0; i < 10; i++ {
  67                 switch resps[i].Code {
  68                 case 200:
  69                         n200++
  70                 case 503:
  71                         n503++
  72                 default:
  73                         t.Fatalf("Unexpected response code %d", resps[i].Code)
  74                 }
  75         }
  76         if n200 != 1 || n503 != 9 {
  77                 t.Fatalf("Got %d 200 responses, %d 503 responses (expected 1, 9)", n200, n503)
  78         }
  79         // Now that all 10 are finished, an 11th request should
  80         // succeed.
  81         go func() {
  82                 <-h.inHandler
  83                 h.okToProceed <- struct{}{}
  84         }()
  85         resp := httptest.NewRecorder()
  86         l.ServeHTTP(resp, &http.Request{})
  87         if resp.Code != 200 {
  88                 t.Errorf("Got status %d on 11th request, want 200", resp.Code)
  89         }
  90 }
  91
  92 func TestRequestLimiter10(t *testing.T) {
  93         h := newTestHandler(10)
  94         l := NewRequestLimiter(10, h)
  95         var wg sync.WaitGroup
  96         for i := 0; i < 10; i++ {
  97                 wg.Add(1)
  98                 go func() {
  99                         l.ServeHTTP(httptest.NewRecorder(), &http.Request{})
 100                         wg.Done()
 101                 }()
 102                 // Make sure the handler starts before we initiate the
 103                 // next request, but don't let it finish yet.
 104                 <-h.inHandler
 105         }
 106         for i := 0; i < 10; i++ {
 107                 h.okToProceed <- struct{}{}
 108         }
 109         wg.Wait()
 110 }