├── .github ├── FUNDING.yml ├── dependabot.yml └── workflows │ ├── ci.yml │ └── codeql-analysis.yml ├── .gitignore ├── LICENSE ├── OSSMETADATA ├── README.md ├── core ├── dep.go ├── doc.go ├── limit.go ├── limiter_registry.go ├── metric_registry.go └── strategy.go ├── doc.go ├── examples ├── doc.go ├── example_blocking_limit │ └── main.go ├── example_concurrent_loading │ └── main.go ├── example_simple_limit │ └── main.go ├── grpc_streaming │ ├── .gitignore │ ├── main.go │ └── pb │ │ ├── pingpong.pb.go │ │ └── pingpong.proto └── grpc_unary │ ├── .gitignore │ ├── main.go │ └── pb │ ├── pingpong.pb.go │ └── pingpong.proto ├── go.mod ├── go.sum ├── grpc ├── doc.go ├── grpc_streaming.go ├── grpc_unary.go ├── option_streaming.go └── option_unary.go ├── limit ├── aimd.go ├── aimd_test.go ├── doc.go ├── fixed.go ├── fixed_test.go ├── functions │ ├── dep.go │ ├── doc.go │ ├── fixed.go │ ├── fixed_test.go │ ├── init.go │ ├── log10_root.go │ ├── log10_root_test.go │ ├── square_root.go │ └── square_root_test.go ├── gradient.go ├── gradient2.go ├── gradient2_test.go ├── gradient_test.go ├── settable.go ├── settable_test.go ├── traced.go ├── traced_test.go ├── vegas.go ├── vegas_test.go ├── windowed.go └── windowed_test.go ├── limiter ├── blocking.go ├── blocking_test.go ├── deadline.go ├── deadline_test.go ├── default.go ├── default_test.go ├── delegate_listener.go ├── delegate_listener_test.go ├── doc.go ├── fifo_blocking.go ├── fifo_blocking_test.go ├── lifo_blocking.go ├── lifo_blocking_test.go ├── queue_blocking.go └── queue_blocking_test.go ├── measurements ├── doc.go ├── exponential_average.go ├── exponential_average_test.go ├── immutable_sample_window.go ├── immutable_sample_window_test.go ├── minimum.go ├── minimum_test.go ├── moving_average.go ├── moving_average_test.go ├── moving_variance.go ├── moving_variance_test.go ├── single.go ├── single_test.go ├── windowless_moving_percentile.go └── windowless_moving_percentile_test.go ├── metric_registry ├── datadog │ └── registry.go └── gometrics │ └── registry.go ├── patterns ├── doc.go └── pool │ ├── doc.go │ ├── example_fixed_pool_test.go │ ├── example_generic_pool_test.go │ ├── fixed_pool.go │ ├── fixed_pool_test.go │ ├── generic_pool.go │ └── generic_pool_test.go └── strategy ├── doc.go ├── lookup_partition.go ├── lookup_partition_test.go ├── matchers ├── doc.go ├── string.go └── string_test.go ├── precise.go ├── precise_test.go ├── predicate_partition.go ├── predicate_partition_test.go ├── simple.go └── simple_test.go /.github/FUNDING.yml: -------------------------------------------------------------------------------- 1 | github: [platinummonkey] 2 | custom: https://www.paypal.com/cgi-bin/webscr?cmd=_donations&business=UV978FNAYLYZE&item_name=Github+Donations¤cy_code=USD&source=url 3 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | # To get started with Dependabot version updates, you'll need to specify which 2 | # package ecosystems to update and where the package manifests are located. 3 | # Please see the documentation for all configuration options: 4 | # https://help.github.com/github/administering-a-repository/configuration-options-for-dependency-updates 5 | 6 | version: 2 7 | updates: 8 | - package-ecosystem: "gomod" # See documentation for possible values 9 | directory: "/" # Location of package manifests 10 | schedule: 11 | interval: "weekly" 12 | day: "monday" 13 | time: "09:00" 14 | timezone: "America/Chicago" 15 | assignees: 16 | - platinummonkey 17 | reviewers: 18 | - platinummonkey 19 | labels: 20 | - dependencies_upgrade 21 | open-pull-requests-limit: 10 22 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | on: 3 | push: 4 | branches: 5 | - master 6 | pull_request: 7 | branches: 8 | - master 9 | 10 | jobs: 11 | build: 12 | # The type of runner that the job will run on 13 | runs-on: ubuntu-latest 14 | 15 | # Steps represent a sequence of tasks that will be executed as part of the job 16 | steps: 17 | # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it 18 | - uses: actions/checkout@v2 19 | 20 | # Setup Go 21 | - name: Setup Go 22 | uses: actions/setup-go@v2 23 | with: 24 | go-version: '1.21.0' # The Go version to download (if necessary) and use. 25 | 26 | # Install all the dependencies 27 | - name: Install dependencies 28 | run: | 29 | go version 30 | go mod download 31 | go install golang.org/x/lint/golint@latest 32 | go install golang.org/x/tools/cmd/cover@latest 33 | 34 | # Run vet & lint on the code 35 | - name: Fmt, Lint, Vet, Test 36 | run: | 37 | go fmt ./... 38 | golint ./... 39 | go vet ./... 40 | 41 | # Run testing on the code 42 | - name: Run testing 43 | run: go test -v -race -covermode=atomic -coverprofile=covprofile ./... 44 | 45 | # todo: fix coverage reports... 46 | 47 | -------------------------------------------------------------------------------- /.github/workflows/codeql-analysis.yml: -------------------------------------------------------------------------------- 1 | name: "CodeQL" 2 | 3 | on: 4 | push: 5 | branches: [master] 6 | pull_request: 7 | # The branches below must be a subset of the branches above 8 | branches: [master] 9 | schedule: 10 | - cron: '0 16 * * 1' 11 | 12 | jobs: 13 | analyze: 14 | name: Analyze 15 | runs-on: ubuntu-latest 16 | 17 | strategy: 18 | fail-fast: false 19 | matrix: 20 | # Override automatic language detection by changing the below list 21 | # Supported options are ['csharp', 'cpp', 'go', 'java', 'javascript', 'python'] 22 | language: ['go'] 23 | # Learn more... 24 | # https://docs.github.com/en/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#overriding-automatic-language-detection 25 | 26 | steps: 27 | - name: Checkout repository 28 | uses: actions/checkout@v2 29 | with: 30 | # We must fetch at least the immediate parents so that if this is 31 | # a pull request then we can checkout the head. 32 | fetch-depth: 2 33 | 34 | # If this run was triggered by a pull request event, then checkout 35 | # the head of the pull request instead of the merge commit. 36 | - run: git checkout HEAD^2 37 | if: ${{ github.event_name == 'pull_request' }} 38 | 39 | # Initializes the CodeQL tools for scanning. 40 | - name: Initialize CodeQL 41 | uses: github/codeql-action/init@v1 42 | with: 43 | languages: ${{ matrix.language }} 44 | 45 | # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). 46 | # If this step fails, then you should remove it and run the build manually (see below) 47 | - name: Autobuild 48 | uses: github/codeql-action/autobuild@v1 49 | 50 | # ℹ️ Command-line programs to run using the OS shell. 51 | # 📚 https://git.io/JvXDl 52 | 53 | # ✏️ If the Autobuild fails above, remove it and uncomment the following three lines 54 | # and modify them (or add more) to build your code if your project 55 | # uses a compiled language 56 | 57 | #- run: | 58 | # make bootstrap 59 | # make release 60 | 61 | - name: Perform CodeQL Analysis 62 | uses: github/codeql-action/analyze@v1 63 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | public/ 3 | vendor/ 4 | -------------------------------------------------------------------------------- /OSSMETADATA: -------------------------------------------------------------------------------- 1 | osslifecycle=active 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![GoDoc](https://godoc.org/github.com/platinummonkey/go-concurrency-limits?status.svg)](https://godoc.org/github.com/platinummonkey/go-concurrency-limits) 2 | [![Build Status](https://travis-ci.org/platinummonkey/go-concurrency-limits.svg?branch=master)](https://travis-ci.org/platinummonkey/go-concurrency-limits) [![Coverage Status](https://img.shields.io/coveralls/github/platinummonkey/go-concurrency-limits/master.svg)](https://coveralls.io/github/platinummonkey/go-concurrency-limits) 3 | [![Releases](https://img.shields.io/github/release/platinummonkey/go-concurrency-limits.svg)](https://github.com/platinummonkey/go-concurrency-limits/releases) [![Releases](https://img.shields.io/github/downloads/platinummonkey/go-concurrency-limits/total.svg)](https://github.com/platinummonkey/go-concurrency-limits/releases) 4 | 5 | # Background 6 | 7 | When thinking of service availability operators traditionally think in terms of RPS (requests per second). Stress tests 8 | are normally performed to determine the RPS at which point the service tips over. RPS limits are then set somewhere 9 | below this tipping point (say 75% of this value) and enforced via a token bucket. However, in large distributed systems 10 | that auto-scale this value quickly goes out of date and the service falls over by becoming non-responsive as it is 11 | unable to gracefully shed excess load. Instead of thinking in terms of RPS, we should be thinking in terms of 12 | concurrent request where we apply queuing theory to determine the number of concurrent requests a service can handle 13 | before a queue starts to build up, latencies increase and the service eventually exhausts a hard limit such as CPU, 14 | memory, disk or network. This relationship is covered very nicely with Little's Law where 15 | Limit = Average RPS * Average Latency. 16 | 17 | Concurrency limits are very easy to enforce but difficult to determine as they would require operators to fully 18 | understand the hardware services run on and coordinate how they scale. Instead we'd prefer to measure or estimate the 19 | concurrency limits at each point in the network. As systems scale and hit limits each node will adjust and enforce its 20 | local view of the limit. To estimate the limit we borrow from common TCP congestion control algorithms by equating a 21 | system's concurrency limit to a TCP congestion window. 22 | 23 | Before applying the algorithm we need to set some ground rules. 24 | 25 | - We accept that every system has an inherent concurrency limit that is determined by a hard resources, such as number of CPU cores. 26 | - We accept that this limit can change as a system auto-scales. 27 | - For large and complex distributed systems it's impossible to know all the hard resources. 28 | - We can use latency measurements to determine when queuing happens. 29 | - We can use timeouts and rejected requests to aggressively back off. 30 | 31 | # Limit Algorithms 32 | 33 | ## Vegas 34 | 35 | Delay based algorithm where the bottleneck queue is estimated as 36 | 37 | ``` 38 | L * (1 - minRTT/sampleRtt) 39 | ``` 40 | 41 | At the end of each sampling window the limit is increased by 1 if the queue is less than alpha (typically a value 42 | between 2-3) or decreased by 1 if the queue is greater than beta (typically a value between 4-6 requests). 43 | 44 | ## Gradient2 45 | 46 | This algorithm attempts to address bias and drift when using minimum latency measurements. To do this the algorithm 47 | tracks uses the measure of divergence between two exponential averages over a long and short time time window. Using 48 | averages the algorithm can smooth out the impact of outliers for bursty traffic. Divergence duration is used as a proxy 49 | to identify a queueing trend at which point the algorithm aggresively reduces the limit. 50 | 51 | # Enforcement Strategies 52 | 53 | ## Simple 54 | 55 | In the simplest use case we don't want to differentiate between requests and so enforce a single gauge of the number of 56 | inflight requests. Requests are rejected immediately once the gauge value equals the limit. 57 | 58 | ## Partitioned 59 | 60 | For a slightly more complex system, it's desirable to partition requests to different backend/services. For example, 61 | you might shard by a customer id modulus 64 and the remainder you use as a unique backend identifier to target the 62 | the request. This allows for specific partitions to begin failing while others are operation normally. 63 | 64 | ## Percentage 65 | 66 | For more complex systems it's desirable to provide certain quality of service guarantees while still making efficient 67 | use of resources. Here we guarantee specific types of requests get a certain percentage of the concurrency limit. For 68 | example, a system that takes both live and batch traffic may want to give live traffic 100% of the limit during heavy 69 | load and is OK with starving batch traffic. Or, a system may want to guarantee that 50% of the limit is given to write 70 | traffic so writes are never starved. 71 | 72 | # Integrations 73 | 74 | ## GRPC 75 | 76 | A concurrency limiter may be installed either on the server or client. The choice of limiter depends on your use case. 77 | For the most part it is recommended to use a dynamic delay based limiter such as the VegasLimit on the server and 78 | either a pure loss based (AIMDLimit) or combined loss and delay based limiter on the client. 79 | 80 | ### Server Limiter 81 | 82 | The purpose of the server limiter is to protect the server from either increased client traffic (batch apps or retry 83 | storms) or latency spikes from a dependent service. With the limiter installed the server can ensure that latencies 84 | remain low by rejecting excess traffic with `Status.UNAVAILABLE` errors. 85 | 86 | In this example a GRPC server is configured with a single adaptive limiter that is shared among batch and live traffic 87 | with live traffic guaranteed 90% of throughput and 10% guaranteed to batch. For simplicity we just expect the client to 88 | send a "group" header identifying it as 'live' or 'batch'. Ideally this should be done using TLS certificates and a 89 | server side lookup of identity to grouping. Any requests not identified as either live or batch may only use excess 90 | capacity. 91 | 92 | ```golang 93 | import ( 94 | gclGrpc "github.com/platnummonkey/go-concurrency-limits/grpc" 95 | ) 96 | 97 | // setup grpc server with this option 98 | serverOption := grpc.UnaryInterceptor( 99 | gclGrpc.UnaryServerInterceptor( 100 | gclGrpc.WithLimiter(...), 101 | gclGrpc.WithServerResponseTypeClassifier(..), 102 | ), 103 | ) 104 | ``` 105 | 106 | ### Client Limiter 107 | 108 | There are two main use cases for client side limiters. A client side limiter can protect the client service from its 109 | dependent services by failing fast and serving a degraded experience to its client instead of having its latency go up 110 | and its resources eventually exhausted. For batch applications that call other services a client side limiter acts as a 111 | backpressure mechanism ensuring that the batch application does not put unnecessary load on dependent services. 112 | 113 | In this example a GRPC client will use a blocking version of the VegasLimit to block the caller when the limit has been 114 | reached. 115 | 116 | ```golang 117 | import ( 118 | gclGrpc "github.com/platnummonkey/go-concurrency-limits/grpc" 119 | ) 120 | 121 | // setup grpc client with this option 122 | dialOption := grpc.WithUnaryInterceptor( 123 | gclGrpc.UnaryClientInterceptor( 124 | gclGrpc.WithLimiter(...), 125 | gclGrpc.WithClientResponseTypeClassifier(...), 126 | ), 127 | ) 128 | ``` 129 | 130 | # References Used 131 | 1. Original Java implementation - Netflix - https://github.com/netflix/concurrency-limits/ 132 | 1. Windowless Moving Percentile - Martin Jambon - https://mjambon.com/2016-07-23-moving-percentile/ 133 | -------------------------------------------------------------------------------- /core/dep.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "strings" 5 | ) 6 | 7 | const ( 8 | // MetricLimit is the name of the metric for current limit 9 | MetricLimit = "limit" 10 | // MetricDropped is the name of the metric for dropped counts 11 | MetricDropped = "dropped" 12 | // MetricInFlight is the name of the metric for current in flight count 13 | MetricInFlight = "inflight" 14 | // MetricPartitionLimit is the name of the metric for a current partition's limit 15 | MetricPartitionLimit = "limit.partition" 16 | // MetricRTT is the name of the metric for the sample Round Trip Time distribution 17 | MetricRTT = "rtt" 18 | // MetricMinRTT is the name of the metric for the Minimum Round Trip Time 19 | MetricMinRTT = "min_rtt" 20 | // MetricWindowMinRTT is the name of the metric for the Window's Minimum Round Trip Time 21 | MetricWindowMinRTT = "window.min_rtt" 22 | // MetricWindowQueueSize represents the name of the metric for the Window's Queue Size 23 | MetricWindowQueueSize = "window.queue_size" 24 | // MetricQueueSize represents the name of the metric for the size of a lifo queue 25 | MetricQueueSize = "queue_size" 26 | // MetricQueueLimit represents the name of the metric for the max size of a lifo queue 27 | MetricQueueLimit = "queue_limit" 28 | ) 29 | 30 | // PrefixMetricWithName will prefix a given name with the metric name in the form "." 31 | var PrefixMetricWithName = func(metric, name string) string { 32 | if name == "" { 33 | name = "default" 34 | } 35 | 36 | if strings.HasSuffix(name, ".") { 37 | return name + metric 38 | } 39 | return name + "." + metric 40 | } 41 | -------------------------------------------------------------------------------- /core/doc.go: -------------------------------------------------------------------------------- 1 | // Package core provides the package interfaces. 2 | package core 3 | -------------------------------------------------------------------------------- /core/limit.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "context" 5 | ) 6 | 7 | // MeasurementInterface defines the contract for tracking a measurement such as a minimum or average of a sample set. 8 | type MeasurementInterface interface { 9 | // Add a single sample and update the internal state. 10 | // returns true if the internal state was updated, also return the current value. 11 | Add(value float64) (float64, bool) 12 | 13 | // Get the current value. 14 | Get() float64 15 | 16 | // Reset the internal state as if no samples were ever added. 17 | Reset() 18 | 19 | // Update will update the value given an operation function 20 | Update(operation func(value float64) float64) 21 | } 22 | 23 | // SampleWindow represents the details of the current sample window 24 | type SampleWindow interface { 25 | // StartTimeNanoseconds returns the epoch start time in nanoseconds. 26 | StartTimeNanoseconds() int64 27 | // CandidateRTTNanoseconds returns the candidate RTT in the sample window. This is traditionally the minimum rtt. 28 | CandidateRTTNanoseconds() int64 29 | // AverageRTTNanoseconds returns the average RTT in the sample window. Excludes timeouts and dropped rtt. 30 | AverageRTTNanoseconds() int64 31 | // MaxInFlight returns the maximum number of in-flight observed during the sample window. 32 | MaxInFlight() int 33 | // SampleCount is the number of observed RTTs in the sample window. 34 | SampleCount() int 35 | // DidDrop returns True if there was a timeout. 36 | DidDrop() bool 37 | } 38 | 39 | // LimitChangeListener is a callback method to receive a notification whenever the limit is updated to a new value. 40 | type LimitChangeListener func(limit int) 41 | 42 | // Limit is a Contract for an algorithm that calculates a concurrency limit based on rtt measurements. 43 | type Limit interface { 44 | // EstimatedLimit returns the current estimated limit. 45 | EstimatedLimit() int 46 | 47 | // NotifyOnChange will register a callback to receive notification whenever the limit is updated to a new value. 48 | // 49 | // consumer - the callback 50 | NotifyOnChange(consumer LimitChangeListener) 51 | 52 | // OnSample the concurrency limit using a new rtt sample. 53 | // 54 | // startTime - in epoch nanoseconds 55 | // rtt - round trip time of sample 56 | // inFlight - in flight observed count during the sample 57 | // didDrop - true if there was a timeout 58 | OnSample(startTime int64, rtt int64, inFlight int, didDrop bool) 59 | } 60 | 61 | // Listener implements token listener for callback to the limiter when and how it should be released. 62 | type Listener interface { 63 | // OnSuccess is called as a notification that the operation succeeded and internally measured latency should be 64 | // used as an RTT sample. 65 | OnSuccess() 66 | // OnIgnore is called to indicate the operation failed before any meaningful RTT measurement could be made and 67 | // should be ignored to not introduce an artificially low RTT. 68 | OnIgnore() 69 | // OnDropped is called to indicate the request failed and was dropped due to being rejected by an external limit or 70 | // hitting a timeout. Loss based Limit implementations will likely do an aggressive reducing in limit when this 71 | // happens. 72 | OnDropped() 73 | } 74 | 75 | // Limiter defines the contract for a concurrency limiter. The caller is expected to call acquire() for each request 76 | // and must also release the returned listener when the operation completes. Releasing the Listener 77 | // may trigger an update to the concurrency limit based on error rate or latency measurement. 78 | type Limiter interface { 79 | // Acquire a token from the limiter. Returns a nil listener if the limit has been exceeded. 80 | // If acquired the caller must call one of the Listener methods when the operation has been completed to release 81 | // the count. 82 | // 83 | // context - Context for the request. The context is used by advanced strategies such as LookupPartitionStrategy. 84 | Acquire(ctx context.Context) (listener Listener, ok bool) 85 | } 86 | -------------------------------------------------------------------------------- /core/limiter_registry.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | // LimiterRegistry lookup for integrations that support multiple Limiters, i.e. one per RPC method. 4 | type LimiterRegistry interface { 5 | // Get a limiter given a key. 6 | Get(key string) Limiter 7 | } 8 | -------------------------------------------------------------------------------- /core/metric_registry.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | // MetricSampleListener is a listener to receive samples for a distribution 4 | type MetricSampleListener interface { 5 | // AddSample will add a sample metric to the listener 6 | AddSample(value float64, tags ...string) 7 | } 8 | 9 | // EmptyMetricSampleListener implements a sample listener that ignores everything. 10 | type EmptyMetricSampleListener struct{} 11 | 12 | // AddSample will add a metric sample to this listener 13 | func (*EmptyMetricSampleListener) AddSample(value float64, tags ...string) { 14 | // noop 15 | } 16 | 17 | // MetricSupplier will return the supplied metric value 18 | type MetricSupplier func() (value float64, ok bool) 19 | 20 | // NewIntMetricSupplierWrapper will wrap a int-return value func to a supplier func 21 | func NewIntMetricSupplierWrapper(s func() int) MetricSupplier { 22 | return MetricSupplier(func() (float64, bool) { 23 | val := s() 24 | return float64(val), true 25 | }) 26 | } 27 | 28 | // NewUint64MetricSupplierWrapper will wrap a uint64-return value func to a supplier func 29 | func NewUint64MetricSupplierWrapper(s func() uint64) MetricSupplier { 30 | return MetricSupplier(func() (float64, bool) { 31 | val := s() 32 | return float64(val), true 33 | }) 34 | } 35 | 36 | // NewFloat64MetricSupplierWrapper will wrap a float64-return value func to a supplier func 37 | func NewFloat64MetricSupplierWrapper(s func() float64) MetricSupplier { 38 | return MetricSupplier(func() (float64, bool) { 39 | val := s() 40 | return val, true 41 | }) 42 | } 43 | 44 | // MetricRegistry is a simple abstraction for tracking metrics in the limiters. 45 | type MetricRegistry interface { 46 | // RegisterDistribution will register a sample distribution. Samples are added to the distribution via the returned 47 | // MetricSampleListener. Will reuse an existing MetricSampleListener if the distribution already exists. 48 | RegisterDistribution(ID string, tags ...string) MetricSampleListener 49 | 50 | // RegisterTiming will register a sample timing distribution. Samples are added to the distribution via the 51 | // returned MetricSampleListener. Will reuse an existing MetricSampleListener if the distribution already exists. 52 | RegisterTiming(ID string, tags ...string) MetricSampleListener 53 | 54 | // RegisterCount will register a sample counter. Samples are added to the counter via the returned 55 | // MetricSampleListener. Will reuse an existing MetricSampleListener if the counter already exists. 56 | RegisterCount(ID string, tags ...string) MetricSampleListener 57 | 58 | // RegisterGauge will register a gauge using the provided supplier. The supplier will be polled whenever the gauge 59 | // value is flushed by the registry. 60 | RegisterGauge(ID string, supplier MetricSupplier, tags ...string) 61 | 62 | // Start will start the metric registry polling 63 | Start() 64 | 65 | // Stop will stop the metric registry polling 66 | Stop() 67 | } 68 | 69 | // EmptyMetricRegistry implements a void reporting metric registry 70 | type EmptyMetricRegistry struct{} 71 | 72 | // EmptyMetricRegistryInstance is a singleton empty metric registry instance. 73 | var EmptyMetricRegistryInstance = &EmptyMetricRegistry{} 74 | 75 | // RegisterDistribution will register a distribution sample to this registry 76 | func (*EmptyMetricRegistry) RegisterDistribution(ID string, tags ...string) MetricSampleListener { 77 | return &EmptyMetricSampleListener{} 78 | } 79 | 80 | // RegisterTiming will register a timing distribution sample to this registry 81 | func (*EmptyMetricRegistry) RegisterTiming(ID string, tags ...string) MetricSampleListener { 82 | return &EmptyMetricSampleListener{} 83 | } 84 | 85 | // RegisterCount will register a count sample to this registry 86 | func (*EmptyMetricRegistry) RegisterCount(ID string, tags ...string) MetricSampleListener { 87 | return &EmptyMetricSampleListener{} 88 | } 89 | 90 | // RegisterGauge will register a gauge sample to this registry 91 | func (*EmptyMetricRegistry) RegisterGauge(ID string, supplier MetricSupplier, tags ...string) {} 92 | 93 | // Start will start the metric registry polling 94 | func (*EmptyMetricRegistry) Start() {} 95 | 96 | // Stop will stop the metric registry polling 97 | func (*EmptyMetricRegistry) Stop() {} 98 | 99 | // CommonMetricSampler is a set of common metrics reported by all Limit implementations 100 | type CommonMetricSampler struct { 101 | RTTListener MetricSampleListener 102 | DropCounterListener MetricSampleListener 103 | InFlightListener MetricSampleListener 104 | } 105 | 106 | // NewCommonMetricSamplerOrNil will only create a new CommonMetricSampler if a valid registry is supplied 107 | func NewCommonMetricSamplerOrNil(registry MetricRegistry, limit Limit, name string, tags ...string) *CommonMetricSampler { 108 | if registry == nil { 109 | return nil 110 | } 111 | if _, ok := registry.(*EmptyMetricRegistry); ok { 112 | return nil 113 | } 114 | return NewCommonMetricSampler(registry, limit, name, tags...) 115 | } 116 | 117 | // NewCommonMetricSampler will create a new CommonMetricSampler that will auto-instrument metrics 118 | func NewCommonMetricSampler(registry MetricRegistry, limit Limit, name string, tags ...string) *CommonMetricSampler { 119 | if registry == nil { 120 | registry = EmptyMetricRegistryInstance 121 | } 122 | 123 | registry.RegisterGauge( 124 | PrefixMetricWithName(MetricLimit, name), 125 | NewIntMetricSupplierWrapper(limit.EstimatedLimit), 126 | tags..., 127 | ) 128 | 129 | return &CommonMetricSampler{ 130 | RTTListener: registry.RegisterTiming(PrefixMetricWithName(MetricRTT, name), tags...), 131 | DropCounterListener: registry.RegisterCount(PrefixMetricWithName(MetricDropped, name), tags...), 132 | InFlightListener: registry.RegisterDistribution(PrefixMetricWithName(MetricInFlight, name), tags...), 133 | } 134 | } 135 | 136 | // Sample will sample the current sample for metric reporting. 137 | func (s *CommonMetricSampler) Sample(rtt int64, inFlight int, didDrop bool) { 138 | // from noop metrics registry 139 | if s == nil { 140 | return 141 | } 142 | 143 | if didDrop { 144 | s.DropCounterListener.AddSample(1.0) 145 | } 146 | s.RTTListener.AddSample(float64(rtt)) 147 | s.InFlightListener.AddSample(float64(inFlight)) 148 | } 149 | -------------------------------------------------------------------------------- /core/strategy.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "context" 5 | ) 6 | 7 | // StrategyToken represents a token from a limiter algorithm 8 | type StrategyToken interface { 9 | // IsAcquired returns true if acquired or false if limit has been reached. 10 | IsAcquired() bool 11 | // InFlightCount will return the number of pending requests. 12 | InFlightCount() int 13 | // Release the acquired token and decrement the current in-flight count. 14 | Release() 15 | } 16 | 17 | // StaticStrategyToken represents a static strategy token, simple but flexible. 18 | type StaticStrategyToken struct { 19 | acquired bool 20 | inFlightCount int 21 | releaseFunc func() 22 | } 23 | 24 | // IsAcquired will return true if the token is acquired 25 | func (t *StaticStrategyToken) IsAcquired() bool { 26 | return t.acquired 27 | } 28 | 29 | // InFlightCount represents the instantaneous snapshot on token creation in-flight count 30 | func (t *StaticStrategyToken) InFlightCount() int { 31 | return t.inFlightCount 32 | } 33 | 34 | // Release will release the current token, it's very important to release all tokens! 35 | func (t *StaticStrategyToken) Release() { 36 | if t.releaseFunc != nil { 37 | t.releaseFunc() 38 | } 39 | } 40 | 41 | // NewNotAcquiredStrategyToken will create a new un-acquired strategy token. 42 | func NewNotAcquiredStrategyToken(inFlightCount int) StrategyToken { 43 | return &StaticStrategyToken{ 44 | acquired: false, 45 | inFlightCount: inFlightCount, 46 | releaseFunc: func() {}, 47 | } 48 | } 49 | 50 | // NewAcquiredStrategyToken will create a new acquired strategy token. 51 | func NewAcquiredStrategyToken(inFlightCount int, releaseFunc func()) StrategyToken { 52 | return &StaticStrategyToken{ 53 | acquired: true, 54 | inFlightCount: inFlightCount, 55 | releaseFunc: releaseFunc, 56 | } 57 | } 58 | 59 | // Strategy defines how the limiter logic should acquire or not acquire tokens. 60 | type Strategy interface { 61 | // TryAcquire will try to acquire a token from the limiter. 62 | // context Context of the request for partitioned limits. 63 | // returns not ok if limit is exceeded, or a StrategyToken that must be released when the operation completes. 64 | TryAcquire(ctx context.Context) (token StrategyToken, ok bool) 65 | 66 | // SetLimit will update the strategy with a new limit. 67 | SetLimit(limit int) 68 | } 69 | -------------------------------------------------------------------------------- /doc.go: -------------------------------------------------------------------------------- 1 | // Package provides primitives for concurrency control in complex systems. 2 | package main 3 | -------------------------------------------------------------------------------- /examples/doc.go: -------------------------------------------------------------------------------- 1 | // Package examples contains examples of using this package to solve concurrency problems. 2 | package examples 3 | -------------------------------------------------------------------------------- /examples/example_blocking_limit/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "math/rand" 8 | "os" 9 | "sync" 10 | "time" 11 | 12 | "golang.org/x/time/rate" 13 | 14 | "github.com/platinummonkey/go-concurrency-limits/core" 15 | "github.com/platinummonkey/go-concurrency-limits/limit" 16 | "github.com/platinummonkey/go-concurrency-limits/limiter" 17 | "github.com/platinummonkey/go-concurrency-limits/strategy" 18 | ) 19 | 20 | type contextKey string 21 | 22 | const testContextKey contextKey = "jobID" 23 | 24 | type resource struct { 25 | limiter *rate.Limiter 26 | } 27 | 28 | func (r *resource) poll(ctx context.Context) (bool, error) { 29 | id := ctx.Value(testContextKey).(int) 30 | log.Printf("request started for id=%d\n", id) 31 | if !r.limiter.Allow() { 32 | time.Sleep(time.Millisecond * 10) 33 | return false, fmt.Errorf("limit exceeded for id=%d", id) 34 | } 35 | // sleep some time 36 | time.Sleep(time.Second * time.Duration(rand.Intn(2))) 37 | log.Printf("request succeeded for id=%d\n", id) 38 | return true, nil 39 | } 40 | 41 | type protectedResource struct { 42 | external *resource 43 | guard core.Limiter 44 | } 45 | 46 | func (r *protectedResource) poll(ctx context.Context) (bool, error) { 47 | id := ctx.Value(testContextKey).(int) 48 | log.Printf("guarded request started for id=%d\n", id) 49 | token, ok := r.guard.Acquire(ctx) 50 | if !ok { 51 | // short circuit no need to try 52 | log.Printf("guarded request short circuited for id=%d\n", id) 53 | if token != nil { 54 | token.OnDropped() 55 | } 56 | return false, fmt.Errorf("short circuited request id=%d", id) 57 | } 58 | 59 | // try to make request 60 | _, err := r.external.poll(ctx) 61 | if err != nil { 62 | token.OnDropped() 63 | log.Printf("guarded request failed for id=%d err=%v\n", id, err) 64 | return false, fmt.Errorf("request failed err=%v", err) 65 | } 66 | token.OnSuccess() 67 | log.Printf("guarded request succeeded for id=%d\n", id) 68 | return true, nil 69 | } 70 | 71 | func main() { 72 | limitStrategy := strategy.NewSimpleStrategy(10) 73 | logger := limit.BuiltinLimitLogger{} 74 | defaultLimiter, err := limiter.NewDefaultLimiterWithDefaults( 75 | "example_blocking_limit", 76 | limitStrategy, 77 | logger, 78 | core.EmptyMetricRegistryInstance, 79 | ) 80 | externalResourceLimiter := limiter.NewBlockingLimiter(defaultLimiter, 0, logger) 81 | 82 | if err != nil { 83 | log.Fatalf("Error creating limiter err=%v\n", err) 84 | os.Exit(-1) 85 | } 86 | 87 | fakeExternalResource := &resource{ 88 | limiter: rate.NewLimiter(5, 15), 89 | } 90 | 91 | guardedResource := &protectedResource{ 92 | external: fakeExternalResource, 93 | guard: externalResourceLimiter, 94 | } 95 | 96 | endOfExampleTimer := time.NewTimer(time.Second * 10) 97 | ticker := time.NewTicker(time.Millisecond * 500) 98 | counter := 0 99 | wg := sync.WaitGroup{} 100 | 101 | for { 102 | select { 103 | case <-endOfExampleTimer.C: 104 | log.Printf("Waiting for go-routines to finish...") 105 | wg.Wait() 106 | return 107 | case <-ticker.C: 108 | // make a few requests 109 | wg.Add(5) 110 | go func(c int) { 111 | for i := 0; i < 5; i++ { 112 | defer wg.Done() 113 | ctx := context.WithValue(context.Background(), testContextKey, c+i) 114 | guardedResource.poll(ctx) 115 | } 116 | }(counter) 117 | } 118 | counter += 5 119 | } 120 | } 121 | -------------------------------------------------------------------------------- /examples/example_concurrent_loading/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "math" 8 | "math/rand" 9 | "os" 10 | "sync" 11 | "sync/atomic" 12 | "time" 13 | 14 | "github.com/platinummonkey/go-concurrency-limits/core" 15 | "github.com/platinummonkey/go-concurrency-limits/limit" 16 | "github.com/platinummonkey/go-concurrency-limits/limiter" 17 | "github.com/platinummonkey/go-concurrency-limits/strategy" 18 | ) 19 | 20 | type contextKey string 21 | 22 | const testContextKey contextKey = "jobID" 23 | 24 | type resource struct { 25 | counter *int64 26 | } 27 | 28 | func (r *resource) poll(ctx context.Context) (bool, error) { 29 | currentCount := atomic.AddInt64(r.counter, 1) 30 | id := ctx.Value(testContextKey).(int) 31 | log.Printf("request started for id=%d currentCount=%d\n", id, currentCount) 32 | numKeys := rand.Int63n(1000) 33 | 34 | // fake a scan, every key lookup additional non-linear time 35 | scanTime := time.Duration(numKeys)*time.Nanosecond + time.Duration(int64(math.Exp(float64(numKeys)/100.0)))*time.Millisecond 36 | scanTimeMillis := scanTime.Milliseconds() 37 | 38 | // sleep some time 39 | time.Sleep(scanTime) 40 | currentCount = atomic.AddInt64(r.counter, -1) 41 | log.Printf("request succeeded for id=%d currentCount=%d scanTime=%d ms\n", id, currentCount, scanTimeMillis) 42 | return true, nil 43 | } 44 | 45 | type protectedResource struct { 46 | external *resource 47 | guard core.Limiter 48 | } 49 | 50 | func (r *protectedResource) poll(ctx context.Context) (bool, error) { 51 | id := ctx.Value(testContextKey).(int) 52 | log.Printf("guarded request started for id=%d\n", id) 53 | token, ok := r.guard.Acquire(ctx) 54 | if !ok { 55 | // short circuit no need to try 56 | log.Printf("guarded request short circuited for id=%d\n", id) 57 | if token != nil { 58 | token.OnDropped() 59 | } 60 | return false, fmt.Errorf("short circuited request id=%d", id) 61 | } 62 | 63 | // try to make request 64 | _, err := r.external.poll(ctx) 65 | if err != nil { 66 | token.OnDropped() 67 | log.Printf("guarded request failed for id=%d err=%v\n", id, err) 68 | return false, fmt.Errorf("request failed err=%v", err) 69 | } 70 | token.OnSuccess() 71 | log.Printf("guarded request succeeded for id=%d\n", id) 72 | return true, nil 73 | } 74 | 75 | func main() { 76 | l := 1000 77 | limitStrategy := strategy.NewSimpleStrategy(l) 78 | logger := limit.BuiltinLimitLogger{} 79 | defaultLimiter, err := limiter.NewDefaultLimiter( 80 | limit.NewFixedLimit( 81 | "initializer_limiter", 82 | l, 83 | core.EmptyMetricRegistryInstance, 84 | ), 85 | int64(time.Millisecond*250), 86 | int64(time.Millisecond*500), 87 | int64(time.Millisecond*10), 88 | 100, 89 | limitStrategy, 90 | logger, 91 | core.EmptyMetricRegistryInstance, 92 | ) 93 | externalResourceLimiter := limiter.NewBlockingLimiter(defaultLimiter, time.Second, logger) 94 | 95 | if err != nil { 96 | log.Fatalf("Error creating limiter err=%v\n", err) 97 | os.Exit(-1) 98 | } 99 | 100 | initialCount := int64(0) 101 | fakeExternalResource := &resource{ 102 | counter: &initialCount, 103 | } 104 | 105 | guardedResource := &protectedResource{ 106 | external: fakeExternalResource, 107 | guard: externalResourceLimiter, 108 | } 109 | atomic.StoreInt64(fakeExternalResource.counter, 0) 110 | 111 | endOfExampleTimer := time.NewTimer(time.Second * 10) 112 | wg := sync.WaitGroup{} 113 | 114 | // spin up 10*l consumers 115 | wg.Add(10 * l) 116 | for i := 0; i < 10*l; i++ { 117 | go func(c int) { 118 | for i := 0; i < 5; i++ { 119 | defer wg.Done() 120 | ctx := context.WithValue(context.Background(), testContextKey, c+i) 121 | guardedResource.poll(ctx) 122 | } 123 | }(i) 124 | } 125 | 126 | <-endOfExampleTimer.C 127 | log.Printf("Waiting for go-routines to finish...") 128 | wg.Wait() 129 | } 130 | -------------------------------------------------------------------------------- /examples/example_simple_limit/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "math/rand" 8 | "os" 9 | "sync" 10 | "time" 11 | 12 | "golang.org/x/time/rate" 13 | 14 | "github.com/platinummonkey/go-concurrency-limits/core" 15 | "github.com/platinummonkey/go-concurrency-limits/limit" 16 | "github.com/platinummonkey/go-concurrency-limits/limiter" 17 | "github.com/platinummonkey/go-concurrency-limits/strategy" 18 | ) 19 | 20 | type contextKey uint8 21 | 22 | const testContextKey = contextKey(1) 23 | 24 | type resource struct { 25 | limiter *rate.Limiter 26 | } 27 | 28 | func (r *resource) poll(ctx context.Context) (bool, error) { 29 | id := ctx.Value(testContextKey).(int) 30 | log.Printf("request started for id=%d\n", id) 31 | if !r.limiter.Allow() { 32 | time.Sleep(time.Millisecond * 10) 33 | return false, fmt.Errorf("limit exceeded for id=%d", id) 34 | } 35 | // sleep some time 36 | time.Sleep(time.Millisecond * time.Duration(rand.Intn(90)+10)) 37 | log.Printf("request succeeded for id=%d\n", id) 38 | return true, nil 39 | } 40 | 41 | type protectedResource struct { 42 | external *resource 43 | guard core.Limiter 44 | } 45 | 46 | func (r *protectedResource) poll(ctx context.Context) (bool, error) { 47 | id := ctx.Value(testContextKey).(int) 48 | log.Printf("guarded request started for id=%d\n", id) 49 | token, ok := r.guard.Acquire(ctx) 50 | if !ok { 51 | // short circuit no need to try 52 | log.Printf("guarded request short circuited for id=%d\n", id) 53 | if token != nil { 54 | token.OnDropped() 55 | } 56 | return false, fmt.Errorf("short circuited request id=%d", id) 57 | } 58 | 59 | // try to make request 60 | _, err := r.external.poll(ctx) 61 | if err != nil { 62 | token.OnDropped() 63 | log.Printf("guarded request failed for id=%d err=%v\n", id, err) 64 | return false, fmt.Errorf("request failed err=%v", err) 65 | } 66 | token.OnSuccess() 67 | log.Printf("guarded request succeeded for id=%d\n", id) 68 | return true, nil 69 | } 70 | 71 | func main() { 72 | limitStrategy := strategy.NewSimpleStrategy(10) 73 | externalResourceLimiter, err := limiter.NewDefaultLimiterWithDefaults( 74 | "example_single_limit", 75 | limitStrategy, 76 | limit.BuiltinLimitLogger{}, 77 | core.EmptyMetricRegistryInstance, 78 | ) 79 | if err != nil { 80 | log.Fatalf("Error creating limiter err=%v\n", err) 81 | os.Exit(-1) 82 | } 83 | 84 | fakeExternalResource := &resource{ 85 | limiter: rate.NewLimiter(5, 15), 86 | } 87 | 88 | guardedResource := &protectedResource{ 89 | external: fakeExternalResource, 90 | guard: externalResourceLimiter, 91 | } 92 | 93 | endOfExampleTimer := time.NewTimer(time.Second * 10) 94 | ticker := time.NewTicker(time.Millisecond * 100) 95 | wg := sync.WaitGroup{} 96 | counter := 0 97 | for { 98 | select { 99 | case <-endOfExampleTimer.C: 100 | log.Printf("Waiting for goroutines to finish...") 101 | wg.Wait() 102 | return 103 | case <-ticker.C: 104 | // make a few requests 105 | wg.Add(5) 106 | go func(c int) { 107 | for i := 0; i < 5; i++ { 108 | defer wg.Done() 109 | ctx := context.WithValue(context.Background(), testContextKey, c+i) 110 | guardedResource.poll(ctx) 111 | } 112 | }(counter) 113 | } 114 | counter += 5 115 | } 116 | } 117 | -------------------------------------------------------------------------------- /examples/grpc_streaming/.gitignore: -------------------------------------------------------------------------------- 1 | grpc_streaming 2 | -------------------------------------------------------------------------------- /examples/grpc_streaming/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | //go:generate protoc -I ./pb --go_out=plugins=grpc:${GOPATH}/src *.proto 4 | 5 | import ( 6 | "context" 7 | "flag" 8 | "fmt" 9 | "log" 10 | "net" 11 | "sync" 12 | "time" 13 | 14 | golangGrpc "google.golang.org/grpc" 15 | 16 | "github.com/platinummonkey/go-concurrency-limits/examples/grpc_streaming/pb" 17 | "github.com/platinummonkey/go-concurrency-limits/grpc" 18 | "github.com/platinummonkey/go-concurrency-limits/limit" 19 | "github.com/platinummonkey/go-concurrency-limits/limiter" 20 | "github.com/platinummonkey/go-concurrency-limits/strategy" 21 | ) 22 | 23 | var options = struct { 24 | mode string 25 | port int 26 | numThreads int 27 | }{ 28 | mode: "server", 29 | port: 8080, 30 | numThreads: 105, 31 | } 32 | 33 | func init() { 34 | flag.StringVar(&options.mode, "mode", options.mode, "choose `client` or `server` mode") 35 | flag.IntVar(&options.port, "port", options.port, "grpc port") 36 | flag.IntVar(&options.numThreads, "threads", options.numThreads, "number of client threads") 37 | } 38 | 39 | func checkOptions() { 40 | switch options.mode { 41 | case "client": 42 | fallthrough 43 | case "server": 44 | // no-op 45 | default: 46 | panic(fmt.Sprintf("invalid mode specified: '%s'", options.mode)) 47 | } 48 | } 49 | 50 | type server struct { 51 | } 52 | 53 | func (s *server) PingPong(ss pb.PingPong_PingPongServer) error { 54 | ping, err := ss.Recv() 55 | if err != nil { 56 | log.Printf("Recv Error: %v", err) 57 | return nil 58 | } 59 | log.Printf("Received: '%s'", ping.GetMessage()) 60 | // pretend to do some work 61 | time.Sleep(time.Millisecond * 10) 62 | err = ss.Send(&pb.Pong{Message: ping.GetMessage()}) 63 | if err != nil { 64 | log.Printf("Send Error: %v", err) 65 | } 66 | return nil 67 | } 68 | 69 | func main() { 70 | flag.Parse() 71 | checkOptions() 72 | switch options.mode { 73 | case "server": 74 | runServer() 75 | case "client": 76 | runClient() 77 | } 78 | } 79 | 80 | func runServer() { 81 | logger := limit.BuiltinLimitLogger{} 82 | lis, err := net.Listen("tcp", fmt.Sprintf(":%d", options.port)) 83 | if err != nil { 84 | panic(err) 85 | } 86 | 87 | serverLimitSend := limit.NewFixedLimit("server-fixed-limit-send", 1000, nil) 88 | serverLimiterSend, err := limiter.NewDefaultLimiter(serverLimitSend, 1000, 10000, 1e5, 1000, strategy.NewSimpleStrategy(1000), logger, nil) 89 | if err != nil { 90 | panic(err) 91 | } 92 | serverLimitRecv := limit.NewFixedLimit("server-fixed-limit-recv", 10, nil) 93 | serverLimiterRecv, err := limiter.NewDefaultLimiter(serverLimitRecv, 1000, 10000, 1e5, 1000, strategy.NewSimpleStrategy(10), logger, nil) 94 | if err != nil { 95 | panic(err) 96 | } 97 | serverOpts := []grpc.StreamInterceptorOption{ 98 | grpc.WithStreamSendName("grpc-stream-server-send"), 99 | grpc.WithStreamRecvName("grpc-stream-server-recv"), 100 | grpc.WithStreamSendLimiter(serverLimiterSend), // outbound guard 101 | grpc.WithStreamRecvLimiter(serverLimiterRecv), // inbound guard 102 | } 103 | serverInterceptor := grpc.StreamServerInterceptor(serverOpts...) 104 | svc := golangGrpc.NewServer(golangGrpc.StreamInterceptor(serverInterceptor)) 105 | s := &server{} 106 | pb.RegisterPingPongServer(svc, s) 107 | if err := svc.Serve(lis); err != nil { 108 | panic(nil) 109 | } 110 | } 111 | 112 | func resetConnection() (*golangGrpc.ClientConn, pb.PingPong_PingPongClient) { 113 | conn, err := golangGrpc.Dial(fmt.Sprintf("localhost:%d", options.port), golangGrpc.WithInsecure()) 114 | if err != nil { 115 | panic(err) 116 | } 117 | clientConn := pb.NewPingPongClient(conn) 118 | ctx := context.Background() 119 | client, err := clientConn.PingPong(ctx) 120 | if err != nil { 121 | panic(err) 122 | } 123 | return conn, client 124 | } 125 | 126 | func runClient() { 127 | wg := sync.WaitGroup{} 128 | wg.Add(options.numThreads) 129 | for i := 0; i < options.numThreads; i++ { 130 | go func(workerID int) { 131 | conn, client := resetConnection() 132 | j := 0 133 | for { 134 | // do this as fast as possible 135 | err := queryServer(client, workerID, j) 136 | if err != nil { 137 | client.CloseSend() 138 | conn.Close() 139 | conn, client = resetConnection() 140 | } 141 | j++ 142 | } 143 | }(i) 144 | } 145 | wg.Wait() 146 | } 147 | 148 | func queryServer(client pb.PingPong_PingPongClient, workerID int, i int) error { 149 | msg := &pb.Ping{ 150 | Message: fmt.Sprintf("hello %d from %d", i, workerID), 151 | } 152 | err := client.Send(msg) 153 | if err != nil { 154 | log.Printf("[failed](%d - %d)\t - %v", workerID, i, err) 155 | return err 156 | } 157 | pong, err := client.Recv() 158 | if err != nil { 159 | log.Printf("[failed](%d - %d)\t - %v", workerID, i, err) 160 | } else { 161 | log.Printf("[pass](%d - %d)\t - %s", workerID, i, pong.GetMessage()) 162 | } 163 | return err 164 | } 165 | -------------------------------------------------------------------------------- /examples/grpc_streaming/pb/pingpong.pb.go: -------------------------------------------------------------------------------- 1 | // Code generated by protoc-gen-go. DO NOT EDIT. 2 | // source: pingpong.proto 3 | 4 | /* 5 | Package pb is a generated protocol buffer package. 6 | 7 | It is generated from these files: 8 | pingpong.proto 9 | 10 | It has these top-level messages: 11 | Ping 12 | Pong 13 | */ 14 | package pb 15 | 16 | import proto "github.com/golang/protobuf/proto" 17 | import fmt "fmt" 18 | import math "math" 19 | 20 | import ( 21 | context "golang.org/x/net/context" 22 | grpc "google.golang.org/grpc" 23 | ) 24 | 25 | // Reference imports to suppress errors if they are not otherwise used. 26 | var _ = proto.Marshal 27 | var _ = fmt.Errorf 28 | var _ = math.Inf 29 | 30 | // This is a compile-time assertion to ensure that this generated file 31 | // is compatible with the proto package it is being compiled against. 32 | // A compilation error at this line likely means your copy of the 33 | // proto package needs to be updated. 34 | const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package 35 | 36 | type Ping struct { 37 | Message string `protobuf:"bytes,1,opt,name=message" json:"message,omitempty"` 38 | } 39 | 40 | func (m *Ping) Reset() { *m = Ping{} } 41 | func (m *Ping) String() string { return proto.CompactTextString(m) } 42 | func (*Ping) ProtoMessage() {} 43 | func (*Ping) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} } 44 | 45 | func (m *Ping) GetMessage() string { 46 | if m != nil { 47 | return m.Message 48 | } 49 | return "" 50 | } 51 | 52 | type Pong struct { 53 | Message string `protobuf:"bytes,1,opt,name=message" json:"message,omitempty"` 54 | } 55 | 56 | func (m *Pong) Reset() { *m = Pong{} } 57 | func (m *Pong) String() string { return proto.CompactTextString(m) } 58 | func (*Pong) ProtoMessage() {} 59 | func (*Pong) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{1} } 60 | 61 | func (m *Pong) GetMessage() string { 62 | if m != nil { 63 | return m.Message 64 | } 65 | return "" 66 | } 67 | 68 | func init() { 69 | proto.RegisterType((*Ping)(nil), "main.Ping") 70 | proto.RegisterType((*Pong)(nil), "main.Pong") 71 | } 72 | 73 | // Reference imports to suppress errors if they are not otherwise used. 74 | var _ context.Context 75 | var _ grpc.ClientConn 76 | 77 | // This is a compile-time assertion to ensure that this generated file 78 | // is compatible with the grpc package it is being compiled against. 79 | const _ = grpc.SupportPackageIsVersion4 80 | 81 | // Client API for PingPong service 82 | 83 | type PingPongClient interface { 84 | PingPong(ctx context.Context, opts ...grpc.CallOption) (PingPong_PingPongClient, error) 85 | } 86 | 87 | type pingPongClient struct { 88 | cc *grpc.ClientConn 89 | } 90 | 91 | func NewPingPongClient(cc *grpc.ClientConn) PingPongClient { 92 | return &pingPongClient{cc} 93 | } 94 | 95 | func (c *pingPongClient) PingPong(ctx context.Context, opts ...grpc.CallOption) (PingPong_PingPongClient, error) { 96 | stream, err := grpc.NewClientStream(ctx, &_PingPong_serviceDesc.Streams[0], c.cc, "/main.PingPong/PingPong", opts...) 97 | if err != nil { 98 | return nil, err 99 | } 100 | x := &pingPongPingPongClient{stream} 101 | return x, nil 102 | } 103 | 104 | type PingPong_PingPongClient interface { 105 | Send(*Ping) error 106 | Recv() (*Pong, error) 107 | grpc.ClientStream 108 | } 109 | 110 | type pingPongPingPongClient struct { 111 | grpc.ClientStream 112 | } 113 | 114 | func (x *pingPongPingPongClient) Send(m *Ping) error { 115 | return x.ClientStream.SendMsg(m) 116 | } 117 | 118 | func (x *pingPongPingPongClient) Recv() (*Pong, error) { 119 | m := new(Pong) 120 | if err := x.ClientStream.RecvMsg(m); err != nil { 121 | return nil, err 122 | } 123 | return m, nil 124 | } 125 | 126 | // Server API for PingPong service 127 | 128 | type PingPongServer interface { 129 | PingPong(PingPong_PingPongServer) error 130 | } 131 | 132 | func RegisterPingPongServer(s *grpc.Server, srv PingPongServer) { 133 | s.RegisterService(&_PingPong_serviceDesc, srv) 134 | } 135 | 136 | func _PingPong_PingPong_Handler(srv interface{}, stream grpc.ServerStream) error { 137 | return srv.(PingPongServer).PingPong(&pingPongPingPongServer{stream}) 138 | } 139 | 140 | type PingPong_PingPongServer interface { 141 | Send(*Pong) error 142 | Recv() (*Ping, error) 143 | grpc.ServerStream 144 | } 145 | 146 | type pingPongPingPongServer struct { 147 | grpc.ServerStream 148 | } 149 | 150 | func (x *pingPongPingPongServer) Send(m *Pong) error { 151 | return x.ServerStream.SendMsg(m) 152 | } 153 | 154 | func (x *pingPongPingPongServer) Recv() (*Ping, error) { 155 | m := new(Ping) 156 | if err := x.ServerStream.RecvMsg(m); err != nil { 157 | return nil, err 158 | } 159 | return m, nil 160 | } 161 | 162 | var _PingPong_serviceDesc = grpc.ServiceDesc{ 163 | ServiceName: "main.PingPong", 164 | HandlerType: (*PingPongServer)(nil), 165 | Methods: []grpc.MethodDesc{}, 166 | Streams: []grpc.StreamDesc{ 167 | { 168 | StreamName: "PingPong", 169 | Handler: _PingPong_PingPong_Handler, 170 | ServerStreams: true, 171 | ClientStreams: true, 172 | }, 173 | }, 174 | Metadata: "pingpong.proto", 175 | } 176 | 177 | func init() { proto.RegisterFile("pingpong.proto", fileDescriptor0) } 178 | 179 | var fileDescriptor0 = []byte{ 180 | // 188 bytes of a gzipped FileDescriptorProto 181 | 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x7c, 0x8f, 0x3f, 0x6b, 0xc3, 0x30, 182 | 0x10, 0x47, 0x2b, 0x30, 0xfd, 0xa3, 0xa1, 0x83, 0x27, 0xd3, 0xc9, 0x78, 0xf2, 0x62, 0xa9, 0xb4, 183 | 0xfd, 0x04, 0x1d, 0x4b, 0x87, 0xd2, 0x31, 0x4b, 0x90, 0x85, 0xb8, 0x1c, 0xf1, 0xdd, 0x09, 0x49, 184 | 0x86, 0xf8, 0xdb, 0x07, 0x1b, 0x0c, 0x99, 0xb2, 0xbd, 0x07, 0x8f, 0x3b, 0x7e, 0xfa, 0x35, 0x22, 185 | 0x43, 0x14, 0x06, 0x13, 0x93, 0x14, 0xa9, 0x2b, 0x72, 0xc8, 0x5d, 0xab, 0xab, 0x3f, 0x64, 0xa8, 186 | 0x1b, 0xfd, 0x44, 0x21, 0x67, 0x07, 0xa1, 0x51, 0xad, 0xea, 0x5f, 0xfe, 0x77, 0xdd, 0x0a, 0xb9, 187 | 0x57, 0x7c, 0x7c, 0xe9, 0xe7, 0xf5, 0xc6, 0x56, 0xf5, 0x37, 0xac, 0xcd, 0xfa, 0xc2, 0xac, 0xfe, 188 | 0xb6, 0xb3, 0x30, 0x74, 0x0f, 0xbd, 0x7a, 0x57, 0xdf, 0xbf, 0x87, 0x1f, 0xc0, 0x72, 0x9a, 0x47, 189 | 0xe3, 0x85, 0x6c, 0x9c, 0x5c, 0x41, 0x9e, 0x89, 0x84, 0xcf, 0x61, 0xb1, 0x20, 0x83, 0x17, 0xf6, 190 | 0x73, 0x4a, 0x81, 0xfd, 0x32, 0x4c, 0x48, 0x58, 0xb2, 0x0d, 0x17, 0x47, 0x71, 0x0a, 0xd9, 0x42, 191 | 0x8a, 0xfe, 0x98, 0x4b, 0x0a, 0x8e, 0x90, 0xc1, 0xc6, 0x71, 0x7c, 0xdc, 0x46, 0x7d, 0x5e, 0x03, 192 | 0x00, 0x00, 0xff, 0xff, 0x16, 0xc5, 0xb0, 0x39, 0xe6, 0x00, 0x00, 0x00, 193 | } 194 | -------------------------------------------------------------------------------- /examples/grpc_streaming/pb/pingpong.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package main; 4 | 5 | option go_package="github.com/platinummonkey/go-concurrency-limits/examples/grpc_streaming/pb"; 6 | 7 | 8 | message Ping { 9 | string message = 1; 10 | } 11 | 12 | message Pong { 13 | string message = 1; 14 | } 15 | 16 | service PingPong { 17 | rpc PingPong (stream Ping) returns (stream Pong) {} 18 | } 19 | -------------------------------------------------------------------------------- /examples/grpc_unary/.gitignore: -------------------------------------------------------------------------------- 1 | grpc_unary 2 | -------------------------------------------------------------------------------- /examples/grpc_unary/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | //go:generate protoc -I ./pb --go_out=plugins=grpc:${GOPATH}/src *.proto 4 | 5 | import ( 6 | "context" 7 | "flag" 8 | "fmt" 9 | "log" 10 | "net" 11 | "sync" 12 | "time" 13 | 14 | golangGrpc "google.golang.org/grpc" 15 | 16 | "github.com/platinummonkey/go-concurrency-limits/examples/grpc_unary/pb" 17 | "github.com/platinummonkey/go-concurrency-limits/grpc" 18 | "github.com/platinummonkey/go-concurrency-limits/limit" 19 | "github.com/platinummonkey/go-concurrency-limits/limiter" 20 | "github.com/platinummonkey/go-concurrency-limits/strategy" 21 | ) 22 | 23 | var options = struct { 24 | mode string 25 | port int 26 | numThreads int 27 | }{ 28 | mode: "server", 29 | port: 8080, 30 | numThreads: 105, 31 | } 32 | 33 | func init() { 34 | flag.StringVar(&options.mode, "mode", options.mode, "choose `client` or `server` mode") 35 | flag.IntVar(&options.port, "port", options.port, "grpc port") 36 | flag.IntVar(&options.numThreads, "threads", options.numThreads, "number of client threads") 37 | } 38 | 39 | func checkOptions() { 40 | switch options.mode { 41 | case "client": 42 | fallthrough 43 | case "server": 44 | // no-op 45 | default: 46 | panic(fmt.Sprintf("invalid mode specified: '%s'", options.mode)) 47 | } 48 | } 49 | 50 | type server struct { 51 | } 52 | 53 | func (s *server) PingPong(ctx context.Context, in *pb.Ping) (*pb.Pong, error) { 54 | log.Printf("Received: '%s'", in.GetMessage()) 55 | // pretend to do some work 56 | time.Sleep(time.Millisecond * 10) 57 | return &pb.Pong{Message: in.GetMessage()}, nil 58 | } 59 | 60 | func main() { 61 | flag.Parse() 62 | checkOptions() 63 | switch options.mode { 64 | case "server": 65 | runServer() 66 | case "client": 67 | runClient() 68 | } 69 | } 70 | 71 | func runServer() { 72 | logger := limit.BuiltinLimitLogger{} 73 | lis, err := net.Listen("tcp", fmt.Sprintf(":%d", options.port)) 74 | if err != nil { 75 | panic(err) 76 | } 77 | 78 | serverLimit := limit.NewFixedLimit("server-fixed-limit", 10, nil) 79 | serverLimiter, err := limiter.NewDefaultLimiter(serverLimit, 1, 1000, 1e6, 100, strategy.NewSimpleStrategy(10), logger, nil) 80 | if err != nil { 81 | panic(err) 82 | } 83 | serverOpts := []grpc.InterceptorOption{grpc.WithName("grpc-unary-server"), grpc.WithLimiter(serverLimiter)} 84 | serverInterceptor := grpc.UnaryServerInterceptor(serverOpts...) 85 | svc := golangGrpc.NewServer(golangGrpc.UnaryInterceptor(serverInterceptor)) 86 | s := &server{} 87 | pb.RegisterPingPongServer(svc, s) 88 | if err := svc.Serve(lis); err != nil { 89 | panic(nil) 90 | } 91 | } 92 | 93 | func runClient() { 94 | conn, err := golangGrpc.Dial(fmt.Sprintf("localhost:%d", options.port), golangGrpc.WithInsecure()) 95 | if err != nil { 96 | panic(err) 97 | } 98 | defer conn.Close() 99 | client := pb.NewPingPongClient(conn) 100 | 101 | wg := sync.WaitGroup{} 102 | wg.Add(options.numThreads) 103 | for i := 0; i < options.numThreads; i++ { 104 | go func(workerID int) { 105 | j := 0 106 | // do this as fast as possible 107 | for { 108 | queryServer(client, workerID, j) 109 | j++ 110 | } 111 | }(i) 112 | } 113 | wg.Wait() 114 | } 115 | 116 | func queryServer(client pb.PingPongClient, workerID int, i int) { 117 | ctx, cancel := context.WithTimeout(context.Background(), time.Second) 118 | defer cancel() 119 | msg := &pb.Ping{ 120 | Message: fmt.Sprintf("hello %d from %d", i, workerID), 121 | } 122 | r, err := client.PingPong(ctx, msg) 123 | if err != nil { 124 | log.Printf("[failed](%d - %d)\t - %v", workerID, i, err) 125 | } else { 126 | log.Printf("[pass](%d - %d)\t - %s", workerID, i, r.GetMessage()) 127 | } 128 | } 129 | -------------------------------------------------------------------------------- /examples/grpc_unary/pb/pingpong.pb.go: -------------------------------------------------------------------------------- 1 | // Code generated by protoc-gen-go. DO NOT EDIT. 2 | // source: pingpong.proto 3 | 4 | /* 5 | Package pb is a generated protocol buffer package. 6 | 7 | It is generated from these files: 8 | pingpong.proto 9 | 10 | It has these top-level messages: 11 | Ping 12 | Pong 13 | */ 14 | package pb 15 | 16 | import proto "github.com/golang/protobuf/proto" 17 | import fmt "fmt" 18 | import math "math" 19 | 20 | import ( 21 | context "golang.org/x/net/context" 22 | grpc "google.golang.org/grpc" 23 | ) 24 | 25 | // Reference imports to suppress errors if they are not otherwise used. 26 | var _ = proto.Marshal 27 | var _ = fmt.Errorf 28 | var _ = math.Inf 29 | 30 | // This is a compile-time assertion to ensure that this generated file 31 | // is compatible with the proto package it is being compiled against. 32 | // A compilation error at this line likely means your copy of the 33 | // proto package needs to be updated. 34 | const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package 35 | 36 | type Ping struct { 37 | Message string `protobuf:"bytes,1,opt,name=message" json:"message,omitempty"` 38 | } 39 | 40 | func (m *Ping) Reset() { *m = Ping{} } 41 | func (m *Ping) String() string { return proto.CompactTextString(m) } 42 | func (*Ping) ProtoMessage() {} 43 | func (*Ping) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} } 44 | 45 | func (m *Ping) GetMessage() string { 46 | if m != nil { 47 | return m.Message 48 | } 49 | return "" 50 | } 51 | 52 | type Pong struct { 53 | Message string `protobuf:"bytes,1,opt,name=message" json:"message,omitempty"` 54 | } 55 | 56 | func (m *Pong) Reset() { *m = Pong{} } 57 | func (m *Pong) String() string { return proto.CompactTextString(m) } 58 | func (*Pong) ProtoMessage() {} 59 | func (*Pong) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{1} } 60 | 61 | func (m *Pong) GetMessage() string { 62 | if m != nil { 63 | return m.Message 64 | } 65 | return "" 66 | } 67 | 68 | func init() { 69 | proto.RegisterType((*Ping)(nil), "main.Ping") 70 | proto.RegisterType((*Pong)(nil), "main.Pong") 71 | } 72 | 73 | // Reference imports to suppress errors if they are not otherwise used. 74 | var _ context.Context 75 | var _ grpc.ClientConn 76 | 77 | // This is a compile-time assertion to ensure that this generated file 78 | // is compatible with the grpc package it is being compiled against. 79 | const _ = grpc.SupportPackageIsVersion4 80 | 81 | // Client API for PingPong service 82 | 83 | type PingPongClient interface { 84 | PingPong(ctx context.Context, in *Ping, opts ...grpc.CallOption) (*Pong, error) 85 | } 86 | 87 | type pingPongClient struct { 88 | cc *grpc.ClientConn 89 | } 90 | 91 | func NewPingPongClient(cc *grpc.ClientConn) PingPongClient { 92 | return &pingPongClient{cc} 93 | } 94 | 95 | func (c *pingPongClient) PingPong(ctx context.Context, in *Ping, opts ...grpc.CallOption) (*Pong, error) { 96 | out := new(Pong) 97 | err := grpc.Invoke(ctx, "/main.PingPong/PingPong", in, out, c.cc, opts...) 98 | if err != nil { 99 | return nil, err 100 | } 101 | return out, nil 102 | } 103 | 104 | // Server API for PingPong service 105 | 106 | type PingPongServer interface { 107 | PingPong(context.Context, *Ping) (*Pong, error) 108 | } 109 | 110 | func RegisterPingPongServer(s *grpc.Server, srv PingPongServer) { 111 | s.RegisterService(&_PingPong_serviceDesc, srv) 112 | } 113 | 114 | func _PingPong_PingPong_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { 115 | in := new(Ping) 116 | if err := dec(in); err != nil { 117 | return nil, err 118 | } 119 | if interceptor == nil { 120 | return srv.(PingPongServer).PingPong(ctx, in) 121 | } 122 | info := &grpc.UnaryServerInfo{ 123 | Server: srv, 124 | FullMethod: "/main.PingPong/PingPong", 125 | } 126 | handler := func(ctx context.Context, req interface{}) (interface{}, error) { 127 | return srv.(PingPongServer).PingPong(ctx, req.(*Ping)) 128 | } 129 | return interceptor(ctx, in, info, handler) 130 | } 131 | 132 | var _PingPong_serviceDesc = grpc.ServiceDesc{ 133 | ServiceName: "main.PingPong", 134 | HandlerType: (*PingPongServer)(nil), 135 | Methods: []grpc.MethodDesc{ 136 | { 137 | MethodName: "PingPong", 138 | Handler: _PingPong_PingPong_Handler, 139 | }, 140 | }, 141 | Streams: []grpc.StreamDesc{}, 142 | Metadata: "pingpong.proto", 143 | } 144 | 145 | func init() { proto.RegisterFile("pingpong.proto", fileDescriptor0) } 146 | 147 | var fileDescriptor0 = []byte{ 148 | // 182 bytes of a gzipped FileDescriptorProto 149 | 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x7c, 0x8f, 0xbd, 0xaa, 0xc2, 0x40, 150 | 0x10, 0x46, 0x6f, 0x20, 0xdc, 0x9f, 0x2d, 0x6e, 0x91, 0x2a, 0x58, 0x85, 0x60, 0x61, 0x93, 0x5d, 151 | 0xd1, 0x37, 0xb0, 0x10, 0x4b, 0xb1, 0xb4, 0x91, 0xcd, 0xb2, 0x8c, 0x8b, 0xd9, 0x99, 0x61, 0x7f, 152 | 0xc0, 0xbc, 0xbd, 0x24, 0x10, 0xb0, 0xb2, 0x3b, 0x07, 0x0e, 0x33, 0x7c, 0xe2, 0x9f, 0x1d, 0x02, 153 | 0x13, 0x82, 0xe4, 0x40, 0x89, 0xaa, 0xd2, 0x6b, 0x87, 0x6d, 0x23, 0xca, 0xb3, 0x43, 0xa8, 0x6a, 154 | 0xf1, 0xe3, 0x6d, 0x8c, 0x1a, 0x6c, 0x5d, 0x34, 0xc5, 0xe6, 0xef, 0xb2, 0xe8, 0x5c, 0xd0, 0xa7, 155 | 0x62, 0xb7, 0x15, 0xbf, 0xd3, 0x8d, 0xb9, 0x5a, 0xbf, 0xb1, 0x90, 0xd3, 0x0b, 0x39, 0xf9, 0x6a, 156 | 0x61, 0x42, 0x68, 0xbf, 0x0e, 0xa7, 0xeb, 0x11, 0x5c, 0xba, 0xe7, 0x5e, 0x1a, 0xf2, 0x8a, 0x07, 157 | 0x9d, 0x1c, 0x66, 0xef, 0x09, 0x1f, 0x76, 0x54, 0x40, 0x9d, 0x21, 0x34, 0x39, 0x04, 0x8b, 0x66, 158 | 0xec, 0x06, 0xe7, 0x5d, 0x8a, 0xca, 0x3e, 0xb5, 0xe7, 0xc1, 0x46, 0x05, 0x81, 0xcd, 0x2d, 0xa3, 159 | 0x0e, 0xa3, 0xe2, 0xbe, 0xff, 0x9e, 0xc7, 0xec, 0x5f, 0x01, 0x00, 0x00, 0xff, 0xff, 0xf6, 0x90, 160 | 0xab, 0x49, 0xde, 0x00, 0x00, 0x00, 161 | } 162 | -------------------------------------------------------------------------------- /examples/grpc_unary/pb/pingpong.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package main; 4 | 5 | option go_package="github.com/platinummonkey/go-concurrency-limits/examples/grpc_unary/pb"; 6 | 7 | 8 | message Ping { 9 | string message = 1; 10 | } 11 | 12 | message Pong { 13 | string message = 1; 14 | } 15 | 16 | service PingPong { 17 | rpc PingPong (Ping) returns (Pong) {} 18 | } 19 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/platinummonkey/go-concurrency-limits 2 | 3 | go 1.23.0 4 | 5 | toolchain go1.24.1 6 | 7 | require ( 8 | github.com/DataDog/datadog-go/v5 v5.6.0 9 | github.com/golang/protobuf v1.5.4 10 | github.com/rcrowley/go-metrics v0.0.0-20180503174638-e2704e165165 11 | github.com/stretchr/testify v1.10.0 12 | golang.org/x/net v0.40.0 13 | golang.org/x/time v0.11.0 14 | google.golang.org/grpc v1.72.2 15 | ) 16 | 17 | require ( 18 | github.com/Microsoft/go-winio v0.5.0 // indirect 19 | github.com/davecgh/go-spew v1.1.1 // indirect 20 | github.com/pmezard/go-difflib v1.0.0 // indirect 21 | golang.org/x/sys v0.33.0 // indirect 22 | golang.org/x/text v0.25.0 // indirect 23 | google.golang.org/genproto/googleapis/rpc v0.0.0-20250218202821-56aae31c358a // indirect 24 | google.golang.org/protobuf v1.36.5 // indirect 25 | gopkg.in/yaml.v3 v3.0.1 // indirect 26 | ) 27 | -------------------------------------------------------------------------------- /grpc/doc.go: -------------------------------------------------------------------------------- 1 | // Package grpc provides GRPC client/server mixins to add concurrency control. 2 | package grpc 3 | -------------------------------------------------------------------------------- /grpc/grpc_streaming.go: -------------------------------------------------------------------------------- 1 | package grpc 2 | 3 | import ( 4 | golangGrpc "google.golang.org/grpc" 5 | "google.golang.org/grpc/status" 6 | ) 7 | 8 | type ssRecvWrapper struct { 9 | golangGrpc.ServerStream 10 | info *golangGrpc.StreamServerInfo 11 | cfg *streamInterceptorConfig 12 | } 13 | 14 | // RecvMsg wrapps the underlying StreamServer RecvMsg with the limiter. 15 | func (s *ssRecvWrapper) RecvMsg(m interface{}) error { 16 | ctx := s.Context() 17 | token, ok := s.cfg.recvLimiter.Acquire(ctx) 18 | if !ok { 19 | _, errCode, err := s.cfg.recvLimitExceededResponseClassifier(ctx, s.info.FullMethod, m, s.cfg.recvLimiter) 20 | return status.Error(errCode, err.Error()) 21 | } 22 | err := s.ServerStream.RecvMsg(m) 23 | if err != nil { 24 | respType := s.cfg.serverResponseClassifer(ctx, m, s.info, err) 25 | switch respType { 26 | case ResponseTypeSuccess: 27 | token.OnSuccess() 28 | case ResponseTypeIgnore: 29 | token.OnIgnore() 30 | case ResponseTypeDropped: 31 | token.OnDropped() 32 | } 33 | return err 34 | } 35 | token.OnSuccess() 36 | return nil 37 | } 38 | 39 | // SendMsg wrapps the underlying StreamServer SendMsg with the limiter. 40 | func (s *ssRecvWrapper) SendMsg(m interface{}) error { 41 | ctx := s.Context() 42 | token, ok := s.cfg.recvLimiter.Acquire(ctx) 43 | if !ok { 44 | _, errCode, err := s.cfg.sendLimitExceededResponseClassifier(ctx, s.info.FullMethod, m, s.cfg.recvLimiter) 45 | return status.Error(errCode, err.Error()) 46 | } 47 | err := s.ServerStream.SendMsg(m) 48 | if err != nil { 49 | respType := s.cfg.clientResponseClassifer(ctx, m, s.info, err) 50 | switch respType { 51 | case ResponseTypeSuccess: 52 | token.OnSuccess() 53 | case ResponseTypeIgnore: 54 | token.OnIgnore() 55 | case ResponseTypeDropped: 56 | token.OnDropped() 57 | } 58 | return err 59 | } 60 | token.OnSuccess() 61 | return nil 62 | } 63 | 64 | // StreamServerInterceptor will add tracing to a gprc streaming client. 65 | func StreamServerInterceptor(opts ...StreamInterceptorOption) golangGrpc.StreamServerInterceptor { 66 | cfg := new(streamInterceptorConfig) 67 | streamDefaults(cfg) 68 | for _, fn := range opts { 69 | fn(cfg) 70 | } 71 | return func(srv interface{}, ss golangGrpc.ServerStream, info *golangGrpc.StreamServerInfo, handler golangGrpc.StreamHandler) error { 72 | wrappedSs := &ssRecvWrapper{ 73 | ServerStream: ss, 74 | info: info, 75 | cfg: cfg, 76 | } 77 | return handler(srv, wrappedSs) 78 | } 79 | } 80 | -------------------------------------------------------------------------------- /grpc/grpc_unary.go: -------------------------------------------------------------------------------- 1 | package grpc 2 | 3 | import ( 4 | "context" 5 | 6 | golangGrpc "google.golang.org/grpc" 7 | "google.golang.org/grpc/status" 8 | ) 9 | 10 | // UnaryServerInterceptor will trace requests to the given grpc server. 11 | func UnaryServerInterceptor(opts ...InterceptorOption) golangGrpc.UnaryServerInterceptor { 12 | cfg := new(interceptorConfig) 13 | defaults(cfg) 14 | for _, fn := range opts { 15 | fn(cfg) 16 | } 17 | return func(ctx context.Context, req interface{}, info *golangGrpc.UnaryServerInfo, handler golangGrpc.UnaryHandler) (interface{}, error) { 18 | token, ok := cfg.limiter.Acquire(ctx) 19 | if !ok { 20 | errResp, errCode, err := cfg.limitExceededResponseClassifier(ctx, info.FullMethod, req, cfg.limiter) 21 | return errResp, status.Error(errCode, err.Error()) 22 | } 23 | resp, err := handler(ctx, req) 24 | respType := cfg.serverResponseClassifer(ctx, req, info, resp, err) 25 | switch respType { 26 | case ResponseTypeSuccess: 27 | token.OnSuccess() 28 | case ResponseTypeIgnore: 29 | token.OnIgnore() 30 | case ResponseTypeDropped: 31 | token.OnDropped() 32 | } 33 | return resp, err 34 | } 35 | } 36 | 37 | // UnaryClientInterceptor will add tracing to a gprc client. 38 | func UnaryClientInterceptor(opts ...InterceptorOption) golangGrpc.UnaryClientInterceptor { 39 | cfg := new(interceptorConfig) 40 | defaults(cfg) 41 | for _, fn := range opts { 42 | fn(cfg) 43 | } 44 | return func(ctx context.Context, method string, req, reply interface{}, cc *golangGrpc.ClientConn, invoker golangGrpc.UnaryInvoker, opts ...golangGrpc.CallOption) error { 45 | token, ok := cfg.limiter.Acquire(ctx) 46 | if !ok { 47 | _, errCode, err := cfg.limitExceededResponseClassifier(ctx, method, req, cfg.limiter) 48 | return status.Error(errCode, err.Error()) 49 | } 50 | err := invoker(ctx, method, req, reply, cc, opts...) 51 | respType := cfg.clientResponseClassifer(ctx, method, req, reply, err) 52 | switch respType { 53 | case ResponseTypeSuccess: 54 | token.OnSuccess() 55 | case ResponseTypeIgnore: 56 | token.OnIgnore() 57 | case ResponseTypeDropped: 58 | token.OnDropped() 59 | } 60 | return err 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /grpc/option_streaming.go: -------------------------------------------------------------------------------- 1 | package grpc 2 | 3 | import ( 4 | "context" 5 | 6 | golangGrpc "google.golang.org/grpc" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | "github.com/platinummonkey/go-concurrency-limits/limit" 10 | "github.com/platinummonkey/go-concurrency-limits/limiter" 11 | "github.com/platinummonkey/go-concurrency-limits/strategy" 12 | ) 13 | 14 | // StreamClientResponseClassifier is a method definition for defining custom response types to the limiter algorithm to 15 | // correctly handle certain types of errors or embedded data. 16 | type StreamClientResponseClassifier func(ctx context.Context, req interface{}, info *golangGrpc.StreamServerInfo, err error) ResponseType 17 | 18 | // StreamServerResponseClassifier is a method definition for defining custom response types to the limiter algorithm to 19 | // correctly handle certain types of errors or embedded data. 20 | type StreamServerResponseClassifier func( 21 | ctx context.Context, req interface{}, info *golangGrpc.StreamServerInfo, err error, 22 | ) ResponseType 23 | 24 | func defaultStreamClientResponseClassifier( 25 | ctx context.Context, 26 | req interface{}, 27 | info *golangGrpc.StreamServerInfo, 28 | err error, 29 | ) ResponseType { 30 | if err != nil { 31 | return ResponseTypeDropped 32 | } 33 | return ResponseTypeSuccess 34 | } 35 | 36 | func defaultStreamServerResponseClassifier( 37 | ctx context.Context, 38 | req interface{}, 39 | info *golangGrpc.StreamServerInfo, 40 | err error, 41 | ) ResponseType { 42 | if err != nil { 43 | return ResponseTypeDropped 44 | } 45 | return ResponseTypeSuccess 46 | } 47 | 48 | type streamInterceptorConfig struct { 49 | recvName string 50 | sendName string 51 | tags []string 52 | recvLimiter core.Limiter 53 | sendLimiter core.Limiter 54 | recvLimitExceededResponseClassifier LimitExceededResponseClassifier 55 | sendLimitExceededResponseClassifier LimitExceededResponseClassifier 56 | serverResponseClassifer StreamServerResponseClassifier 57 | clientResponseClassifer StreamClientResponseClassifier 58 | } 59 | 60 | // StreamInterceptorOption represents an option that can be passed to the stream 61 | // client and server interceptors. 62 | type StreamInterceptorOption func(*streamInterceptorConfig) 63 | 64 | func streamDefaults(cfg *streamInterceptorConfig) { 65 | recvName := cfg.recvName 66 | if recvName == "" { 67 | recvName = "default-recv" 68 | } 69 | sendName := cfg.sendName 70 | if sendName == "" { 71 | sendName = "default-send" 72 | } 73 | tags := cfg.tags 74 | if tags == nil { 75 | tags = make([]string, 0) 76 | } 77 | cfg.recvLimiter, _ = limiter.NewDefaultLimiterWithDefaults( 78 | recvName, 79 | strategy.NewSimpleStrategy(1000), 80 | limit.NoopLimitLogger{}, 81 | core.EmptyMetricRegistryInstance, 82 | tags..., 83 | ) 84 | cfg.sendLimiter, _ = limiter.NewDefaultLimiterWithDefaults( 85 | sendName, 86 | strategy.NewSimpleStrategy(1000), 87 | limit.NoopLimitLogger{}, 88 | core.EmptyMetricRegistryInstance, 89 | tags..., 90 | ) 91 | cfg.recvLimitExceededResponseClassifier = defaultLimitExceededResponseClassifier 92 | cfg.sendLimitExceededResponseClassifier = defaultLimitExceededResponseClassifier 93 | cfg.clientResponseClassifer = defaultStreamClientResponseClassifier 94 | cfg.serverResponseClassifer = defaultStreamServerResponseClassifier 95 | } 96 | 97 | // WithStreamSendName sets the default SendMsg limiter name if the default limiter is used, otherwise unused. 98 | func WithStreamSendName(name string) StreamInterceptorOption { 99 | return func(cfg *streamInterceptorConfig) { 100 | cfg.sendName = name 101 | } 102 | } 103 | 104 | // WithStreamRecvName sets the default RecvMsg limiter name if the default limiter is used, otherwise unused. 105 | func WithStreamRecvName(name string) StreamInterceptorOption { 106 | return func(cfg *streamInterceptorConfig) { 107 | cfg.recvName = name 108 | } 109 | } 110 | 111 | // WithStreamTags sets the default limiter tags if the default limiter is used, otherwise unused. 112 | func WithStreamTags(tags []string) InterceptorOption { 113 | return func(cfg *interceptorConfig) { 114 | cfg.tags = tags 115 | } 116 | } 117 | 118 | // WithStreamSendLimiter sets the given limiter for the intercepted stream client for SendMsg. 119 | func WithStreamSendLimiter(limiter core.Limiter) StreamInterceptorOption { 120 | return func(cfg *streamInterceptorConfig) { 121 | cfg.sendLimiter = limiter 122 | } 123 | } 124 | 125 | // WithStreamRecvLimiter sets the given limiter for the intercepted stream client for RecvMsg. 126 | func WithStreamRecvLimiter(limiter core.Limiter) StreamInterceptorOption { 127 | return func(cfg *streamInterceptorConfig) { 128 | cfg.recvLimiter = limiter 129 | } 130 | } 131 | 132 | // WithStreamSendLimitExceededResponseClassifier sets the response classifier for the intercepted stream client on SendMsg. 133 | func WithStreamSendLimitExceededResponseClassifier(classifier LimitExceededResponseClassifier) StreamInterceptorOption { 134 | return func(cfg *streamInterceptorConfig) { 135 | cfg.sendLimitExceededResponseClassifier = classifier 136 | } 137 | } 138 | 139 | // WithStreamRecvLimitExceededResponseClassifier sets the response classifier for the intercepted stream client on RecvMsg. 140 | func WithStreamRecvLimitExceededResponseClassifier(classifier LimitExceededResponseClassifier) StreamInterceptorOption { 141 | return func(cfg *streamInterceptorConfig) { 142 | cfg.recvLimitExceededResponseClassifier = classifier 143 | } 144 | } 145 | 146 | // WithStreamClientResponseTypeClassifier sets the response classifier for the intercepted client response 147 | func WithStreamClientResponseTypeClassifier(classifier StreamClientResponseClassifier) StreamInterceptorOption { 148 | return func(cfg *streamInterceptorConfig) { 149 | cfg.clientResponseClassifer = classifier 150 | } 151 | } 152 | 153 | // WithStreamServerResponseTypeClassifier sets the response classifier for the intercepted server response 154 | func WithStreamServerResponseTypeClassifier(classifier StreamServerResponseClassifier) StreamInterceptorOption { 155 | return func(cfg *streamInterceptorConfig) { 156 | cfg.serverResponseClassifer = classifier 157 | } 158 | } 159 | -------------------------------------------------------------------------------- /grpc/option_unary.go: -------------------------------------------------------------------------------- 1 | package grpc 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | golangGrpc "google.golang.org/grpc" 8 | "google.golang.org/grpc/codes" 9 | 10 | "github.com/platinummonkey/go-concurrency-limits/core" 11 | "github.com/platinummonkey/go-concurrency-limits/limit" 12 | "github.com/platinummonkey/go-concurrency-limits/limiter" 13 | "github.com/platinummonkey/go-concurrency-limits/strategy" 14 | ) 15 | 16 | // ResponseType is the type of token release that should be specified to the limiter algorithm. 17 | type ResponseType int 18 | 19 | const ( 20 | // ResponseTypeSuccess represents a successful response for the limiter algorithm 21 | ResponseTypeSuccess ResponseType = iota 22 | // ResponseTypeIgnore represents an ignorable error or response for the limiter algorithm 23 | ResponseTypeIgnore 24 | // ResponseTypeDropped represents a dropped request type for the limiter algorithm 25 | ResponseTypeDropped 26 | ) 27 | 28 | // LimitExceededResponseClassifier is a method definition for defining the error response type when the limit is exceeded 29 | // and a token is not able to be acquired. By default the RESOURCE_EXHUASTED type is returend. 30 | type LimitExceededResponseClassifier func(ctx context.Context, method string, req interface{}, l core.Limiter) (interface{}, codes.Code, error) 31 | 32 | // ClientResponseClassifier is a method definition for defining custom response types to the limiter algorithm to 33 | // correctly handle certain types of errors or embedded data. 34 | type ClientResponseClassifier func(ctx context.Context, method string, req, reply interface{}, err error) ResponseType 35 | 36 | // ServerResponseClassifier is a method definition for defining custom response types to the limiter algorithm to 37 | // correctly handle certain types of errors or embedded data. 38 | type ServerResponseClassifier func( 39 | ctx context.Context, req interface{}, info *golangGrpc.UnaryServerInfo, resp interface{}, err error, 40 | ) ResponseType 41 | 42 | func defaultLimitExceededResponseClassifier( 43 | ctx context.Context, 44 | method string, 45 | req interface{}, 46 | l core.Limiter, 47 | ) (interface{}, codes.Code, error) { 48 | return nil, codes.ResourceExhausted, fmt.Errorf("limit exceeded for limiter=%v", l) 49 | } 50 | 51 | func defaultClientResponseClassifier( 52 | ctx context.Context, 53 | method string, 54 | req interface{}, 55 | reply interface{}, 56 | err error, 57 | ) ResponseType { 58 | if err != nil { 59 | return ResponseTypeDropped 60 | } 61 | return ResponseTypeSuccess 62 | } 63 | 64 | func defaultServerResponseClassifier( 65 | ctx context.Context, 66 | req interface{}, 67 | info *golangGrpc.UnaryServerInfo, 68 | resp interface{}, 69 | err error, 70 | ) ResponseType { 71 | if err != nil { 72 | return ResponseTypeDropped 73 | } 74 | return ResponseTypeSuccess 75 | } 76 | 77 | type interceptorConfig struct { 78 | name string 79 | tags []string 80 | limiter core.Limiter 81 | limitExceededResponseClassifier LimitExceededResponseClassifier 82 | serverResponseClassifer ServerResponseClassifier 83 | clientResponseClassifer ClientResponseClassifier 84 | } 85 | 86 | // InterceptorOption represents an option that can be passed to the grpc unary 87 | // client and server interceptors. 88 | type InterceptorOption func(*interceptorConfig) 89 | 90 | func defaults(cfg *interceptorConfig) { 91 | name := cfg.name 92 | if name == "" { 93 | name = "default" 94 | } 95 | tags := cfg.tags 96 | if tags == nil { 97 | tags = make([]string, 0) 98 | } 99 | cfg.limiter, _ = limiter.NewDefaultLimiterWithDefaults( 100 | name, 101 | strategy.NewSimpleStrategy(1000), 102 | limit.NoopLimitLogger{}, 103 | core.EmptyMetricRegistryInstance, 104 | tags..., 105 | ) 106 | cfg.limitExceededResponseClassifier = defaultLimitExceededResponseClassifier 107 | cfg.clientResponseClassifer = defaultClientResponseClassifier 108 | cfg.serverResponseClassifer = defaultServerResponseClassifier 109 | } 110 | 111 | // WithName sets the default limiter name if the default limiter is used, otherwise unused. 112 | func WithName(name string) InterceptorOption { 113 | return func(cfg *interceptorConfig) { 114 | cfg.name = name 115 | } 116 | } 117 | 118 | // WithTags sets the default limiter tags if the default limiter is used, otherwise unused. 119 | func WithTags(tags []string) InterceptorOption { 120 | return func(cfg *interceptorConfig) { 121 | cfg.tags = tags 122 | } 123 | } 124 | 125 | // WithLimiter sets the given limiter for the intercepted client. 126 | func WithLimiter(limiter core.Limiter) InterceptorOption { 127 | return func(cfg *interceptorConfig) { 128 | cfg.limiter = limiter 129 | } 130 | } 131 | 132 | // WithLimitExceededResponseClassifier sets the response classifier for the intercepted client 133 | func WithLimitExceededResponseClassifier(classifier LimitExceededResponseClassifier) InterceptorOption { 134 | return func(cfg *interceptorConfig) { 135 | cfg.limitExceededResponseClassifier = classifier 136 | } 137 | } 138 | 139 | // WithClientResponseTypeClassifier sets the response classifier for the intercepted client 140 | func WithClientResponseTypeClassifier(classifier ClientResponseClassifier) InterceptorOption { 141 | return func(cfg *interceptorConfig) { 142 | cfg.clientResponseClassifer = classifier 143 | } 144 | } 145 | 146 | // WithServerResponseTypeClassifier sets the response classifier for the intercepted client 147 | func WithServerResponseTypeClassifier(classifier ServerResponseClassifier) InterceptorOption { 148 | return func(cfg *interceptorConfig) { 149 | cfg.serverResponseClassifer = classifier 150 | } 151 | } 152 | -------------------------------------------------------------------------------- /limit/aimd.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "fmt" 5 | "math" 6 | "sync" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | ) 10 | 11 | // AIMDLimit implements a Loss based dynamic Limit that does an additive increment as long as there are no errors and a 12 | // multiplicative decrement when there is an error. 13 | type AIMDLimit struct { 14 | name string 15 | limit int 16 | increaseBy int 17 | backOffRatio float64 18 | 19 | listeners []core.LimitChangeListener 20 | registry core.MetricRegistry 21 | commonSampler *core.CommonMetricSampler 22 | 23 | mu sync.RWMutex 24 | } 25 | 26 | // NewDefaultAIMDLimit will create a default AIMDLimit. 27 | func NewDefaultAIMDLimit( 28 | name string, 29 | registry core.MetricRegistry, 30 | tags ...string, 31 | ) *AIMDLimit { 32 | return NewAIMDLimit(name, 10, 0.9, 1, registry, tags...) 33 | } 34 | 35 | // NewAIMDLimit will create a new AIMDLimit. 36 | func NewAIMDLimit( 37 | name string, 38 | initialLimit int, 39 | backOffRatio float64, 40 | increaseBy int, 41 | registry core.MetricRegistry, 42 | tags ...string, 43 | ) *AIMDLimit { 44 | if registry == nil { 45 | registry = core.EmptyMetricRegistryInstance 46 | } 47 | if increaseBy <= 0 { 48 | increaseBy = 1 49 | } 50 | 51 | l := &AIMDLimit{ 52 | name: name, 53 | limit: initialLimit, 54 | backOffRatio: backOffRatio, 55 | increaseBy: increaseBy, 56 | listeners: make([]core.LimitChangeListener, 0), 57 | registry: registry, 58 | } 59 | l.commonSampler = core.NewCommonMetricSamplerOrNil(registry, l, name, tags...) 60 | return l 61 | } 62 | 63 | // EstimatedLimit returns the current estimated limit. 64 | func (l *AIMDLimit) EstimatedLimit() int { 65 | l.mu.RLock() 66 | defer l.mu.RUnlock() 67 | return l.limit 68 | } 69 | 70 | // NotifyOnChange will register a callback to receive notification whenever the limit is updated to a new value. 71 | func (l *AIMDLimit) NotifyOnChange(consumer core.LimitChangeListener) { 72 | l.mu.Lock() 73 | l.listeners = append(l.listeners, consumer) 74 | l.mu.Unlock() 75 | } 76 | 77 | // notifyListeners will call the callbacks on limit changes 78 | func (l *AIMDLimit) notifyListeners(newLimit int) { 79 | for _, listener := range l.listeners { 80 | listener(newLimit) 81 | } 82 | } 83 | 84 | // OnSample the concurrency limit using a new rtt sample. 85 | func (l *AIMDLimit) OnSample(startTime int64, rtt int64, inFlight int, didDrop bool) { 86 | l.mu.Lock() 87 | defer l.mu.Unlock() 88 | 89 | l.commonSampler.Sample(rtt, inFlight, didDrop) 90 | 91 | if didDrop { 92 | l.limit = int(math.Max(1, math.Min(float64(l.limit-1), float64(l.limit)*l.backOffRatio))) 93 | l.notifyListeners(l.limit) 94 | } else if inFlight >= l.limit { 95 | l.limit += l.increaseBy 96 | l.notifyListeners(l.limit) 97 | } 98 | return 99 | } 100 | 101 | // BackOffRatio return the current back-off-ratio for the AIMDLimit 102 | func (l *AIMDLimit) BackOffRatio() float64 { 103 | l.mu.RLock() 104 | defer l.mu.RUnlock() 105 | return l.backOffRatio 106 | } 107 | 108 | func (l *AIMDLimit) String() string { 109 | return fmt.Sprintf("AIMDLimit{limit=%d, backOffRatio=%0.4f}", l.EstimatedLimit(), l.BackOffRatio()) 110 | } 111 | -------------------------------------------------------------------------------- /limit/aimd_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestAIMDLimit(t *testing.T) { 11 | t.Parallel() 12 | 13 | t.Run("DefaultAIMD", func(t2 *testing.T) { 14 | t2.Parallel() 15 | asrt := assert.New(t2) 16 | l := NewDefaultAIMDLimit("test", nil) 17 | asrt.Equal(10, l.EstimatedLimit()) 18 | }) 19 | 20 | t.Run("Default", func(t2 *testing.T) { 21 | t2.Parallel() 22 | asrt := assert.New(t2) 23 | l := NewAIMDLimit("test", 10, 0.9, 1, nil) 24 | asrt.Equal(10, l.EstimatedLimit()) 25 | asrt.Equal(0.9, l.BackOffRatio()) 26 | }) 27 | 28 | t.Run("IncreaseOnSuccess", func(t2 *testing.T) { 29 | t2.Parallel() 30 | asrt := assert.New(t2) 31 | l := NewAIMDLimit("test", 10, 0.9, 1, nil) 32 | listener := testNotifyListener{changes: make([]int, 0)} 33 | l.NotifyOnChange(listener.updater()) 34 | l.OnSample(-1, (time.Millisecond * 1).Nanoseconds(), 10, false) 35 | asrt.Equal(11, l.EstimatedLimit()) 36 | asrt.Equal(11, listener.changes[0]) 37 | }) 38 | 39 | t.Run("DecreaseOnDrops", func(t2 *testing.T) { 40 | t2.Parallel() 41 | asrt := assert.New(t2) 42 | l := NewAIMDLimit("test", 10, 0.9, 1, nil) 43 | l.OnSample(-1, 1, 1, true) 44 | asrt.Equal(9, l.EstimatedLimit()) 45 | }) 46 | 47 | t.Run("String", func(t2 *testing.T) { 48 | t2.Parallel() 49 | asrt := assert.New(t2) 50 | l := NewAIMDLimit("test", 10, 0.9, 1, nil) 51 | asrt.Equal("AIMDLimit{limit=10, backOffRatio=0.9000}", l.String()) 52 | }) 53 | } 54 | -------------------------------------------------------------------------------- /limit/doc.go: -------------------------------------------------------------------------------- 1 | // Package limit provides several useful limit implementations. 2 | package limit 3 | -------------------------------------------------------------------------------- /limit/fixed.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/platinummonkey/go-concurrency-limits/core" 7 | ) 8 | 9 | // FixedLimit is a non dynamic limit with fixed value. 10 | type FixedLimit struct { 11 | limit int 12 | registry core.MetricRegistry 13 | commonSampler *core.CommonMetricSampler 14 | } 15 | 16 | // NewFixedLimit will return a new FixedLimit 17 | func NewFixedLimit(name string, limit int, registry core.MetricRegistry, tags ...string) *FixedLimit { 18 | if limit < 0 { 19 | // force to a positive value 20 | limit = 10 21 | } 22 | if registry == nil { 23 | registry = core.EmptyMetricRegistryInstance 24 | } 25 | 26 | l := &FixedLimit{ 27 | limit: limit, 28 | registry: registry, 29 | } 30 | l.commonSampler = core.NewCommonMetricSamplerOrNil(registry, l, name, tags...) 31 | return l 32 | } 33 | 34 | // EstimatedLimit will return the current limit. 35 | func (l *FixedLimit) EstimatedLimit() int { 36 | return l.limit 37 | } 38 | 39 | // NotifyOnChange will register a callback to receive notification whenever the limit is updated to a new value. 40 | func (l *FixedLimit) NotifyOnChange(consumer core.LimitChangeListener) { 41 | // noop for fixed limit 42 | } 43 | 44 | // OnSample will update the limit with the sample. 45 | func (l *FixedLimit) OnSample(startTime int64, rtt int64, inFlight int, didDrop bool) { 46 | // noop for fixed limit, just record metrics 47 | l.commonSampler.Sample(rtt, inFlight, didDrop) 48 | } 49 | 50 | func (l FixedLimit) String() string { 51 | return fmt.Sprintf("FixedLimit{limit=%d}", l.limit) 52 | } 53 | -------------------------------------------------------------------------------- /limit/fixed_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestFixedLimit(t *testing.T) { 11 | t.Parallel() 12 | asrt := assert.New(t) 13 | l := NewFixedLimit("test", 10, nil) 14 | asrt.Equal(10, l.EstimatedLimit()) 15 | 16 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 10, false) 17 | asrt.Equal(10, l.EstimatedLimit()) 18 | 19 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 100, false) 20 | asrt.Equal(10, l.EstimatedLimit()) 21 | 22 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 100, true) 23 | asrt.Equal(10, l.EstimatedLimit()) 24 | 25 | // NOOP 26 | listener := testNotifyListener{} 27 | l.NotifyOnChange(listener.updater()) 28 | 29 | asrt.Equal("FixedLimit{limit=10}", l.String()) 30 | } 31 | -------------------------------------------------------------------------------- /limit/functions/dep.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | func max(a, b int) int { 4 | if a > b { 5 | return a 6 | } 7 | return b 8 | } 9 | -------------------------------------------------------------------------------- /limit/functions/doc.go: -------------------------------------------------------------------------------- 1 | // Package functions provides additional helper functions to the limit package. 2 | package functions 3 | -------------------------------------------------------------------------------- /limit/functions/fixed.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | // FixedQueueSizeFunc implements a fixed amount the estimated limit can grow while latencies remain low 4 | func FixedQueueSizeFunc(queueSize int) func(int) int { 5 | return func(_ int) int { 6 | return queueSize 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /limit/functions/fixed_test.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestFixedQueueSizeFunc(t *testing.T) { 10 | t.Parallel() 11 | 12 | f := FixedQueueSizeFunc(4) 13 | for i := -4; i < 5; i++ { 14 | assert.Equal(t, 4, f(0)) 15 | } 16 | } 17 | -------------------------------------------------------------------------------- /limit/functions/init.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | import ( 4 | "math" 5 | "os" 6 | "strconv" 7 | ) 8 | 9 | func init() { 10 | initSqrtFunc() 11 | initLog10RootFunc() 12 | } 13 | 14 | func initSqrtFunc() { 15 | // compute square root numbers up to Max(1000||env['GO_CONCURRENCY_LIMIT_SQRT_PRE_COMPUTE']) 16 | defaultVal := 1000 17 | valStr, ok := os.LookupEnv("GO_CONCURRENCY_LIMIT_SQRT_PRE_COMPUTE") 18 | if !ok { 19 | valStr = "1000" 20 | } 21 | val, err := strconv.Atoi(valStr) 22 | if err != nil { 23 | val = defaultVal 24 | } 25 | if val < defaultVal { 26 | val = defaultVal 27 | } 28 | for i := 0; i < val; i++ { 29 | // pre-compute the integer value with a min of 1 up to val samples. 30 | sqrtRootLookup = append(sqrtRootLookup, int(math.Max(1, float64(int(math.Sqrt(float64(i))))))) 31 | } 32 | } 33 | 34 | func initLog10RootFunc() { 35 | // compute log10 root numbers up to Max(1000||env['GO_CONCURRENCY_LIMIT_LOG10ROOT_PRE_COMPUTE']) 36 | defaultVal := 1000 37 | valStr, ok := os.LookupEnv("GO_CONCURRENCY_LIMIT_LOG10ROOT_PRE_COMPUTE") 38 | if !ok { 39 | valStr = "1000" 40 | } 41 | val, err := strconv.Atoi(valStr) 42 | if err != nil { 43 | val = defaultVal 44 | } 45 | if val < defaultVal { 46 | val = defaultVal 47 | } 48 | for i := 0; i < val; i++ { 49 | // pre-compute the integer value with a min of 1 up to val samples. 50 | log10RootLookup = append(log10RootLookup, int(math.Max(1, float64(int(math.Log10(float64(i))))))) 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /limit/functions/log10_root.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | import ( 4 | "math" 5 | ) 6 | 7 | var log10RootLookup []int 8 | 9 | // Log10RootFunction is a specialized utility function used by limiters to calculate thresholds using log10 of the 10 | // current limit. Here we pre-compute the log10 root of numbers up to 1000 (or more) because the log10 root 11 | // operation can be slow. 12 | func Log10RootFunction(baseline int) func(estimatedLimit int) int { 13 | return func(estimatedLimit int) int { 14 | if estimatedLimit < len(log10RootLookup) { 15 | return baseline + log10RootLookup[estimatedLimit] 16 | } 17 | return baseline + int(math.Log10(float64(estimatedLimit))) 18 | } 19 | } 20 | 21 | // Log10RootFloatFunction is a specialized utility function used by limiters to calculate thresholds using log10 of the 22 | // current limit. Here we pre-compute the log10 root of numbers up to 1000 (or more) because the log10 root 23 | // operation can be slow. 24 | func Log10RootFloatFunction(baseline float64) func(estimatedLimit float64) float64 { 25 | return func(estimatedLimit float64) float64 { 26 | if int(estimatedLimit) < len(log10RootLookup) { 27 | return baseline + float64(log10RootLookup[int(estimatedLimit)]) 28 | } 29 | return baseline + math.Log10(estimatedLimit) 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /limit/functions/log10_root_test.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestLog10RootFunction(t *testing.T) { 10 | t.Parallel() 11 | 12 | t.Run("ZeroIndex", func(t2 *testing.T) { 13 | t2.Parallel() 14 | f := Log10RootFunction(4) 15 | assert.Equal(t2, 5, f(0)) 16 | }) 17 | 18 | t.Run("MaxIndex", func(t2 *testing.T) { 19 | t2.Parallel() 20 | f := Log10RootFunction(4) 21 | assert.Equal(t2, 7, f(1000)) 22 | }) 23 | 24 | t.Run("OutOfLookupRange", func(t2 *testing.T) { 25 | t2.Parallel() 26 | f := Log10RootFunction(4) 27 | assert.Equal(t2, 7, f(2500)) 28 | }) 29 | } 30 | 31 | func TestLog10RootFloatFunction(t *testing.T) { 32 | t.Parallel() 33 | 34 | t.Run("ZeroIndex", func(t2 *testing.T) { 35 | t2.Parallel() 36 | f := Log10RootFloatFunction(4) 37 | assert.Equal(t2, 5.0, f(0)) 38 | }) 39 | 40 | t.Run("MaxIndex", func(t2 *testing.T) { 41 | t2.Parallel() 42 | f := Log10RootFloatFunction(4) 43 | assert.InDelta(t2, 7.0, f(1000), 0.001) 44 | }) 45 | 46 | t.Run("OutOfLookupRange", func(t2 *testing.T) { 47 | t2.Parallel() 48 | f := Log10RootFloatFunction(4) 49 | assert.Equal(t2, 7.3979400086720375, f(2500)) 50 | }) 51 | } 52 | -------------------------------------------------------------------------------- /limit/functions/square_root.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | import ( 4 | "math" 5 | ) 6 | 7 | var sqrtRootLookup []int 8 | 9 | // SqrtRootFunction is a specialized utility function used by limiters to calculate thresholds using square root of the 10 | // current limit. Here we pre-compute the square root of numbers up to 1000 (or more) because the square root 11 | // operation can be slow. 12 | func SqrtRootFunction(baseline int) func(estimatedLimit int) int { 13 | return func(estimatedLimit int) int { 14 | if estimatedLimit < len(sqrtRootLookup) { 15 | return max(baseline, sqrtRootLookup[estimatedLimit]) 16 | } 17 | return max(baseline, int(math.Sqrt(float64(estimatedLimit)))) 18 | } 19 | } 20 | -------------------------------------------------------------------------------- /limit/functions/square_root_test.go: -------------------------------------------------------------------------------- 1 | package functions 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestSqrtRootFunction(t *testing.T) { 10 | t.Parallel() 11 | 12 | t.Run("ZeroIndex", func(t2 *testing.T) { 13 | t2.Parallel() 14 | f := SqrtRootFunction(4) 15 | assert.Equal(t2, 4, f(0)) 16 | }) 17 | 18 | t.Run("MaxIndex", func(t2 *testing.T) { 19 | t2.Parallel() 20 | f := SqrtRootFunction(4) 21 | assert.Equal(t2, 31, f(999)) 22 | }) 23 | 24 | t.Run("OutOfLookupRange", func(t2 *testing.T) { 25 | t2.Parallel() 26 | f := SqrtRootFunction(4) 27 | assert.Equal(t2, 50, f(2500)) 28 | }) 29 | } 30 | -------------------------------------------------------------------------------- /limit/gradient2_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | ) 10 | 11 | func TestGradient2Limit(t *testing.T) { 12 | t.Parallel() 13 | 14 | t.Run("Default", func(t2 *testing.T) { 15 | t2.Parallel() 16 | asrt := assert.New(t2) 17 | l := NewDefaultGradient2Limit("test", nil, nil) 18 | asrt.NotNil(l) 19 | 20 | asrt.Equal(20, l.EstimatedLimit()) 21 | asrt.Equal("Gradient2Limit{limit=20}", l.String()) 22 | }) 23 | 24 | t.Run("OnSample", func(t2 *testing.T) { 25 | t2.Parallel() 26 | asrt := assert.New(t2) 27 | l, err := NewGradient2Limit( 28 | "test", 29 | 50, 30 | 0, 31 | 0, 32 | nil, 33 | -1, 34 | -1, 35 | NoopLimitLogger{}, 36 | core.EmptyMetricRegistryInstance, 37 | ) 38 | asrt.NoError(err) 39 | asrt.NotNil(l) 40 | listener := testNotifyListener{} 41 | l.NotifyOnChange(listener.updater()) 42 | 43 | // nothing should change 44 | l.OnSample(0, 10, 1, false) 45 | asrt.Equal(50, l.EstimatedLimit()) 46 | 47 | for i := 0; i < 25; i++ { 48 | l.OnSample(int64(i), 10*int64(i), i, false) 49 | asrt.Equal(50, l.EstimatedLimit()) 50 | } 51 | 52 | // dropped samples cut off limit, smoothed down 53 | l.OnSample(25, 2500, 25, true) 54 | asrt.Equal(45, l.EstimatedLimit()) 55 | asrt.Equal(45, listener.changes[0]) 56 | 57 | // test new sample shouldn't grow too fast 58 | l.OnSample(26, 10, 5, false) 59 | asrt.Equal(45, l.EstimatedLimit()) 60 | 61 | // drain down again 62 | for i := 0; i < 100; i++ { 63 | l.OnSample(int64(i+27), 1000, i, true) 64 | } 65 | asrt.Equal(21, l.EstimatedLimit()) 66 | 67 | // slowly grow back up 68 | for i := 0; i < 100; i++ { 69 | l.OnSample(int64(i+127), 1, 1, false) 70 | } 71 | asrt.Equal(21, l.EstimatedLimit()) 72 | }) 73 | } 74 | -------------------------------------------------------------------------------- /limit/gradient_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | ) 10 | 11 | func TestGradientLimit(t *testing.T) { 12 | t.Parallel() 13 | t.Run("nextProbeInterval", func(t2 *testing.T) { 14 | t2.Parallel() 15 | asrt := assert.New(t2) 16 | asrt.Equal(ProbeDisabled, nextProbeCountdown(ProbeDisabled)) 17 | asrt.True(nextProbeCountdown(1) > 0) 18 | }) 19 | 20 | t.Run("Default", func(t2 *testing.T) { 21 | t2.Parallel() 22 | asrt := assert.New(t2) 23 | l := NewGradientLimitWithRegistry( 24 | "test", 25 | 0, 26 | 0, 27 | 0, 28 | -1, 29 | nil, 30 | -1, 31 | 0, 32 | NoopLimitLogger{}, 33 | core.EmptyMetricRegistryInstance, 34 | ) 35 | 36 | asrt.Equal(50, l.EstimatedLimit()) 37 | asrt.Equal(int64(0), l.RTTNoLoad()) 38 | asrt.Equal("GradientLimit{limit=50, rttNoLoad=0 ms}", l.String()) 39 | }) 40 | 41 | t.Run("OnSample", func(t2 *testing.T) { 42 | t2.Parallel() 43 | asrt := assert.New(t2) 44 | l := NewGradientLimitWithRegistry( 45 | "test", 46 | 0, 47 | 0, 48 | 0, 49 | -1, 50 | nil, 51 | -1, 52 | 0, 53 | NoopLimitLogger{}, 54 | core.EmptyMetricRegistryInstance, 55 | ) 56 | listener := testNotifyListener{} 57 | l.NotifyOnChange(listener.updater()) 58 | // nothing should change 59 | l.OnSample(0, 10, 1, false) 60 | asrt.Equal(50, l.EstimatedLimit()) 61 | 62 | // dropped samples cut off limit, smoothed down 63 | l.OnSample(10, 1, 1, true) 64 | asrt.Equal(45, l.EstimatedLimit()) 65 | asrt.Equal(45, listener.changes[0]) 66 | 67 | // test new sample shouldn't grow with current conditions 68 | l.OnSample(20, 10, 5, false) 69 | asrt.Equal(45, l.EstimatedLimit()) 70 | 71 | // drain down pretty far 72 | for i := 0; i < 100; i++ { 73 | l.OnSample(int64(i*10+30), 0, 0, true) 74 | } 75 | asrt.Equal(4, l.EstimatedLimit()) 76 | 77 | // slowly grow back up 78 | for i := 0; i < 100; i++ { 79 | l.OnSample(int64(i*10+3030), 1, 5, false) 80 | } 81 | asrt.Equal(12, l.EstimatedLimit()) 82 | }) 83 | } 84 | -------------------------------------------------------------------------------- /limit/settable.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | "sync/atomic" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | ) 10 | 11 | // SettableLimit is a fixed limit that can be changed. 12 | // Note: to be used mostly for testing where the limit can be manually adjusted. 13 | type SettableLimit struct { 14 | limit int32 15 | 16 | listeners []core.LimitChangeListener 17 | commonSampler *core.CommonMetricSampler 18 | mu sync.RWMutex 19 | } 20 | 21 | // NewSettableLimit will create a new SettableLimit. 22 | func NewSettableLimit(name string, limit int, registry core.MetricRegistry, tags ...string) *SettableLimit { 23 | if limit < 0 { 24 | limit = 10 25 | } 26 | if registry == nil { 27 | registry = core.EmptyMetricRegistryInstance 28 | } 29 | 30 | l := &SettableLimit{ 31 | limit: int32(limit), 32 | listeners: make([]core.LimitChangeListener, 0), 33 | } 34 | l.commonSampler = core.NewCommonMetricSamplerOrNil(registry, l, name, tags...) 35 | return l 36 | } 37 | 38 | // EstimatedLimit will return the estimated limit. 39 | func (l *SettableLimit) EstimatedLimit() int { 40 | return int(atomic.LoadInt32(&l.limit)) 41 | } 42 | 43 | // NotifyOnChange will register a callback to receive notification whenever the limit is updated to a new value. 44 | func (l *SettableLimit) NotifyOnChange(consumer core.LimitChangeListener) { 45 | l.mu.Lock() 46 | l.listeners = append(l.listeners, consumer) 47 | l.mu.Unlock() 48 | } 49 | 50 | // notifyListeners will call the callbacks on limit changes 51 | func (l *SettableLimit) notifyListeners(newLimit int) { 52 | l.mu.Lock() 53 | defer l.mu.Unlock() 54 | for _, listener := range l.listeners { 55 | listener(newLimit) 56 | } 57 | } 58 | 59 | // OnSample will update the limit with the given sample. 60 | func (l *SettableLimit) OnSample(startTime int64, rtt int64, inFlight int, didDrop bool) { 61 | // noop for SettableLimit, just record metrics 62 | l.commonSampler.Sample(rtt, inFlight, didDrop) 63 | } 64 | 65 | // SetLimit will update the current limit. 66 | func (l *SettableLimit) SetLimit(limit int) { 67 | atomic.StoreInt32(&l.limit, int32(limit)) 68 | l.notifyListeners(limit) 69 | } 70 | 71 | func (l *SettableLimit) String() string { 72 | return fmt.Sprintf("SettableLimit{limit=%d}", atomic.LoadInt32(&l.limit)) 73 | } 74 | -------------------------------------------------------------------------------- /limit/settable_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "sync" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | 9 | "github.com/platinummonkey/go-concurrency-limits/core" 10 | ) 11 | 12 | type testNotifyListener struct { 13 | changes []int 14 | mu sync.Mutex 15 | } 16 | 17 | func (l *testNotifyListener) updater() core.LimitChangeListener { 18 | return func(limit int) { 19 | l.mu.Lock() 20 | l.changes = append(l.changes, limit) 21 | l.mu.Unlock() 22 | } 23 | } 24 | 25 | func TestSettableLimit(t *testing.T) { 26 | t.Parallel() 27 | asrt := assert.New(t) 28 | l := NewSettableLimit("test", 10, nil) 29 | asrt.Equal(10, l.EstimatedLimit()) 30 | 31 | l.SetLimit(5) 32 | asrt.Equal(5, l.EstimatedLimit()) 33 | 34 | // should be a noop 35 | l.OnSample(0, 10, 1, true) 36 | asrt.Equal(5, l.EstimatedLimit()) 37 | 38 | asrt.Equal("SettableLimit{limit=5}", l.String()) 39 | 40 | // notify on change 41 | listener := testNotifyListener{changes: make([]int, 0)} 42 | l.NotifyOnChange(listener.updater()) 43 | l.SetLimit(10) 44 | l.SetLimit(5) 45 | l.SetLimit(1) 46 | asrt.Equal(10, listener.changes[0]) 47 | asrt.Equal(5, listener.changes[1]) 48 | asrt.Equal(1, listener.changes[2]) 49 | } 50 | -------------------------------------------------------------------------------- /limit/traced.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | 7 | "github.com/platinummonkey/go-concurrency-limits/core" 8 | ) 9 | 10 | // Logger implements a basic dependency to log. Feel free to report stats as well. 11 | type Logger interface { 12 | // Log a Debug statement already formatted. 13 | Debugf(msg string, params ...interface{}) 14 | // Check if debug is enabled 15 | IsDebugEnabled() bool 16 | } 17 | 18 | // NoopLimitLogger implements a NO-OP logger, it does nothing. 19 | type NoopLimitLogger struct{} 20 | 21 | // Debugf debug formatted log 22 | func (l NoopLimitLogger) Debugf(msg string, params ...interface{}) {} 23 | 24 | // IsDebugEnabled will return true if debug is enabled. NoopLimitLogger is always `false` 25 | func (l NoopLimitLogger) IsDebugEnabled() bool { 26 | return false 27 | } 28 | 29 | func (l NoopLimitLogger) String() string { 30 | return "NoopLimitLogger{}" 31 | } 32 | 33 | // BuiltinLimitLogger implements a STDOUT limit logger. 34 | type BuiltinLimitLogger struct{} 35 | 36 | // Debugf debug formatted log 37 | func (l BuiltinLimitLogger) Debugf(msg string, params ...interface{}) { 38 | log.Println(fmt.Sprintf(msg, params...)) 39 | } 40 | 41 | // IsDebugEnabled will return true if debug is enabled. BuiltinLimitLogger is always `true` 42 | func (l BuiltinLimitLogger) IsDebugEnabled() bool { 43 | return true 44 | } 45 | 46 | func (l BuiltinLimitLogger) String() string { 47 | return "BuiltinLimitLogger{}" 48 | } 49 | 50 | // TracedLimit implements core.Limit but adds some additional logging 51 | type TracedLimit struct { 52 | limit core.Limit 53 | logger Logger 54 | } 55 | 56 | // NewTracedLimit returns a new wrapped Limit with TracedLimit. 57 | func NewTracedLimit(limit core.Limit, logger Logger) *TracedLimit { 58 | return &TracedLimit{ 59 | limit: limit, 60 | logger: logger, 61 | } 62 | } 63 | 64 | // EstimatedLimit returns the estimated limit. 65 | func (l *TracedLimit) EstimatedLimit() int { 66 | estimatedLimit := l.limit.EstimatedLimit() 67 | l.logger.Debugf("estimatedLimit=%d\n", estimatedLimit) 68 | return estimatedLimit 69 | } 70 | 71 | // NotifyOnChange will register a callback to receive notification whenever the limit is updated to a new value. 72 | func (l *TracedLimit) NotifyOnChange(consumer core.LimitChangeListener) { 73 | l.limit.NotifyOnChange(consumer) 74 | } 75 | 76 | // OnSample will log and deleate the update of the sample. 77 | func (l *TracedLimit) OnSample(startTime int64, rtt int64, inFlight int, didDrop bool) { 78 | l.logger.Debugf("startTime=%d, rtt=%d ms, inFlight=%d, didDrop=%t", startTime, rtt/1e6, inFlight, didDrop) 79 | l.limit.OnSample(startTime, rtt, inFlight, didDrop) 80 | } 81 | 82 | func (l *TracedLimit) String() string { 83 | return fmt.Sprintf("TracedLimit{limit=%v, logger=%v}", l.limit, l.logger) 84 | } 85 | -------------------------------------------------------------------------------- /limit/traced_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestNoopLimitLogger(t *testing.T) { 10 | t.Parallel() 11 | asrt := assert.New(t) 12 | l := NoopLimitLogger{} 13 | asrt.NotPanics(func() { l.Debugf("") }) 14 | asrt.False(l.IsDebugEnabled()) 15 | asrt.Equal("NoopLimitLogger{}", l.String()) 16 | } 17 | 18 | func TestBuiltinLimitLogger(t *testing.T) { 19 | t.Parallel() 20 | asrt := assert.New(t) 21 | l := BuiltinLimitLogger{} 22 | asrt.NotPanics(func() { l.Debugf("") }) 23 | asrt.True(l.IsDebugEnabled()) 24 | asrt.Equal("BuiltinLimitLogger{}", l.String()) 25 | } 26 | 27 | func TestTracedLimit(t *testing.T) { 28 | t.Parallel() 29 | asrt := assert.New(t) 30 | delegate := NewSettableLimit("test", 10, nil) 31 | l := NewTracedLimit(delegate, NoopLimitLogger{}) 32 | listener := testNotifyListener{} 33 | l.NotifyOnChange(listener.updater()) 34 | 35 | asrt.Equal(10, l.EstimatedLimit()) 36 | 37 | l.OnSample(0, 0, 1, true) 38 | asrt.Equal(10, l.EstimatedLimit()) 39 | 40 | asrt.Equal("TracedLimit{limit=SettableLimit{limit=10}, logger=NoopLimitLogger{}}", l.String()) 41 | } 42 | -------------------------------------------------------------------------------- /limit/vegas_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/stretchr/testify/assert" 8 | 9 | "github.com/platinummonkey/go-concurrency-limits/core" 10 | "github.com/platinummonkey/go-concurrency-limits/limit/functions" 11 | ) 12 | 13 | func createVegasLimit() *VegasLimit { 14 | return NewVegasLimitWithRegistry( 15 | "test", 16 | 10, 17 | nil, 18 | 20, 19 | 1.0, 20 | functions.FixedQueueSizeFunc(3), 21 | functions.FixedQueueSizeFunc(6), 22 | nil, 23 | nil, 24 | nil, 25 | 0, 26 | NoopLimitLogger{}, 27 | core.EmptyMetricRegistryInstance) 28 | } 29 | 30 | func TestVegasLimit(t *testing.T) { 31 | t.Parallel() 32 | 33 | t.Run("NewDefaultVegasLimit", func(t2 *testing.T) { 34 | t2.Parallel() 35 | l := NewDefaultVegasLimit("test", NoopLimitLogger{}, nil) 36 | assert.Equal(t2, 20, l.EstimatedLimit()) 37 | }) 38 | 39 | t.Run("NewDefaultVegasLimitWithLimit", func(t2 *testing.T) { 40 | t2.Parallel() 41 | l := NewDefaultVegasLimitWithLimit("test", 5, NoopLimitLogger{}, nil) 42 | assert.Equal(t2, 5, l.EstimatedLimit()) 43 | }) 44 | 45 | t.Run("InitialLimit", func(t2 *testing.T) { 46 | t2.Parallel() 47 | l := createVegasLimit() 48 | assert.Equal(t2, l.EstimatedLimit(), 10) 49 | assert.Equal(t2, l.RTTNoLoad(), int64(0)) 50 | assert.Equal(t2, "VegasLimit{limit=10, rttNoLoad=0 ms}", l.String()) 51 | }) 52 | 53 | t.Run("IncreaseLimit", func(t2 *testing.T) { 54 | t2.Parallel() 55 | asrt := assert.New(t2) 56 | l := createVegasLimit() 57 | listener := testNotifyListener{} 58 | l.NotifyOnChange(listener.updater()) 59 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 10, false) 60 | asrt.Equal(10, l.EstimatedLimit()) 61 | l.OnSample(10, (time.Millisecond * 10).Nanoseconds(), 11, false) 62 | asrt.Equal(16, l.EstimatedLimit()) 63 | asrt.Equal(16, listener.changes[0]) 64 | }) 65 | 66 | t.Run("DecreaseLimit", func(t2 *testing.T) { 67 | t2.Parallel() 68 | asrt := assert.New(t2) 69 | l := createVegasLimit() 70 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 10, false) 71 | asrt.Equal(10, l.EstimatedLimit()) 72 | l.OnSample(10, (time.Millisecond * 50).Nanoseconds(), 11, false) 73 | asrt.Equal(9, l.EstimatedLimit()) 74 | }) 75 | 76 | t.Run("NoChangeIfWithinThresholds", func(t2 *testing.T) { 77 | t2.Parallel() 78 | asrt := assert.New(t2) 79 | l := createVegasLimit() 80 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 10, false) 81 | asrt.Equal(10, l.EstimatedLimit()) 82 | l.OnSample(10, (time.Millisecond * 14).Nanoseconds(), 14, false) 83 | asrt.Equal(10, l.EstimatedLimit()) 84 | }) 85 | 86 | t.Run("DecreaseSmoothing", func(t2 *testing.T) { 87 | t2.Parallel() 88 | asrt := assert.New(t2) 89 | l := NewVegasLimitWithRegistry( 90 | "test", 91 | 100, 92 | nil, 93 | 200, 94 | 0.5, 95 | nil, 96 | nil, 97 | nil, 98 | nil, 99 | func(estimatedLimit float64) float64 { 100 | return estimatedLimit / 2.0 101 | }, 102 | 0, 103 | NoopLimitLogger{}, 104 | core.EmptyMetricRegistryInstance) 105 | 106 | // Pick up first min-rtt 107 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 100, false) 108 | asrt.Equal(100, l.EstimatedLimit()) 109 | 110 | // First decrease 111 | l.OnSample(10, (time.Millisecond * 20).Nanoseconds(), 100, false) 112 | asrt.Equal(75, l.EstimatedLimit()) 113 | 114 | // Second decrease 115 | l.OnSample(20, (time.Millisecond * 20).Nanoseconds(), 100, false) 116 | asrt.Equal(56, l.EstimatedLimit()) 117 | }) 118 | 119 | t.Run("DecreaseWithoutSmoothing", func(t2 *testing.T) { 120 | t2.Parallel() 121 | asrt := assert.New(t2) 122 | l := NewVegasLimitWithRegistry( 123 | "test", 124 | 100, 125 | nil, 126 | 200, 127 | -1, 128 | nil, 129 | nil, 130 | nil, 131 | nil, 132 | func(estimatedLimit float64) float64 { 133 | return estimatedLimit / 2.0 134 | }, 135 | 0, 136 | NoopLimitLogger{}, 137 | core.EmptyMetricRegistryInstance) 138 | 139 | // Pick up first min-rtt 140 | l.OnSample(0, (time.Millisecond * 10).Nanoseconds(), 100, false) 141 | asrt.Equal(100, l.EstimatedLimit()) 142 | 143 | // First decrease 144 | l.OnSample(10, (time.Millisecond * 20).Nanoseconds(), 100, false) 145 | asrt.Equal(50, l.EstimatedLimit()) 146 | 147 | // Second decrease 148 | l.OnSample(20, (time.Millisecond * 20).Nanoseconds(), 100, false) 149 | asrt.Equal(25, l.EstimatedLimit()) 150 | }) 151 | } 152 | -------------------------------------------------------------------------------- /limit/windowed.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "fmt" 5 | "math" 6 | "sync" 7 | "time" 8 | 9 | "github.com/platinummonkey/go-concurrency-limits/core" 10 | "github.com/platinummonkey/go-concurrency-limits/measurements" 11 | ) 12 | 13 | // WindowedLimit implements a windowed limit 14 | type WindowedLimit struct { 15 | minWindowTime int64 // Minimum window duration for sampling a new minRtt 16 | maxWindowTime int64 // Maximum window duration for sampling a new minRtt 17 | nextUpdateTime int64 // End time for the sampling window at which point the limit should be updated 18 | windowSize int32 // Minimum sampling window size for finding a new minimum rtt 19 | minRTTThreshold int64 20 | 21 | delegate core.Limit 22 | sample *measurements.ImmutableSampleWindow 23 | listeners []core.LimitChangeListener 24 | registry core.MetricRegistry 25 | commonSampler *core.CommonMetricSampler 26 | 27 | mu sync.RWMutex 28 | } 29 | 30 | const ( 31 | defaultWindowedMinWindowTime = 1e9 // 1 second 32 | defaultWindowedMaxWindowTime = 1e9 // 1 second 33 | defaultWindowedMinRTTThreshold = 1e8 // 100 microseconds 34 | defaultWindowedWindowSize = 10 35 | ) 36 | 37 | // NewDefaultWindowedLimit will create a new default WindowedLimit 38 | func NewDefaultWindowedLimit( 39 | name string, 40 | delegate core.Limit, 41 | registry core.MetricRegistry, 42 | tags ...string, 43 | ) *WindowedLimit { 44 | l, _ := NewWindowedLimit( 45 | name, 46 | defaultWindowedMinWindowTime, 47 | defaultWindowedMaxWindowTime, 48 | defaultWindowedWindowSize, 49 | defaultWindowedMinRTTThreshold, 50 | delegate, 51 | registry, 52 | tags..., 53 | ) 54 | return l 55 | } 56 | 57 | // NewWindowedLimit will create a new WindowedLimit 58 | func NewWindowedLimit( 59 | name string, 60 | minWindowTime int64, 61 | maxWindowTime int64, 62 | windowSize int32, 63 | minRTTThreshold int64, 64 | delegate core.Limit, 65 | registry core.MetricRegistry, 66 | tags ...string, 67 | ) (*WindowedLimit, error) { 68 | if minWindowTime < (time.Duration(100) * time.Millisecond).Nanoseconds() { 69 | return nil, fmt.Errorf("minWindowTime must be >= 100 ms") 70 | } 71 | 72 | if maxWindowTime < (time.Duration(100) * time.Millisecond).Nanoseconds() { 73 | return nil, fmt.Errorf("maxWindowTime must be >= 100 ms") 74 | } 75 | 76 | if windowSize < 10 { 77 | return nil, fmt.Errorf("windowSize must be >= 10 ms") 78 | } 79 | 80 | if delegate == nil { 81 | return nil, fmt.Errorf("delegate must be specified") 82 | } 83 | 84 | if registry == nil { 85 | registry = core.EmptyMetricRegistryInstance 86 | } 87 | 88 | l := &WindowedLimit{ 89 | minWindowTime: minWindowTime, 90 | maxWindowTime: maxWindowTime, 91 | nextUpdateTime: 0, 92 | windowSize: windowSize, 93 | minRTTThreshold: minRTTThreshold, 94 | delegate: delegate, 95 | sample: measurements.NewDefaultImmutableSampleWindow(), 96 | listeners: make([]core.LimitChangeListener, 0), 97 | registry: registry, 98 | } 99 | l.commonSampler = core.NewCommonMetricSamplerOrNil(registry, l, name, tags...) 100 | return l, nil 101 | 102 | } 103 | 104 | // EstimatedLimit returns the current estimated limit. 105 | func (l *WindowedLimit) EstimatedLimit() int { 106 | l.mu.RLock() 107 | defer l.mu.RUnlock() 108 | return l.delegate.EstimatedLimit() 109 | } 110 | 111 | // NotifyOnChange will register a callback to receive notification whenever the limit is updated to a new value. 112 | func (l *WindowedLimit) NotifyOnChange(consumer core.LimitChangeListener) { 113 | l.mu.Lock() 114 | l.listeners = append(l.listeners, consumer) 115 | l.delegate.NotifyOnChange(consumer) 116 | l.mu.Unlock() 117 | } 118 | 119 | // notifyListeners will call the callbacks on limit changes 120 | func (l *WindowedLimit) notifyListeners(newLimit int) { 121 | for _, listener := range l.listeners { 122 | listener(newLimit) 123 | } 124 | } 125 | 126 | // OnSample the concurrency limit using a new rtt sample. 127 | func (l *WindowedLimit) OnSample(startTime int64, rtt int64, inFlight int, didDrop bool) { 128 | l.mu.Lock() 129 | defer l.mu.Unlock() 130 | l.commonSampler.Sample(rtt, inFlight, didDrop) 131 | 132 | endTime := startTime + rtt 133 | if rtt < l.minRTTThreshold { 134 | return 135 | } 136 | 137 | if didDrop { 138 | l.sample = l.sample.AddDroppedSample(-1, inFlight) 139 | } else { 140 | l.sample = l.sample.AddSample(-1, rtt, inFlight) 141 | } 142 | 143 | if endTime > l.nextUpdateTime && l.isWindowReady(rtt, inFlight) { 144 | current := l.sample 145 | l.sample = measurements.NewDefaultImmutableSampleWindow() 146 | l.nextUpdateTime = endTime + minInt64(maxInt64(current.CandidateRTTNanoseconds()*2, l.minWindowTime), l.maxWindowTime) 147 | l.delegate.OnSample(startTime, current.AverageRTTNanoseconds(), current.MaxInFlight(), didDrop) 148 | } 149 | } 150 | 151 | func (l *WindowedLimit) String() string { 152 | l.mu.RLock() 153 | defer l.mu.RUnlock() 154 | return fmt.Sprintf("WindowedLimit{minWindowTime=%d, maxWindowTime=%d, minRTTThreshold=%d, windowSize=%d,"+ 155 | " delegate=%v", l.minWindowTime, l.maxWindowTime, l.minRTTThreshold, l.windowSize, l.delegate) 156 | } 157 | 158 | func (l *WindowedLimit) isWindowReady(rtt int64, inFlight int) bool { 159 | return rtt < int64(math.MaxInt64) && int32(inFlight) > l.windowSize 160 | } 161 | 162 | func minInt64(a, b int64) int64 { 163 | if a < b { 164 | return a 165 | } 166 | return b 167 | } 168 | 169 | func maxInt64(a, b int64) int64 { 170 | if a > b { 171 | return a 172 | } 173 | return b 174 | } 175 | -------------------------------------------------------------------------------- /limit/windowed_test.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestWindowedLimit(t *testing.T) { 11 | t.Parallel() 12 | 13 | t.Run("DefaultWindowedLimit", func(t2 *testing.T) { 14 | t2.Parallel() 15 | asrt := assert.New(t2) 16 | delegate := NewSettableLimit("test", 10, nil) 17 | l := NewDefaultWindowedLimit("test", delegate, nil) 18 | asrt.Equal(10, l.EstimatedLimit()) 19 | }) 20 | 21 | t.Run("NewWindowedLimit", func(t2 *testing.T) { 22 | t2.Parallel() 23 | asrt := assert.New(t2) 24 | delegate := NewSettableLimit("test", 10, nil) 25 | minWindowTime := (time.Millisecond * 100).Nanoseconds() 26 | l, err := NewWindowedLimit("test", minWindowTime, minWindowTime*2, 10, 10, delegate, nil) 27 | asrt.NoError(err) 28 | asrt.NotNil(l) 29 | asrt.Equal(10, l.EstimatedLimit()) 30 | }) 31 | 32 | t.Run("DecreaseOnDrops", func(t2 *testing.T) { 33 | t2.Parallel() 34 | asrt := assert.New(t2) 35 | delegate := NewDefaultAIMDLimit("test", nil) 36 | minWindowTime := (time.Millisecond * 100).Nanoseconds() 37 | l, err := NewWindowedLimit("test", minWindowTime, minWindowTime*2, 10, 10, delegate, nil) 38 | asrt.NoError(err) 39 | asrt.NotNil(l) 40 | asrt.Equal(10, l.EstimatedLimit()) 41 | 42 | l.OnSample(0, 10, 1, false) 43 | asrt.Equal(10, l.EstimatedLimit()) 44 | for i := 0; i < 10; i++ { 45 | l.OnSample(0, minWindowTime*1000, 15, true) 46 | } 47 | asrt.Equal(9, l.EstimatedLimit()) 48 | }) 49 | 50 | t.Run("IncreaseOnSuccess", func(t2 *testing.T) { 51 | t2.Parallel() 52 | asrt := assert.New(t2) 53 | delegate := NewDefaultAIMDLimit("test", nil) 54 | minWindowTime := (time.Millisecond * 100).Nanoseconds() 55 | l, err := NewWindowedLimit("test", minWindowTime, minWindowTime*2, 10, 10, delegate, nil) 56 | asrt.NoError(err) 57 | asrt.NotNil(l) 58 | asrt.Equal(10, l.EstimatedLimit()) 59 | 60 | listener := testNotifyListener{} 61 | l.NotifyOnChange(listener.updater()) 62 | 63 | for i := 0; i < 40; i++ { 64 | l.OnSample(l.minWindowTime*int64(i*i), minWindowTime+10, 15, false) 65 | } 66 | asrt.Equal(16, l.EstimatedLimit()) 67 | asrt.Equal([]int{11, 12, 13, 14, 15, 16}, listener.changes) 68 | }) 69 | 70 | t.Run("String", func(t2 *testing.T) { 71 | t2.Parallel() 72 | asrt := assert.New(t2) 73 | delegate := NewSettableLimit("test", 10, nil) 74 | minWindowTime := (time.Millisecond * 100).Nanoseconds() 75 | l, err := NewWindowedLimit("test", minWindowTime, minWindowTime*2, 10, 10, delegate, nil) 76 | asrt.NoError(err) 77 | asrt.NotNil(l) 78 | asrt.Equal("WindowedLimit{minWindowTime=100000000, maxWindowTime=200000000, minRTTThreshold=10, "+ 79 | "windowSize=10, delegate=SettableLimit{limit=10}", l.String()) 80 | }) 81 | } 82 | -------------------------------------------------------------------------------- /limiter/blocking.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sync" 7 | "time" 8 | 9 | "github.com/platinummonkey/go-concurrency-limits/core" 10 | "github.com/platinummonkey/go-concurrency-limits/limit" 11 | ) 12 | 13 | // blockUntilSignaled will wait for context cancellation, an unblock signal or timeout 14 | // This method will return true if we were successfully signalled. 15 | func blockUntilSignaled(ctx context.Context, c *sync.Cond, timeout time.Duration) bool { 16 | ready := make(chan struct{}) 17 | 18 | go func() { 19 | c.L.Lock() 20 | defer c.L.Unlock() 21 | c.Wait() 22 | close(ready) 23 | }() 24 | 25 | if timeout > 0 { 26 | // use NewTimer over time.After so that we don't have to 27 | // wait for the timeout to elapse in order to release memory 28 | timer := time.NewTimer(timeout) 29 | defer timer.Stop() 30 | 31 | select { 32 | case <-ctx.Done(): 33 | return false 34 | case <-ready: 35 | return true 36 | case <-timer.C: 37 | return false 38 | } 39 | } 40 | 41 | select { 42 | case <-ctx.Done(): 43 | return false 44 | case <-ready: 45 | return true 46 | } 47 | } 48 | 49 | // BlockingLimiter implements a Limiter that blocks the caller when the limit has been reached. The caller is 50 | // blocked until the limiter has been released. This limiter is commonly used in batch clients that use the limiter 51 | // as a back-pressure mechanism. 52 | type BlockingLimiter struct { 53 | logger limit.Logger 54 | delegate core.Limiter 55 | c *sync.Cond 56 | timeout time.Duration 57 | } 58 | 59 | // NewBlockingLimiter will create a new blocking limiter 60 | func NewBlockingLimiter( 61 | delegate core.Limiter, 62 | timeout time.Duration, 63 | logger limit.Logger, 64 | ) *BlockingLimiter { 65 | mu := sync.Mutex{} 66 | if timeout < 0 { 67 | timeout = 0 68 | } 69 | if logger == nil { 70 | logger = limit.NoopLimitLogger{} 71 | } 72 | return &BlockingLimiter{ 73 | logger: logger, 74 | delegate: delegate, 75 | c: sync.NewCond(&mu), 76 | timeout: timeout, 77 | } 78 | } 79 | 80 | // tryAcquire will block when attempting to acquire a token 81 | func (l *BlockingLimiter) tryAcquire(ctx context.Context) (core.Listener, bool) { 82 | 83 | for { 84 | // if the context has already been cancelled, fail quickly 85 | if err := ctx.Err(); err != nil { 86 | l.logger.Debugf("context cancelled ctx=%v", ctx) 87 | return nil, false 88 | } 89 | 90 | // try to acquire a new token and return immediately if successful 91 | listener, ok := l.delegate.Acquire(ctx) 92 | if ok && listener != nil { 93 | l.logger.Debugf("delegate returned a listener ctx=%v", ctx) 94 | return listener, true 95 | } 96 | 97 | // We have reached the limit so block until: 98 | // - A token is released 99 | // - A timeout 100 | // - The context is cancelled 101 | l.logger.Debugf("Blocking waiting for release or timeout ctx=%v", ctx) 102 | if shouldAcquire := blockUntilSignaled(ctx, l.c, l.timeout); shouldAcquire { 103 | listener, ok := l.delegate.Acquire(ctx) 104 | if ok && listener != nil { 105 | l.logger.Debugf("delegate returned a listener ctx=%v", ctx) 106 | return listener, true 107 | } 108 | } 109 | l.logger.Debugf("blocking released, trying again to acquire ctx=%v", ctx) 110 | } 111 | } 112 | 113 | // Acquire a token from the limiter. Returns `nil, false` if the limit has been exceeded. 114 | // If acquired the caller must call one of the Listener methods when the operation has been completed to release 115 | // the count. 116 | // 117 | // context Context for the request. The context is used by advanced strategies such as LookupPartitionStrategy. 118 | func (l *BlockingLimiter) Acquire(ctx context.Context) (core.Listener, bool) { 119 | delegateListener, ok := l.tryAcquire(ctx) 120 | if !ok && delegateListener == nil { 121 | l.logger.Debugf("did not acquire ctx=%v", ctx) 122 | return nil, false 123 | } 124 | l.logger.Debugf("acquired, returning listener ctx=%v", ctx) 125 | return &DelegateListener{ 126 | delegateListener: delegateListener, 127 | c: l.c, 128 | }, true 129 | } 130 | 131 | func (l BlockingLimiter) String() string { 132 | return fmt.Sprintf("BlockingLimiter{delegate=%v}", l.delegate) 133 | } 134 | -------------------------------------------------------------------------------- /limiter/blocking_test.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | "sync" 7 | "testing" 8 | "time" 9 | 10 | "github.com/stretchr/testify/assert" 11 | 12 | "github.com/platinummonkey/go-concurrency-limits/core" 13 | "github.com/platinummonkey/go-concurrency-limits/limit" 14 | "github.com/platinummonkey/go-concurrency-limits/strategy" 15 | ) 16 | 17 | type testListener struct { 18 | successCount int 19 | ignoreCount int 20 | dropCount int 21 | } 22 | 23 | func (l *testListener) OnSuccess() { 24 | l.successCount++ 25 | } 26 | 27 | func (l *testListener) OnIgnore() { 28 | l.ignoreCount++ 29 | } 30 | 31 | func (l *testListener) OnDropped() { 32 | l.dropCount++ 33 | } 34 | 35 | type contextKey string 36 | 37 | var testContextKey contextKey = "jobID" 38 | 39 | func BenchmarkBlockingLimiter(b *testing.B) { 40 | 41 | benchLimiter := func(b *testing.B, limitCount, acquireCount int, timeout time.Duration) { 42 | b.StopTimer() 43 | 44 | l := limit.NewSettableLimit("test", limitCount, nil) 45 | defaultLimiter, err := NewDefaultLimiter( 46 | l, 47 | defaultMinWindowTime, 48 | defaultMaxWindowTime, 49 | defaultMinRTTThreshold, 50 | defaultWindowSize, 51 | strategy.NewSimpleStrategy(limitCount), 52 | limit.NoopLimitLogger{}, 53 | core.EmptyMetricRegistryInstance, 54 | ) 55 | if err != nil { 56 | b.Fatal(err.Error()) 57 | } 58 | 59 | blockingLimiter := NewBlockingLimiter(defaultLimiter, timeout, limit.NoopLimitLogger{}) 60 | 61 | wg := sync.WaitGroup{} 62 | wg.Add(acquireCount) 63 | 64 | startCount := sync.WaitGroup{} 65 | startCount.Add(acquireCount) 66 | 67 | released := make(chan int, acquireCount) 68 | startAcquire := make(chan struct{}) 69 | for i := 0; i < acquireCount; i++ { 70 | go func(j int) { 71 | startCount.Done() 72 | defer wg.Done() 73 | <-startAcquire 74 | 75 | listener, ok := blockingLimiter.Acquire(context.Background()) 76 | if ok && listener != nil { 77 | listener.OnSuccess() 78 | released <- 1 79 | return 80 | } 81 | released <- 0 82 | }(i) 83 | } 84 | 85 | startCount.Wait() 86 | b.StartTimer() 87 | close(startAcquire) 88 | 89 | wg.Wait() 90 | 91 | b.StopTimer() 92 | sumReleased := 0 93 | for i := 0; i < acquireCount; i++ { 94 | sumReleased += <-released 95 | } 96 | if sumReleased != acquireCount { 97 | b.Fatal("not enough released") 98 | } 99 | b.StartTimer() 100 | } 101 | 102 | b.Run("limiter_contention", func(b *testing.B) { 103 | b.ReportAllocs() 104 | for i := 0; i < b.N; i++ { 105 | benchLimiter(b, 2, 10, 0) 106 | } 107 | }) 108 | 109 | b.Run("limiter_no_contention", func(b *testing.B) { 110 | b.ReportAllocs() 111 | for i := 0; i < b.N; i++ { 112 | benchLimiter(b, 100, 10, 0) 113 | } 114 | }) 115 | 116 | b.Run("limiter_contention_w_timeout", func(b *testing.B) { 117 | b.ReportAllocs() 118 | for i := 0; i < b.N; i++ { 119 | benchLimiter(b, 2, 10, 10*time.Minute) 120 | } 121 | }) 122 | 123 | b.Run("limiter_no_contention_w_timeout", func(b *testing.B) { 124 | b.ReportAllocs() 125 | for i := 0; i < b.N; i++ { 126 | benchLimiter(b, 100, 10, 10*time.Minute) 127 | } 128 | }) 129 | 130 | } 131 | 132 | func TestBlockingLimiter(t *testing.T) { 133 | t.Run("Unblocked", func(t2 *testing.T) { 134 | asrt := assert.New(t2) 135 | l := limit.NewSettableLimit("test", 10, nil) 136 | noopLogger := limit.NoopLimitLogger{} 137 | defaultLimiter, err := NewDefaultLimiter( 138 | l, 139 | defaultMinWindowTime, 140 | defaultMaxWindowTime, 141 | defaultMinRTTThreshold, 142 | defaultWindowSize, 143 | strategy.NewSimpleStrategy(10), 144 | noopLogger, 145 | core.EmptyMetricRegistryInstance, 146 | ) 147 | if !asrt.NoError(err) { 148 | asrt.FailNow("") 149 | } 150 | asrt.NotNil(defaultLimiter) 151 | blockingLimiter := NewBlockingLimiter(defaultLimiter, 0, noopLogger) 152 | // stringer 153 | asrt.True(strings.Contains(blockingLimiter.String(), "BlockingLimiter{delegate=DefaultLimiter{")) 154 | 155 | var listeners []core.Listener 156 | for i := 0; i < 10; i++ { 157 | listener, ok := blockingLimiter.Acquire(context.Background()) 158 | if ok && listener != nil { 159 | listeners = append(listeners, listener) 160 | } 161 | } 162 | 163 | l.SetLimit(1) 164 | 165 | for _, listener := range listeners { 166 | listener.OnSuccess() 167 | } 168 | 169 | blockingLimiter.Acquire(context.Background()) 170 | }) 171 | 172 | t.Run("MultipleBlocked", func(t2 *testing.T) { 173 | asrt := assert.New(t2) 174 | l := limit.NewSettableLimit("test", 1, nil) 175 | noopLogger := limit.NoopLimitLogger{} 176 | defaultLimiter, err := NewDefaultLimiter( 177 | l, 178 | defaultMinWindowTime, 179 | defaultMaxWindowTime, 180 | defaultMinRTTThreshold, 181 | defaultWindowSize, 182 | strategy.NewSimpleStrategy(1), 183 | noopLogger, 184 | core.EmptyMetricRegistryInstance, 185 | ) 186 | if !asrt.NoError(err) { 187 | asrt.FailNow("") 188 | } 189 | asrt.NotNil(defaultLimiter) 190 | blockingLimiter := NewBlockingLimiter(defaultLimiter, 0, noopLogger) 191 | 192 | wg := sync.WaitGroup{} 193 | wg.Add(8) 194 | 195 | released := make(chan int, 8) 196 | 197 | for i := 0; i < 8; i++ { 198 | go func(j int) { 199 | defer wg.Done() 200 | listener, ok := blockingLimiter.Acquire(context.Background()) 201 | if ok && listener != nil { 202 | listener.OnSuccess() 203 | released <- 1 204 | return 205 | } 206 | released <- 0 207 | }(i) 208 | } 209 | 210 | wg.Wait() 211 | 212 | sumReleased := 0 213 | for i := 0; i < 8; i++ { 214 | sumReleased += <-released 215 | } 216 | asrt.Equal(8, sumReleased) 217 | }) 218 | 219 | t.Run("BlockingLimiterTimeout", func(t2 *testing.T) { 220 | asrt := assert.New(t2) 221 | l := limit.NewSettableLimit("test", 1, nil) 222 | noopLogger := limit.NoopLimitLogger{} 223 | defaultLimiter, err := NewDefaultLimiter( 224 | l, 225 | defaultMinWindowTime, 226 | defaultMaxWindowTime, 227 | defaultMinRTTThreshold, 228 | defaultWindowSize, 229 | strategy.NewSimpleStrategy(1), 230 | noopLogger, 231 | core.EmptyMetricRegistryInstance, 232 | ) 233 | if !asrt.NoError(err) { 234 | asrt.FailNow("") 235 | } 236 | asrt.NotNil(defaultLimiter) 237 | blockingLimiter := NewBlockingLimiter(defaultLimiter, time.Millisecond*25, noopLogger) 238 | 239 | wg := sync.WaitGroup{} 240 | wg.Add(8) 241 | 242 | released := make(chan int, 8) 243 | 244 | for i := 0; i < 8; i++ { 245 | go func(j int) { 246 | defer wg.Done() 247 | ctx, cancel := context.WithTimeout(context.WithValue(context.Background(), testContextKey, j), time.Millisecond*400) 248 | defer cancel() 249 | listener, ok := blockingLimiter.Acquire(ctx) 250 | if ok && listener != nil { 251 | time.Sleep(time.Millisecond * 100) 252 | listener.OnSuccess() 253 | released <- 1 254 | return 255 | } 256 | released <- 0 257 | }(i) 258 | time.Sleep(time.Nanosecond * 50) 259 | } 260 | 261 | wg.Wait() 262 | 263 | sumReleased := 0 264 | for i := 0; i < 8; i++ { 265 | sumReleased += <-released 266 | } 267 | // we only expect half of them to complete before their deadlines 268 | asrt.InDelta(4, sumReleased, 1.0, "expected roughly half to succeed") 269 | }) 270 | } 271 | 272 | func TestBlockUntilSignaled(t *testing.T) { 273 | asrt := assert.New(t) 274 | var mu sync.Mutex 275 | cond := sync.NewCond(&mu) 276 | 277 | // Use a channel to control test completion 278 | done := make(chan struct{}) 279 | 280 | // Create a context for testing cancellation 281 | ctx, cancel := context.WithCancel(context.Background()) 282 | 283 | const waitTime = 50 * time.Millisecond 284 | go func() { 285 | time.Sleep(waitTime * 2) 286 | cond.Signal() 287 | close(done) 288 | }() 289 | 290 | timeout := 500 * time.Millisecond 291 | now := time.Now() 292 | signalled := blockUntilSignaled(ctx, cond, timeout) 293 | asrt.True(signalled) 294 | asrt.True(time.Since(now) >= waitTime, "expected to wait at least 50ms") 295 | 296 | <-done 297 | cancel() 298 | } 299 | -------------------------------------------------------------------------------- /limiter/deadline.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sync" 7 | "time" 8 | 9 | "github.com/platinummonkey/go-concurrency-limits/core" 10 | "github.com/platinummonkey/go-concurrency-limits/limit" 11 | ) 12 | 13 | // DeadlineLimiter that blocks the caller when the limit has been reached. The caller is 14 | // blocked until the limiter has been released, or a deadline has been passed. 15 | type DeadlineLimiter struct { 16 | logger limit.Logger 17 | delegate core.Limiter 18 | deadline time.Time 19 | c *sync.Cond 20 | } 21 | 22 | // NewDeadlineLimiter will create a new DeadlineLimiter that will wrap a limiter such that acquire will block until a 23 | // provided deadline if the limit was reached instead of returning an empty listener immediately. 24 | func NewDeadlineLimiter( 25 | delegate core.Limiter, 26 | deadline time.Time, 27 | logger limit.Logger, 28 | ) *DeadlineLimiter { 29 | mu := sync.Mutex{} 30 | if logger == nil { 31 | logger = limit.NoopLimitLogger{} 32 | } 33 | return &DeadlineLimiter{ 34 | logger: logger, 35 | delegate: delegate, 36 | c: sync.NewCond(&mu), 37 | deadline: deadline, 38 | } 39 | } 40 | 41 | // tryAcquire will block when attempting to acquire a token 42 | func (l *DeadlineLimiter) tryAcquire(ctx context.Context) (listener core.Listener, ok bool) { 43 | 44 | for { 45 | // if the context has already been cancelled, fail quickly 46 | if err := ctx.Err(); err != nil { 47 | l.logger.Debugf("context cancelled ctx=%v", ctx) 48 | return nil, false 49 | } 50 | 51 | // if the deadline has passed, fail quickly 52 | if time.Now().UTC().After(l.deadline) { 53 | return nil, false 54 | } 55 | 56 | // try to acquire a new token and return immediately if successful 57 | listener, ok := l.delegate.Acquire(ctx) 58 | if ok && listener != nil { 59 | l.logger.Debugf("delegate returned a listener ctx=%v", ctx) 60 | return listener, true 61 | } 62 | 63 | // We have reached the limit so block until a token is released 64 | timeout := l.deadline.Sub(time.Now().UTC()) 65 | 66 | // We have reached the limit so block until: 67 | // - A token is released 68 | // - A timeout 69 | // - The context is cancelled 70 | l.logger.Debugf("Blocking waiting for release or timeout ctx=%v", ctx) 71 | if shouldAcquire := blockUntilSignaled(ctx, l.c, timeout); shouldAcquire { 72 | listener, ok := l.delegate.Acquire(ctx) 73 | if ok && listener != nil { 74 | l.logger.Debugf("delegate returned a listener ctx=%v", ctx) 75 | return listener, true 76 | } 77 | } 78 | l.logger.Debugf("blocking released, trying again to acquire ctx=%v", ctx) 79 | } 80 | } 81 | 82 | // Acquire a token from the limiter. Returns `nil, false` if the limit has been exceeded. 83 | // If acquired the caller must call one of the Listener methods when the operation has been completed to release 84 | // the count. 85 | // 86 | // context Context for the request. The context is used by advanced strategies such as LookupPartitionStrategy. 87 | func (l *DeadlineLimiter) Acquire(ctx context.Context) (listener core.Listener, ok bool) { 88 | delegateListener, ok := l.tryAcquire(ctx) 89 | if !ok && delegateListener == nil { 90 | l.logger.Debugf("did not acquire ctx=%v", ctx) 91 | return nil, false 92 | } 93 | l.logger.Debugf("acquired, returning listener ctx=%v", ctx) 94 | return &DelegateListener{ 95 | delegateListener: delegateListener, 96 | c: l.c, 97 | }, true 98 | } 99 | 100 | // String implements Stringer for easy debugging. 101 | func (l DeadlineLimiter) String() string { 102 | return fmt.Sprintf("DeadlineLimiter{delegate=%v}", l.delegate) 103 | } 104 | -------------------------------------------------------------------------------- /limiter/deadline_test.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | "sync" 7 | "sync/atomic" 8 | "testing" 9 | "time" 10 | 11 | "github.com/platinummonkey/go-concurrency-limits/core" 12 | "github.com/platinummonkey/go-concurrency-limits/limit" 13 | "github.com/platinummonkey/go-concurrency-limits/strategy" 14 | "github.com/stretchr/testify/assert" 15 | ) 16 | 17 | func TestDeadlineLimiter(t *testing.T) { 18 | t.Run("Unblocked", func(t2 *testing.T) { 19 | asrt := assert.New(t2) 20 | l := limit.NewSettableLimit("test", 10, nil) 21 | defaultLimiter, err := NewDefaultLimiter( 22 | l, 23 | defaultMinWindowTime, 24 | defaultMaxWindowTime, 25 | defaultMinRTTThreshold, 26 | defaultWindowSize, 27 | strategy.NewSimpleStrategy(10), 28 | nil, 29 | core.EmptyMetricRegistryInstance, 30 | ) 31 | if !asrt.NoError(err) { 32 | asrt.FailNow("") 33 | } 34 | asrt.NotNil(defaultLimiter) 35 | deadline := time.Now().Add(time.Second * 15) 36 | deadlineLimiter := NewDeadlineLimiter(defaultLimiter, deadline, nil) 37 | // stringer 38 | asrt.True(strings.Contains(deadlineLimiter.String(), "DeadlineLimiter{delegate=DefaultLimiter{")) 39 | 40 | var listeners []core.Listener 41 | for i := 0; i < 10; i++ { 42 | listener, ok := deadlineLimiter.Acquire(context.Background()) 43 | if ok && listener != nil { 44 | listeners = append(listeners, listener) 45 | } 46 | } 47 | 48 | l.SetLimit(1) 49 | 50 | for _, listener := range listeners { 51 | listener.OnSuccess() 52 | } 53 | 54 | deadlineLimiter.Acquire(context.Background()) 55 | }) 56 | 57 | t.Run("Deadline passed", func(t2 *testing.T) { 58 | asrt := assert.New(t2) 59 | l := limit.NewSettableLimit("test", 1, nil) 60 | noopLogger := limit.NoopLimitLogger{} 61 | defaultLimiter, err := NewDefaultLimiter( 62 | l, 63 | defaultMinWindowTime, 64 | defaultMaxWindowTime, 65 | defaultMinRTTThreshold, 66 | defaultWindowSize, 67 | strategy.NewSimpleStrategy(1), 68 | noopLogger, 69 | core.EmptyMetricRegistryInstance, 70 | ) 71 | if !asrt.NoError(err) { 72 | asrt.FailNow("") 73 | } 74 | asrt.NotNil(defaultLimiter) 75 | deadline := time.Now().Add(time.Second * 1) 76 | deadlineLimiter := NewDeadlineLimiter(defaultLimiter, deadline, noopLogger) 77 | 78 | i := 1 79 | deadlineFound := false 80 | 81 | for { 82 | listener, ok := deadlineLimiter.Acquire(context.Background()) 83 | if ok && listener != nil { 84 | listener.OnSuccess() 85 | time.Sleep(time.Second) 86 | } else if i > 3 { 87 | break 88 | } else { 89 | deadlineFound = true 90 | break 91 | } 92 | i++ 93 | } 94 | 95 | asrt.True(deadlineFound, "expected deadline to be reached but not after %d attempts", i) 96 | asrt.Equal(2, i, "expected deadline to be exceeded on second attempt") 97 | }) 98 | 99 | t.Run("limit reached", func(t2 *testing.T) { 100 | asrt := assert.New(t2) 101 | l := limit.NewFixedLimit("test", 1, nil) 102 | noopLogger := limit.BuiltinLimitLogger{} 103 | defaultLimiter, err := NewDefaultLimiter( 104 | l, 105 | defaultMinWindowTime, 106 | defaultMaxWindowTime, 107 | defaultMinRTTThreshold, 108 | defaultWindowSize, 109 | strategy.NewSimpleStrategy(1), 110 | noopLogger, 111 | core.EmptyMetricRegistryInstance, 112 | ) 113 | if !asrt.NoError(err) { 114 | asrt.FailNow("") 115 | } 116 | asrt.NotNil(defaultLimiter) 117 | deadline := time.Now().Add(time.Second * 1) 118 | deadlineLimiter := NewDeadlineLimiter(defaultLimiter, deadline, noopLogger) 119 | 120 | var deadlineFound atomic.Value 121 | deadlineFound.Store(false) 122 | var wg sync.WaitGroup 123 | wg.Add(3) 124 | 125 | for i := 0; i < 3; i++ { 126 | go func(c int) { 127 | defer wg.Done() 128 | listener, ok := deadlineLimiter.Acquire(context.Background()) 129 | if ok && listener != nil { 130 | time.Sleep(time.Second) 131 | listener.OnSuccess() 132 | } else { 133 | deadlineFound.Store(true) 134 | } 135 | }(i) 136 | time.Sleep(time.Millisecond * 5) 137 | } 138 | wg.Wait() 139 | 140 | asrt.True(deadlineFound.Load().(bool), "expected deadline limit to be reached but not after 2 attempts") 141 | }) 142 | } 143 | -------------------------------------------------------------------------------- /limiter/default_test.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | "testing" 7 | "time" 8 | 9 | "github.com/stretchr/testify/assert" 10 | 11 | "github.com/platinummonkey/go-concurrency-limits/core" 12 | "github.com/platinummonkey/go-concurrency-limits/limit" 13 | "github.com/platinummonkey/go-concurrency-limits/measurements" 14 | "github.com/platinummonkey/go-concurrency-limits/strategy" 15 | ) 16 | 17 | func TestDefaultListener(t *testing.T) { 18 | t.Parallel() 19 | asrt := assert.New(t) 20 | inFlight := int64(3) 21 | releaseCount := int64(0) 22 | f := func() { 23 | releaseCount++ 24 | } 25 | limiter, _ := NewDefaultLimiterWithDefaults( 26 | "", 27 | strategy.NewSimpleStrategy(10), 28 | limit.NoopLimitLogger{}, 29 | core.EmptyMetricRegistryInstance, 30 | ) 31 | limiter.sample = measurements.NewDefaultImmutableSampleWindow() 32 | listener := DefaultListener{ 33 | currentMaxInFlight: 1, 34 | inFlight: &inFlight, 35 | token: core.NewAcquiredStrategyToken(1, f), 36 | startTime: time.Now().Unix(), 37 | minRTTThreshold: 10, 38 | limiter: limiter, 39 | nextUpdateTime: time.Now().Add(time.Minute * 10).Unix(), 40 | } 41 | 42 | // On Success 43 | listener.OnSuccess() 44 | asrt.Equal(int64(2), inFlight) 45 | asrt.Equal(int64(1), releaseCount) 46 | 47 | // On Ignore 48 | listener.OnIgnore() 49 | asrt.Equal(int64(1), inFlight) 50 | asrt.Equal(int64(2), releaseCount) 51 | 52 | // On Dropped 53 | listener.OnDropped() 54 | asrt.Equal(int64(0), inFlight) 55 | asrt.Equal(int64(3), releaseCount) 56 | } 57 | 58 | func TestDefaultLimiter(t *testing.T) { 59 | t.Parallel() 60 | 61 | t.Run("NewDefaultLimiterWithDefaults", func(t2 *testing.T) { 62 | t2.Parallel() 63 | asrt := assert.New(t2) 64 | l, err := NewDefaultLimiterWithDefaults("", strategy.NewSimpleStrategy(10), limit.NoopLimitLogger{}, core.EmptyMetricRegistryInstance) 65 | asrt.NoError(err) 66 | asrt.NotNil(l) 67 | asrt.Equal(20, l.EstimatedLimit()) 68 | asrt.True(strings.Contains(l.String(), "DefaultLimiter{RTTCandidate=")) 69 | }) 70 | 71 | t.Run("NewDefaultLimiter", func(t2 *testing.T) { 72 | t2.Parallel() 73 | asrt := assert.New(t2) 74 | l, err := NewDefaultLimiter( 75 | limit.NewFixedLimit("test", 10, nil), 76 | defaultMinWindowTime, 77 | defaultMaxWindowTime, 78 | defaultMinRTTThreshold, 79 | defaultWindowSize, 80 | strategy.NewSimpleStrategy(10), 81 | limit.NoopLimitLogger{}, 82 | core.EmptyMetricRegistryInstance, 83 | ) 84 | asrt.NoError(err) 85 | asrt.NotNil(l) 86 | asrt.Equal(10, l.EstimatedLimit()) 87 | }) 88 | 89 | t.Run("Acquire", func(t2 *testing.T) { 90 | t2.Parallel() 91 | asrt := assert.New(t2) 92 | l, err := NewDefaultLimiter( 93 | limit.NewFixedLimit("test", 10, nil), 94 | defaultMinWindowTime, 95 | defaultMaxWindowTime, 96 | defaultMinRTTThreshold, 97 | defaultWindowSize, 98 | strategy.NewSimpleStrategy(10), 99 | limit.NoopLimitLogger{}, 100 | core.EmptyMetricRegistryInstance, 101 | ) 102 | asrt.NoError(err) 103 | asrt.NotNil(l) 104 | asrt.Equal(10, l.EstimatedLimit()) 105 | 106 | listeners := make([]core.Listener, 0) 107 | 108 | // Acquire tokens 109 | for i := 0; i < 10; i++ { 110 | listener, ok := l.Acquire(context.Background()) 111 | asrt.True(ok) 112 | asrt.NotNil(listener) 113 | listeners = append(listeners, listener) 114 | } 115 | 116 | // try to acquire one more 117 | listener, ok := l.Acquire(context.Background()) 118 | asrt.False(ok) 119 | asrt.Nil(listener) 120 | 121 | // release all 122 | for _, listener = range listeners { 123 | listener.OnSuccess() 124 | } 125 | 126 | listener, ok = l.Acquire(context.Background()) 127 | asrt.True(ok) 128 | asrt.NotNil(listener) 129 | listener.OnSuccess() 130 | }) 131 | } 132 | -------------------------------------------------------------------------------- /limiter/delegate_listener.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "sync" 5 | 6 | "github.com/platinummonkey/go-concurrency-limits/core" 7 | ) 8 | 9 | // DelegateListener wraps the wrapped Limiter's Listener to simply delegate as a wrapper. 10 | type DelegateListener struct { 11 | delegateListener core.Listener 12 | c *sync.Cond 13 | } 14 | 15 | // NewDelegateListener creates a new wrapped listener. 16 | func NewDelegateListener(delegateListener core.Listener) *DelegateListener { 17 | mu := sync.Mutex{} 18 | return &DelegateListener{ 19 | delegateListener: delegateListener, 20 | c: sync.NewCond(&mu), 21 | } 22 | } 23 | 24 | // OnDropped is called to indicate the request failed and was dropped due to being rejected by an external limit or 25 | // hitting a timeout. Loss based Limit implementations will likely do an aggressive reducing in limit when this 26 | // happens. 27 | func (l *DelegateListener) OnDropped() { 28 | l.delegateListener.OnDropped() 29 | // unblock 30 | l.c.Broadcast() 31 | } 32 | 33 | // OnIgnore is called to indicate the operation failed before any meaningful RTT measurement could be made and 34 | // should be ignored to not introduce an artificially low RTT. 35 | func (l *DelegateListener) OnIgnore() { 36 | l.delegateListener.OnIgnore() 37 | // unblock 38 | l.c.Broadcast() 39 | } 40 | 41 | // OnSuccess is called as a notification that the operation succeeded and internally measured latency should be 42 | // used as an RTT sample. 43 | func (l *DelegateListener) OnSuccess() { 44 | l.delegateListener.OnSuccess() 45 | // unblock 46 | l.c.Broadcast() 47 | } 48 | -------------------------------------------------------------------------------- /limiter/delegate_listener_test.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestDelegateListener(t *testing.T) { 10 | asrt := assert.New(t) 11 | delegateListener := testListener{} 12 | listener := NewDelegateListener(&delegateListener) 13 | listener.OnSuccess() 14 | asrt.Equal(1, delegateListener.successCount) 15 | listener.OnIgnore() 16 | asrt.Equal(1, delegateListener.ignoreCount) 17 | listener.OnDropped() 18 | asrt.Equal(1, delegateListener.dropCount) 19 | } 20 | -------------------------------------------------------------------------------- /limiter/doc.go: -------------------------------------------------------------------------------- 1 | // Package limiter provides common limiter implementations that are useful. 2 | package limiter 3 | -------------------------------------------------------------------------------- /limiter/fifo_blocking.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "time" 5 | 6 | "github.com/platinummonkey/go-concurrency-limits/core" 7 | ) 8 | 9 | // FifoBlockingLimiter implements a Limiter that blocks the caller when the limit has been reached. This strategy 10 | // ensures the resource is properly protected but favors availability over latency by not fast failing requests when 11 | // the limit has been reached. To help keep success latencies low and minimize timeouts any blocked requests are 12 | // processed in last in/first out order. 13 | // 14 | // Use this limiter only when the concurrency model allows the limiter to be blocked. 15 | // Deprecated in favor of QueueBlockingLimiter 16 | type FifoBlockingLimiter struct { 17 | *QueueBlockingLimiter 18 | } 19 | 20 | // NewFifoBlockingLimiter will create a new FifoBlockingLimiter 21 | // Deprecated, use NewQueueBlockingLimiterFromConfig instead 22 | func NewFifoBlockingLimiter( 23 | delegate core.Limiter, 24 | maxBacklogSize int, 25 | maxBacklogTimeout time.Duration, 26 | ) *FifoBlockingLimiter { 27 | 28 | return &FifoBlockingLimiter{ 29 | NewQueueBlockingLimiterFromConfig(delegate, QueueLimiterConfig{ 30 | Ordering: OrderingFIFO, 31 | MaxBacklogSize: maxBacklogSize, 32 | MaxBacklogTimeout: maxBacklogTimeout, 33 | }), 34 | } 35 | } 36 | 37 | // NewFifoBlockingLimiterWithDefaults will create a new FifoBlockingLimiter with default values. 38 | // Deprecated, use NewQueueBlockingLimiterWithDefaults instead 39 | func NewFifoBlockingLimiterWithDefaults( 40 | delegate core.Limiter, 41 | ) *FifoBlockingLimiter { 42 | return NewFifoBlockingLimiter(delegate, 100, time.Millisecond*1000) 43 | } 44 | -------------------------------------------------------------------------------- /limiter/fifo_blocking_test.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | "sync" 7 | "testing" 8 | 9 | "github.com/stretchr/testify/assert" 10 | 11 | "github.com/platinummonkey/go-concurrency-limits/core" 12 | "github.com/platinummonkey/go-concurrency-limits/limit" 13 | "github.com/platinummonkey/go-concurrency-limits/strategy" 14 | ) 15 | 16 | func TestFifoBlockingLimiter(t *testing.T) { 17 | t.Parallel() 18 | 19 | t.Run("NewFifoBlockingLimiterWithDefaults", func(t2 *testing.T) { 20 | t2.Parallel() 21 | asrt := assert.New(t2) 22 | delegateLimiter, _ := NewDefaultLimiterWithDefaults( 23 | "", 24 | strategy.NewSimpleStrategy(20), 25 | limit.NoopLimitLogger{}, 26 | core.EmptyMetricRegistryInstance, 27 | ) 28 | limiter := NewFifoBlockingLimiterWithDefaults(delegateLimiter) 29 | asrt.NotNil(limiter) 30 | asrt.True(strings.Contains(limiter.String(), "QueueBlockingLimiter{delegate=DefaultLimiter{")) 31 | }) 32 | 33 | t.Run("NewFifoBlockingLimiter", func(t2 *testing.T) { 34 | t2.Parallel() 35 | asrt := assert.New(t2) 36 | delegateLimiter, _ := NewDefaultLimiterWithDefaults( 37 | "", 38 | strategy.NewSimpleStrategy(20), 39 | limit.NoopLimitLogger{}, 40 | core.EmptyMetricRegistryInstance, 41 | ) 42 | limiter := NewFifoBlockingLimiter(delegateLimiter, -1, 0) 43 | asrt.NotNil(limiter) 44 | asrt.True(strings.Contains(limiter.String(), "QueueBlockingLimiter{delegate=DefaultLimiter{")) 45 | }) 46 | 47 | t.Run("Acquire", func(t2 *testing.T) { 48 | t2.Parallel() 49 | asrt := assert.New(t2) 50 | delegateLimiter, _ := NewDefaultLimiter( 51 | limit.NewFixedLimit("test", 10, nil), 52 | defaultMinWindowTime, 53 | defaultMaxWindowTime, 54 | defaultMinRTTThreshold, 55 | defaultWindowSize, 56 | strategy.NewSimpleStrategy(10), 57 | limit.NoopLimitLogger{}, 58 | core.EmptyMetricRegistryInstance, 59 | ) 60 | limiter := NewFifoBlockingLimiterWithDefaults(delegateLimiter) 61 | asrt.NotNil(limiter) 62 | 63 | // acquire all tokens first 64 | listeners := make([]core.Listener, 0) 65 | for i := 0; i < 10; i++ { 66 | listener, ok := limiter.Acquire(context.Background()) 67 | asrt.True(ok) 68 | asrt.NotNil(listener) 69 | listeners = append(listeners, listener) 70 | } 71 | 72 | // queue up 10 more waiting 73 | waitingListeners := make([]acquiredListenerFifo, 0) 74 | mu := sync.Mutex{} 75 | startupReady := make(chan bool, 1) 76 | wg := sync.WaitGroup{} 77 | wg.Add(10) 78 | for i := 0; i < 10; i++ { 79 | if i > 0 { 80 | select { 81 | case <-startupReady: 82 | // proceed 83 | } 84 | } 85 | go func(j int) { 86 | startupReady <- true 87 | listener, ok := limiter.Acquire(context.Background()) 88 | asrt.True(ok) 89 | asrt.NotNil(listener) 90 | mu.Lock() 91 | waitingListeners = append(waitingListeners, acquiredListenerFifo{id: j, listener: listener}) 92 | mu.Unlock() 93 | wg.Done() 94 | }(i) 95 | } 96 | 97 | // release all other listeners, so we can continue 98 | for _, listener := range listeners { 99 | listener.OnSuccess() 100 | } 101 | 102 | // wait for others 103 | wg.Wait() 104 | 105 | // check all eventually required. Note: due to scheduling, it's not entirely LIFO as scheduling will allow 106 | // some non-determinism 107 | asrt.Len(waitingListeners, 10) 108 | // release all 109 | for _, acquired := range waitingListeners { 110 | if acquired.listener != nil { 111 | acquired.listener.OnSuccess() 112 | } 113 | } 114 | }) 115 | } 116 | -------------------------------------------------------------------------------- /limiter/lifo_blocking.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "time" 5 | 6 | "github.com/platinummonkey/go-concurrency-limits/core" 7 | ) 8 | 9 | // LifoBlockingLimiter implements a Limiter that blocks the caller when the limit has been reached. This strategy 10 | // ensures the resource is properly protected but favors availability over latency by not fast failing requests when 11 | // the limit has been reached. To help keep success latencies low and minimize timeouts any blocked requests are 12 | // processed in last in/first out order. 13 | // 14 | // Use this limiter only when the concurrency model allows the limiter to be blocked. 15 | // Deprecated in favor of QueueBlockingLimiter 16 | type LifoBlockingLimiter struct { 17 | *QueueBlockingLimiter 18 | } 19 | 20 | // NewLifoBlockingLimiter will create a new LifoBlockingLimiter 21 | // Deprecated, use NewQueueBlockingLimiterFromConfig instead 22 | func NewLifoBlockingLimiter( 23 | delegate core.Limiter, 24 | maxBacklogSize int, 25 | maxBacklogTimeout time.Duration, 26 | registry core.MetricRegistry, 27 | tags ...string, 28 | ) *LifoBlockingLimiter { 29 | 30 | return &LifoBlockingLimiter{ 31 | NewQueueBlockingLimiterFromConfig( 32 | delegate, 33 | QueueLimiterConfig{ 34 | MaxBacklogSize: maxBacklogSize, 35 | MaxBacklogTimeout: maxBacklogTimeout, 36 | MetricRegistry: registry, 37 | Tags: tags, 38 | }, 39 | ), 40 | } 41 | } 42 | 43 | // NewLifoBlockingLimiterWithDefaults will create a new LifoBlockingLimiter with default values. 44 | // Deprecated, use NewQueueBlockingLimiterFromConfig 45 | func NewLifoBlockingLimiterWithDefaults( 46 | delegate core.Limiter, 47 | ) *LifoBlockingLimiter { 48 | return NewLifoBlockingLimiter(delegate, 0, 0, nil) 49 | } 50 | -------------------------------------------------------------------------------- /limiter/lifo_blocking_test.go: -------------------------------------------------------------------------------- 1 | package limiter 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | "sync" 7 | "testing" 8 | 9 | "github.com/stretchr/testify/assert" 10 | 11 | "github.com/platinummonkey/go-concurrency-limits/core" 12 | "github.com/platinummonkey/go-concurrency-limits/limit" 13 | "github.com/platinummonkey/go-concurrency-limits/strategy" 14 | ) 15 | 16 | type acquiredListenerLifo struct { 17 | id int 18 | listener core.Listener 19 | } 20 | 21 | func TestLifoBlockingLimiter(t *testing.T) { 22 | t.Parallel() 23 | 24 | t.Run("NewLifoBlockingLimiterWithDefaults", func(t2 *testing.T) { 25 | t2.Parallel() 26 | asrt := assert.New(t2) 27 | delegateLimiter, _ := NewDefaultLimiterWithDefaults( 28 | "", 29 | strategy.NewSimpleStrategy(20), 30 | limit.NoopLimitLogger{}, 31 | core.EmptyMetricRegistryInstance, 32 | ) 33 | limiter := NewLifoBlockingLimiterWithDefaults(delegateLimiter) 34 | asrt.NotNil(limiter) 35 | asrt.True(strings.Contains(limiter.String(), "QueueBlockingLimiter{delegate=DefaultLimiter{")) 36 | }) 37 | 38 | t.Run("NewLifoBlockingLimiter", func(t2 *testing.T) { 39 | t2.Parallel() 40 | asrt := assert.New(t2) 41 | delegateLimiter, _ := NewDefaultLimiterWithDefaults( 42 | "", 43 | strategy.NewSimpleStrategy(20), 44 | limit.NoopLimitLogger{}, 45 | core.EmptyMetricRegistryInstance, 46 | ) 47 | limiter := NewLifoBlockingLimiter(delegateLimiter, -1, 0, nil) 48 | asrt.NotNil(limiter) 49 | asrt.True(strings.Contains(limiter.String(), "QueueBlockingLimiter{delegate=DefaultLimiter{")) 50 | }) 51 | 52 | t.Run("Acquire", func(t2 *testing.T) { 53 | t2.Parallel() 54 | asrt := assert.New(t2) 55 | delegateLimiter, _ := NewDefaultLimiter( 56 | limit.NewFixedLimit("test", 10, nil), 57 | defaultMinWindowTime, 58 | defaultMaxWindowTime, 59 | defaultMinRTTThreshold, 60 | defaultWindowSize, 61 | strategy.NewSimpleStrategy(10), 62 | limit.NoopLimitLogger{}, 63 | core.EmptyMetricRegistryInstance, 64 | ) 65 | limiter := NewLifoBlockingLimiterWithDefaults(delegateLimiter) 66 | asrt.NotNil(limiter) 67 | 68 | // acquire all tokens first 69 | listeners := make([]core.Listener, 0) 70 | for i := 0; i < 10; i++ { 71 | listener, ok := limiter.Acquire(context.Background()) 72 | asrt.True(ok) 73 | asrt.NotNil(listener) 74 | listeners = append(listeners, listener) 75 | } 76 | 77 | // queue up 10 more waiting 78 | waitingListeners := make([]acquiredListenerLifo, 0) 79 | mu := sync.Mutex{} 80 | startupReady := make(chan bool, 1) 81 | wg := sync.WaitGroup{} 82 | wg.Add(10) 83 | for i := 0; i < 10; i++ { 84 | if i > 0 { 85 | select { 86 | case <-startupReady: 87 | // proceed 88 | } 89 | } 90 | go func(j int) { 91 | startupReady <- true 92 | listener, ok := limiter.Acquire(context.Background()) 93 | asrt.True(ok, "must be true for j %d", j) 94 | asrt.NotNil(listener, "must be not be nil for j %d", j) 95 | mu.Lock() 96 | waitingListeners = append(waitingListeners, acquiredListenerLifo{id: j, listener: listener}) 97 | mu.Unlock() 98 | wg.Done() 99 | }(i) 100 | } 101 | 102 | // release all other listeners, so we can continue 103 | for _, listener := range listeners { 104 | listener.OnSuccess() 105 | } 106 | 107 | // wait for others 108 | wg.Wait() 109 | 110 | // check all eventually required. Note: due to scheduling, it's not entirely LIFO as scheduling will allow 111 | // some non-determinism 112 | asrt.Len(waitingListeners, 10) 113 | // release all 114 | for _, acquired := range waitingListeners { 115 | if acquired.listener != nil { 116 | acquired.listener.OnSuccess() 117 | } 118 | } 119 | }) 120 | } 121 | -------------------------------------------------------------------------------- /measurements/doc.go: -------------------------------------------------------------------------------- 1 | // Package measurements provides measurement reading implementations 2 | package measurements 3 | -------------------------------------------------------------------------------- /measurements/exponential_average.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | ) 7 | 8 | // ExponentialAverageMeasurement is an exponential average measurement implementation. 9 | type ExponentialAverageMeasurement struct { 10 | value float64 11 | sum float64 12 | window int 13 | warmupWindow int 14 | count int 15 | 16 | mu sync.RWMutex 17 | } 18 | 19 | // NewExponentialAverageMeasurement will create a new ExponentialAverageMeasurement 20 | func NewExponentialAverageMeasurement( 21 | window int, 22 | warmupWindow int, 23 | ) *ExponentialAverageMeasurement { 24 | return &ExponentialAverageMeasurement{ 25 | window: window, 26 | warmupWindow: warmupWindow, 27 | } 28 | } 29 | 30 | // Add a single sample and update the internal state. 31 | func (m *ExponentialAverageMeasurement) Add(value float64) (float64, bool) { 32 | m.mu.Lock() 33 | defer m.mu.Unlock() 34 | if m.count < m.warmupWindow { 35 | m.count++ 36 | m.sum += value 37 | m.value = m.sum / float64(m.count) 38 | } else { 39 | f := factor(m.window) 40 | m.value = m.value*(1-f) + value*f 41 | } 42 | return m.value, true 43 | } 44 | 45 | // Get the current value. 46 | func (m *ExponentialAverageMeasurement) Get() float64 { 47 | m.mu.RLock() 48 | defer m.mu.RUnlock() 49 | return m.value 50 | } 51 | 52 | // Reset the internal state as if no samples were ever added. 53 | func (m *ExponentialAverageMeasurement) Reset() { 54 | m.mu.Lock() 55 | defer m.mu.Unlock() 56 | m.value = 0 57 | m.count = 0 58 | m.sum = 0 59 | } 60 | 61 | // Update will update the value given an operation function 62 | func (m *ExponentialAverageMeasurement) Update(operation func(value float64) float64) { 63 | m.mu.Lock() 64 | defer m.mu.Unlock() 65 | m.value = operation(m.value) 66 | } 67 | 68 | func factor(n int) float64 { 69 | return 2.0 / float64(n+1) 70 | } 71 | 72 | func (m *ExponentialAverageMeasurement) String() string { 73 | m.mu.RLock() 74 | defer m.mu.RUnlock() 75 | return fmt.Sprintf( 76 | "ExponentialAverageMeasurement{value=%0.5f, count=%d, window=%d, warmupWindow=%d}", 77 | m.value, m.count, m.window, m.warmupWindow, 78 | ) 79 | } 80 | -------------------------------------------------------------------------------- /measurements/exponential_average_test.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestExponentialAverageMeasurement(t *testing.T) { 10 | t.Parallel() 11 | asrt := assert.New(t) 12 | m := NewExponentialAverageMeasurement(100, 10) 13 | 14 | expected := []float64{10, 10.5, 11, 11.5, 12, 12.5, 13, 13.5, 14, 14.5} 15 | for i := 0; i < 10; i++ { 16 | result, _ := m.Add(float64(i + 10)) 17 | asrt.Equal(expected[i], result) 18 | } 19 | 20 | m.Add(100) 21 | asrt.InDelta(float64(16.2), m.Get(), 0.01) 22 | 23 | m.Update(func(value float64) float64 { 24 | return value - 1 25 | }) 26 | asrt.InDelta(float64(15.19), m.Get(), 0.01) 27 | 28 | m.Reset() 29 | asrt.Equal(float64(0), m.Get()) 30 | 31 | asrt.Equal( 32 | "ExponentialAverageMeasurement{value=0.00000, count=0, window=100, warmupWindow=10}", 33 | m.String(), 34 | ) 35 | } 36 | -------------------------------------------------------------------------------- /measurements/immutable_sample_window.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "fmt" 5 | "math" 6 | "time" 7 | ) 8 | 9 | // ImmutableSampleWindow is used to track immutable samples atomically. 10 | type ImmutableSampleWindow struct { 11 | startTime int64 12 | minRTT int64 13 | maxInFlight int 14 | sampleCount int 15 | sum int64 16 | didDrop bool 17 | } 18 | 19 | // NewDefaultImmutableSampleWindow will create a new ImmutableSampleWindow with defaults 20 | func NewDefaultImmutableSampleWindow() *ImmutableSampleWindow { 21 | return NewImmutableSampleWindow( 22 | time.Now().UnixNano(), 23 | math.MaxInt64, 24 | 0, 25 | 0, 26 | 0, 27 | false, 28 | ) 29 | } 30 | 31 | // NewImmutableSampleWindow will create a new ImmutableSampleWindow with defaults 32 | func NewImmutableSampleWindow( 33 | startTime int64, 34 | minRTT int64, 35 | sum int64, 36 | maxInFlight int, 37 | sampleCount int, 38 | didDrop bool, 39 | ) *ImmutableSampleWindow { 40 | if minRTT == 0 { 41 | minRTT = math.MaxInt64 42 | } 43 | return &ImmutableSampleWindow{ 44 | startTime: startTime, 45 | minRTT: minRTT, 46 | sum: sum, 47 | maxInFlight: maxInFlight, 48 | sampleCount: sampleCount, 49 | didDrop: didDrop, 50 | } 51 | } 52 | 53 | // AddSample will create a new immutable sample for which to use. 54 | func (s *ImmutableSampleWindow) AddSample(startTime int64, rtt int64, maxInFlight int) *ImmutableSampleWindow { 55 | minRTT := s.minRTT 56 | if rtt < s.minRTT { 57 | minRTT = rtt 58 | } 59 | if maxInFlight < s.maxInFlight { 60 | maxInFlight = s.maxInFlight 61 | } 62 | if startTime < 0 { 63 | startTime = time.Now().UnixNano() 64 | } 65 | return NewImmutableSampleWindow(startTime, minRTT, s.sum+rtt, maxInFlight, s.sampleCount+1, s.didDrop) 66 | } 67 | 68 | // AddDroppedSample will create a new immutable sample that was dropped. 69 | func (s *ImmutableSampleWindow) AddDroppedSample(startTime int64, maxInFlight int) *ImmutableSampleWindow { 70 | if maxInFlight < s.maxInFlight { 71 | maxInFlight = s.maxInFlight 72 | } 73 | if startTime < 0 { 74 | startTime = time.Now().UnixNano() 75 | } 76 | return NewImmutableSampleWindow(startTime, s.minRTT, s.sum, maxInFlight, s.sampleCount, true) 77 | } 78 | 79 | // StartTimeNanoseconds returns the epoch start time in nanoseconds. 80 | func (s *ImmutableSampleWindow) StartTimeNanoseconds() int64 { 81 | return s.startTime 82 | } 83 | 84 | // CandidateRTTNanoseconds returns the candidate RTT in the sample window. This is traditionally the minimum rtt. 85 | func (s *ImmutableSampleWindow) CandidateRTTNanoseconds() int64 { 86 | return s.minRTT 87 | } 88 | 89 | // AverageRTTNanoseconds returns the average RTT in the sample window. Excludes timeouts and dropped rtt. 90 | func (s *ImmutableSampleWindow) AverageRTTNanoseconds() int64 { 91 | if s.sampleCount == 0 { 92 | return 0 93 | } 94 | return s.sum / int64(s.sampleCount) 95 | } 96 | 97 | // MaxInFlight returns the maximum number of in-flight observed during the sample window. 98 | func (s *ImmutableSampleWindow) MaxInFlight() int { 99 | return s.maxInFlight 100 | } 101 | 102 | // SampleCount is the number of observed RTTs in the sample window. 103 | func (s *ImmutableSampleWindow) SampleCount() int { 104 | return s.sampleCount 105 | } 106 | 107 | // DidDrop returns True if there was a timeout. 108 | func (s *ImmutableSampleWindow) DidDrop() bool { 109 | return s.didDrop 110 | } 111 | 112 | func (s *ImmutableSampleWindow) String() string { 113 | return fmt.Sprintf( 114 | "ImmutableSampleWindow{minRTT=%d, averageRTT=%d, maxInFlight=%d, sampleCount=%d, didDrop=%t}", 115 | s.minRTT, s.AverageRTTNanoseconds(), s.maxInFlight, s.sampleCount, s.didDrop) 116 | } 117 | -------------------------------------------------------------------------------- /measurements/immutable_sample_window_test.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestImmutableSampleWindow(t *testing.T) { 10 | t.Parallel() 11 | asrt := assert.New(t) 12 | w := NewDefaultImmutableSampleWindow() 13 | asrt.False(w.DidDrop()) 14 | w2 := w.AddSample(10, 10, 5) 15 | asrt.NotEqual(w, w2) 16 | asrt.Equal(int64(w2.StartTimeNanoseconds()), w2.StartTimeNanoseconds()) 17 | asrt.Equal(5, w2.MaxInFlight()) 18 | asrt.Equal(1, w2.SampleCount()) 19 | asrt.Equal(int64(10), w2.CandidateRTTNanoseconds()) 20 | asrt.Equal(int64(10), w2.AverageRTTNanoseconds()) 21 | asrt.Equal( 22 | "ImmutableSampleWindow{minRTT=10, averageRTT=10, maxInFlight=5, sampleCount=1, didDrop=false}", 23 | w2.String(), 24 | ) 25 | 26 | // Adding a dropped sample should mark the window as having contained dropped tokens 27 | w3 := w2.AddDroppedSample(-10, 500) 28 | asrt.True(w3.DidDrop()) 29 | 30 | // Adding a successful sample should not void the dropped marker on the window 31 | w4 := w3.AddSample(10, 10, 5) 32 | asrt.True(w4.DidDrop()) 33 | } 34 | -------------------------------------------------------------------------------- /measurements/minimum.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | ) 7 | 8 | // MinimumMeasurement implements a minimum value measurement 9 | type MinimumMeasurement struct { 10 | value float64 11 | mu sync.RWMutex 12 | } 13 | 14 | // Add will compare the sample and save if it's the minimum value. 15 | func (m *MinimumMeasurement) Add(sample float64) (float64, bool) { 16 | m.mu.Lock() 17 | defer m.mu.Unlock() 18 | oldValue := float64(m.value) 19 | if oldValue == 0.0 || sample < oldValue { 20 | m.value = sample 21 | } 22 | return m.value, oldValue == m.value 23 | } 24 | 25 | // Get will return the current minimum value 26 | func (m *MinimumMeasurement) Get() float64 { 27 | m.mu.RLock() 28 | defer m.mu.RUnlock() 29 | return m.value 30 | } 31 | 32 | // Reset will reset the minimum value to 0.0 33 | func (m *MinimumMeasurement) Reset() { 34 | m.mu.Lock() 35 | defer m.mu.Unlock() 36 | m.value = 0.0 37 | } 38 | 39 | // Update will update the value given an operation function 40 | func (m *MinimumMeasurement) Update(operation func(value float64) float64) { 41 | m.mu.RLock() 42 | current := m.value 43 | m.mu.RUnlock() 44 | m.Add(operation(current)) 45 | } 46 | 47 | func (m *MinimumMeasurement) String() string { 48 | m.mu.RLock() 49 | defer m.mu.RUnlock() 50 | return fmt.Sprintf("MinimumMeasurement{value=%0.5f}", m.value) 51 | } 52 | -------------------------------------------------------------------------------- /measurements/minimum_test.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestMinimumMeasurement(t *testing.T) { 10 | t.Parallel() 11 | asrt := assert.New(t) 12 | m := MinimumMeasurement{} 13 | asrt.Equal(float64(0.0), m.Get()) 14 | m.Add(-1) 15 | asrt.Equal(float64(-1.0), m.Get()) 16 | m.Add(0) 17 | asrt.Equal(float64(-1.0), m.Get()) 18 | m.Reset() 19 | asrt.Equal(float64(0.0), m.Get()) 20 | m.Update(func(value float64) float64 { 21 | return value - 1 22 | }) 23 | asrt.Equal(float64(-1.0), m.Get()) 24 | 25 | asrt.Equal("MinimumMeasurement{value=-1.00000}", m.String()) 26 | } 27 | -------------------------------------------------------------------------------- /measurements/moving_average.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "fmt" 5 | "math" 6 | "sync" 7 | ) 8 | 9 | // SimpleExponentialMovingAverage implements a simple exponential moving average 10 | // this implementation only uses a single alpha value to determine warm-up time and provides a mean 11 | // approximation 12 | type SimpleExponentialMovingAverage struct { 13 | alpha float64 14 | initialAlpha float64 15 | minSamples int 16 | seenSamples int 17 | 18 | value float64 19 | 20 | mu sync.RWMutex 21 | } 22 | 23 | // NewSimpleExponentialMovingAverage creates a new simple moving average 24 | func NewSimpleExponentialMovingAverage( 25 | alpha float64, 26 | ) (*SimpleExponentialMovingAverage, error) { 27 | if alpha < 0 || alpha > 1 { 28 | return nil, fmt.Errorf("alpha must be [0, 1]") 29 | } 30 | minSamples := int(math.Trunc(math.Ceil(1 / alpha))) 31 | return &SimpleExponentialMovingAverage{ 32 | alpha: alpha, 33 | initialAlpha: alpha, 34 | minSamples: minSamples, 35 | }, nil 36 | } 37 | 38 | // Add a single sample and update the internal state. 39 | // returns true if the internal state was updated, also return the current value. 40 | func (m *SimpleExponentialMovingAverage) Add(value float64) (float64, bool) { 41 | m.mu.Lock() 42 | defer m.mu.Unlock() 43 | return m.add(value) 44 | } 45 | 46 | func (m *SimpleExponentialMovingAverage) add(value float64) (float64, bool) { 47 | changed := false 48 | if m.seenSamples < m.minSamples { 49 | m.seenSamples++ 50 | } 51 | var alpha float64 52 | if m.seenSamples >= m.minSamples { 53 | alpha = m.alpha 54 | } else { 55 | alpha = 1 / float64(m.seenSamples) 56 | } 57 | newValue := (1-alpha)*m.value + alpha*value 58 | if newValue != m.value { 59 | changed = true 60 | } 61 | m.value = newValue 62 | return m.value, changed 63 | } 64 | 65 | // Get the current value. 66 | func (m *SimpleExponentialMovingAverage) Get() float64 { 67 | m.mu.RLock() 68 | defer m.mu.RUnlock() 69 | return m.value 70 | } 71 | 72 | // Reset the internal state as if no samples were ever added. 73 | func (m *SimpleExponentialMovingAverage) Reset() { 74 | m.mu.Lock() 75 | m.seenSamples = 0 76 | m.value = 0 77 | m.alpha = m.initialAlpha 78 | m.mu.Unlock() 79 | } 80 | 81 | // Update will update the value given an operation function 82 | func (m *SimpleExponentialMovingAverage) Update(operation func(value float64) float64) { 83 | m.mu.Lock() 84 | defer m.mu.Unlock() 85 | newValue, _ := m.add(m.value) 86 | m.value = operation(newValue) 87 | } 88 | -------------------------------------------------------------------------------- /measurements/moving_average_test.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestSimpleExponentialMovingAverage(t *testing.T) { 10 | t.Parallel() 11 | asrt := assert.New(t) 12 | m, err := NewSimpleExponentialMovingAverage(0.05) 13 | asrt.NoError(err) 14 | asrt.NotNil(m) 15 | 16 | asrt.Equal(float64(0), m.Get()) 17 | m.Add(10) 18 | asrt.Equal(float64(10), m.Get()) 19 | m.Add(11) 20 | asrt.Equal(float64(10.5), m.Get()) 21 | m.Add(11) 22 | m.Add(11) 23 | asrt.Equal(float64(10.75), m.Get()) 24 | 25 | m.Reset() 26 | asrt.Equal(float64(0), m.Get()) 27 | m.Update(func(value float64) float64 { 28 | return 1.0 29 | }) 30 | asrt.Equal(float64(1.0), m.Get()) 31 | } 32 | -------------------------------------------------------------------------------- /measurements/moving_variance.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "math" 5 | "sync" 6 | ) 7 | 8 | // SimpleMovingVariance implements a simple moving variance calculation based on the simple moving average. 9 | type SimpleMovingVariance struct { 10 | average *SimpleExponentialMovingAverage 11 | variance *SimpleExponentialMovingAverage 12 | 13 | stdev float64 // square root of the estimated variance 14 | normalized float64 // (input - mean) / stdev 15 | 16 | mu sync.RWMutex 17 | } 18 | 19 | // NewSimpleMovingVariance will create a new exponential moving variance approximation based on the SimpleMovingAverage 20 | func NewSimpleMovingVariance( 21 | alphaAverage float64, 22 | alphaVariance float64, 23 | ) (*SimpleMovingVariance, error) { 24 | movingAverage, err := NewSimpleExponentialMovingAverage(alphaAverage) 25 | if err != nil { 26 | return nil, err 27 | } 28 | variance, err := NewSimpleExponentialMovingAverage(alphaVariance) 29 | if err != nil { 30 | return nil, err 31 | } 32 | return &SimpleMovingVariance{ 33 | average: movingAverage, 34 | variance: variance, 35 | }, nil 36 | } 37 | 38 | // Add a single sample and update the internal state. 39 | // returns true if the internal state was updated, also return the current value. 40 | func (m *SimpleMovingVariance) Add(value float64) (float64, bool) { 41 | m.mu.Lock() 42 | defer m.mu.Unlock() 43 | changed := false 44 | if m.average.seenSamples > 0 { 45 | m.variance.Add(math.Pow(value-m.average.Get(), 2)) 46 | } 47 | m.average.Add(value) 48 | 49 | mean := m.average.Get() 50 | variance := m.variance.Get() 51 | stdev := math.Sqrt(variance) 52 | normalized := m.normalized 53 | if stdev != 0 { 54 | // edge case 55 | normalized = (value - mean) / stdev 56 | } 57 | 58 | if stdev != m.stdev || normalized != m.normalized { 59 | changed = true 60 | } 61 | m.stdev = stdev 62 | m.normalized = normalized 63 | return stdev, changed 64 | } 65 | 66 | // Get the current value. 67 | func (m *SimpleMovingVariance) Get() float64 { 68 | m.mu.RLock() 69 | defer m.mu.RUnlock() 70 | return m.variance.Get() 71 | } 72 | 73 | // Reset the internal state as if no samples were ever added. 74 | func (m *SimpleMovingVariance) Reset() { 75 | m.mu.Lock() 76 | m.average.Reset() 77 | m.variance.Reset() 78 | m.stdev = 0 79 | m.normalized = 0 80 | m.mu.Unlock() 81 | } 82 | 83 | // Update will update the value given an operation function 84 | func (m *SimpleMovingVariance) Update(operation func(value float64) float64) { 85 | m.mu.Lock() 86 | defer m.mu.Unlock() 87 | m.stdev = operation(m.variance.Get()) 88 | } 89 | -------------------------------------------------------------------------------- /measurements/moving_variance_test.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestSimpleMovingVariance(t *testing.T) { 10 | t.Parallel() 11 | asrt := assert.New(t) 12 | m, err := NewSimpleMovingVariance(0.05, 0.05) 13 | asrt.NoError(err) 14 | asrt.NotNil(m) 15 | 16 | // initial condition 17 | asrt.Equal(float64(0), m.Get()) 18 | // warmup first sample 19 | m.Add(10) 20 | asrt.Equal(float64(0), m.Get()) 21 | // first variance reading is expected to be 1 here 22 | m.Add(11) 23 | asrt.Equal(float64(1), m.Get()) 24 | m.Add(10) 25 | m.Add(11) 26 | asrt.InDelta(float64(0.5648), m.Get(), 0.00005) 27 | m.Add(20) 28 | m.Add(100) 29 | m.Add(30) 30 | asrt.InDelta(float64(1295.7841), m.Get(), 0.00005) 31 | 32 | m.Reset() 33 | asrt.Equal(float64(0), m.Get()) 34 | m.Update(func(value float64) float64 { 35 | return 1.0 36 | }) 37 | // this is ignored 38 | asrt.Equal(float64(0), m.Get()) 39 | } 40 | -------------------------------------------------------------------------------- /measurements/single.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | ) 7 | 8 | // SingleMeasurement only keeps the latest value used. 9 | type SingleMeasurement struct { 10 | value float64 11 | mu sync.RWMutex 12 | } 13 | 14 | // Add a single sample and update the internal state. 15 | func (m *SingleMeasurement) Add(value float64) (float64, bool) { 16 | m.mu.Lock() 17 | defer m.mu.Unlock() 18 | m.value = value 19 | return m.value, true 20 | } 21 | 22 | // Get the current value. 23 | func (m *SingleMeasurement) Get() float64 { 24 | m.mu.RLock() 25 | defer m.mu.RUnlock() 26 | return m.value 27 | } 28 | 29 | // Reset the internal state as if no samples were ever added. 30 | func (m *SingleMeasurement) Reset() { 31 | m.mu.Lock() 32 | m.value = 0 33 | m.mu.Unlock() 34 | } 35 | 36 | // Update will update the value given an operation function 37 | func (m *SingleMeasurement) Update(operation func(value float64) float64) { 38 | m.mu.Lock() 39 | m.value = operation(m.value) 40 | m.mu.Unlock() 41 | } 42 | 43 | func (m *SingleMeasurement) String() string { 44 | m.mu.RLock() 45 | defer m.mu.RUnlock() 46 | return fmt.Sprintf("SingleMeasurement{value=%0.5f}", m.value) 47 | } 48 | -------------------------------------------------------------------------------- /measurements/single_test.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestSingleMeasurement(t *testing.T) { 10 | asrt := assert.New(t) 11 | m := SingleMeasurement{} 12 | asrt.Equal(float64(0.0), m.Get()) 13 | 14 | m.Add(1) 15 | m.Add(2) 16 | asrt.Equal(float64(2.0), m.Get()) 17 | 18 | m.Update(func(value float64) float64 { 19 | return value * 2 20 | }) 21 | asrt.Equal(float64(4.0), m.Get()) 22 | asrt.Equal("SingleMeasurement{value=4.00000}", m.String()) 23 | 24 | m.Reset() 25 | asrt.Equal(float64(0), m.Get()) 26 | } 27 | -------------------------------------------------------------------------------- /measurements/windowless_moving_percentile.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | ) 7 | 8 | // WindowlessMovingPercentile implements a moving percentile. 9 | // This implementation uses a windowless calculation that while not strictly always accurate, 10 | // provides a very close estimation in O(1) time and space. 11 | // Much credit goes to Martin Jambon here: https://mjambon.com/2016-07-23-moving-percentile/ 12 | // a copy can be found in github.com/platinummonkey/go-concurrency-limits/docs/assets/moving_percentile_reference.pdf 13 | // and this is a port of the OCaml implementation provided in that reference. 14 | type WindowlessMovingPercentile struct { 15 | p float64 16 | deltaInitial float64 17 | 18 | value float64 19 | delta float64 20 | deltaState *SimpleMovingVariance 21 | 22 | seenCount int 23 | 24 | mu sync.RWMutex 25 | } 26 | 27 | // NewWindowlessMovingPercentile creates a new Windowless Moving Percentile 28 | // p - percentile requested, accepts (0,1) 29 | // deltaInitial - the initial delta value, here 0 is acceptable if you expect it to be rather stable at start, otherwise 30 | // choose a larger value. This would be estimated: `delta := stdev * r` where `r` is a user chosen 31 | // constant. Good values are generally from 0.001 to 0.01 32 | // movingAvgAlphaAvg - this is the alpha value for the simple moving average. A good start is 0.05. Accepts [0,1] 33 | // movingVarianceAlphaAvg - this is the alpha value for the simple moving variance. A good start is 0.05. Accepts [0,1] 34 | func NewWindowlessMovingPercentile( 35 | p float64, // percentile requested 36 | deltaInitial float64, 37 | movingAvgAlphaAvg float64, 38 | movingVarianceAlphaVar float64, 39 | ) (*WindowlessMovingPercentile, error) { 40 | if p <= 0 || p >= 1 { 41 | return nil, fmt.Errorf("p must be between (0,1)") 42 | } 43 | variance, err := NewSimpleMovingVariance(movingAvgAlphaAvg, movingVarianceAlphaVar) 44 | if err != nil { 45 | return nil, err 46 | } 47 | 48 | return &WindowlessMovingPercentile{ 49 | p: p, 50 | deltaInitial: deltaInitial, 51 | delta: deltaInitial, 52 | deltaState: variance, 53 | }, nil 54 | } 55 | 56 | // Add a single sample and update the internal state. 57 | // returns true if the internal state was updated, also return the current value. 58 | func (m *WindowlessMovingPercentile) Add(value float64) (float64, bool) { 59 | m.mu.Lock() 60 | defer m.mu.Unlock() 61 | return m.add(value) 62 | } 63 | 64 | func (m *WindowlessMovingPercentile) add(value float64) (float64, bool) { 65 | changed := false 66 | if m.seenCount < 2 { 67 | // we only need 2 samples to continue 68 | m.seenCount++ 69 | } 70 | originalDelta := m.delta 71 | stdev, _ := m.deltaState.Add(value) 72 | if m.seenCount >= 2 { 73 | m.delta = m.deltaInitial * stdev 74 | if m.delta != originalDelta { 75 | changed = true 76 | } 77 | } 78 | newValue := float64(m.value) 79 | if m.seenCount == 1 { 80 | newValue = value 81 | changed = true 82 | } else if value < m.value { 83 | newValue = m.value - m.delta/m.p 84 | } else if value > m.value { 85 | newValue = m.value + m.delta/(1-m.p) 86 | } 87 | // else the same 88 | if newValue != m.value { 89 | changed = true 90 | } 91 | m.value = newValue 92 | return m.value, changed 93 | } 94 | 95 | // Get the current value. 96 | func (m *WindowlessMovingPercentile) Get() float64 { 97 | m.mu.RLock() 98 | defer m.mu.RUnlock() 99 | return m.value 100 | } 101 | 102 | // Reset the internal state as if no samples were ever added. 103 | func (m *WindowlessMovingPercentile) Reset() { 104 | m.mu.Lock() 105 | m.value = 0 106 | m.seenCount = 0 107 | m.mu.Unlock() 108 | } 109 | 110 | // Update will update the value given an operation function 111 | func (m *WindowlessMovingPercentile) Update(operation func(value float64) float64) { 112 | m.mu.Lock() 113 | defer m.mu.Unlock() 114 | newValue, _ := m.add(m.value) 115 | m.value = operation(newValue) 116 | } 117 | -------------------------------------------------------------------------------- /measurements/windowless_moving_percentile_test.go: -------------------------------------------------------------------------------- 1 | package measurements 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestWindowlessMovingPercentile(t *testing.T) { 10 | t.Parallel() 11 | asrt := assert.New(t) 12 | m, err := NewWindowlessMovingPercentile(0.9, 0.01, 0.05, 0.05) 13 | asrt.NoError(err) 14 | asrt.NotNil(m) 15 | asrt.Equal(float64(0.0), m.Get()) 16 | for i := 0; i < 10; i++ { 17 | m.Add(100) 18 | } 19 | asrt.Equal(float64(100), m.Get()) 20 | m.Add(99) 21 | for i := 0; i < 10; i++ { 22 | m.Add(1000) 23 | } 24 | m.Add(0.1) 25 | asrt.Equal(520, int(m.Get())) 26 | 27 | m.Reset() 28 | asrt.Equal(float64(0), m.Get()) 29 | m.Update(func(value float64) float64 { 30 | return 1.0 31 | }) 32 | asrt.Equal(float64(1.0), m.Get()) 33 | } 34 | -------------------------------------------------------------------------------- /metric_registry/datadog/registry.go: -------------------------------------------------------------------------------- 1 | // Package datadog implements the metric registry interface for a Datadog provider. 2 | package datadog 3 | 4 | import ( 5 | "fmt" 6 | "strings" 7 | "sync" 8 | "time" 9 | 10 | dogstatsd "github.com/DataDog/datadog-go/v5/statsd" 11 | 12 | "github.com/platinummonkey/go-concurrency-limits/core" 13 | ) 14 | 15 | const defaultMetricPrefix = "limiter." 16 | const defaultPollFrequency = time.Second * 5 17 | 18 | type metricSampleListener struct { 19 | client *dogstatsd.Client 20 | id string 21 | metricType uint8 22 | } 23 | 24 | // AddSample will add a sample metric to the listener 25 | func (l *metricSampleListener) AddSample(value float64, tags ...string) { 26 | switch l.metricType { 27 | case 0: // distribution 28 | l.client.Distribution(l.id, value, tags, 1.0) 29 | case 1: // timing 30 | l.client.TimeInMilliseconds(l.id, value, tags, 1.0) 31 | case 2: // count 32 | l.client.Count(l.id, int64(value), tags, 1.0) 33 | default: 34 | // unsupported 35 | } 36 | } 37 | 38 | type metricPoller struct { 39 | supplier core.MetricSupplier 40 | id string 41 | tags []string 42 | } 43 | 44 | func (p *metricPoller) poll() (string, float64, []string, bool) { 45 | val, ok := p.supplier() 46 | return p.id, val, p.tags, ok 47 | } 48 | 49 | // MetricRegistry will implements a MetricRegistry for sending metrics to Datadog via dogstatsd. 50 | type MetricRegistry struct { 51 | client *dogstatsd.Client 52 | prefix string 53 | pollFrequency time.Duration 54 | registeredGauges map[string]*metricPoller 55 | registeredListeners map[string]*metricSampleListener 56 | 57 | mu sync.Mutex 58 | wg sync.WaitGroup 59 | 60 | started bool 61 | stopper chan bool 62 | } 63 | 64 | // NewMetricRegistry will create a new Datadog MetricRegistry. 65 | // This registry reports metrics to datadog using the datadog dogstatsd forwarding. 66 | func NewMetricRegistry(addr string, prefix string, pollFrequency time.Duration) (*MetricRegistry, error) { 67 | if prefix == "" { 68 | prefix = defaultMetricPrefix 69 | } 70 | if !strings.HasSuffix(prefix, ".") { 71 | prefix = prefix + "." 72 | } 73 | 74 | if pollFrequency == 0 { 75 | pollFrequency = defaultPollFrequency 76 | } 77 | 78 | client, err := dogstatsd.New(addr) 79 | if err != nil { 80 | return nil, err 81 | } 82 | return &MetricRegistry{ 83 | client: client, 84 | prefix: prefix, 85 | pollFrequency: pollFrequency, 86 | stopper: make(chan bool, 1), 87 | registeredGauges: make(map[string]*metricPoller, 0), 88 | registeredListeners: make(map[string]*metricSampleListener, 0), 89 | }, nil 90 | } 91 | 92 | // NewMetricRegistryWithClient will create a new Datadog MetricRegistry with the provided client instead. 93 | // This registry reports metrics to datadog using the datadog dogstatsd forwarding. 94 | func NewMetricRegistryWithClient( 95 | client *dogstatsd.Client, 96 | prefix string, 97 | pollFrequency time.Duration, 98 | ) (*MetricRegistry, error) { 99 | if client == nil { 100 | return nil, fmt.Errorf("client is nil") 101 | } 102 | 103 | if !strings.HasSuffix(prefix, ".") { 104 | prefix = prefix + "." 105 | } 106 | 107 | if pollFrequency == 0 { 108 | pollFrequency = defaultPollFrequency 109 | } 110 | 111 | return &MetricRegistry{ 112 | client: client, 113 | prefix: prefix, 114 | pollFrequency: pollFrequency, 115 | stopper: make(chan bool, 1), 116 | registeredGauges: make(map[string]*metricPoller, 0), 117 | registeredListeners: make(map[string]*metricSampleListener, 0), 118 | }, nil 119 | } 120 | 121 | // Start will start the metric registry polling 122 | func (r *MetricRegistry) Start() { 123 | r.mu.Lock() 124 | if !r.started { 125 | r.wg.Add(1) 126 | go func() { 127 | defer r.wg.Done() 128 | r.run() 129 | }() 130 | } 131 | r.mu.Unlock() 132 | } 133 | 134 | func (r *MetricRegistry) run() { 135 | ticker := time.NewTicker(r.pollFrequency) 136 | for { 137 | select { 138 | case <-r.stopper: 139 | return 140 | case <-ticker.C: 141 | // poll the gauges 142 | r.mu.Lock() 143 | for _, g := range r.registeredGauges { 144 | metricSuffix, value, tags, ok := g.poll() 145 | if ok { 146 | r.client.Gauge(r.prefix+metricSuffix, value, tags, 1.0) 147 | } 148 | } 149 | r.mu.Unlock() 150 | } 151 | } 152 | } 153 | 154 | // Stop will gracefully stop the registry 155 | func (r *MetricRegistry) Stop() { 156 | r.mu.Lock() 157 | if !r.started { 158 | r.mu.Unlock() 159 | return 160 | } 161 | r.stopper <- true 162 | r.wg.Wait() 163 | r.started = false 164 | r.mu.Unlock() 165 | } 166 | 167 | // RegisterDistribution will register a distribution sample to this registry 168 | func (r *MetricRegistry) RegisterDistribution( 169 | ID string, 170 | tags ...string, 171 | ) core.MetricSampleListener { 172 | if strings.HasPrefix(ID, ".") { 173 | ID = strings.TrimPrefix(ID, ".") 174 | } 175 | 176 | // only add once 177 | if l, ok := r.registeredListeners[ID]; ok { 178 | return l 179 | } 180 | 181 | r.registeredListeners[ID] = &metricSampleListener{ 182 | client: r.client, 183 | metricType: 0, 184 | id: r.prefix + ID, 185 | } 186 | 187 | return r.registeredListeners[ID] 188 | } 189 | 190 | // RegisterTiming will register a timing distribution sample to this registry 191 | func (r *MetricRegistry) RegisterTiming( 192 | ID string, 193 | tags ...string, 194 | ) core.MetricSampleListener { 195 | if strings.HasPrefix(ID, ".") { 196 | ID = strings.TrimPrefix(ID, ".") 197 | } 198 | 199 | // only add once 200 | if l, ok := r.registeredListeners[ID]; ok { 201 | return l 202 | } 203 | 204 | r.registeredListeners[ID] = &metricSampleListener{ 205 | client: r.client, 206 | metricType: 1, 207 | id: r.prefix + ID, 208 | } 209 | 210 | return r.registeredListeners[ID] 211 | } 212 | 213 | // RegisterCount will register a count sample to this registry 214 | func (r *MetricRegistry) RegisterCount( 215 | ID string, 216 | tags ...string, 217 | ) core.MetricSampleListener { 218 | if strings.HasPrefix(ID, ".") { 219 | ID = strings.TrimPrefix(ID, ".") 220 | } 221 | 222 | // only add once 223 | if l, ok := r.registeredListeners[ID]; ok { 224 | return l 225 | } 226 | 227 | r.registeredListeners[ID] = &metricSampleListener{ 228 | client: r.client, 229 | metricType: 2, 230 | id: r.prefix + ID, 231 | } 232 | 233 | return r.registeredListeners[ID] 234 | } 235 | 236 | // RegisterGauge will register a gauge sample to this registry 237 | func (r *MetricRegistry) RegisterGauge( 238 | ID string, 239 | supplier core.MetricSupplier, 240 | tags ...string, 241 | ) { 242 | if strings.HasPrefix(ID, ".") { 243 | ID = strings.TrimPrefix(ID, ".") 244 | } 245 | 246 | r.mu.Lock() 247 | defer r.mu.Unlock() 248 | 249 | // only add once 250 | if _, ok := r.registeredGauges[ID]; ok { 251 | return 252 | } 253 | 254 | r.registeredGauges[ID] = &metricPoller{ 255 | supplier: supplier, 256 | id: ID, 257 | tags: tags, 258 | } 259 | } 260 | -------------------------------------------------------------------------------- /metric_registry/gometrics/registry.go: -------------------------------------------------------------------------------- 1 | // Package gometrics implements the metric registry interface for a gometrics provider. 2 | package gometrics 3 | 4 | import ( 5 | "fmt" 6 | "strings" 7 | "sync" 8 | "time" 9 | 10 | gometrics "github.com/rcrowley/go-metrics" 11 | 12 | "github.com/platinummonkey/go-concurrency-limits/core" 13 | ) 14 | 15 | const defaultMetricPrefix = "limiter." 16 | const defaultPollFrequency = time.Second * 5 17 | 18 | type metricSampleListener struct { 19 | distribution gometrics.Histogram 20 | timer gometrics.Timer 21 | counter gometrics.Counter 22 | id string 23 | metricType uint8 24 | } 25 | 26 | // AddSample will add a sample metric to the listener 27 | func (l *metricSampleListener) AddSample(value float64, tags ...string) { 28 | switch l.metricType { 29 | case 0: // distribution 30 | l.distribution.Update(int64(value)) 31 | case 1: // timing 32 | // value is in epoch milliseconds 33 | l.timer.Update(time.Duration(value) * time.Millisecond) 34 | case 2: // count 35 | l.counter.Inc(int64(value)) 36 | default: 37 | // unsupported 38 | } 39 | } 40 | 41 | type gometricsMetricPoller struct { 42 | supplier core.MetricSupplier 43 | id string 44 | tags []string 45 | } 46 | 47 | func (p *gometricsMetricPoller) poll() (string, float64, []string, bool) { 48 | val, ok := p.supplier() 49 | return p.id, val, p.tags, ok 50 | } 51 | 52 | // MetricRegistry will implements a MetricRegistry for sending metrics via go-metrics with any reporter. 53 | type MetricRegistry struct { 54 | registry gometrics.Registry 55 | prefix string 56 | pollFrequency time.Duration 57 | registeredGauges map[string]*gometricsMetricPoller 58 | registeredListeners map[string]*metricSampleListener 59 | 60 | mu sync.Mutex 61 | wg sync.WaitGroup 62 | 63 | started bool 64 | stopper chan bool 65 | } 66 | 67 | // NewGoMetricsMetricRegistry will create a new Datadog MetricRegistry. 68 | // This registry reports metrics to datadog using the datadog dogstatsd forwarding. 69 | func NewGoMetricsMetricRegistry( 70 | registry gometrics.Registry, 71 | addr string, 72 | prefix string, 73 | pollFrequency time.Duration, 74 | ) (*MetricRegistry, error) { 75 | if prefix == "" { 76 | prefix = defaultMetricPrefix 77 | } 78 | if !strings.HasSuffix(prefix, ".") { 79 | prefix = prefix + "." 80 | } 81 | 82 | if pollFrequency == 0 { 83 | pollFrequency = defaultPollFrequency 84 | } 85 | 86 | if registry == nil { 87 | return nil, fmt.Errorf("registry required") 88 | } 89 | 90 | return &MetricRegistry{ 91 | registry: registry, 92 | prefix: prefix, 93 | pollFrequency: pollFrequency, 94 | stopper: make(chan bool, 1), 95 | registeredGauges: make(map[string]*gometricsMetricPoller, 0), 96 | registeredListeners: make(map[string]*metricSampleListener, 0), 97 | }, nil 98 | } 99 | 100 | // Start will start the metric registry polling 101 | func (r *MetricRegistry) Start() { 102 | r.mu.Lock() 103 | if !r.started { 104 | r.wg.Add(1) 105 | go func() { 106 | defer r.wg.Done() 107 | r.run() 108 | }() 109 | } 110 | r.mu.Unlock() 111 | } 112 | 113 | func (r *MetricRegistry) run() { 114 | ticker := time.NewTicker(r.pollFrequency) 115 | for { 116 | select { 117 | case <-r.stopper: 118 | return 119 | case <-ticker.C: 120 | // poll the gauges 121 | r.mu.Lock() 122 | for _, g := range r.registeredGauges { 123 | metricSuffix, value, _, ok := g.poll() 124 | if ok { 125 | m := gometrics.GetOrRegisterGaugeFloat64(r.prefix+metricSuffix, r.registry) 126 | m.Update(value) 127 | } 128 | } 129 | r.mu.Unlock() 130 | } 131 | } 132 | } 133 | 134 | // Stop will gracefully stop the registry 135 | func (r *MetricRegistry) Stop() { 136 | r.mu.Lock() 137 | if !r.started { 138 | r.mu.Unlock() 139 | return 140 | } 141 | r.stopper <- true 142 | r.wg.Wait() 143 | r.started = false 144 | r.mu.Unlock() 145 | } 146 | 147 | // RegisterDistribution will register a distribution sample to this registry 148 | func (r *MetricRegistry) RegisterDistribution( 149 | ID string, 150 | tags ...string, 151 | ) core.MetricSampleListener { 152 | if strings.HasPrefix(ID, ".") { 153 | ID = strings.TrimPrefix(ID, ".") 154 | } 155 | 156 | // only add once 157 | if l, ok := r.registeredListeners[ID]; ok { 158 | return l 159 | } 160 | 161 | r.registeredListeners[ID] = &metricSampleListener{ 162 | distribution: gometrics.GetOrRegisterHistogram( 163 | r.prefix+ID, 164 | r.registry, 165 | gometrics.NewUniformSample(100), 166 | ), 167 | metricType: 0, 168 | id: r.prefix + ID, 169 | } 170 | 171 | return r.registeredListeners[ID] 172 | } 173 | 174 | // RegisterTiming will register a timing distribution sample to this registry 175 | func (r *MetricRegistry) RegisterTiming( 176 | ID string, 177 | tags ...string, 178 | ) core.MetricSampleListener { 179 | if strings.HasPrefix(ID, ".") { 180 | ID = strings.TrimPrefix(ID, ".") 181 | } 182 | 183 | // only add once 184 | if l, ok := r.registeredListeners[ID]; ok { 185 | return l 186 | } 187 | 188 | r.registeredListeners[ID] = &metricSampleListener{ 189 | timer: gometrics.GetOrRegisterTimer( 190 | r.prefix+ID, 191 | r.registry, 192 | ), 193 | metricType: 1, 194 | id: r.prefix + ID, 195 | } 196 | 197 | return r.registeredListeners[ID] 198 | } 199 | 200 | // RegisterCount will register a count sample to this registry 201 | func (r *MetricRegistry) RegisterCount( 202 | ID string, 203 | tags ...string, 204 | ) core.MetricSampleListener { 205 | if strings.HasPrefix(ID, ".") { 206 | ID = strings.TrimPrefix(ID, ".") 207 | } 208 | 209 | // only add once 210 | if l, ok := r.registeredListeners[ID]; ok { 211 | return l 212 | } 213 | 214 | r.registeredListeners[ID] = &metricSampleListener{ 215 | counter: gometrics.GetOrRegisterCounter( 216 | r.prefix+ID, 217 | r.registry, 218 | ), 219 | metricType: 2, 220 | id: r.prefix + ID, 221 | } 222 | 223 | return r.registeredListeners[ID] 224 | } 225 | 226 | // RegisterGauge will register a gauge sample to this registry 227 | func (r *MetricRegistry) RegisterGauge( 228 | ID string, 229 | supplier core.MetricSupplier, 230 | tags ...string, 231 | ) { 232 | if strings.HasPrefix(ID, ".") { 233 | ID = strings.TrimPrefix(ID, ".") 234 | } 235 | 236 | r.mu.Lock() 237 | defer r.mu.Unlock() 238 | 239 | // only add once 240 | if _, ok := r.registeredGauges[ID]; ok { 241 | return 242 | } 243 | 244 | r.registeredGauges[ID] = &gometricsMetricPoller{ 245 | supplier: supplier, 246 | id: ID, 247 | tags: tags, 248 | } 249 | } 250 | -------------------------------------------------------------------------------- /patterns/doc.go: -------------------------------------------------------------------------------- 1 | // Package patterns provides common patterns as higher level abstractions from the library building blocks. 2 | package patterns 3 | -------------------------------------------------------------------------------- /patterns/pool/doc.go: -------------------------------------------------------------------------------- 1 | // Package pool provides common pool patterns for concurrency control. 2 | package pool 3 | -------------------------------------------------------------------------------- /patterns/pool/example_fixed_pool_test.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "log" 6 | "sync" 7 | "time" 8 | 9 | "github.com/platinummonkey/go-concurrency-limits/limit" 10 | ) 11 | 12 | func ExampleFixedPool() { 13 | type JobKey string 14 | var JobKeyID = JobKey("job_id") 15 | 16 | l := 1000 // limit to 1000 concurrent requests. 17 | // create a new pool 18 | pool, err := NewFixedPool( 19 | "protected_resource_pool", 20 | OrderingRandom, 21 | l, 22 | 100, 23 | time.Millisecond*250, 24 | time.Millisecond*500, 25 | time.Millisecond*10, 26 | 0, 27 | time.Second, 28 | limit.BuiltinLimitLogger{}, 29 | nil, 30 | ) 31 | if err != nil { 32 | panic(err) 33 | } 34 | 35 | wg := sync.WaitGroup{} 36 | wg.Add(l * 3) 37 | // spawn 3000 concurrent requests that would normally be too much load for the protected resource. 38 | for i := 0; i <= l*3; i++ { 39 | go func(c int) { 40 | defer wg.Done() 41 | ctx := context.WithValue(context.Background(), JobKeyID, c) 42 | // this will block until timeout or token was acquired. 43 | listener, ok := pool.Acquire(ctx) 44 | if !ok { 45 | log.Printf("was not able to acquire lock for id=%d\n", c) 46 | return 47 | } 48 | log.Printf("acquired lock for id=%d\n", c) 49 | // do something... 50 | time.Sleep(time.Millisecond * 10) 51 | listener.OnSuccess() 52 | log.Printf("released lock for id=%d\n", c) 53 | }(i) 54 | } 55 | 56 | // wait for completion 57 | wg.Wait() 58 | log.Println("Finished") 59 | } 60 | -------------------------------------------------------------------------------- /patterns/pool/example_generic_pool_test.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "log" 6 | "sync" 7 | "time" 8 | 9 | "github.com/platinummonkey/go-concurrency-limits/limit" 10 | "github.com/platinummonkey/go-concurrency-limits/limiter" 11 | "github.com/platinummonkey/go-concurrency-limits/strategy" 12 | ) 13 | 14 | func ExamplePool() { 15 | type JobKey string 16 | var JobKeyID = JobKey("job_id") 17 | 18 | l := 1000 // limit to adjustable 1000 concurrent requests. 19 | delegateLimit := limit.NewDefaultAIMDLimit( 20 | "aimd_limiter", 21 | nil, 22 | ) 23 | // wrap with a default limiter and simple strategy 24 | // you could of course get very complicated with this. 25 | delegateLimiter, err := limiter.NewDefaultLimiter( 26 | delegateLimit, 27 | (time.Millisecond * 250).Nanoseconds(), 28 | (time.Millisecond * 500).Nanoseconds(), 29 | (time.Millisecond * 10).Nanoseconds(), 30 | 100, 31 | strategy.NewSimpleStrategy(l), 32 | limit.BuiltinLimitLogger{}, 33 | nil, 34 | ) 35 | if err != nil { 36 | panic(err) 37 | } 38 | 39 | // create a new pool 40 | pool, err := NewPool( 41 | delegateLimiter, 42 | OrderingRandom, 43 | 0, 44 | time.Second, 45 | limit.BuiltinLimitLogger{}, 46 | nil, 47 | ) 48 | if err != nil { 49 | panic(err) 50 | } 51 | 52 | wg := sync.WaitGroup{} 53 | wg.Add(l * 3) 54 | // spawn 3000 concurrent requests that would normally be too much load for the protected resource. 55 | for i := 0; i <= l*3; i++ { 56 | go func(c int) { 57 | defer wg.Done() 58 | ctx := context.WithValue(context.Background(), JobKeyID, c) 59 | // this will block until timeout or token was acquired. 60 | listener, ok := pool.Acquire(ctx) 61 | if !ok { 62 | log.Printf("was not able to acquire lock for id=%d\n", c) 63 | return 64 | } 65 | log.Printf("acquired lock for id=%d\n", c) 66 | // do something... 67 | time.Sleep(time.Millisecond * 10) 68 | listener.OnSuccess() 69 | log.Printf("released lock for id=%d\n", c) 70 | }(i) 71 | } 72 | 73 | // wait for completion 74 | wg.Wait() 75 | log.Println("Finished") 76 | } 77 | -------------------------------------------------------------------------------- /patterns/pool/fixed_pool.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/platinummonkey/go-concurrency-limits/core" 8 | "github.com/platinummonkey/go-concurrency-limits/limit" 9 | "github.com/platinummonkey/go-concurrency-limits/limiter" 10 | "github.com/platinummonkey/go-concurrency-limits/strategy" 11 | ) 12 | 13 | // FixedPool implements a fixed size blocking pool. 14 | type FixedPool struct { 15 | limit int 16 | ordering Ordering 17 | limiter core.Limiter 18 | } 19 | 20 | // NewFixedPool creates a named fixed pool resource. You can use this to guard another resource from too many concurrent 21 | // requests. 22 | // 23 | // use < 0 values for defaults, but fixedLimit and name are required. 24 | func NewFixedPool( 25 | name string, 26 | ordering Ordering, 27 | fixedLimit int, 28 | windowSize int, 29 | minWindowTime time.Duration, 30 | maxWindowTime time.Duration, 31 | minRTTThreshold time.Duration, 32 | maxBacklog int, 33 | timeout time.Duration, 34 | logger limit.Logger, 35 | metricRegistry core.MetricRegistry, 36 | ) (*FixedPool, error) { 37 | if minWindowTime < 0 { 38 | minWindowTime = time.Millisecond * 250 39 | } 40 | if maxWindowTime < 0 { 41 | maxWindowTime = time.Millisecond * 500 42 | } 43 | if minRTTThreshold < 0 { 44 | minRTTThreshold = time.Millisecond * 10 45 | } 46 | if windowSize <= 0 { 47 | windowSize = 100 48 | } 49 | if timeout < 0 { 50 | timeout = 0 51 | } 52 | if logger == nil { 53 | logger = limit.NoopLimitLogger{} 54 | } 55 | if metricRegistry == nil { 56 | metricRegistry = core.EmptyMetricRegistryInstance 57 | } 58 | 59 | limitStrategy := strategy.NewPreciseStrategy(fixedLimit) 60 | defaultLimiter, err := limiter.NewDefaultLimiter( 61 | limit.NewFixedLimit( 62 | name, 63 | fixedLimit, 64 | metricRegistry, 65 | ), 66 | minWindowTime.Nanoseconds(), 67 | maxWindowTime.Nanoseconds(), 68 | minRTTThreshold.Nanoseconds(), 69 | windowSize, 70 | limitStrategy, 71 | logger, 72 | metricRegistry, 73 | ) 74 | if err != nil { 75 | return nil, err 76 | } 77 | 78 | var fp FixedPool 79 | switch ordering { 80 | case OrderingFIFO: 81 | fp = FixedPool{ 82 | limit: fixedLimit, 83 | limiter: limiter.NewQueueBlockingLimiterFromConfig(defaultLimiter, limiter.QueueLimiterConfig{ 84 | Ordering: limiter.OrderingFIFO, 85 | MaxBacklogSize: maxBacklog, 86 | MaxBacklogTimeout: timeout, 87 | MetricRegistry: metricRegistry, 88 | }), 89 | ordering: ordering, 90 | } 91 | case OrderingLIFO: 92 | fp = FixedPool{ 93 | limit: fixedLimit, 94 | limiter: limiter.NewQueueBlockingLimiterFromConfig(defaultLimiter, limiter.QueueLimiterConfig{ 95 | Ordering: limiter.OrderingLIFO, 96 | MaxBacklogSize: maxBacklog, 97 | MaxBacklogTimeout: timeout, 98 | MetricRegistry: metricRegistry, 99 | }), 100 | ordering: ordering, 101 | } 102 | default: 103 | fp = FixedPool{ 104 | limit: fixedLimit, 105 | limiter: limiter.NewBlockingLimiter(defaultLimiter, timeout, logger), 106 | ordering: ordering, 107 | } 108 | } 109 | return &fp, nil 110 | } 111 | 112 | // Limit will return the configured limit 113 | func (p *FixedPool) Limit() int { 114 | return p.limit 115 | } 116 | 117 | // Ordering the ordering strategy configured for this pool. 118 | func (p *FixedPool) Ordering() Ordering { 119 | return p.ordering 120 | } 121 | 122 | // Acquire a token for the protected resource. This method will block until acquisition or the configured timeout 123 | // has expired. 124 | func (p *FixedPool) Acquire(ctx context.Context) (core.Listener, bool) { 125 | return p.limiter.Acquire(ctx) 126 | } 127 | -------------------------------------------------------------------------------- /patterns/pool/fixed_pool_test.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "sync" 8 | "testing" 9 | "time" 10 | 11 | "github.com/stretchr/testify/assert" 12 | ) 13 | 14 | type testKey string 15 | 16 | const testKeyID = testKey("id") 17 | 18 | func TestFixedPool(t *testing.T) { 19 | asrt := assert.New(t) 20 | p, err := NewFixedPool( 21 | "test-fixed-pool", 22 | OrderingRandom, 23 | 10, 24 | -1, 25 | -1, 26 | -1, 27 | -1, 28 | 0, 29 | 0, 30 | nil, 31 | nil, 32 | ) 33 | asrt.NoError(err) 34 | 35 | asrt.Equal(10, p.Limit()) 36 | 37 | var wg sync.WaitGroup 38 | for i := 0; i < 20; i++ { 39 | wg.Add(1) 40 | go func(c int) { 41 | defer wg.Done() 42 | l, _ := p.Acquire(context.WithValue(context.Background(), testKeyID, fmt.Sprint(c))) 43 | log.Printf("acquired now, sleeping - %d\n", c) 44 | time.Sleep(time.Millisecond * 100) 45 | l.OnSuccess() 46 | log.Printf("no longer acquired, released - %d\n", c) 47 | }(i) 48 | time.Sleep(time.Millisecond * 5) 49 | } 50 | wg.Wait() 51 | } 52 | 53 | func TestFIFOFixedPool(t *testing.T) { 54 | asrt := assert.New(t) 55 | p, err := NewFixedPool( 56 | "test-fifo-fixed-pool", 57 | OrderingFIFO, 58 | 10, 59 | -1, 60 | -1, 61 | -1, 62 | -1, 63 | 0, 64 | 0, 65 | nil, 66 | nil, 67 | ) 68 | asrt.NoError(err) 69 | 70 | asrt.Equal(10, p.Limit()) 71 | 72 | var wg sync.WaitGroup 73 | for i := 0; i < 20; i++ { 74 | wg.Add(1) 75 | go func(c int) { 76 | defer wg.Done() 77 | l, _ := p.Acquire(context.WithValue(context.Background(), testKeyID, fmt.Sprint(c))) 78 | log.Printf("acquired now, sleeping - %d\n", c) 79 | time.Sleep(time.Millisecond * 100) 80 | l.OnSuccess() 81 | log.Printf("no longer acquired, released - %d\n", c) 82 | }(i) 83 | time.Sleep(time.Millisecond * 5) 84 | } 85 | wg.Wait() 86 | } 87 | 88 | func TestLIFOFixedPool(t *testing.T) { 89 | asrt := assert.New(t) 90 | p, err := NewFixedPool( 91 | "test-lifo-fixed-pool", 92 | OrderingLIFO, 93 | 10, 94 | -1, 95 | -1, 96 | -1, 97 | -1, 98 | 0, 99 | 0, 100 | nil, 101 | nil, 102 | ) 103 | asrt.NoError(err) 104 | 105 | asrt.Equal(10, p.Limit()) 106 | 107 | var wg sync.WaitGroup 108 | for i := 0; i < 20; i++ { 109 | wg.Add(1) 110 | go func(c int) { 111 | defer wg.Done() 112 | l, _ := p.Acquire(context.WithValue(context.Background(), testKeyID, fmt.Sprint(c))) 113 | log.Printf("acquired now, sleeping - %d\n", c) 114 | time.Sleep(time.Millisecond * 100) 115 | l.OnSuccess() 116 | log.Printf("no longer acquired, released - %d\n", c) 117 | }(i) 118 | time.Sleep(time.Millisecond * 5) 119 | } 120 | wg.Wait() 121 | } 122 | -------------------------------------------------------------------------------- /patterns/pool/generic_pool.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "time" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | "github.com/platinummonkey/go-concurrency-limits/limit" 10 | "github.com/platinummonkey/go-concurrency-limits/limiter" 11 | ) 12 | 13 | // Ordering define the pattern for ordering requests on Pool 14 | type Ordering int 15 | 16 | // The available options 17 | const ( 18 | OrderingRandom Ordering = iota 19 | OrderingFIFO 20 | OrderingLIFO 21 | ) 22 | 23 | // Pool implements a generic blocking pool pattern. 24 | type Pool struct { 25 | limiter core.Limiter 26 | } 27 | 28 | // NewPool creates a pool resource. You can use this to guard another resource from too many concurrent requests. 29 | // 30 | // use < 0 values for defaults, but delegateLimiter and name are required. 31 | func NewPool( 32 | delegateLimiter core.Limiter, 33 | ordering Ordering, 34 | maxBacklog int, 35 | timeout time.Duration, 36 | logger limit.Logger, 37 | metricRegistry core.MetricRegistry, 38 | ) (*Pool, error) { 39 | if delegateLimiter == nil { 40 | return nil, fmt.Errorf("must specify a delegateLimiter") 41 | } 42 | 43 | if timeout < 0 { 44 | timeout = 0 45 | } 46 | if logger == nil { 47 | logger = limit.NoopLimitLogger{} 48 | } 49 | if metricRegistry == nil { 50 | metricRegistry = core.EmptyMetricRegistryInstance 51 | } 52 | 53 | var p Pool 54 | switch ordering { 55 | case OrderingFIFO: 56 | p = Pool{ 57 | limiter: limiter.NewQueueBlockingLimiterFromConfig(delegateLimiter, limiter.QueueLimiterConfig{ 58 | Ordering: limiter.OrderingFIFO, 59 | MaxBacklogSize: maxBacklog, 60 | MaxBacklogTimeout: timeout, 61 | MetricRegistry: metricRegistry, 62 | }), 63 | } 64 | case OrderingLIFO: 65 | p = Pool{ 66 | limiter: limiter.NewQueueBlockingLimiterFromConfig(delegateLimiter, limiter.QueueLimiterConfig{ 67 | Ordering: limiter.OrderingLIFO, 68 | MaxBacklogSize: maxBacklog, 69 | MaxBacklogTimeout: timeout, 70 | MetricRegistry: metricRegistry, 71 | }), 72 | } 73 | default: 74 | p = Pool{ 75 | limiter: limiter.NewBlockingLimiter(delegateLimiter, timeout, logger), 76 | } 77 | } 78 | 79 | return &p, nil 80 | } 81 | 82 | // Acquire a token for the protected resource. This method will block until acquisition or the configured timeout 83 | // has expired. 84 | func (p *Pool) Acquire(ctx context.Context) (core.Listener, bool) { 85 | return p.limiter.Acquire(ctx) 86 | } 87 | -------------------------------------------------------------------------------- /patterns/pool/generic_pool_test.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "sync" 8 | "testing" 9 | "time" 10 | 11 | "github.com/platinummonkey/go-concurrency-limits/limit" 12 | "github.com/platinummonkey/go-concurrency-limits/limiter" 13 | "github.com/platinummonkey/go-concurrency-limits/strategy" 14 | "github.com/stretchr/testify/assert" 15 | ) 16 | 17 | func TestGenericPool(t *testing.T) { 18 | asrt := assert.New(t) 19 | delegateLimiter, err := limiter.NewDefaultLimiter( 20 | limit.NewFixedLimit("test-generic-pool", 10, nil), 21 | (time.Millisecond * 250).Nanoseconds(), 22 | (time.Millisecond * 500).Nanoseconds(), 23 | (time.Millisecond * 10).Nanoseconds(), 24 | 100, 25 | strategy.NewPreciseStrategy(10), 26 | nil, 27 | nil, 28 | ) 29 | asrt.NoError(err) 30 | 31 | p, err := NewPool( 32 | delegateLimiter, 33 | OrderingRandom, 34 | 10, 35 | -1, 36 | nil, 37 | nil, 38 | ) 39 | asrt.NoError(err) 40 | 41 | var wg sync.WaitGroup 42 | for i := 0; i < 20; i++ { 43 | wg.Add(1) 44 | go func(c int) { 45 | defer wg.Done() 46 | l, _ := p.Acquire(context.WithValue(context.Background(), testKeyID, fmt.Sprint(c))) 47 | log.Printf("acquired now, sleeping - %d\n", c) 48 | time.Sleep(time.Millisecond * 100) 49 | l.OnSuccess() 50 | log.Printf("no longer acquired, released - %d\n", c) 51 | }(i) 52 | time.Sleep(time.Millisecond * 5) 53 | } 54 | wg.Wait() 55 | } 56 | -------------------------------------------------------------------------------- /strategy/doc.go: -------------------------------------------------------------------------------- 1 | // Package strategy provides common strategy implementations. 2 | package strategy 3 | -------------------------------------------------------------------------------- /strategy/matchers/doc.go: -------------------------------------------------------------------------------- 1 | // Package matchers provides basic matchers for partitioned type strategies. 2 | package matchers 3 | -------------------------------------------------------------------------------- /strategy/matchers/string.go: -------------------------------------------------------------------------------- 1 | package matchers 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | ) 7 | 8 | type strategyContextKey string 9 | 10 | func (c strategyContextKey) String() string { 11 | return "go-concurrency-limits|strategy|" + string(c) 12 | } 13 | 14 | // StringPredicateContextKey is the StringPredicate context key 15 | // use this in your context.Context 16 | const StringPredicateContextKey = strategyContextKey("stringPredicate") 17 | 18 | // StringPredicateMatcher implements the string predicate matcher. 19 | func StringPredicateMatcher(matchString string, caseInsensitive bool) func(ctx context.Context) bool { 20 | return func(ctx context.Context) bool { 21 | val := ctx.Value(StringPredicateContextKey) 22 | if val != nil { 23 | strVal, ok := val.(string) 24 | if ok { 25 | if caseInsensitive { 26 | return strings.ToLower(strVal) == strings.ToLower(matchString) 27 | } 28 | return strVal == matchString 29 | } 30 | } 31 | return false 32 | } 33 | } 34 | 35 | // LookupPartitionContextKey is the StringLookup context key 36 | // use this in your context.Context 37 | const LookupPartitionContextKey = strategyContextKey("stringLookup") 38 | 39 | // DefaultStringLookupFunc implements the default string lookup partition based on single-value matching. 40 | func DefaultStringLookupFunc(ctx context.Context) string { 41 | val := ctx.Value(LookupPartitionContextKey) 42 | if val != nil { 43 | strVal, ok := val.(string) 44 | if ok { 45 | return strVal 46 | } 47 | } 48 | return "" 49 | } 50 | -------------------------------------------------------------------------------- /strategy/matchers/string_test.go: -------------------------------------------------------------------------------- 1 | package matchers 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestStringPredicateMatcher(t *testing.T) { 11 | t.Parallel() 12 | 13 | asrt := assert.New(t) 14 | matcher := StringPredicateMatcher("foo", false) 15 | ctx1 := context.WithValue(context.Background(), StringPredicateContextKey, "foo") 16 | ctx2 := context.WithValue(context.Background(), StringPredicateContextKey, "Foo") 17 | ctx3 := context.WithValue(context.Background(), StringPredicateContextKey, "bar") 18 | 19 | // stringer test 20 | asrt.Equal("go-concurrency-limits|strategy|stringPredicate", StringPredicateContextKey.String()) 21 | 22 | asrt.True(matcher(ctx1), "expected case match") 23 | asrt.False(matcher(ctx2), "expected case sensitive failure here") 24 | asrt.False(matcher(ctx3), "this shouldn't match") 25 | asrt.False(matcher(context.Background()), "expect default false") 26 | 27 | matcher = StringPredicateMatcher("foo", true) 28 | asrt.True(matcher(ctx1), "expected case match") 29 | asrt.True(matcher(ctx2), "expected case insensitive match") 30 | asrt.False(matcher(ctx3), "this shouldn't match") 31 | } 32 | 33 | func TestDefaultStringLookupFunc(t *testing.T) { 34 | t.Parallel() 35 | 36 | asrt := assert.New(t) 37 | f := DefaultStringLookupFunc 38 | ctx1 := context.WithValue(context.Background(), LookupPartitionContextKey, "foo") 39 | ctx2 := context.WithValue(context.Background(), LookupPartitionContextKey, "bar") 40 | 41 | // stringer test 42 | asrt.Equal("go-concurrency-limits|strategy|stringLookup", LookupPartitionContextKey.String()) 43 | 44 | asrt.Equal("foo", f(ctx1)) 45 | asrt.Equal("bar", f(ctx2)) 46 | asrt.Equal("", f(context.Background()), "expected default value") 47 | } 48 | -------------------------------------------------------------------------------- /strategy/precise.go: -------------------------------------------------------------------------------- 1 | package strategy 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sync" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | ) 10 | 11 | // PreciseStrategy strategy is much more strict than Simple strategy. It uses a sync.Mutex to keep track of 12 | // of the number of inFlight requests instead of atomic accounting. The major difference is this adds additional 13 | // contention for precise limiting, vs SimpleStrategy which trades off exact precision for speed. 14 | type PreciseStrategy struct { 15 | mu sync.Mutex 16 | inFlight int32 17 | limit int32 18 | metricListener core.MetricSampleListener 19 | } 20 | 21 | // NewPreciseStrategy will return a new PreciseStrategy. 22 | func NewPreciseStrategy(limit int) *PreciseStrategy { 23 | return NewPreciseStrategyWithMetricRegistry(limit, core.EmptyMetricRegistryInstance) 24 | } 25 | 26 | // NewPreciseStrategyWithMetricRegistry will create a new PreciseStrategy 27 | func NewPreciseStrategyWithMetricRegistry(limit int, registry core.MetricRegistry, tags ...string) *PreciseStrategy { 28 | if limit < 1 { 29 | limit = 1 30 | } 31 | listener := registry.RegisterDistribution(core.MetricInFlight, tags...) 32 | strategy := &PreciseStrategy{ 33 | limit: int32(limit), 34 | metricListener: listener, 35 | } 36 | 37 | registry.RegisterGauge(core.MetricLimit, core.NewIntMetricSupplierWrapper(strategy.GetLimit), tags...) 38 | return strategy 39 | } 40 | 41 | // TryAcquire will try to acquire a token from the limiter. 42 | // context Context of the request for partitioned limits. 43 | // returns not ok if limit is exceeded, or a StrategyToken that must be released when the operation completes. 44 | func (s *PreciseStrategy) TryAcquire(ctx context.Context) (token core.StrategyToken, ok bool) { 45 | s.mu.Lock() 46 | defer s.mu.Unlock() 47 | if s.inFlight >= s.limit { 48 | s.metricListener.AddSample(float64(s.inFlight)) 49 | return core.NewNotAcquiredStrategyToken(int(s.inFlight)), false 50 | } 51 | s.inFlight++ 52 | s.metricListener.AddSample(float64(s.inFlight)) 53 | return core.NewAcquiredStrategyToken(int(s.inFlight), s.releaseHandler), true 54 | } 55 | 56 | func (s *PreciseStrategy) releaseHandler() { 57 | s.mu.Lock() 58 | s.inFlight-- 59 | s.mu.Unlock() 60 | } 61 | 62 | // SetLimit will update the strategy with a new limit. 63 | func (s *PreciseStrategy) SetLimit(limit int) { 64 | if limit < 1 { 65 | limit = 1 66 | } 67 | s.mu.Lock() 68 | s.limit = int32(limit) 69 | s.mu.Unlock() 70 | } 71 | 72 | // GetLimit will get the current limit 73 | func (s *PreciseStrategy) GetLimit() int { 74 | s.mu.Lock() 75 | defer s.mu.Unlock() 76 | return int(s.limit) 77 | } 78 | 79 | // GetBusyCount will get the current busy count 80 | func (s *PreciseStrategy) GetBusyCount() int { 81 | s.mu.Lock() 82 | defer s.mu.Unlock() 83 | return int(s.inFlight) 84 | } 85 | 86 | func (s *PreciseStrategy) String() string { 87 | s.mu.Lock() 88 | defer s.mu.Unlock() 89 | return fmt.Sprintf("PreciseStrategy{inFlight=%d, limit=%d}", s.inFlight, s.limit) 90 | } 91 | -------------------------------------------------------------------------------- /strategy/precise_test.go: -------------------------------------------------------------------------------- 1 | package strategy 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestPreciseStrategy(t *testing.T) { 11 | t.Parallel() 12 | 13 | t.Run("LimitLessThanOneSetAsOne", func(t2 *testing.T) { 14 | t2.Parallel() 15 | asrt := assert.New(t2) 16 | strategy := NewPreciseStrategy(-10) 17 | asrt.Equal(1, strategy.GetLimit(), "expected a default limit of 1") 18 | }) 19 | 20 | t.Run("InitialState", func(t2 *testing.T) { 21 | t2.Parallel() 22 | asrt := assert.New(t2) 23 | strategy := NewPreciseStrategy(1) 24 | asrt.Equal(1, strategy.GetLimit(), "expected a default limit of 1") 25 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 26 | asrt.Contains(strategy.String(), "PreciseStrategy{inFlight=0, ") 27 | }) 28 | 29 | t.Run("SetLimit", func(t2 *testing.T) { 30 | t2.Parallel() 31 | asrt := assert.New(t2) 32 | strategy := NewPreciseStrategy(0) 33 | asrt.Equal(1, strategy.GetLimit(), "expected a default limit of 3") 34 | strategy.SetLimit(2) 35 | asrt.Equal(2, strategy.GetLimit()) 36 | // negative limits result in 1 37 | strategy.SetLimit(-10) 38 | asrt.Equal(1, strategy.GetLimit()) 39 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 40 | }) 41 | 42 | t.Run("AcquireIncrementsBusy", func(t2 *testing.T) { 43 | t2.Parallel() 44 | asrt := assert.New(t2) 45 | strategy := NewPreciseStrategy(1) 46 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 47 | token, ok := strategy.TryAcquire(context.Background()) 48 | asrt.True(ok && token != nil, "expected token") 49 | asrt.True(token.IsAcquired(), "expected acquired token") 50 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 51 | }) 52 | 53 | t.Run("ExceedingLimitReturnsFalse", func(t2 *testing.T) { 54 | t2.Parallel() 55 | asrt := assert.New(t2) 56 | strategy := NewPreciseStrategy(1) 57 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 58 | token, ok := strategy.TryAcquire(context.Background()) 59 | asrt.True(ok && token != nil, "expected token") 60 | asrt.True(token.IsAcquired(), "expected acquired token") 61 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 62 | 63 | // try again but we expect this to fail 64 | token2, ok2 := strategy.TryAcquire(context.Background()) 65 | asrt.False(ok2, "expected token fail") 66 | if token2 != nil { 67 | asrt.False(token2.IsAcquired(), "token should not be acquired") 68 | } 69 | asrt.Equal(1, strategy.GetBusyCount(), "expected only 1 resource taken") 70 | }) 71 | 72 | t.Run("AcquireAndRelease", func(t2 *testing.T) { 73 | t2.Parallel() 74 | asrt := assert.New(t2) 75 | strategy := NewPreciseStrategy(1) 76 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 77 | token, ok := strategy.TryAcquire(context.Background()) 78 | asrt.True(ok && token != nil, "expected token") 79 | asrt.True(token.IsAcquired(), "expected acquired token") 80 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 81 | 82 | token.Release() 83 | 84 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 85 | 86 | token, ok = strategy.TryAcquire(context.Background()) 87 | asrt.True(ok && token != nil, "expected token") 88 | asrt.True(token.IsAcquired(), "expected acquired token") 89 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 90 | }) 91 | } 92 | -------------------------------------------------------------------------------- /strategy/simple.go: -------------------------------------------------------------------------------- 1 | package strategy 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sync/atomic" 7 | 8 | "github.com/platinummonkey/go-concurrency-limits/core" 9 | ) 10 | 11 | // SimpleStrategy is the simplest strategy for enforcing a concurrency limit that has a single counter for tracking 12 | // total usage. 13 | type SimpleStrategy struct { 14 | inFlight *int32 15 | limit *int32 16 | metricListener core.MetricSampleListener 17 | } 18 | 19 | // NewSimpleStrategy will create a new SimpleStrategy 20 | func NewSimpleStrategy(limit int) *SimpleStrategy { 21 | return NewSimpleStrategyWithMetricRegistry(limit, core.EmptyMetricRegistryInstance) 22 | } 23 | 24 | // NewSimpleStrategyWithMetricRegistry will create a new SimpleStrategy 25 | func NewSimpleStrategyWithMetricRegistry(limit int, registry core.MetricRegistry, tags ...string) *SimpleStrategy { 26 | if limit < 1 { 27 | limit = 1 28 | } 29 | currentLimit := int32(limit) 30 | inFlight := int32(0) 31 | listener := registry.RegisterDistribution(core.MetricInFlight, tags...) 32 | strategy := &SimpleStrategy{ 33 | limit: ¤tLimit, 34 | inFlight: &inFlight, 35 | metricListener: listener, 36 | } 37 | 38 | registry.RegisterGauge(core.MetricLimit, core.NewIntMetricSupplierWrapper(strategy.GetLimit), tags...) 39 | return strategy 40 | } 41 | 42 | // TryAcquire will try to acquire a token from the limiter. 43 | // context Context of the request for partitioned limits. 44 | // returns not ok if limit is exceeded, or a StrategyToken that must be released when the operation completes. 45 | func (s *SimpleStrategy) TryAcquire(ctx context.Context) (token core.StrategyToken, ok bool) { 46 | inFlight := atomic.LoadInt32(s.inFlight) 47 | if inFlight >= atomic.LoadInt32(s.limit) { 48 | s.metricListener.AddSample(float64(inFlight)) 49 | return core.NewNotAcquiredStrategyToken(int(inFlight)), false 50 | } 51 | inFlight = atomic.AddInt32(s.inFlight, 1) 52 | s.metricListener.AddSample(float64(inFlight)) 53 | f := func(ref *int32) func() { 54 | return func() { 55 | atomic.AddInt32(ref, -1) 56 | } 57 | } 58 | return core.NewAcquiredStrategyToken(int(inFlight), f(s.inFlight)), true 59 | } 60 | 61 | // SetLimit will update the strategy with a new limit. 62 | func (s *SimpleStrategy) SetLimit(limit int) { 63 | if limit < 1 { 64 | limit = 1 65 | } 66 | atomic.StoreInt32(s.limit, int32(limit)) 67 | } 68 | 69 | // GetLimit will get the current limit 70 | func (s *SimpleStrategy) GetLimit() int { 71 | return int(atomic.LoadInt32(s.limit)) 72 | } 73 | 74 | // GetBusyCount will get the current busy count 75 | func (s *SimpleStrategy) GetBusyCount() int { 76 | return int(atomic.LoadInt32(s.inFlight)) 77 | } 78 | 79 | func (s *SimpleStrategy) String() string { 80 | return fmt.Sprintf("SimpleStrategy{inFlight=%d, limit=%d}", atomic.LoadInt32(s.inFlight), s.limit) 81 | } 82 | -------------------------------------------------------------------------------- /strategy/simple_test.go: -------------------------------------------------------------------------------- 1 | package strategy 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestSimpleStrategy(t *testing.T) { 11 | t.Parallel() 12 | 13 | t.Run("LimitLessThanOneSetAsOne", func(t2 *testing.T) { 14 | t2.Parallel() 15 | asrt := assert.New(t2) 16 | strategy := NewSimpleStrategy(-10) 17 | asrt.Equal(1, strategy.GetLimit(), "expected a default limit of 1") 18 | }) 19 | 20 | t.Run("InitialState", func(t2 *testing.T) { 21 | t2.Parallel() 22 | asrt := assert.New(t2) 23 | strategy := NewSimpleStrategy(1) 24 | asrt.Equal(1, strategy.GetLimit(), "expected a default limit of 1") 25 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 26 | asrt.Contains(strategy.String(), "SimpleStrategy{inFlight=0, ") 27 | }) 28 | 29 | t.Run("SetLimit", func(t2 *testing.T) { 30 | t2.Parallel() 31 | asrt := assert.New(t2) 32 | strategy := NewSimpleStrategy(0) 33 | asrt.Equal(1, strategy.GetLimit(), "expected a default limit of 3") 34 | strategy.SetLimit(2) 35 | asrt.Equal(2, strategy.GetLimit()) 36 | // negative limits result in 1 37 | strategy.SetLimit(-10) 38 | asrt.Equal(1, strategy.GetLimit()) 39 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 40 | }) 41 | 42 | t.Run("AcquireIncrementsBusy", func(t2 *testing.T) { 43 | t2.Parallel() 44 | asrt := assert.New(t2) 45 | strategy := NewSimpleStrategy(1) 46 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 47 | token, ok := strategy.TryAcquire(context.Background()) 48 | asrt.True(ok && token != nil, "expected token") 49 | asrt.True(token.IsAcquired(), "expected acquired token") 50 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 51 | }) 52 | 53 | t.Run("ExceedingLimitReturnsFalse", func(t2 *testing.T) { 54 | t2.Parallel() 55 | asrt := assert.New(t2) 56 | strategy := NewSimpleStrategy(1) 57 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 58 | token, ok := strategy.TryAcquire(context.Background()) 59 | asrt.True(ok && token != nil, "expected token") 60 | asrt.True(token.IsAcquired(), "expected acquired token") 61 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 62 | 63 | // try again but we expect this to fail 64 | token2, ok2 := strategy.TryAcquire(context.Background()) 65 | asrt.False(ok2, "expected token fail") 66 | if token2 != nil { 67 | asrt.False(token2.IsAcquired(), "token should not be acquired") 68 | } 69 | asrt.Equal(1, strategy.GetBusyCount(), "expected only 1 resource taken") 70 | }) 71 | 72 | t.Run("AcquireAndRelease", func(t2 *testing.T) { 73 | t2.Parallel() 74 | asrt := assert.New(t2) 75 | strategy := NewSimpleStrategy(1) 76 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 77 | token, ok := strategy.TryAcquire(context.Background()) 78 | asrt.True(ok && token != nil, "expected token") 79 | asrt.True(token.IsAcquired(), "expected acquired token") 80 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 81 | 82 | token.Release() 83 | 84 | asrt.Equal(0, strategy.GetBusyCount(), "expected all resources free") 85 | 86 | token, ok = strategy.TryAcquire(context.Background()) 87 | asrt.True(ok && token != nil, "expected token") 88 | asrt.True(token.IsAcquired(), "expected acquired token") 89 | asrt.Equal(1, strategy.GetBusyCount(), "expected 1 resource taken") 90 | }) 91 | } 92 | --------------------------------------------------------------------------------