Merge cockroachdb#26637

26637: distsql: move tableReader stat collection setup out of Start r=asubiotto a=asubiotto

It makes more sense to have the tableReader instrument itself for stat
collection at the same time other setup is performed in newTableReader.

Release note: None

Co-authored-by: Alfonso Subiotto Marqués <[email protected]>

pkg/sql/distsqlrun/tablereader.go

@@ -109,6 +109,12 @@ func newTableReader(
 		tr.spans[i] = s.Span
 	}
 	tr.input = &rowFetcherWrapper{RowFetcher: &tr.fetcher}
+
+	if sp := opentracing.SpanFromContext(flowCtx.EvalCtx.Ctx()); sp != nil && tracing.IsRecording(sp) {
+		tr.input = NewInputStatCollector(tr.input)
+		tr.finishTrace = tr.outputStatsToTrace
+	}
+
 	return tr, nil
 }
 
@@ -207,11 +213,6 @@ func (tr *tableReader) Start(ctx context.Context) context.Context {
 		log.Fatalf(ctx, "tableReader outside of txn")
 	}
 
-	if sp := opentracing.SpanFromContext(ctx); sp != nil && tracing.IsRecording(sp) {
-		tr.input = NewInputStatCollector(tr.input)
-		tr.finishTrace = tr.outputStatsToTrace
-	}
-
 	// Like every processor, the tableReader will have a context with a log tag
 	// and a span. The underlying fetcher inherits the proc's span, but not the
 	// log tag.

pkg/sql/distsqlrun/tablereader_test.go

@@ -329,16 +329,17 @@ func TestLimitScans(t *testing.T) {
 	const limit = 3
 	post := PostProcessSpec{Limit: limit}
 
-	tr, err := newTableReader(&flowCtx, 0 /* processorID */, &spec, &post, nil /* output */)
-	if err != nil {
-		t.Fatal(err)
-	}
-
 	// Now we're going to run the tableReader and trace it.
 	tracer := tracing.NewTracer()
 	sp := tracer.StartSpan("root", tracing.Recordable)
 	tracing.StartRecording(sp, tracing.SnowballRecording)
 	ctx := opentracing.ContextWithSpan(context.Background(), sp)
+	flowCtx.EvalCtx.CtxProvider = tree.FixedCtxProvider{Context: ctx}
+
+	tr, err := newTableReader(&flowCtx, 0 /* processorID */, &spec, &post, nil /* output */)
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	tr.Start(ctx)
 	rows := 0