github.com/MetalBlockchain/metalgo@v1.11.9/snow/networking/router/traced_router.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package router
     5  
     6  import (
     7  	"context"
     8  	"time"
     9  
    10  	"github.com/prometheus/client_golang/prometheus"
    11  	"go.opentelemetry.io/otel/attribute"
    12  
    13  	"github.com/MetalBlockchain/metalgo/ids"
    14  	"github.com/MetalBlockchain/metalgo/message"
    15  	"github.com/MetalBlockchain/metalgo/proto/pb/p2p"
    16  	"github.com/MetalBlockchain/metalgo/snow/networking/handler"
    17  	"github.com/MetalBlockchain/metalgo/snow/networking/timeout"
    18  	"github.com/MetalBlockchain/metalgo/trace"
    19  	"github.com/MetalBlockchain/metalgo/utils/logging"
    20  	"github.com/MetalBlockchain/metalgo/utils/set"
    21  	"github.com/MetalBlockchain/metalgo/version"
    22  
    23  	oteltrace "go.opentelemetry.io/otel/trace"
    24  )
    25  
    26  var _ Router = (*tracedRouter)(nil)
    27  
    28  type tracedRouter struct {
    29  	router Router
    30  	tracer trace.Tracer
    31  }
    32  
    33  func Trace(router Router, tracer trace.Tracer) Router {
    34  	return &tracedRouter{
    35  		router: router,
    36  		tracer: tracer,
    37  	}
    38  }
    39  
    40  func (r *tracedRouter) Initialize(
    41  	nodeID ids.NodeID,
    42  	log logging.Logger,
    43  	timeoutManager timeout.Manager,
    44  	closeTimeout time.Duration,
    45  	criticalChains set.Set[ids.ID],
    46  	sybilProtectionEnabled bool,
    47  	trackedSubnets set.Set[ids.ID],
    48  	onFatal func(exitCode int),
    49  	healthConfig HealthConfig,
    50  	reg prometheus.Registerer,
    51  ) error {
    52  	return r.router.Initialize(
    53  		nodeID,
    54  		log,
    55  		timeoutManager,
    56  		closeTimeout,
    57  		criticalChains,
    58  		sybilProtectionEnabled,
    59  		trackedSubnets,
    60  		onFatal,
    61  		healthConfig,
    62  		reg,
    63  	)
    64  }
    65  
    66  func (r *tracedRouter) RegisterRequest(
    67  	ctx context.Context,
    68  	nodeID ids.NodeID,
    69  	requestingChainID ids.ID,
    70  	respondingChainID ids.ID,
    71  	requestID uint32,
    72  	op message.Op,
    73  	failedMsg message.InboundMessage,
    74  	engineType p2p.EngineType,
    75  ) {
    76  	r.router.RegisterRequest(
    77  		ctx,
    78  		nodeID,
    79  		requestingChainID,
    80  		respondingChainID,
    81  		requestID,
    82  		op,
    83  		failedMsg,
    84  		engineType,
    85  	)
    86  }
    87  
    88  func (r *tracedRouter) HandleInbound(ctx context.Context, msg message.InboundMessage) {
    89  	m := msg.Message()
    90  	destinationChainID, err := message.GetChainID(m)
    91  	if err != nil {
    92  		r.router.HandleInbound(ctx, msg)
    93  		return
    94  	}
    95  
    96  	sourceChainID, err := message.GetSourceChainID(m)
    97  	if err != nil {
    98  		r.router.HandleInbound(ctx, msg)
    99  		return
   100  	}
   101  
   102  	ctx, span := r.tracer.Start(ctx, "tracedRouter.HandleInbound", oteltrace.WithAttributes(
   103  		attribute.Stringer("nodeID", msg.NodeID()),
   104  		attribute.Stringer("messageOp", msg.Op()),
   105  		attribute.Stringer("chainID", destinationChainID),
   106  		attribute.Stringer("sourceChainID", sourceChainID),
   107  	))
   108  	defer span.End()
   109  
   110  	r.router.HandleInbound(ctx, msg)
   111  }
   112  
   113  func (r *tracedRouter) Shutdown(ctx context.Context) {
   114  	ctx, span := r.tracer.Start(ctx, "tracedRouter.Shutdown")
   115  	defer span.End()
   116  
   117  	r.router.Shutdown(ctx)
   118  }
   119  
   120  func (r *tracedRouter) AddChain(ctx context.Context, chain handler.Handler) {
   121  	chainCtx := chain.Context()
   122  	ctx, span := r.tracer.Start(ctx, "tracedRouter.AddChain", oteltrace.WithAttributes(
   123  		attribute.Stringer("subnetID", chainCtx.SubnetID),
   124  		attribute.Stringer("chainID", chainCtx.ChainID),
   125  	))
   126  	defer span.End()
   127  
   128  	r.router.AddChain(ctx, chain)
   129  }
   130  
   131  func (r *tracedRouter) Connected(nodeID ids.NodeID, nodeVersion *version.Application, subnetID ids.ID) {
   132  	r.router.Connected(nodeID, nodeVersion, subnetID)
   133  }
   134  
   135  func (r *tracedRouter) Disconnected(nodeID ids.NodeID) {
   136  	r.router.Disconnected(nodeID)
   137  }
   138  
   139  func (r *tracedRouter) Benched(chainID ids.ID, nodeID ids.NodeID) {
   140  	r.router.Benched(chainID, nodeID)
   141  }
   142  
   143  func (r *tracedRouter) Unbenched(chainID ids.ID, nodeID ids.NodeID) {
   144  	r.router.Unbenched(chainID, nodeID)
   145  }
   146  
   147  func (r *tracedRouter) HealthCheck(ctx context.Context) (interface{}, error) {
   148  	ctx, span := r.tracer.Start(ctx, "tracedRouter.HealthCheck")
   149  	defer span.End()
   150  
   151  	return r.router.HealthCheck(ctx)
   152  }