github.com/ava-labs/avalanchego@v1.11.11/network/p2p/throttler_handler.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package p2p
     5  
     6  import (
     7  	"context"
     8  	"time"
     9  
    10  	"go.uber.org/zap"
    11  
    12  	"github.com/ava-labs/avalanchego/ids"
    13  	"github.com/ava-labs/avalanchego/snow/engine/common"
    14  	"github.com/ava-labs/avalanchego/utils/logging"
    15  )
    16  
    17  var _ Handler = (*ThrottlerHandler)(nil)
    18  
    19  func NewThrottlerHandler(handler Handler, throttler Throttler, log logging.Logger) *ThrottlerHandler {
    20  	return &ThrottlerHandler{
    21  		handler:   handler,
    22  		throttler: throttler,
    23  		log:       log,
    24  	}
    25  }
    26  
    27  type ThrottlerHandler struct {
    28  	handler   Handler
    29  	throttler Throttler
    30  	log       logging.Logger
    31  }
    32  
    33  func (t ThrottlerHandler) AppGossip(ctx context.Context, nodeID ids.NodeID, gossipBytes []byte) {
    34  	if !t.throttler.Handle(nodeID) {
    35  		t.log.Debug("dropping message",
    36  			zap.Stringer("nodeID", nodeID),
    37  			zap.String("reason", "throttled"),
    38  		)
    39  		return
    40  	}
    41  
    42  	t.handler.AppGossip(ctx, nodeID, gossipBytes)
    43  }
    44  
    45  func (t ThrottlerHandler) AppRequest(ctx context.Context, nodeID ids.NodeID, deadline time.Time, requestBytes []byte) ([]byte, *common.AppError) {
    46  	if !t.throttler.Handle(nodeID) {
    47  		return nil, ErrThrottled
    48  	}
    49  
    50  	return t.handler.AppRequest(ctx, nodeID, deadline, requestBytes)
    51  }