implement review feedback from lightclient

internal/telemetry/telemetry.go

@@ -27,6 +27,14 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
+// RPCInfo contains information about the RPC request.
+type RPCInfo struct {
+	System    string
+	Service   string
+	Method    string
+	RequestID string
+}
+
 // Attribute is an alias for attribute.KeyValue.
 type Attribute = attribute.KeyValue
 
@@ -43,16 +51,8 @@ func Int64Attribute(key string, val int64) Attribute {
 // StartServerSpan creates a SpanKind=SERVER span at the JSON-RPC boundary.
 // The span name is formatted as $rpcSystem.$rpcService/$rpcMethod
 // (e.g. "jsonrpc.engine/newPayloadV4").
-func StartServerSpan(
+func StartServerSpan(ctx context.Context, tracer trace.Tracer, rpc RPCInfo, others ...Attribute) (context.Context, func(error)) {
-	ctx context.Context,
+	spanName := fmt.Sprintf("%s.%s/%s", rpc.System, rpc.Service, rpc.Method)
-	tracer trace.Tracer,
-	rpcSystem string,
-	rpcService string,
-	rpcMethod string,
-	requestID string,
-	additionalAttributes ...Attribute,
-) (context.Context, func(error)) {
-	spanName := fmt.Sprintf("%s.%s/%s", rpcSystem, rpcService, rpcMethod)
 	ctx, span := tracer.Start(
 		ctx,
 		spanName,
@@ -66,46 +66,46 @@ func StartServerSpan(
 
 	// Define required attributes
 	attrs := []Attribute{
-		semconv.RPCSystemKey.String(rpcSystem),
+		semconv.RPCSystemKey.String(rpc.System),
-		semconv.RPCServiceKey.String(rpcService),
+		semconv.RPCServiceKey.String(rpc.Service),
-		semconv.RPCMethodKey.String(rpcMethod),
+		semconv.RPCMethodKey.String(rpc.Method),
-		semconv.RPCJSONRPCRequestID(requestID),
+		semconv.RPCJSONRPCRequestID(rpc.RequestID),
 	}
 
 	// Add any additional attributes provided
-	if len(additionalAttributes) > 0 {
+	if len(others) > 0 {
-		attrs = append(attrs, additionalAttributes...)
+		attrs = append(attrs, others...)
 	}
 	span.SetAttributes(attrs...)
 	return ctx, endSpan(span)
 }
 
-// StartInternalSpan creates a SpanKind=INTERNAL span.
+// StartSpan creates a SpanKind=INTERNAL span.
-func StartInternalSpan(
+func StartSpan(
 	ctx context.Context,
 	spanName string,
 	attributes ...Attribute,
-) (context.Context, func(error)) {
+) (context.Context, trace.Span, func(error)) {
-	return StartInternalSpanWithTracer(ctx, otel.Tracer(""), spanName, attributes...)
+	return StartSpanWithTracer(ctx, otel.Tracer(""), spanName, attributes...)
 }
 
-// StartInternalSpanWithTracer requires a tracer to be passed in and creates a SpanKind=INTERNAL span.
+// StartSpanWithTracer requires a tracer to be passed in and creates a SpanKind=INTERNAL span.
-func StartInternalSpanWithTracer(
+func StartSpanWithTracer(
 	ctx context.Context,
 	tracer trace.Tracer,
 	spanName string,
 	attributes ...Attribute,
-) (context.Context, func(error)) {
+) (context.Context, trace.Span, func(error)) {
-	return startInternalSpan(ctx, tracer, spanName, attributes...)
+	return startSpan(ctx, tracer, spanName, attributes...)
 }
 
-// startInternalSpan creates a SpanKind=INTERNAL span.
+// startSpan creates a SpanKind=INTERNAL span.
-func startInternalSpan(
+func startSpan(
 	ctx context.Context,
 	tracer trace.Tracer,
 	spanName string,
 	attributes ...Attribute,
-) (context.Context, func(error)) {
+) (context.Context, trace.Span, func(error)) {
 	ctx, span := tracer.Start(
 		ctx,
 		spanName,
@@ -114,13 +114,13 @@ func startInternalSpan(
 
 	// Fast path
 	if !span.IsRecording() {
-		return ctx, func(error) { span.End() }
+		return ctx, span, func(error) { span.End() }
 	}
 
 	if len(attributes) > 0 {
 		span.SetAttributes(attributes...)
 	}
-	return ctx, endSpan(span)
+	return ctx, span, endSpan(span)
 }
 
 // endSpan ends the span and handles error recording.

rpc/handler.go

@@ -514,7 +514,7 @@ func (h *handler) handleCall(cp *callProc, msg *jsonrpcMessage) *jsonrpcMessage
 	if len(msg.Method) > maxMethodNameLength {
 		return msg.errorResponse(&invalidRequestError{fmt.Sprintf("method name too long: %d > %d", len(msg.Method), maxMethodNameLength)})
 	}
-	callb, rpcService, rpcMethod := h.reg.callback(msg.Method)
+	callb, service, method := h.reg.callback(msg.Method)
 
 	// If the method is not found, return an error.
 	if callb == nil {
@@ -523,11 +523,17 @@ func (h *handler) handleCall(cp *callProc, msg *jsonrpcMessage) *jsonrpcMessage
 
 	// Start root span for the request.
 	var err error
-	ctx, spanEnd := telemetry.StartServerSpan(cp.ctx, h.tracer(), "jsonrpc", rpcService, rpcMethod, string(msg.ID))
+	rpcInfo := telemetry.RPCInfo{
+		System:    "jsonrpc",
+		Service:   service,
+		Method:    method,
+		RequestID: string(msg.ID),
+	}
+	ctx, spanEnd := telemetry.StartServerSpan(cp.ctx, h.tracer(), rpcInfo)
 	defer spanEnd(err)
 
 	// Start tracing span before parsing arguments.
-	_, pSpanEnd := telemetry.StartInternalSpanWithTracer(ctx, h.tracer(), "rpc.parsePositionalArguments")
+	_, _, pSpanEnd := telemetry.StartSpanWithTracer(ctx, h.tracer(), "rpc.parsePositionalArguments")
 	args, err := parsePositionalArguments(msg.Params, callb.argTypes)
 	pSpanEnd(err)
 	if err != nil {
@@ -536,7 +542,7 @@ func (h *handler) handleCall(cp *callProc, msg *jsonrpcMessage) *jsonrpcMessage
 	start := time.Now()
 
 	// Start tracing span before running the method.
-	rctx, rSpanEnd := telemetry.StartInternalSpanWithTracer(ctx, h.tracer(), "rpc.runMethod")
+	rctx, _, rSpanEnd := telemetry.StartSpanWithTracer(ctx, h.tracer(), "rpc.runMethod")
 	answer := h.runMethod(rctx, msg, callb, args)
 	if answer.Error != nil {
 		err = errors.New(answer.Error.Message)
@@ -609,7 +615,7 @@ func (h *handler) runMethod(ctx context.Context, msg *jsonrpcMessage, callb *cal
 	if err != nil {
 		return msg.errorResponse(err)
 	}
-	_, spanEnd := telemetry.StartInternalSpanWithTracer(ctx, h.tracer(), "rpc.encodeJSONResponse", attributes...)
+	_, _, spanEnd := telemetry.StartSpanWithTracer(ctx, h.tracer(), "rpc.encodeJSONResponse", attributes...)
 	response := msg.response(result)
 	if response.Error != nil {
 		err = errors.New(response.Error.Message)

rpc/server.go

@@ -101,15 +101,6 @@ func (s *Server) SetWebsocketReadLimit(limit int64) {
 	s.wsReadLimit = limit
 }
 
-// SetTracerProvider configures the OpenTelemetry TracerProvider for RPC call tracing.
-// Note: This method (and the TracerProvider field in the Server/Handler struct) is
-// primarily intended for testing. In particular, it allows tests to configure an
-// isolated TracerProvider without changing the global provider, avoiding
-// interference between tests running in parallel.
-func (s *Server) SetTracerProvider(tp trace.TracerProvider) {
-	s.tracerProvider = tp
-}
-
 // RegisterName creates a service for the given receiver type under the given name. When no
 // methods on the given receiver match the criteria to be either an RPC method or a
 // subscription an error is returned. Otherwise a new service is created and added to the
@@ -141,6 +132,15 @@ func (s *Server) ServeCodec(codec ServerCodec, options CodecOption) {
 	c.Close()
 }
 
+// setTracerProvider configures the OpenTelemetry TracerProvider for RPC call tracing.
+// Note: This method (and the TracerProvider field in the Server/Handler struct) is
+// primarily intended for testing. In particular, it allows tests to configure an
+// isolated TracerProvider without changing the global provider, avoiding
+// interference between tests running in parallel.
+func (s *Server) setTracerProvider(tp trace.TracerProvider) {
+	s.tracerProvider = tp
+}
+
 func (s *Server) trackCodec(codec ServerCodec) bool {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()

rpc/tracing_test.go

@@ -51,7 +51,7 @@ func newTracingServer(t *testing.T) (*Server, *sdktrace.TracerProvider, *tracete
 	tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))
 	t.Cleanup(func() { _ = tp.Shutdown(context.Background()) })
 	server := newTestServer()
-	server.SetTracerProvider(tp)
+	server.setTracerProvider(tp)
 	t.Cleanup(server.Stop)
 	return server, tp, exporter
 }