gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/eventchannel/rate.go (about)

     1  // Copyright 2019 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package eventchannel
    16  
    17  import (
    18  	"golang.org/x/time/rate"
    19  	"google.golang.org/protobuf/proto"
    20  )
    21  
    22  // rateLimitedEmitter wraps an emitter and limits events to the given limits.
    23  // Events that would exceed the limit are discarded.
    24  type rateLimitedEmitter struct {
    25  	inner   Emitter
    26  	limiter *rate.Limiter
    27  }
    28  
    29  // RateLimitedEmitterFrom creates a new event channel emitter that wraps the
    30  // existing emitter and enforces rate limits. The limits are imposed via a
    31  // token bucket, with `maxRate` events per second, with burst size of `burst`
    32  // events. See the golang.org/x/time/rate package and
    33  // https://en.wikipedia.org/wiki/Token_bucket for more information about token
    34  // buckets generally.
    35  func RateLimitedEmitterFrom(inner Emitter, maxRate float64, burst int) Emitter {
    36  	return &rateLimitedEmitter{
    37  		inner:   inner,
    38  		limiter: rate.NewLimiter(rate.Limit(maxRate), burst),
    39  	}
    40  }
    41  
    42  // Emit implements EventEmitter.Emit.
    43  func (rle *rateLimitedEmitter) Emit(msg proto.Message) (bool, error) {
    44  	if !rle.limiter.Allow() {
    45  		// Drop event.
    46  		return false, nil
    47  	}
    48  	return rle.inner.Emit(msg)
    49  }
    50  
    51  // Close implements EventEmitter.Close.
    52  func (rle *rateLimitedEmitter) Close() error {
    53  	return rle.inner.Close()
    54  }