go.temporal.io/server@v1.23.0/common/quotas/map_request_rate_limiter_impl.go (about)

     1  // The MIT License
     2  //
     3  // Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
     4  //
     5  // Copyright (c) 2020 Uber Technologies, Inc.
     6  //
     7  // Permission is hereby granted, free of charge, to any person obtaining a copy
     8  // of this software and associated documentation files (the "Software"), to deal
     9  // in the Software without restriction, including without limitation the rights
    10  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    11  // copies of the Software, and to permit persons to whom the Software is
    12  // furnished to do so, subject to the following conditions:
    13  //
    14  // The above copyright notice and this permission notice shall be included in
    15  // all copies or substantial portions of the Software.
    16  //
    17  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    18  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    19  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    20  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    21  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    22  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    23  // THE SOFTWARE.
    24  
    25  package quotas
    26  
    27  import (
    28  	"context"
    29  	"sync"
    30  	"time"
    31  )
    32  
    33  type (
    34  	// RequestRateLimiterKeyFn extracts the map key from the request
    35  	RequestRateLimiterKeyFn[K comparable] func(req Request) K
    36  
    37  	// MapRequestRateLimiterImpl is a generic wrapper rate limiter for a set of rate limiters
    38  	// identified by a key
    39  	MapRequestRateLimiterImpl[K comparable] struct {
    40  		rateLimiterGenFn RequestRateLimiterFn
    41  		rateLimiterKeyFn RequestRateLimiterKeyFn[K]
    42  
    43  		sync.RWMutex
    44  		rateLimiters map[K]RequestRateLimiter
    45  	}
    46  )
    47  
    48  func NewMapRequestRateLimiter[K comparable](
    49  	rateLimiterGenFn RequestRateLimiterFn,
    50  	rateLimiterKeyFn RequestRateLimiterKeyFn[K],
    51  ) *MapRequestRateLimiterImpl[K] {
    52  	return &MapRequestRateLimiterImpl[K]{
    53  		rateLimiterGenFn: rateLimiterGenFn,
    54  		rateLimiterKeyFn: rateLimiterKeyFn,
    55  		rateLimiters:     make(map[K]RequestRateLimiter),
    56  	}
    57  }
    58  
    59  func namespaceRequestRateLimiterKeyFn(req Request) string {
    60  	return req.Caller
    61  }
    62  
    63  func NewNamespaceRequestRateLimiter(
    64  	rateLimiterGenFn RequestRateLimiterFn,
    65  ) *MapRequestRateLimiterImpl[string] {
    66  	return NewMapRequestRateLimiter[string](rateLimiterGenFn, namespaceRequestRateLimiterKeyFn)
    67  }
    68  
    69  // Allow attempts to allow a request to go through. The method returns
    70  // immediately with a true or false indicating if the request can make
    71  // progress
    72  func (r *MapRequestRateLimiterImpl[_]) Allow(
    73  	now time.Time,
    74  	request Request,
    75  ) bool {
    76  	rateLimiter := r.getOrInitRateLimiter(request)
    77  	return rateLimiter.Allow(now, request)
    78  }
    79  
    80  // Reserve returns a Reservation that indicates how long the caller
    81  // must wait before event happen.
    82  func (r *MapRequestRateLimiterImpl[_]) Reserve(
    83  	now time.Time,
    84  	request Request,
    85  ) Reservation {
    86  	rateLimiter := r.getOrInitRateLimiter(request)
    87  	return rateLimiter.Reserve(now, request)
    88  }
    89  
    90  // Wait waits till the deadline for a rate limit token to allow the request
    91  // to go through.
    92  func (r *MapRequestRateLimiterImpl[_]) Wait(
    93  	ctx context.Context,
    94  	request Request,
    95  ) error {
    96  	rateLimiter := r.getOrInitRateLimiter(request)
    97  	return rateLimiter.Wait(ctx, request)
    98  }
    99  
   100  func (r *MapRequestRateLimiterImpl[_]) getOrInitRateLimiter(
   101  	req Request,
   102  ) RequestRateLimiter {
   103  	key := r.rateLimiterKeyFn(req)
   104  
   105  	r.RLock()
   106  	rateLimiter, ok := r.rateLimiters[key]
   107  	r.RUnlock()
   108  	if ok {
   109  		return rateLimiter
   110  	}
   111  
   112  	newRateLimiter := r.rateLimiterGenFn(req)
   113  	r.Lock()
   114  	defer r.Unlock()
   115  
   116  	rateLimiter, ok = r.rateLimiters[key]
   117  	if ok {
   118  		return rateLimiter
   119  	}
   120  
   121  	r.rateLimiters[key] = newRateLimiter
   122  	return newRateLimiter
   123  }