2022-09-09 09:16:22 +01:00
|
|
|
// Copyright (C) 2022 Storj Labs, Inc.
|
|
|
|
// See LICENSE for copying information.
|
|
|
|
|
|
|
|
package main
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2023-01-05 18:00:00 +00:00
|
|
|
"errors"
|
2022-09-16 16:35:19 +01:00
|
|
|
"io"
|
2022-09-12 20:22:01 +01:00
|
|
|
"time"
|
2022-09-09 09:16:22 +01:00
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
"github.com/blang/semver"
|
2022-09-09 09:16:22 +01:00
|
|
|
"github.com/zeebo/errs"
|
2022-09-12 20:22:01 +01:00
|
|
|
"go.uber.org/zap"
|
|
|
|
|
|
|
|
"storj.io/common/errs2"
|
2022-12-21 15:14:53 +00:00
|
|
|
"storj.io/common/pb"
|
2022-09-12 20:22:01 +01:00
|
|
|
"storj.io/common/rpc"
|
2022-10-11 07:06:44 +01:00
|
|
|
"storj.io/common/rpc/rpcpool"
|
2022-09-12 20:22:01 +01:00
|
|
|
"storj.io/common/rpc/rpcstatus"
|
|
|
|
"storj.io/common/storj"
|
|
|
|
"storj.io/common/sync2"
|
2022-12-13 22:52:36 +00:00
|
|
|
"storj.io/storj/satellite/audit"
|
2022-10-06 13:26:57 +01:00
|
|
|
"storj.io/storj/satellite/metabase"
|
2022-09-12 20:22:01 +01:00
|
|
|
"storj.io/storj/satellite/orders"
|
|
|
|
"storj.io/uplink/private/piecestore"
|
|
|
|
)
|
|
|
|
|
2022-09-15 15:23:10 +01:00
|
|
|
// ErrNodeOffline is returned when it was not possible to contact a node or the node was not responding.
|
|
|
|
var ErrNodeOffline = errs.Class("node offline")
|
2022-09-09 09:16:22 +01:00
|
|
|
|
2023-01-04 16:33:07 +00:00
|
|
|
// ErrNoSuchNode is returned when a node has been disqualified or gracefully exited, or no known node
|
|
|
|
// has the specified alias. Pieces associated with the given alias should be considered lost.
|
|
|
|
var ErrNoSuchNode = errs.Class("no such node")
|
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
var errWrongNodeVersion = errs.Class("wrong node version")
|
|
|
|
|
2022-09-15 15:23:10 +01:00
|
|
|
// VerifierConfig contains configurations for operation.
|
|
|
|
type VerifierConfig struct {
|
2022-09-19 13:40:22 +01:00
|
|
|
DialTimeout time.Duration `help:"how long to wait for a successful dial" default:"2s"`
|
2022-09-15 15:23:10 +01:00
|
|
|
PerPieceTimeout time.Duration `help:"duration to wait per piece download" default:"800ms"`
|
|
|
|
OrderRetryThrottle time.Duration `help:"how much to wait before retrying order creation" default:"50ms"`
|
2022-09-12 20:22:01 +01:00
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
RequestThrottle time.Duration `help:"minimum interval for sending out each request" default:"150ms"`
|
|
|
|
VersionWithExists string `help:"minimum storage node version with implemented Exists method" default:"v1.69.2"`
|
2022-09-15 15:23:10 +01:00
|
|
|
}
|
2022-09-12 20:22:01 +01:00
|
|
|
|
|
|
|
// NodeVerifier implements segment verification by dialing nodes.
|
|
|
|
type NodeVerifier struct {
|
2022-09-15 15:23:10 +01:00
|
|
|
log *zap.Logger
|
|
|
|
|
|
|
|
config VerifierConfig
|
|
|
|
|
2022-09-12 20:22:01 +01:00
|
|
|
dialer rpc.Dialer
|
|
|
|
orders *orders.Service
|
2022-12-13 22:52:36 +00:00
|
|
|
|
2023-01-05 17:32:18 +00:00
|
|
|
// this is a callback so that problematic pieces can be reported as they are found,
|
|
|
|
// rather than being kept in a list which might grow unreasonably large.
|
2022-12-13 22:52:36 +00:00
|
|
|
reportPiece pieceReporterFunc
|
2022-12-21 15:14:53 +00:00
|
|
|
|
|
|
|
versionWithExists semver.Version
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
|
|
|
|
|
|
|
var _ Verifier = (*NodeVerifier)(nil)
|
|
|
|
|
|
|
|
// NewVerifier creates a new segment verifier using the specified dialer.
|
2022-09-15 15:23:10 +01:00
|
|
|
func NewVerifier(log *zap.Logger, dialer rpc.Dialer, orders *orders.Service, config VerifierConfig) *NodeVerifier {
|
2022-09-19 13:40:22 +01:00
|
|
|
configuredDialer := dialer
|
|
|
|
if config.DialTimeout > 0 {
|
|
|
|
configuredDialer.DialTimeout = config.DialTimeout
|
|
|
|
}
|
|
|
|
|
2022-10-11 07:06:44 +01:00
|
|
|
configuredDialer.Pool = rpcpool.New(rpcpool.Options{
|
|
|
|
Capacity: 1000,
|
|
|
|
KeyCapacity: 5,
|
|
|
|
IdleExpiration: 10 * time.Minute,
|
|
|
|
})
|
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
version, err := semver.ParseTolerant(config.VersionWithExists)
|
|
|
|
if err != nil {
|
|
|
|
log.Warn("invalid VersionWithExists", zap.String("VersionWithExists", config.VersionWithExists), zap.Error(err))
|
|
|
|
}
|
|
|
|
|
2022-09-12 20:22:01 +01:00
|
|
|
return &NodeVerifier{
|
2022-12-21 15:14:53 +00:00
|
|
|
log: log,
|
|
|
|
config: config,
|
|
|
|
dialer: configuredDialer,
|
|
|
|
orders: orders,
|
|
|
|
versionWithExists: version,
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// Verify a collection of segments by attempting to download a byte from each segment from the target node.
|
2022-12-21 15:14:53 +00:00
|
|
|
func (service *NodeVerifier) Verify(ctx context.Context, alias metabase.NodeAlias, target storj.NodeURL, targetVersion string, segments []*Segment, ignoreThrottle bool) (verifiedCount int, err error) {
|
|
|
|
verifiedCount, err = service.VerifyWithExists(ctx, alias, target, targetVersion, segments)
|
|
|
|
// if Exists method is unimplemented or it is wrong node version fallback to download verification
|
2022-12-22 14:14:20 +00:00
|
|
|
if !methodUnimplemented(err) && !errWrongNodeVersion.Has(err) {
|
2022-12-21 15:14:53 +00:00
|
|
|
return verifiedCount, err
|
|
|
|
}
|
|
|
|
if err != nil {
|
|
|
|
service.log.Debug("fallback to download method", zap.Error(err))
|
|
|
|
}
|
|
|
|
|
|
|
|
service.log.Debug("verify segments by downloading pieces")
|
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
var client *piecestore.Client
|
2022-10-13 13:35:31 +01:00
|
|
|
defer func() {
|
|
|
|
if client != nil {
|
|
|
|
_ = client.Close()
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
const maxDials = 2
|
|
|
|
dialCount := 0
|
2022-09-12 20:22:01 +01:00
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
rateLimiter := newRateLimiter(0, 0)
|
|
|
|
if !ignoreThrottle {
|
|
|
|
rateLimiter = newRateLimiter(service.config.RequestThrottle, service.config.RequestThrottle/4)
|
|
|
|
}
|
2022-10-13 13:35:31 +01:00
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
nextRequest := time.Now()
|
|
|
|
for i, segment := range segments {
|
|
|
|
nextRequest, err = rateLimiter.next(ctx, nextRequest)
|
|
|
|
if err != nil {
|
|
|
|
return i, Error.Wrap(err)
|
|
|
|
}
|
2022-10-13 13:35:31 +01:00
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
for client == nil {
|
|
|
|
dialCount++
|
|
|
|
if dialCount > maxDials {
|
|
|
|
return i, ErrNodeOffline.New("too many redials")
|
2022-10-13 13:35:31 +01:00
|
|
|
}
|
2022-10-19 14:32:24 +01:00
|
|
|
client, err = piecestore.Dial(rpcpool.WithForceDial(ctx), service.dialer, target, piecestore.DefaultConfig)
|
2022-10-13 13:35:31 +01:00
|
|
|
if err != nil {
|
2022-12-13 22:52:36 +00:00
|
|
|
service.log.Info("failed to dial node",
|
|
|
|
zap.Stringer("node-id", target.ID),
|
|
|
|
zap.Error(err))
|
|
|
|
client = nil
|
|
|
|
nextRequest, err = rateLimiter.next(ctx, nextRequest)
|
|
|
|
if err != nil {
|
|
|
|
return i, Error.Wrap(err)
|
|
|
|
}
|
2022-10-13 13:35:31 +01:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
outcome, err := service.verifySegment(ctx, client, alias, target, segment)
|
2022-09-12 20:22:01 +01:00
|
|
|
if err != nil {
|
2022-12-13 22:52:36 +00:00
|
|
|
// we could not do the verification, for a reason that implies we won't be able
|
|
|
|
// to do any more
|
2022-10-13 14:25:12 +01:00
|
|
|
return i, Error.Wrap(err)
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
2022-12-13 22:52:36 +00:00
|
|
|
switch outcome {
|
2023-01-05 18:00:00 +00:00
|
|
|
case audit.OutcomeNodeOffline, audit.OutcomeTimedOut:
|
2022-12-13 22:52:36 +00:00
|
|
|
_ = client.Close()
|
|
|
|
client = nil
|
|
|
|
case audit.OutcomeFailure:
|
|
|
|
segment.Status.MarkNotFound()
|
|
|
|
case audit.OutcomeSuccess:
|
|
|
|
segment.Status.MarkFound()
|
2022-09-19 13:40:22 +01:00
|
|
|
}
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
2022-10-13 14:25:12 +01:00
|
|
|
return len(segments), nil
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
// verifySegment tries to verify the segment by downloading a single byte from the piece of the segment
|
|
|
|
// on the specified target node.
|
|
|
|
func (service *NodeVerifier) verifySegment(ctx context.Context, client *piecestore.Client, alias metabase.NodeAlias, target storj.NodeURL, segment *Segment) (outcome audit.Outcome, err error) {
|
2022-10-06 13:26:57 +01:00
|
|
|
pieceNum := findPieceNum(segment, alias)
|
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
logger := service.log.With(
|
|
|
|
zap.Stringer("stream-id", segment.StreamID),
|
|
|
|
zap.Stringer("node-id", target.ID),
|
|
|
|
zap.Uint64("position", segment.Position.Encode()),
|
|
|
|
zap.Uint16("piece-num", pieceNum))
|
|
|
|
|
|
|
|
defer func() {
|
|
|
|
// report the outcome of the piece check, if required
|
|
|
|
if outcome != audit.OutcomeSuccess && service.reportPiece != nil {
|
|
|
|
reportErr := service.reportPiece(ctx, &segment.VerifySegment, target.ID, int(pieceNum), outcome)
|
|
|
|
err = errs.Combine(err, reportErr)
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
|
2022-10-06 13:26:57 +01:00
|
|
|
limit, piecePrivateKey, _, err := service.orders.CreateAuditOrderLimit(ctx, target.ID, pieceNum, segment.RootPieceID, segment.Redundancy.ShareSize)
|
2022-09-12 20:22:01 +01:00
|
|
|
if err != nil {
|
2022-12-13 22:52:36 +00:00
|
|
|
logger.Error("failed to create order limit",
|
2022-09-15 15:23:10 +01:00
|
|
|
zap.Stringer("retrying in", service.config.OrderRetryThrottle),
|
2022-09-12 20:22:01 +01:00
|
|
|
zap.Error(err))
|
|
|
|
|
2022-09-15 15:23:10 +01:00
|
|
|
if !sync2.Sleep(ctx, service.config.OrderRetryThrottle) {
|
2022-12-13 22:52:36 +00:00
|
|
|
return audit.OutcomeNotPerformed, Error.Wrap(ctx.Err())
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
|
|
|
|
2022-10-06 13:26:57 +01:00
|
|
|
limit, piecePrivateKey, _, err = service.orders.CreateAuditOrderLimit(ctx, target.ID, pieceNum, segment.RootPieceID, segment.Redundancy.ShareSize)
|
2022-09-12 20:22:01 +01:00
|
|
|
if err != nil {
|
2022-12-13 22:52:36 +00:00
|
|
|
return audit.OutcomeNotPerformed, Error.Wrap(err)
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-09-15 15:23:10 +01:00
|
|
|
timedCtx, cancel := context.WithTimeout(ctx, service.config.PerPieceTimeout)
|
2022-09-12 20:22:01 +01:00
|
|
|
defer cancel()
|
|
|
|
|
|
|
|
downloader, err := client.Download(timedCtx, limit.GetLimit(), piecePrivateKey, 0, 1)
|
|
|
|
if err != nil {
|
2022-12-13 22:52:36 +00:00
|
|
|
logger.Error("download failed", zap.Error(err))
|
2022-09-12 20:22:01 +01:00
|
|
|
if errs2.IsRPC(err, rpcstatus.NotFound) {
|
2022-12-13 22:52:36 +00:00
|
|
|
return audit.OutcomeFailure, nil
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
2022-12-13 22:52:36 +00:00
|
|
|
if errs2.IsRPC(err, rpcstatus.Unknown) {
|
|
|
|
// dial failed -- offline node
|
|
|
|
return audit.OutcomeNodeOffline, nil
|
|
|
|
}
|
|
|
|
return audit.OutcomeUnknownError, nil
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
|
|
|
|
|
|
|
buf := [1]byte{}
|
2022-09-21 15:55:46 +01:00
|
|
|
_, errRead := io.ReadFull(downloader, buf[:])
|
|
|
|
errClose := downloader.Close()
|
|
|
|
|
|
|
|
err = errs.Combine(errClose, errRead)
|
2022-09-12 20:22:01 +01:00
|
|
|
if err != nil {
|
|
|
|
if errs2.IsRPC(err, rpcstatus.NotFound) {
|
2022-12-13 22:52:36 +00:00
|
|
|
logger.Info("segment not found", zap.Error(err))
|
|
|
|
return audit.OutcomeFailure, nil
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
2023-01-05 18:00:00 +00:00
|
|
|
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
|
|
|
|
logger.Info("failed to get timely response when asking for piece", zap.Error(err))
|
|
|
|
return audit.OutcomeTimedOut, nil
|
|
|
|
}
|
2022-09-12 20:22:01 +01:00
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
logger.Error("read/close failed", zap.Error(err))
|
|
|
|
return audit.OutcomeUnknownError, nil
|
2022-09-12 20:22:01 +01:00
|
|
|
}
|
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
logger.Info("download succeeded")
|
|
|
|
return audit.OutcomeSuccess, nil
|
2022-09-09 09:16:22 +01:00
|
|
|
}
|
2022-10-06 13:26:57 +01:00
|
|
|
|
|
|
|
func findPieceNum(segment *Segment, alias metabase.NodeAlias) uint16 {
|
|
|
|
for _, p := range segment.AliasPieces {
|
|
|
|
if p.Alias == alias {
|
|
|
|
return p.Number
|
|
|
|
}
|
|
|
|
}
|
|
|
|
panic("piece number not found")
|
|
|
|
}
|
2022-12-13 22:52:36 +00:00
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
// VerifyWithExists verifies that the segments exist on the specified node by calling the piecestore Exists
|
|
|
|
// endpoint if the node version supports it.
|
|
|
|
func (service *NodeVerifier) VerifyWithExists(ctx context.Context, alias metabase.NodeAlias, target storj.NodeURL, targetVersion string, segments []*Segment) (verifiedCount int, err error) {
|
2022-12-27 20:55:43 +00:00
|
|
|
defer mon.Task()(&ctx)(&err)
|
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
if service.versionWithExists.String() == "" || targetVersion == "" {
|
|
|
|
return 0, errWrongNodeVersion.New("missing node version or no base version defined")
|
|
|
|
}
|
|
|
|
|
|
|
|
nodeVersion, err := semver.ParseTolerant(targetVersion)
|
|
|
|
if err != nil {
|
|
|
|
return 0, errWrongNodeVersion.Wrap(err)
|
|
|
|
}
|
|
|
|
|
|
|
|
if !nodeVersion.GE(service.versionWithExists) {
|
|
|
|
return 0, errWrongNodeVersion.New("too old version")
|
|
|
|
}
|
|
|
|
|
|
|
|
service.log.Debug("verify segments using Exists method", zap.Stringer("node-id", target.ID))
|
|
|
|
|
|
|
|
var conn *rpc.Conn
|
|
|
|
var client pb.DRPCPiecestoreClient
|
|
|
|
defer func() {
|
|
|
|
if conn != nil {
|
|
|
|
_ = conn.Close()
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
|
|
|
|
const maxDials = 2
|
|
|
|
dialCount := 0
|
|
|
|
|
|
|
|
for client == nil {
|
|
|
|
dialCount++
|
|
|
|
if dialCount > maxDials {
|
|
|
|
return 0, ErrNodeOffline.New("too many redials")
|
|
|
|
}
|
|
|
|
|
|
|
|
conn, err := service.dialer.DialNodeURL(rpcpool.WithForceDial(ctx), target)
|
|
|
|
if err != nil {
|
|
|
|
service.log.Info("failed to dial node",
|
|
|
|
zap.Stringer("node-id", target.ID),
|
|
|
|
zap.Error(err))
|
|
|
|
} else {
|
|
|
|
client = pb.NewDRPCPiecestoreClient(conn)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-01-05 18:02:06 +00:00
|
|
|
// this is maybe too generous, since Exists should be faster than asking for pieces one by one,
|
|
|
|
// but seems like a good first try
|
|
|
|
ctx, cancel := context.WithTimeout(ctx, service.config.PerPieceTimeout*time.Duration(len(segments)))
|
|
|
|
defer cancel()
|
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
err = service.verifySegmentsWithExists(ctx, client, alias, target, segments)
|
|
|
|
if err != nil {
|
|
|
|
// we could not do the verification, for a reason that implies we won't be able
|
|
|
|
// to do any more
|
|
|
|
return 0, Error.Wrap(err)
|
|
|
|
}
|
|
|
|
|
|
|
|
return len(segments), nil
|
|
|
|
}
|
|
|
|
|
2022-12-22 14:14:20 +00:00
|
|
|
func methodUnimplemented(err error) bool {
|
|
|
|
return errs2.IsRPC(err, rpcstatus.Unimplemented) || errs2.IsRPC(err, rpcstatus.Unknown)
|
|
|
|
}
|
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
// verifySegmentsWithExists calls the Exists endpoint on the specified target node for each segment.
|
|
|
|
func (service *NodeVerifier) verifySegmentsWithExists(ctx context.Context, client pb.DRPCPiecestoreClient, alias metabase.NodeAlias, target storj.NodeURL, segments []*Segment) (err error) {
|
2022-12-27 20:55:43 +00:00
|
|
|
defer mon.Task()(&ctx)(&err)
|
|
|
|
|
2022-12-21 15:14:53 +00:00
|
|
|
pieceIds := make([]storj.PieceID, 0, len(segments))
|
2022-12-27 20:57:28 +00:00
|
|
|
pieceNums := make([]uint16, 0, len(segments))
|
2022-12-21 15:14:53 +00:00
|
|
|
|
|
|
|
for _, segment := range segments {
|
|
|
|
pieceNum := findPieceNum(segment, alias)
|
2022-12-27 20:57:28 +00:00
|
|
|
pieceNums = append(pieceNums, pieceNum)
|
2022-12-21 15:14:53 +00:00
|
|
|
pieceId := segment.RootPieceID.Derive(target.ID, int32(pieceNum))
|
|
|
|
pieceIds = append(pieceIds, pieceId)
|
|
|
|
}
|
|
|
|
|
|
|
|
response, err := client.Exists(ctx, &pb.ExistsRequest{
|
|
|
|
PieceIds: pieceIds,
|
|
|
|
})
|
|
|
|
if err != nil {
|
|
|
|
return Error.Wrap(err)
|
|
|
|
}
|
|
|
|
|
2023-01-19 23:44:06 +00:00
|
|
|
service.log.Debug("Exists response",
|
|
|
|
zap.Int("request-count", len(pieceIds)),
|
|
|
|
zap.Int("missing-count", len(response.Missing)))
|
|
|
|
|
2022-12-22 14:14:20 +00:00
|
|
|
if len(response.Missing) == 0 {
|
|
|
|
for i := range segments {
|
|
|
|
segments[i].Status.MarkFound()
|
2022-12-21 15:14:53 +00:00
|
|
|
}
|
2022-12-22 14:14:20 +00:00
|
|
|
return nil
|
2022-12-21 15:14:53 +00:00
|
|
|
}
|
|
|
|
|
2022-12-22 14:14:20 +00:00
|
|
|
missing := make(map[uint32]struct{}, len(segments))
|
|
|
|
for _, m := range response.Missing {
|
|
|
|
missing[m] = struct{}{}
|
|
|
|
}
|
2022-12-21 15:14:53 +00:00
|
|
|
|
2022-12-22 14:14:20 +00:00
|
|
|
for i := range segments {
|
|
|
|
_, miss := missing[uint32(i)]
|
|
|
|
if miss {
|
|
|
|
segments[i].Status.MarkNotFound()
|
2022-12-27 20:57:28 +00:00
|
|
|
if service.reportPiece != nil {
|
|
|
|
reportErr := service.reportPiece(ctx, &segments[i].VerifySegment, target.ID, int(pieceNums[i]), audit.OutcomeFailure)
|
|
|
|
err = errs.Combine(err, reportErr)
|
|
|
|
}
|
2022-12-22 14:14:20 +00:00
|
|
|
} else {
|
|
|
|
segments[i].Status.MarkFound()
|
2022-12-21 15:14:53 +00:00
|
|
|
}
|
|
|
|
}
|
2022-12-22 14:14:20 +00:00
|
|
|
|
2022-12-27 20:57:28 +00:00
|
|
|
return err
|
2022-12-21 15:14:53 +00:00
|
|
|
}
|
|
|
|
|
2022-12-13 22:52:36 +00:00
|
|
|
// rateLimiter limits the rate of some type of event. It acts like a token
|
|
|
|
// bucket, allowing for bursting, as long as the _average_ interval between
|
|
|
|
// events over the lifetime of the rateLimiter is less than or equal to the
|
|
|
|
// specified averageInterval.
|
|
|
|
//
|
|
|
|
// The wait time between events will be at least minInterval, even if an
|
|
|
|
// event would otherwise have been allowed with bursting.
|
|
|
|
type rateLimiter struct {
|
|
|
|
nowFn func() time.Time
|
|
|
|
sleepFn func(ctx context.Context, t time.Duration) bool
|
|
|
|
averageInterval time.Duration
|
|
|
|
minInterval time.Duration
|
|
|
|
}
|
|
|
|
|
|
|
|
// newRateLimiter creates a new rateLimiter. If both arguments are specified
|
|
|
|
// as 0, the rate limiter will have no effect (next() will always return
|
|
|
|
// immediately).
|
|
|
|
func newRateLimiter(averageInterval, minInterval time.Duration) rateLimiter {
|
|
|
|
return rateLimiter{
|
|
|
|
nowFn: time.Now,
|
|
|
|
sleepFn: sync2.Sleep,
|
|
|
|
averageInterval: averageInterval,
|
|
|
|
minInterval: minInterval,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// next() sleeps until the time when the next event should be allowed.
|
|
|
|
// It should be passed the current time for the first call. After that,
|
|
|
|
// the value returned by the last call to next() should be passed in
|
|
|
|
// as its argument.
|
|
|
|
func (r rateLimiter) next(ctx context.Context, nextTime time.Time) (time.Time, error) {
|
|
|
|
sleepTime := nextTime.Sub(r.nowFn())
|
|
|
|
if sleepTime < r.minInterval {
|
|
|
|
sleepTime = r.minInterval
|
|
|
|
}
|
|
|
|
if !r.sleepFn(ctx, sleepTime) {
|
|
|
|
return time.Time{}, ctx.Err()
|
|
|
|
}
|
|
|
|
return nextTime.Add(r.averageInterval), nil
|
|
|
|
}
|