Function TestConcurrencyLimit

coderd/httpmw/ratelimit_test.go:251–357

(t *testing.T)

}

func TestConcurrencyLimit(t *testing.T) {
	t.Parallel()

	t.Run("LimitsConcurrentRequests", func(t *testing.T) {
		t.Parallel()

		const maxConcurrency = 2
		rtr := chi.NewRouter()
		rtr.Use(httpmw.ConcurrencyLimit(maxConcurrency, "Test"))

		// Use a WaitGroup as a barrier to ensure all requests are in the handler
		// before any of them proceed.
		var handlersReady sync.WaitGroup
		handlersReady.Add(maxConcurrency)
		releaseHandler := make(chan struct{})

		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
			handlersReady.Done()
			// Wait until released.
			<-releaseHandler
			rw.WriteHeader(http.StatusOK)
		})

		server := httptest.NewServer(rtr)
		defer server.Close()

		ctx := testutil.Context(t, testutil.WaitShort)

		// Start maxConcurrency requests that will block.
		// We use channels to collect errors instead of require in goroutines.
		type result struct {
			statusCode int
			err        error
		}
		results := make(chan result, maxConcurrency)

		var wg sync.WaitGroup
		for i := 0; i < maxConcurrency; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+"/", nil)
				if err != nil {
					results <- result{err: err}
					return
				}
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					results <- result{err: err}
					return
				}
				defer resp.Body.Close()
				results <- result{statusCode: resp.StatusCode}
			}()
		}

		// Wait for all requests to enter the handler with a timeout.
		handlersReadyCh := make(chan struct{})