github.com/ava-labs/avalanchego@v1.11.11/network/p2p/client.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  	"errors"
     9  	"fmt"
    10  
    11  	"go.uber.org/zap"
    12  
    13  	"github.com/ava-labs/avalanchego/ids"
    14  	"github.com/ava-labs/avalanchego/message"
    15  	"github.com/ava-labs/avalanchego/snow/engine/common"
    16  	"github.com/ava-labs/avalanchego/utils/set"
    17  )
    18  
    19  var (
    20  	ErrRequestPending = errors.New("request pending")
    21  	ErrNoPeers        = errors.New("no peers")
    22  )
    23  
    24  // AppResponseCallback is called upon receiving an AppResponse for an AppRequest
    25  // issued by Client.
    26  // Callers should check [err] to see whether the AppRequest failed or not.
    27  type AppResponseCallback func(
    28  	ctx context.Context,
    29  	nodeID ids.NodeID,
    30  	responseBytes []byte,
    31  	err error,
    32  )
    33  
    34  type Client struct {
    35  	handlerID     uint64
    36  	handlerIDStr  string
    37  	handlerPrefix []byte
    38  	router        *router
    39  	sender        common.AppSender
    40  	options       *clientOptions
    41  }
    42  
    43  // AppRequestAny issues an AppRequest to an arbitrary node decided by Client.
    44  // If a specific node needs to be requested, use AppRequest instead.
    45  // See AppRequest for more docs.
    46  func (c *Client) AppRequestAny(
    47  	ctx context.Context,
    48  	appRequestBytes []byte,
    49  	onResponse AppResponseCallback,
    50  ) error {
    51  	sampled := c.options.nodeSampler.Sample(ctx, 1)
    52  	if len(sampled) != 1 {
    53  		return ErrNoPeers
    54  	}
    55  
    56  	nodeIDs := set.Of(sampled...)
    57  	return c.AppRequest(ctx, nodeIDs, appRequestBytes, onResponse)
    58  }
    59  
    60  // AppRequest issues an arbitrary request to a node.
    61  // [onResponse] is invoked upon an error or a response.
    62  func (c *Client) AppRequest(
    63  	ctx context.Context,
    64  	nodeIDs set.Set[ids.NodeID],
    65  	appRequestBytes []byte,
    66  	onResponse AppResponseCallback,
    67  ) error {
    68  	// Cancellation is removed from this context to avoid erroring unexpectedly.
    69  	// SendAppRequest should be non-blocking and any error other than context
    70  	// cancellation is unexpected.
    71  	//
    72  	// This guarantees that the router should never receive an unexpected
    73  	// AppResponse.
    74  	ctxWithoutCancel := context.WithoutCancel(ctx)
    75  
    76  	c.router.lock.Lock()
    77  	defer c.router.lock.Unlock()
    78  
    79  	appRequestBytes = PrefixMessage(c.handlerPrefix, appRequestBytes)
    80  	for nodeID := range nodeIDs {
    81  		requestID := c.router.requestID
    82  		if _, ok := c.router.pendingAppRequests[requestID]; ok {
    83  			return fmt.Errorf(
    84  				"failed to issue request with request id %d: %w",
    85  				requestID,
    86  				ErrRequestPending,
    87  			)
    88  		}
    89  
    90  		if err := c.sender.SendAppRequest(
    91  			ctxWithoutCancel,
    92  			set.Of(nodeID),
    93  			requestID,
    94  			appRequestBytes,
    95  		); err != nil {
    96  			c.router.log.Error("unexpected error when sending message",
    97  				zap.Stringer("op", message.AppRequestOp),
    98  				zap.Stringer("nodeID", nodeID),
    99  				zap.Uint32("requestID", requestID),
   100  				zap.Error(err),
   101  			)
   102  			return err
   103  		}
   104  
   105  		c.router.pendingAppRequests[requestID] = pendingAppRequest{
   106  			handlerID: c.handlerIDStr,
   107  			callback:  onResponse,
   108  		}
   109  		c.router.requestID += 2
   110  	}
   111  
   112  	return nil
   113  }
   114  
   115  // AppGossip sends a gossip message to a random set of peers.
   116  func (c *Client) AppGossip(
   117  	ctx context.Context,
   118  	config common.SendConfig,
   119  	appGossipBytes []byte,
   120  ) error {
   121  	// Cancellation is removed from this context to avoid erroring unexpectedly.
   122  	// SendAppGossip should be non-blocking and any error other than context
   123  	// cancellation is unexpected.
   124  	ctxWithoutCancel := context.WithoutCancel(ctx)
   125  
   126  	return c.sender.SendAppGossip(
   127  		ctxWithoutCancel,
   128  		config,
   129  		PrefixMessage(c.handlerPrefix, appGossipBytes),
   130  	)
   131  }
   132  
   133  // PrefixMessage prefixes the original message with the protocol identifier.
   134  //
   135  // Only gossip and request messages need to be prefixed.
   136  // Response messages don't need to be prefixed because request ids are tracked
   137  // which map to the expected response handler.
   138  func PrefixMessage(prefix, msg []byte) []byte {
   139  	messageBytes := make([]byte, len(prefix)+len(msg))
   140  	copy(messageBytes, prefix)
   141  	copy(messageBytes[len(prefix):], msg)
   142  	return messageBytes
   143  }