Merge #148072

148072: sql: adopt stopper handle in two more places r=tbg a=tbg

Noticed during heap profiling.

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

Author: craig[bot]

pkg/sql/internal.go

@@ -237,36 +237,35 @@ func (ie *InternalExecutor) runWithEx(
 		ex.close(ctx, closeMode)
 		wg.Done()
 	}
-	if err = ie.s.cfg.Stopper.RunAsyncTaskEx(
-		ctx,
-		stop.TaskOpts{
-			TaskName: opName.StripMarkers(),
-			SpanOpt:  stop.ChildSpan,
-		},
-		func(ctx context.Context) {
-			defer cleanup(ctx)
-			// TODO(yuzefovich): benchmark whether we should be growing the
-			// stack size unconditionally.
-			if growStackSize {
-				growstack.Grow()
-			}
-			if err := ex.run(
-				ctx,
-				ie.mon,
-				&mon.BoundAccount{}, /*reserved*/
-				nil,                 /* cancel */
-			); err != nil {
-				sqltelemetry.RecordError(ctx, err, &ex.server.cfg.Settings.SV)
-				errCallback(err)
-			}
-			w.finish()
-		},
-	); err != nil {
+	ctx, hdl, err := ie.s.cfg.Stopper.GetHandle(ctx, stop.TaskOpts{
+		TaskName: opName.StripMarkers(),
+		SpanOpt:  stop.ChildSpan,
+	})
+	if err != nil {
 		// The goroutine wasn't started, so we need to perform the cleanup
 		// ourselves.
 		cleanup(ctx)
 		return err
 	}
+	go func() {
+		defer hdl.Activate(ctx).Release(ctx)
+		defer cleanup(ctx)
+		// TODO(yuzefovich): benchmark whether we should be growing the
+		// stack size unconditionally.
+		if growStackSize {
+			growstack.Grow()
+		}
+		if err := ex.run(
+			ctx,
+			ie.mon,
+			&mon.BoundAccount{}, /*reserved*/
+			nil,                 /* cancel */
+		); err != nil {
+			sqltelemetry.RecordError(ctx, err, &ex.server.cfg.Settings.SV)
+			errCallback(err)
+		}
+		w.finish()
+	}()
 	return nil
 }
 

pkg/ts/db.go

@@ -141,13 +141,18 @@ func (db *DB) PollSource(
 // start begins the goroutine for this poller, which will periodically request
 // time series data from the DataSource and store it.
 func (p *poller) start() (firstDone <-chan struct{}) {
-	ch := make(chan struct{})
-	// Poll once immediately and synchronously.
-	bgCtx := p.AnnotateCtx(context.Background())
-	if p.stopper.RunAsyncTask(bgCtx, "ts-poller", func(ctx context.Context) {
-		ch := ch // goroutine-local copy
+	ch := make(chan struct{}) // closed on completion of first poll
+	ctx, hdl, err := p.stopper.GetHandle(
+		p.AnnotateCtx(context.Background()), stop.TaskOpts{TaskName: "ts-poller"},
+	)
+	if err != nil {
+		close(ch)
+		return ch
+	}
+	go func(ctx context.Context, ch chan struct{}) {
+		defer hdl.Activate(ctx).Release(ctx)
 		var ticker timeutil.Timer
-		ticker.Reset(0)
+		ticker.Reset(0) // poll immediately
 		defer ticker.Stop()
 		for {
 			select {
@@ -162,9 +167,7 @@ func (p *poller) start() (firstDone <-chan struct{}) {
 				return
 			}
 		}
-	}) != nil {
-		close(ch)
-	}
+	}(ctx, ch)
 	return ch
 }