stop, kvclient: include DistSender partial batches in traces

Before this patch, since fairly recently, the spans created for
Stopper.RunAsyncTask[Ex]() were not included in the recording of the
caller's span because these spans were created with the ForkSpan(). This
patch adds an option to RunAsyncEx to make the task's span a child of
the caller's, and thus to be included in the caller's recording. This
option is used by the DistSender when sending partial batch requests,
therefore re-including these requests in higher-level traces (as they
used to be).

The reason why async tasks were detached from the caller's span
in cockroachdb#59815 was because of concerns about parent and children spans with
very different lifetimes becoming tied through a trace, delaying the
garbage collection of the whole trace until the last span in it
finishes. This patch gives the caller of RunAsyncTaskEx control over
this behavior; in the particular case of the DistSender, the traces are
not too long lived and the DistSender waits for the async tasks it
spawns to finish. In such circumstances, tying the allocations of the
child spans to the whole trace should be fine.

Release note (general change): A recent release removed parts of some
queries from the debugging traces of those queries. This information
(i.e. the execution of some low-level RPCs) has been re-included in the
traces.

Author: andreimatei

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -1347,6 +1347,7 @@ func (ds *DistSender) sendPartialBatchAsync(
 		ctx,
 		stop.TaskOpts{
 			TaskName:   "kv.DistSender: sending partial batch",
+			SameTrace:  true,
 			Sem:        ds.asyncSenderSem,
 			WaitForSem: false,
 		},

pkg/util/stop/BUILD.bazel

@@ -33,6 +33,7 @@ go_test(
         "//pkg/util/log",
         "//pkg/util/quotapool",
         "//pkg/util/syncutil",
+        "//pkg/util/tracing",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],

pkg/util/stop/stopper.go

@@ -336,6 +336,7 @@ func (s *Stopper) RunAsyncTask(
 	return s.RunAsyncTaskEx(ctx,
 		TaskOpts{
 			TaskName:   taskName,
+			SameTrace:  false,
 			Sem:        nil,
 			WaitForSem: false,
 		},
@@ -347,6 +348,23 @@ type TaskOpts struct {
 	// TaskName is a human-readable name for the operation. Used as the name of
 	// the tracing span.
 	TaskName string
+
+	// SameTrace, if set, makes the tracing span created for the task be a child
+	// of the caller's span (if the caller has a span in the passed-in ctx). If
+	// not set, the task's span is not a child span (technically, the task's span
+	// will have a FollowsFrom relationship with the parent instead of ChildOf
+	// relationship). When set, the recording of the parent span (if the parent is
+	// recording) will include the child.
+	//
+	// Regardless of this setting, if the caller doesn't have a span, the task
+	// doesn't get a span either.
+	//
+	// Setting SameTrace can have consequences on memory usage. The lifetime of
+	// the task's span becomes tied to the lifetime of the parent. Generally
+	// SameTrace should be used when the parent waits for the task to complete,
+	// and the parent is not a long-running process.
+	SameTrace bool
+
 	// If set, Sem is used as a semaphore limiting the concurrency (each task has
 	// weight 1).
 	//
@@ -403,7 +421,13 @@ func (s *Stopper) RunAsyncTaskEx(ctx context.Context, opt TaskOpts, f func(conte
 		return ErrUnavailable
 	}
 
-	ctx, span := tracing.ForkSpan(ctx, taskName)
+	// If the caller has a span, the task gets a child span.
+	var span *tracing.Span
+	if opt.SameTrace {
+		ctx, span = tracing.ChildSpan(ctx, taskName)
+	} else {
+		ctx, span = tracing.ForkSpan(ctx, taskName)
+	}
 
 	// Call f on another goroutine.
 	taskStarted = true // Another goroutine now takes ownership of the alloc, if any.

pkg/util/stop/stopper_test.go

@@ -22,10 +22,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	_ "github.com/cockroachdb/cockroach/pkg/util/log" // for flags
 	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 )
@@ -681,3 +683,39 @@ func TestCancelInCloser(t *testing.T) {
 type closerFunc func()
 
 func (cf closerFunc) Close() { cf() }
+
+// Test that task spans are included or not in the parent's recording based on
+// the SameTrace option.
+func TestStopperRunAsyncTaskTracing(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	s := stop.NewStopper()
+
+	ctx, getRecording, finish := tracing.ContextWithRecordingSpan(
+		context.Background(), tracing.NewTracer(), "parent")
+
+	// Start two child tasks. Only the one with SameTrace:true is expected to be
+	// present in the parent's recording.
+	s.RunAsyncTaskEx(ctx, stop.TaskOpts{
+		TaskName:  "async child",
+		SameTrace: false,
+	},
+		func(ctx context.Context) {
+			log.Event(ctx, "async 1")
+		},
+	)
+	s.RunAsyncTaskEx(ctx, stop.TaskOpts{
+		TaskName:  "async child same trace",
+		SameTrace: true,
+	},
+		func(ctx context.Context) {
+			log.Event(ctx, "async 2")
+		},
+	)
+
+	s.Stop(ctx)
+	finish()
+	require.NoError(t, tracing.TestingCheckRecordedSpans(getRecording(), `
+		span: parent
+			span: [async] async child same trace
+				event: async 2`))
+}