chore(p2p): extend server tracing

celestiaorg · Sep 18, 2024 · e8b2052 · e8b2052
1 parent 8f53979
commit e8b2052
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/p2p/server.go b/p2p/server.go
@@ -89,6 +89,14 @@ func (serv *ExchangeServer[H]) Stop(context.Context) error {
 
 // requestHandler handles inbound HeaderRequests.
 func (serv *ExchangeServer[H]) requestHandler(stream network.Stream) {
+	ctx, cancel := context.WithTimeout(serv.ctx, serv.Params.RequestTimeout)
+	defer cancel()
+
+	ctx, span := tracerServ.Start(ctx, "request", trace.WithAttributes(
+		attribute.String("peerID", stream.Conn().RemotePeer().String()),
+	))
+	defer span.End()
+
 	err := stream.SetReadDeadline(time.Now().Add(serv.Params.ReadDeadline))
 	if err != nil {
 		log.Debugf("error setting deadline: %s", err)
@@ -99,15 +107,14 @@ func (serv *ExchangeServer[H]) requestHandler(stream network.Stream) {
 	if err != nil {
 		log.Warnw("server: reading header request from stream", "err", err)
 		stream.Reset() //nolint:errcheck
+		span.SetStatus(codes.Error, err.Error())
 		return
 	}
 	if err = stream.CloseRead(); err != nil {
 		log.Error(err)
+		span.SetStatus(codes.Error, err.Error())
 	}
 
-	ctx, cancel := context.WithTimeout(serv.ctx, serv.Params.RequestTimeout)
-	defer cancel()
-
 	var headers []H
 	// retrieve and write Headers
 	switch pbreq.Data.(type) {
@@ -126,7 +133,9 @@ func (serv *ExchangeServer[H]) requestHandler(stream network.Stream) {
 		code = p2p_pb.StatusCode_OK
 	case header.ErrNotFound:
 		code = p2p_pb.StatusCode_NOT_FOUND
+		span.SetStatus(codes.Ok, err.Error())
 	default:
+		span.SetStatus(codes.Error, err.Error())
 		stream.Reset() //nolint:errcheck
 		return
 	}
@@ -150,13 +159,15 @@ func (serv *ExchangeServer[H]) requestHandler(stream network.Stream) {
 			if err != nil {
 				log.Warnw("server: marshaling header to proto", "height", h.Height, "err", err)
 				stream.Reset() //nolint:errcheck
+				span.SetStatus(codes.Error, err.Error())
 				return
 			}
 		}
 		_, err = serde.Write(stream, &p2p_pb.HeaderResponse{Body: bin, StatusCode: code})
 		if err != nil {
 			log.Warnw("server: writing header to stream", "err", err)
 			stream.Reset() //nolint:errcheck
+			span.SetStatus(codes.Error, err.Error())
 			return
 		}
 	}