sqlstats: stats flush should honor stopper quiesce

pkg/sql/sqlstats/sslocal/sslocal_provider.go

@@ -8,7 +8,6 @@ package sslocal
 import (
 	"context"
 	"sort"
-	"sync"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
@@ -162,11 +161,16 @@ func (s *SQLStats) ConsumeStats(
 		// the system statement|transaction_statistics tables.
 		// In-memory stats storage is not locked here and it is safe to call stmtVisitor or txnVisitor functions
 		// that might be time consuming operations.
-		var wg sync.WaitGroup
-		wg.Add(2)
+
+		// We use a done channel instead of waitGroup because we need to
+		// fail if the stopper is quiescing and the tasks are never going
+		// to run.
+		doneCh := make(chan struct{}, 2)
 
 		err := stopper.RunAsyncTask(ctx, "sql-stmt-stats-flush", func(ctx context.Context) {
-			defer wg.Done()
+			defer func() {
+				doneCh <- struct{}{}
+			}()
 			for _, stat := range stmtStats {
 				stat := stat
 				if err := stmtVisitor(ctx, stat); err != nil {
@@ -176,12 +180,13 @@ func (s *SQLStats) ConsumeStats(
 		})
 		if err != nil {
 			log.Warningf(ctx, "failed to execute sql-stmt-stats-flush task, %s", err.Error())
-			wg.Done()
 			return
 		}
 
 		err = stopper.RunAsyncTask(ctx, "sql-txn-stats-flush", func(ctx context.Context) {
-			defer wg.Done()
+			defer func() {
+				doneCh <- struct{}{}
+			}()
 			for _, stat := range txnStats {
 				stat := stat
 				if err := txnVisitor(ctx, stat); err != nil {
@@ -191,10 +196,24 @@ func (s *SQLStats) ConsumeStats(
 		})
 		if err != nil {
 			log.Warningf(ctx, "failed to execute sql-txn-stats-flush task, %s", err.Error())
-			wg.Done()
 			return
 		}
-		wg.Wait()
+
+		numComplete := 0
+		for {
+			select {
+			case <-doneCh:
+				numComplete++
+				if numComplete >= 2 {
+					close(doneCh)
+					return
+				}
+			case <-stopper.ShouldQuiesce():
+				// Don't close `doneCh` because the tasks could still be
+				// running.
+				return
+			}
+		}
 	}
 }