Merge #147607

147607: kv: increase adoption of Stopper Handle API r=tbg a=tbg

Continues #142059.

I opened a recent TPCC execution trace in https://github.com/felixge/sqlprof and worked through the output of the below query, which roughly shows which stopper tasks were active in the execution. I then added more based on looking at more execution traces from the 300 node scale test.

```
select * from (select count(*) as hits,  list_slice(s.funcs, -3, -2) as root
  from stack_samples ss
  join stacks s on ss.src_stack_id = s.stack_id
  where root[2] LIKE '%Stopper%' group by root) order by hits desc;
```

I recommend reviewing with whitespace changes suppressed, and to watch out for ctx assignments, as the new API is more brittle with regards to them.

Epic: CRDB-42584

Co-authored-by: Tobias Grieger <[email protected]>

Author: craig[bot]

pkg/internal/client/requestbatcher/batcher.go

@@ -242,9 +242,16 @@ func New(cfg Config) *RequestBatcher {
 	}
 	b.sendBatchOpName = redact.Sprintf("%s.sendBatch", b.cfg.Name)
 	bgCtx := cfg.AmbientCtx.AnnotateCtx(context.Background())
-	if err := cfg.Stopper.RunAsyncTask(bgCtx, b.cfg.Name.StripMarkers(), b.run); err != nil {
+	bgCtx, hdl, err := cfg.Stopper.GetHandle(bgCtx, stop.TaskOpts{
+		TaskName: b.cfg.Name.StripMarkers(),
+	})
+	if err != nil {
 		panic(err)
 	}
+	go func(ctx context.Context) {
+		defer hdl.Activate(ctx).Release(ctx)
+		b.run(ctx)
+	}(bgCtx)
 	return b
 }
 
