14360: Merge branch 'master'
[arvados.git] / 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 }