get.pme.sh/pnats@v0.0.0-20240304004023-26bb5a137ed0/server/rate_counter.go (about)

     1  // Copyright 2021-2022 The NATS Authors
     2  // Licensed under the Apache License, Version 2.0 (the "License");
     3  // you may not use this file except in compliance with the License.
     4  // You may obtain a copy of the License at
     5  //
     6  // http://www.apache.org/licenses/LICENSE-2.0
     7  //
     8  // Unless required by applicable law or agreed to in writing, software
     9  // distributed under the License is distributed on an "AS IS" BASIS,
    10  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package server
    15  
    16  import (
    17  	"sync"
    18  	"time"
    19  )
    20  
    21  type rateCounter struct {
    22  	limit    int64
    23  	count    int64
    24  	blocked  uint64
    25  	end      time.Time
    26  	interval time.Duration
    27  	mu       sync.Mutex
    28  }
    29  
    30  func newRateCounter(limit int64) *rateCounter {
    31  	return &rateCounter{
    32  		limit:    limit,
    33  		interval: time.Second,
    34  	}
    35  }
    36  
    37  func (r *rateCounter) allow() bool {
    38  	now := time.Now()
    39  
    40  	r.mu.Lock()
    41  
    42  	if now.After(r.end) {
    43  		r.count = 0
    44  		r.end = now.Add(r.interval)
    45  	} else {
    46  		r.count++
    47  	}
    48  	allow := r.count < r.limit
    49  	if !allow {
    50  		r.blocked++
    51  	}
    52  
    53  	r.mu.Unlock()
    54  
    55  	return allow
    56  }
    57  
    58  func (r *rateCounter) countBlocked() uint64 {
    59  	r.mu.Lock()
    60  	blocked := r.blocked
    61  	r.blocked = 0
    62  	r.mu.Unlock()
    63  
    64  	return blocked
    65  }