@@ -342,7 +349,7 @@ func (b *RequestBatcher) sendDone(ba *batch) {
 }
 
 func (b *RequestBatcher) sendBatch(ctx context.Context, ba *batch) {
-	if err := b.cfg.Stopper.RunAsyncTask(ctx, "send-batch", func(ctx context.Context) {
+	work := func(ctx context.Context) {
 		defer b.sendDone(ba)
 		var batchRequest *kvpb.BatchRequest
 		var br *kvpb.BatchResponse
@@ -426,9 +433,19 @@ func (b *RequestBatcher) sendBatch(ctx context.Context, ba *batch) {
 			}
 			ba.reqs, prevResps = nextReqs, nextPrevResps
 		}
-	}); err != nil {
+	}
+
+	ctx, hdl, err := b.cfg.Stopper.GetHandle(ctx, stop.TaskOpts{
+		TaskName: "send-batch",
+	})
+	if err != nil {
 		b.sendDone(ba)
+		return
 	}
+	go func(ctx context.Context) {
+		defer hdl.Activate(ctx).Release(ctx)
+		work(ctx)
+	}(ctx)
 }
 
 func (b *RequestBatcher) sendResponse(req *request, resp Response) {

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -1503,11 +1503,7 @@ func (ds *DistSender) divideAndSendParallelCommit(
 	qiBatchIdx := batchIdx + 1
 	qiResponseCh := make(chan response, 1)
 
-	runTask := ds.stopper.RunAsyncTask
-	if ds.disableParallelBatches {
-		runTask = ds.stopper.RunTask
-	}
-	if err := runTask(ctx, "kv.DistSender: sending pre-commit query intents", func(ctx context.Context) {
+	sendPreCommit := func(ctx context.Context) {
 		// Map response index to the original un-swapped batch index.
 		// Remember that we moved the last QueryIntent in this batch
 		// from swapIdx to the end.
@@ -1530,8 +1526,24 @@ func (ds *DistSender) divideAndSendParallelCommit(
 		// concurrently with the EndTxn batch below.
 		reply, pErr := ds.divideAndSendBatchToRanges(ctx, qiBa, qiRS, qiIsReverse, true /* withCommit */, qiBatchIdx)
 		qiResponseCh <- response{reply: reply, positions: positions, pErr: pErr}
-	}); err != nil {
-		return nil, kvpb.NewError(err)
+	}
+
+	const taskName = "kv.DistSender: sending pre-commit query intents"
+	if ds.disableParallelBatches {
+		if err := ds.stopper.RunTask(ctx, taskName, sendPreCommit); err != nil {
+			return nil, kvpb.NewError(err)
+		}
+	} else {
+		ctx, hdl, err := ds.stopper.GetHandle(ctx, stop.TaskOpts{
+			TaskName: taskName,
+		})
+		if err != nil {
+			return nil, kvpb.NewError(err)
+		}
+		go func(ctx context.Context) {
+			defer hdl.Activate(ctx).Release(ctx)
+			sendPreCommit(ctx)
+		}(ctx)
 	}
 
 	// Adjust the original batch request to ignore the pre-commit
@@ -2069,26 +2081,25 @@ func (ds *DistSender) sendPartialBatchAsync(
 	responseCh chan response,
 	positions []int,
 ) bool {
-	if err := ds.stopper.RunAsyncTaskEx(
-		ctx,
-		stop.TaskOpts{
-			TaskName:   "kv.DistSender: sending partial batch",
-			SpanOpt:    stop.ChildSpan,
-			Sem:        ds.asyncSenderSem,
-			WaitForSem: false,
-		},
-		func(ctx context.Context) {
-			ds.metrics.AsyncSentCount.Inc(1)
-			ds.metrics.AsyncInProgress.Inc(1)
-			defer ds.metrics.AsyncInProgress.Dec(1)
-			resp := ds.sendPartialBatch(ctx, ba, rs, isReverse, withCommit, batchIdx, routing)
-			resp.positions = positions
-			responseCh <- resp
-		},
-	); err != nil {
+	ctx, hdl, err := ds.stopper.GetHandle(ctx, stop.TaskOpts{
+		TaskName:   "kv.DistSender: sending partial batch",
+		SpanOpt:    stop.ChildSpan,
+		Sem:        ds.asyncSenderSem,
+		WaitForSem: false,
+	})
+	if err != nil {
 		ds.metrics.AsyncThrottledCount.Inc(1)
 		return false
 	}
+	go func(ctx context.Context) {
+		defer hdl.Activate(ctx).Release(ctx)
+		ds.metrics.AsyncSentCount.Inc(1)
+		ds.metrics.AsyncInProgress.Inc(1)
+		defer ds.metrics.AsyncInProgress.Dec(1)
+		resp := ds.sendPartialBatch(ctx, ba, rs, isReverse, withCommit, batchIdx, routing)
+		resp.positions = positions
+		responseCh <- resp
+	}(ctx)
 	return true
 }
 

pkg/kv/kvclient/kvcoord/txn_interceptor_committer.go

@@ -499,17 +499,26 @@ func (tc *txnCommitter) makeTxnCommitExplicitAsync(
 	if multitenant.HasTenantCostControlExemption(ctx) {
 		asyncCtx = multitenant.WithTenantCostControlExemption(asyncCtx)
 	}
-	if err := tc.stopper.RunAsyncTask(
-		asyncCtx, "txnCommitter: making txn commit explicit", func(ctx context.Context) {
-			tc.mu.Lock()
-			defer tc.mu.Unlock()
-			if err := makeTxnCommitExplicitLocked(ctx, tc.wrapped, txn, lockSpans); err != nil {
-				log.Errorf(ctx, "making txn commit explicit failed for %s: %v", txn, err)
-			}
-		},
-	); err != nil {
+
+	work := func(ctx context.Context) {
+		tc.mu.Lock()
+		defer tc.mu.Unlock()
+		if err := makeTxnCommitExplicitLocked(ctx, tc.wrapped, txn, lockSpans); err != nil {
+			log.Errorf(ctx, "making txn commit explicit failed for %s: %v", txn, err)
+		}
+	}
+
+	asyncCtx, hdl, err := tc.stopper.GetHandle(asyncCtx, stop.TaskOpts{
+		TaskName: "txnCommitter: making txn commit explicit",
+	})
+	if err != nil {
 		log.VErrEventf(ctx, 1, "failed to make txn commit explicit: %v", err)
+		return
 	}
+	go func(ctx context.Context) {
+		defer hdl.Activate(ctx).Release(ctx)
+		work(ctx)
+	}(asyncCtx)
 }
 
 func makeTxnCommitExplicitLocked(

pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go

@@ -566,67 +566,76 @@ func (h *txnHeartbeater) abortTxnAsyncLocked(ctx context.Context) {
 
 	const taskName = "txnHeartbeater: aborting txn"
 	log.VEventf(ctx, 2, "async abort for txn: %s", txn)
-	if err := h.stopper.RunAsyncTask(h.AnnotateCtx(context.Background()), taskName,
-		func(ctx context.Context) {
-			if err := timeutil.RunWithTimeout(ctx, taskName, abortTxnAsyncTimeout,
-				func(ctx context.Context) error {
-					h.mu.Lock()
-					defer h.mu.Unlock()
-
-					// If we find an abortTxnAsyncResultC, that means an async
-					// rollback request is already in flight, so there's no
-					// point in us running another. This can happen because the
-					// TxnCoordSender also calls abortTxnAsyncLocked()
-					// independently of the heartbeat loop.
-					if h.mu.abortTxnAsyncResultC != nil {
-						log.VEventf(ctx, 2,
-							"skipping async abort due to concurrent async abort for %s", txn)
-						return nil
-					}
-
-					// TxnCoordSender allows EndTxn(commit=false) through even
-					// after we set finalObservedStatus, and that request can
-					// race with us for the mutex. Thus, if we find an in-flight
-					// request here, after checking ifReqs=0 before being spawned,
-					// we deduce that it must have been a rollback and there's no
-					// point in sending another rollback.
-					if h.mu.ifReqs > 0 {
-						log.VEventf(ctx, 2,
-							"skipping async abort due to client rollback for %s", txn)
-						return nil
-					}
-
-					// Set up a result channel to signal to an incoming client
-					// rollback that an async rollback is already in progress,
-					// and pass it the result. The buffer allows storing the
-					// result even when no client rollback arrives. Recall that
-					// the SendLocked() call below releases the mutex while
-					// running, allowing concurrent incoming requests.
-					h.mu.abortTxnAsyncResultC = make(chan abortTxnAsyncResult, 1)
-
-					// Send the abort request through the interceptor stack. This is
-					// important because we need the txnPipeliner to append lock spans
-					// to the EndTxn request.
-					br, pErr := h.wrapped.SendLocked(ctx, ba)
-					if pErr != nil {
-						log.VErrEventf(ctx, 1, "async abort failed for %s: %s ", txn, pErr)
-						h.metrics.AsyncRollbacksFailed.Inc(1)
-					}
-
-					// Pass the result to a waiting client rollback, if any, and
-					// remove the channel since we're no longer in flight.
-					h.mu.abortTxnAsyncResultC <- abortTxnAsyncResult{br: br, pErr: pErr}
-					h.mu.abortTxnAsyncResultC = nil
+
+	work := func(ctx context.Context) {
+		if err := timeutil.RunWithTimeout(ctx, taskName, abortTxnAsyncTimeout,
+			func(ctx context.Context) error {
+				h.mu.Lock()
+				defer h.mu.Unlock()
+
+				// If we find an abortTxnAsyncResultC, that means an async
+				// rollback request is already in flight, so there's no
+				// point in us running another. This can happen because the
+				// TxnCoordSender also calls abortTxnAsyncLocked()
+				// independently of the heartbeat loop.
+				if h.mu.abortTxnAsyncResultC != nil {
+					log.VEventf(ctx, 2,
+						"skipping async abort due to concurrent async abort for %s", txn)
 					return nil
-				},
-			); err != nil {
-				log.VEventf(ctx, 1, "async abort failed for %s: %s", txn, err)
-			}
-		},
-	); err != nil {
+				}
+
+				// TxnCoordSender allows EndTxn(commit=false) through even
+				// after we set finalObservedStatus, and that request can
+				// race with us for the mutex. Thus, if we find an in-flight
+				// request here, after checking ifReqs=0 before being spawned,
+				// we deduce that it must have been a rollback and there's no
+				// point in sending another rollback.
+				if h.mu.ifReqs > 0 {
+					log.VEventf(ctx, 2,
+						"skipping async abort due to client rollback for %s", txn)
+					return nil
+				}
+
+				// Set up a result channel to signal to an incoming client
+				// rollback that an async rollback is already in progress,
+				// and pass it the result. The buffer allows storing the
+				// result even when no client rollback arrives. Recall that
+				// the SendLocked() call below releases the mutex while
+				// running, allowing concurrent incoming requests.
+				h.mu.abortTxnAsyncResultC = make(chan abortTxnAsyncResult, 1)
+
+				// Send the abort request through the interceptor stack. This is
+				// important because we need the txnPipeliner to append lock spans
+				// to the EndTxn request.
+				br, pErr := h.wrapped.SendLocked(ctx, ba)
+				if pErr != nil {
+					log.VErrEventf(ctx, 1, "async abort failed for %s: %s ", txn, pErr)
+					h.metrics.AsyncRollbacksFailed.Inc(1)
+				}
+
+				// Pass the result to a waiting client rollback, if any, and
+				// remove the channel since we're no longer in flight.
+				h.mu.abortTxnAsyncResultC <- abortTxnAsyncResult{br: br, pErr: pErr}
+				h.mu.abortTxnAsyncResultC = nil
+				return nil
+			},
+		); err != nil {
+			log.VEventf(ctx, 1, "async abort failed for %s: %s", txn, err)
+		}
+	}
+
+	asyncCtx, hdl, err := h.stopper.GetHandle(h.AnnotateCtx(context.Background()), stop.TaskOpts{
+		TaskName: taskName,
+	})
+	if err != nil {
 		log.Warningf(ctx, "%v", err)
 		h.metrics.AsyncRollbacksFailed.Inc(1)
+		return
 	}
+	go func(ctx context.Context) {
+		defer hdl.Activate(ctx).Release(ctx)
+		work(ctx)
+	}(asyncCtx)
 }
 
 // randLockingIndex returns the index of the first request that acquires locks