consolidate span creation and cleanup into startSpan, which now returnsa function to end the span and set its status

This commit is contained in:
jonny rhea 2026-01-04 17:36:57 -06:00
parent afb4f91258
commit 4cd89d4bab
2 changed files with 32 additions and 33 deletions

View file

@ -524,30 +524,26 @@ func (h *handler) handleCall(cp *callProc, msg *jsonrpcMessage) *jsonrpcMessage
}
// Start root span for the request.
ctx, rootSpan := h.startSpan(cp.ctx, msg, "rpc.handleCall", cp.isBatch)
defer rootSpan.End()
var err error
ctx, done := h.startSpan(cp.ctx, msg, "rpc.handleCall", cp.isBatch)
defer done(&err)
// Start tracing span before parsing arguments.
_, pspan := h.startSpan(ctx, msg, "rpc.parsePositionalArguments", cp.isBatch)
_, pDone := h.startSpan(ctx, msg, "rpc.parsePositionalArguments", cp.isBatch)
args, err := parsePositionalArguments(msg.Params, callb.argTypes)
pDone(&err)
if err != nil {
pspan.RecordError(err)
pspan.SetStatus(codes.Error, err.Error())
rootSpan.SetStatus(codes.Error, err.Error())
pspan.End()
return msg.errorResponse(&invalidParamsError{err.Error()})
}
pspan.End()
start := time.Now()
// Start tracing span before running the method.
rctx, rspan := h.startSpan(ctx, msg, "rpc.runMethod", cp.isBatch)
rctx, rDone := h.startSpan(ctx, msg, "rpc.runMethod", cp.isBatch)
answer := h.runMethod(rctx, msg, callb, args, cp.isBatch)
if answer.Error != nil {
err := errors.New(answer.Error.Message)
rootSpan.SetStatus(codes.Error, err.Error())
err = errors.New(answer.Error.Message)
}
rspan.End()
rDone(&err)
// Collect the statistics for RPC calls if metrics is enabled.
rpcRequestGauge.Inc(1)
@ -599,18 +595,20 @@ func (h *handler) handleSubscribe(cp *callProc, msg *jsonrpcMessage) *jsonrpcMes
return h.runMethod(ctx, msg, callb, args, cp.isBatch)
}
// startSpan starts a tracing span for an RPC call.
func (h *handler) startSpan(ctx context.Context, msg *jsonrpcMessage, spanName string, isBatch bool) (context.Context, trace.Span) {
// startSpan starts a tracing span for an RPC call and returns a function to
// end the span. The function will record errors and set span status based on
// the error value.
func (h *handler) startSpan(ctx context.Context, msg *jsonrpcMessage, spanName string, isBatch bool) (context.Context, func(*error)) {
parentSpan := trace.SpanFromContext(ctx)
ctx, span := h.tracer().Start(ctx, spanName)
// Fast path: noop provider or span not sampled
if !span.IsRecording() {
return ctx, span
return ctx, func(*error) { span.End() }
}
// Set span attributes.
span.SetAttributes(
semconv.RPCSystemKey.String("jsonrpc"),
semconv.RPCMethodKey.String(msg.Method),
attribute.Bool("rpc.batch", isBatch),
)
@ -618,7 +616,19 @@ func (h *handler) startSpan(ctx context.Context, msg *jsonrpcMessage, spanName s
if id != "" && id != "null" {
span.SetAttributes(attribute.String("rpc.id", id))
}
return ctx, span
// Define the function to end the span and handle error recording
done := func(err *error) {
if *err != nil {
span.RecordError(*err)
span.SetStatus(codes.Error, (*err).Error())
parentSpan.SetStatus(codes.Error, (*err).Error())
} else {
span.SetStatus(codes.Ok, "")
}
span.End()
}
return ctx, done
}
// tracer returns the OpenTelemetry Tracer for RPC call tracing.
@ -634,26 +644,22 @@ func (h *handler) tracer() trace.Tracer {
// runMethod runs the Go callback for an RPC method.
func (h *handler) runMethod(ctx context.Context, msg *jsonrpcMessage, callb *callback, args []reflect.Value, isBatch bool) *jsonrpcMessage {
parentSpan := trace.SpanFromContext(ctx)
result, err := callb.call(ctx, msg.Method, args)
if err != nil {
parentSpan.SetStatus(codes.Error, err.Error())
return msg.errorResponse(err)
}
// If parent span is recording, start a span for the response.
// Note: This prevents msg.response spans from being created when
// the parent span is not recording (e.g. subscription tracing disabled).
parentSpan := trace.SpanFromContext(ctx)
if parentSpan.IsRecording() {
_, span := h.startSpan(ctx, msg, "rpc.msg.response", isBatch)
defer span.End()
_, done := h.startSpan(ctx, msg, "rpc.msg.response", isBatch)
response := msg.response(result)
if response.Error != nil {
err := errors.New(response.Error.Message)
span.RecordError(errors.New(response.Error.Message))
span.SetStatus(codes.Error, err.Error())
parentSpan.SetStatus(codes.Error, err.Error())
err = errors.New(response.Error.Message)
}
done(&err)
return response
}
return msg.response(result)

View file

@ -47,15 +47,12 @@ func attributeMap(attrs []attribute.KeyValue) map[string]string {
func newTracingServer(t *testing.T) (*Server, *sdktrace.TracerProvider, *tracetest.InMemoryExporter) {
t.Helper()
exporter := tracetest.NewInMemoryExporter()
tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))
t.Cleanup(func() { _ = tp.Shutdown(context.Background()) })
server := newTestServer()
server.SetTracerProvider(tp)
t.Cleanup(server.Stop)
return server, tp, exporter
}
@ -97,9 +94,6 @@ func TestTracingHTTP(t *testing.T) {
}
attrs := attributeMap(rpcSpan.Attributes)
if attrs["rpc.system"] != "jsonrpc" {
t.Errorf("expected rpc.system=jsonrpc, got %v", attrs["rpc.system"])
}
if attrs["rpc.method"] != "test_echo" {
t.Errorf("expected rpc.method=test_echo, got %v", attrs["rpc.method"])
}
@ -149,8 +143,7 @@ func TestTracingBatchHTTP(t *testing.T) {
for i := range spans {
if spans[i].Name == "rpc.handleCall" {
attrs := attributeMap(spans[i].Attributes)
if attrs["rpc.system"] == "jsonrpc" &&
attrs["rpc.method"] == "test_echo" &&
if attrs["rpc.method"] == "test_echo" &&
attrs["rpc.batch"] == "true" {
found++
}