go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/limiter/limiter_test.go (about)

     1  // Copyright 2020 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package limiter
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"testing"
    21  
    22  	"go.chromium.org/luci/common/tsmon"
    23  
    24  	. "github.com/smartystreets/goconvey/convey"
    25  	. "go.chromium.org/luci/common/testing/assertions"
    26  )
    27  
    28  func TestMaxConcurrencyLimit(t *testing.T) {
    29  	t.Parallel()
    30  
    31  	Convey("Works", t, func() {
    32  		const limiterName = "test-limiter"
    33  		const maxConcurrent = 5
    34  		const allConcurrent = 12
    35  
    36  		ctx, _ := tsmon.WithDummyInMemory(context.Background())
    37  		block := make(chan struct{}) // stalls all requests
    38  		wg := &sync.WaitGroup{}      // waits until all requests are done
    39  
    40  		Convey("In enforcing mode", func() {
    41  			l, _ := New(Options{
    42  				Name:                  limiterName,
    43  				MaxConcurrentRequests: maxConcurrent,
    44  			})
    45  
    46  			accepted, rejected := makeConcurrentRequests(ctx, l, allConcurrent, block, wg)
    47  			So(accepted, ShouldEqual, maxConcurrent)
    48  			So(rejected, ShouldEqual, allConcurrent-maxConcurrent)
    49  
    50  			// There are still maxConcurrent requests blocked. Also the rest of the
    51  			// requests were already rejected. Verify metrics reflect all that.
    52  			l.ReportMetrics(ctx)
    53  			So(concurrencyCurGauge.Get(ctx, limiterName), ShouldEqual, maxConcurrent)
    54  			So(concurrencyMaxGauge.Get(ctx, limiterName), ShouldEqual, maxConcurrent)
    55  			So(rejectedCounter.Get(ctx, limiterName, "call", "peer", "max concurrency"), ShouldEqual, allConcurrent-maxConcurrent)
    56  
    57  			// Unblock pending requests.
    58  			close(block)
    59  			wg.Wait()
    60  
    61  			// Metrics show there are no concurrent requests anymore.
    62  			l.ReportMetrics(ctx)
    63  			So(concurrencyCurGauge.Get(ctx, limiterName), ShouldEqual, 0)
    64  		})
    65  
    66  		Convey("In advisory mode", func() {
    67  			l, _ := New(Options{
    68  				Name:                  limiterName,
    69  				MaxConcurrentRequests: maxConcurrent,
    70  				AdvisoryMode:          true,
    71  			})
    72  
    73  			// All requests are actually accepted.
    74  			accepted, rejected := makeConcurrentRequests(ctx, l, allConcurrent, block, wg)
    75  			So(accepted, ShouldEqual, allConcurrent)
    76  			So(rejected, ShouldEqual, 0)
    77  
    78  			// But metrics reflect that some requests should have been rejected if
    79  			// not running in the advisory mode. Also concurrencyCurGauge reflects
    80  			// the reality (all allConcurrent requests are executing now).
    81  			l.ReportMetrics(ctx)
    82  			So(concurrencyCurGauge.Get(ctx, limiterName), ShouldEqual, allConcurrent)
    83  			So(concurrencyMaxGauge.Get(ctx, limiterName), ShouldEqual, maxConcurrent)
    84  			So(rejectedCounter.Get(ctx, limiterName, "call", "peer", "max concurrency"), ShouldEqual, allConcurrent-maxConcurrent)
    85  
    86  			// Unblock pending requests.
    87  			close(block)
    88  			wg.Wait()
    89  		})
    90  	})
    91  }
    92  
    93  func makeConcurrentRequests(ctx context.Context, l *Limiter, count int, block chan struct{}, wg *sync.WaitGroup) (accepted, rejected int) {
    94  	verdicts := make(chan error) // nil if accepted, non-nil if rejected
    95  
    96  	// Note: this test tries to simulate real server environment where calls to
    97  	// CheckRequest happen from multiple goroutines.
    98  	for i := 0; i < count; i++ {
    99  		wg.Add(1)
   100  		go func() {
   101  			defer wg.Done()
   102  			done, err := l.CheckRequest(ctx, &RequestInfo{
   103  				CallLabel: "call",
   104  				PeerLabel: "peer",
   105  			})
   106  			if err != nil {
   107  				verdicts <- err
   108  				return
   109  			}
   110  			verdicts <- nil
   111  			<-block
   112  			defer done()
   113  		}()
   114  	}
   115  
   116  	// Collect the verdicts.
   117  	for i := 0; i < count; i++ {
   118  		err := <-verdicts
   119  		if err == nil {
   120  			accepted++
   121  		} else {
   122  			So(err, ShouldErrLike, "max concurrency limit: the server limit reached")
   123  			rejected++
   124  		}
   125  	}
   126  	return
   127  }