github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/limit/limiter.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package limit
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/util/tracing"
    17  	"github.com/marusama/semaphore"
    18  )
    19  
    20  // ConcurrentRequestLimiter wraps a simple semaphore, adding a tracing span when
    21  // a request is forced to wait.
    22  type ConcurrentRequestLimiter struct {
    23  	spanName string
    24  	sem      semaphore.Semaphore
    25  }
    26  
    27  // MakeConcurrentRequestLimiter creates a ConcurrentRequestLimiter.
    28  func MakeConcurrentRequestLimiter(spanName string, limit int) ConcurrentRequestLimiter {
    29  	return ConcurrentRequestLimiter{spanName: spanName, sem: semaphore.New(limit)}
    30  }
    31  
    32  // Begin attempts to reserve a spot in the pool, blocking if needed until the
    33  // one is available or the context is canceled and adding a tracing span if it
    34  // is forced to block.
    35  func (l *ConcurrentRequestLimiter) Begin(ctx context.Context) error {
    36  	select {
    37  	case <-ctx.Done():
    38  		return ctx.Err()
    39  	default:
    40  	}
    41  
    42  	if l.sem.TryAcquire(1) {
    43  		return nil
    44  	}
    45  	// If not, start a span and begin waiting.
    46  	ctx, span := tracing.ChildSpan(ctx, l.spanName)
    47  	defer tracing.FinishSpan(span)
    48  	return l.sem.Acquire(ctx, 1)
    49  }
    50  
    51  // Finish indicates a concurrent request has completed and its reservation can
    52  // be returned to the pool.
    53  func (l *ConcurrentRequestLimiter) Finish() {
    54  	l.sem.Release(1)
    55  }
    56  
    57  // SetLimit adjusts the size of the pool.
    58  func (l *ConcurrentRequestLimiter) SetLimit(newLimit int) {
    59  	l.sem.SetLimit(newLimit)
    60  }