storage: propagate tracing info through Raft proposals

pkg/storage/replica_application_cmd.go

@@ -54,10 +54,13 @@ type replicatedCmd struct {
 
 	// ctx is a context that follows from the proposal's context if it was
 	// proposed locally. Otherwise, it will follow from the context passed to
-	// ApplyCommittedEntries. sp is the corresponding tracing span, which is
-	// closed in FinishAndAckOutcome.
+	// ApplyCommittedEntries.
 	ctx context.Context
-	sp  opentracing.Span
+	// sp is the tracing span corresponding to ctx. It is closed in
+	// FinishAndAckOutcome. This span "follows from" the proposer's span (even
+	// when the proposer is remote; we marshall tracing info through the
+	// proposal).
+	sp opentracing.Span
 
 	// The following fields are set in shouldApplyCommand when we validate that
 	// a command applies given the current lease and GC threshold. The process

pkg/storage/replica_application_decoder.go

@@ -14,8 +14,10 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/storage/apply"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/opentracing/opentracing-go"
 	"go.etcd.io/etcd/raft/raftpb"
 )
 
@@ -133,14 +135,27 @@ func (d *replicaDecoder) retrieveLocalProposals(ctx context.Context) (anyLocal b
 // command. If a command was proposed locally, it will be given a tracing span
 // that follows from its proposal's span.
 func (d *replicaDecoder) createTracingSpans(ctx context.Context) {
+	const opName = "raft application"
 	var it replicatedCmdBufSlice
 	for it.init(&d.cmdBuf); it.Valid(); it.Next() {
 		cmd := it.cur()
-		parentCtx := ctx
 		if cmd.IsLocal() {
-			parentCtx = cmd.proposal.ctx
+			cmd.ctx, cmd.sp = tracing.ForkCtxSpan(cmd.proposal.ctx, opName)
+		} else if cmd.raftCmd.TraceData != nil {
+			// The proposal isn't local, and trace data is available. Extract
+			// the span context and start a server-side span.
+			spanCtx, err := d.r.AmbientContext.Tracer.Extract(
+				opentracing.TextMap, opentracing.TextMapCarrier(cmd.raftCmd.TraceData))
+			if err != nil {
+				log.Errorf(ctx, "unable to extract trace data from raft command: %s", err)
+			} else {
+				cmd.sp = d.r.AmbientContext.Tracer.StartSpan(
+					"raft application", opentracing.FollowsFrom(spanCtx))
+				cmd.ctx = opentracing.ContextWithSpan(ctx, cmd.sp)
+			}
+		} else {
+			cmd.ctx, cmd.sp = tracing.ForkCtxSpan(ctx, opName)
 		}
-		cmd.ctx, cmd.sp = tracing.ForkCtxSpan(parentCtx, "raft application")
 	}
 }
 

pkg/storage/replica_proposal.go

@@ -778,8 +778,28 @@ func (r *Replica) requestToProposal(
 			ReplicatedEvalResult: res.Replicated,
 			WriteBatch:           res.WriteBatch,
 			LogicalOpLog:         res.LogicalOpLog,
+			TraceData:            r.getTraceData(ctx),
 		}
 	}
 
 	return proposal, pErr
 }
+
+// getTraceData extracts the SpanContext of the current span.
+func (r *Replica) getTraceData(ctx context.Context) opentracing.TextMapCarrier {
+	sp := opentracing.SpanFromContext(ctx)
+	if sp == nil {
+		return nil
+	}
+	if tracing.IsBlackHoleSpan(sp) {
+		return nil
+	}
+	traceData := opentracing.TextMapCarrier{}
+	if err := r.AmbientContext.Tracer.Inject(
+		sp.Context(), opentracing.TextMap, traceData,
+	); err != nil {
+		log.Errorf(ctx, "failed to inject sp context (%+v) as trace data: %s", sp.Context(), err)
+		return nil
+	}
+	return traceData
+}