github.com/ethersphere/bee/v2@v2.2.0/pkg/pingpong/pingpong.go (about)

     1  // Copyright 2020 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package pingpong exposes the simple ping-pong protocol
     6  // which measures round-trip-time with other peers.
     7  package pingpong
     8  
     9  import (
    10  	"context"
    11  	"errors"
    12  	"fmt"
    13  	"io"
    14  	"time"
    15  
    16  	"github.com/ethersphere/bee/v2/pkg/log"
    17  	"github.com/ethersphere/bee/v2/pkg/p2p"
    18  	"github.com/ethersphere/bee/v2/pkg/p2p/protobuf"
    19  	"github.com/ethersphere/bee/v2/pkg/pingpong/pb"
    20  	"github.com/ethersphere/bee/v2/pkg/swarm"
    21  	"github.com/ethersphere/bee/v2/pkg/tracing"
    22  )
    23  
    24  // loggerName is the tree path name of the logger for this package.
    25  const loggerName = "pinpong"
    26  
    27  const (
    28  	protocolName    = "pingpong"
    29  	protocolVersion = "1.0.0"
    30  	streamName      = "pingpong"
    31  )
    32  
    33  type Interface interface {
    34  	Ping(ctx context.Context, address swarm.Address, msgs ...string) (rtt time.Duration, err error)
    35  }
    36  
    37  type Service struct {
    38  	streamer p2p.Streamer
    39  	logger   log.Logger
    40  	tracer   *tracing.Tracer
    41  	metrics  metrics
    42  }
    43  
    44  func New(streamer p2p.Streamer, logger log.Logger, tracer *tracing.Tracer) *Service {
    45  	return &Service{
    46  		streamer: streamer,
    47  		logger:   logger.WithName(loggerName).Register(),
    48  		tracer:   tracer,
    49  		metrics:  newMetrics(),
    50  	}
    51  }
    52  
    53  func (s *Service) Protocol() p2p.ProtocolSpec {
    54  	return p2p.ProtocolSpec{
    55  		Name:    protocolName,
    56  		Version: protocolVersion,
    57  		StreamSpecs: []p2p.StreamSpec{
    58  			{
    59  				Name:    streamName,
    60  				Handler: s.handler,
    61  			},
    62  		},
    63  	}
    64  }
    65  
    66  func (s *Service) Ping(ctx context.Context, address swarm.Address, msgs ...string) (rtt time.Duration, err error) {
    67  	span, _, ctx := s.tracer.StartSpanFromContext(ctx, "pingpong-p2p-ping", s.logger)
    68  	defer span.Finish()
    69  
    70  	start := time.Now()
    71  	stream, err := s.streamer.NewStream(ctx, address, nil, protocolName, protocolVersion, streamName)
    72  	if err != nil {
    73  		return 0, fmt.Errorf("new stream: %w", err)
    74  	}
    75  	defer func() {
    76  		go stream.FullClose()
    77  	}()
    78  
    79  	w, r := protobuf.NewWriterAndReader(stream)
    80  
    81  	var pong pb.Pong
    82  	for _, msg := range msgs {
    83  		if err := w.WriteMsgWithContext(ctx, &pb.Ping{
    84  			Greeting: msg,
    85  		}); err != nil {
    86  			return 0, fmt.Errorf("write message: %w", err)
    87  		}
    88  		s.metrics.PingSentCount.Inc()
    89  
    90  		if err := r.ReadMsgWithContext(ctx, &pong); err != nil {
    91  			if errors.Is(err, io.EOF) {
    92  				break
    93  			}
    94  			return 0, fmt.Errorf("read message: %w", err)
    95  		}
    96  
    97  		s.metrics.PongReceivedCount.Inc()
    98  	}
    99  	return time.Since(start), nil
   100  }
   101  
   102  func (s *Service) handler(ctx context.Context, p p2p.Peer, stream p2p.Stream) error {
   103  	w, r := protobuf.NewWriterAndReader(stream)
   104  	defer stream.FullClose()
   105  
   106  	span, _, ctx := s.tracer.StartSpanFromContext(ctx, "pingpong-p2p-handler", s.logger)
   107  	defer span.Finish()
   108  
   109  	var ping pb.Ping
   110  	for {
   111  		if err := r.ReadMsgWithContext(ctx, &ping); err != nil {
   112  			if errors.Is(err, io.EOF) {
   113  				break
   114  			}
   115  			return fmt.Errorf("read message: %w", err)
   116  		}
   117  		s.metrics.PingReceivedCount.Inc()
   118  
   119  		if err := w.WriteMsgWithContext(ctx, &pb.Pong{
   120  			Response: "{" + ping.Greeting + "}",
   121  		}); err != nil {
   122  			return fmt.Errorf("write message: %w", err)
   123  		}
   124  		s.metrics.PongSentCount.Inc()
   125  	}
   126  	return nil
   127  }