tikv, txnkv: unify the way background goroutines are spawned

tikv/kv.go

@@ -120,6 +120,10 @@ func createEtcdKV(addrs []string, tlsConfig *tls.Config) (*clientv3.Client, erro
 // or at runtime by calling SetLowResolutionTimestampUpdateInterval on the oracle
 var defaultOracleUpdateInterval = 2 * time.Second
 
+// ErrStoreClosed is returned by KVStore.Go when the store is closing, which
+// means the submitted function is rejected and will never be executed.
+var ErrStoreClosed = errors.New("the store is closed")
+
 // KVStore contains methods to interact with a TiKV cluster.
 type KVStore struct {
 	clusterID uint64
@@ -160,16 +164,57 @@ type KVStore struct {
 
 	ctx    context.Context
 	cancel context.CancelFunc
-	wg     sync.WaitGroup
 	close  atomicutil.Bool
 	gP     Pool
+
+	// backgroundJobs guards the spawning of background goroutines. Every
+	// background goroutine owned by the store is started by Go and tracked by
+	// wg, and no new one is admitted once the store starts closing. The lock is
+	// what makes `wg.Add` in Go and `wg.Wait` in stopBackgroundJobs mutually
+	// exclusive, which a plain atomic flag cannot guarantee.
+	backgroundJobs struct {
+		sync.RWMutex
+		closed bool
+	}
+	wg sync.WaitGroup
 }
 
 var _ Storage = (*KVStore)(nil)
 
-// Go run the function in a separate goroutine.
+// Go runs f in a background goroutine owned by the store.
+//
+// This is the only supported way to spawn a background goroutine on behalf of a
+// store: f runs in the store's goroutine pool and is tracked by the store, so
+// that Close waits for it instead of racing with it.
+//
+// A non-nil error means f is rejected and will never run, which happens when
+// the store is closing. The caller then owns the cleanup of whatever it has
+// prepared for f, for example releasing a token or marking a WaitGroup done.
 func (s *KVStore) Go(f func()) error {
-	return s.gP.Run(f)
+	s.backgroundJobs.RLock()
+	if s.backgroundJobs.closed {
+		s.backgroundJobs.RUnlock()
+		return ErrStoreClosed
+	}
+	s.wg.Add(1)
+	err := s.gP.Run(func() {
+		defer s.wg.Done()
+		f()
+	})
+	s.backgroundJobs.RUnlock()
+	if err != nil {
+		s.wg.Done()
+	}
+	return err
+}
+
+// stopBackgroundJobs rejects any new background job and waits for the running
+// ones to finish.
+func (s *KVStore) stopBackgroundJobs() {
+	s.backgroundJobs.Lock()
+	s.backgroundJobs.closed = true
+	s.backgroundJobs.Unlock()
+	s.wg.Wait()
 }
 
 // UpdateTxnSafePointCache updates the cached txn safe point, which is used for safety check of data access
@@ -359,6 +404,9 @@ func NewKVStore(
 		if retErr != nil {
 			regionCache.Close()
 			cancel()
+			// Background jobs may already have been started below, stop them
+			// before reporting the failure.
+			store.stopBackgroundJobs()
 		}
 	}()
 
@@ -378,9 +426,11 @@ func NewKVStore(
 	store.lockResolver = txnlock.NewLockResolver(store)
 	loadOption(store, opt...)
 
-	store.wg.Add(2)
-	go store.runTxnSafePointUpdater()
-	go store.safeTSUpdater()
+	for _, updater := range []func(){store.runTxnSafePointUpdater, store.safeTSUpdater} {
+		if err := store.Go(updater); err != nil {
+			return nil, errors.WithStack(err)
+		}
+	}
 
 	return store, nil
 }
@@ -425,7 +475,6 @@ func (s *KVStore) IsLatchEnabled() bool {
 }
 
 func (s *KVStore) runTxnSafePointUpdater() {
-	defer s.wg.Done()
 	if _, e := util.EvalFailpoint("noBuiltInTxnSafePointUpdater"); e == nil {
 		return
 	}
@@ -502,7 +551,7 @@ func (s *KVStore) Close() error {
 	defer s.gP.Close()
 	s.close.Store(true)
 	s.cancel()
-	s.wg.Wait()
+	s.stopBackgroundJobs()
 
 	s.oracle.Close()
 	if s.txnLatches != nil {
@@ -715,7 +764,12 @@ func (s *KVStore) IsClose() bool {
 	return s.close.Load()
 }
 
-// WaitGroup returns wg
+// WaitGroup returns the WaitGroup that tracks the store's background
+// goroutines.
+//
+// Deprecated: use Go to spawn a background goroutine, which registers it with
+// this WaitGroup for you. Adding to the WaitGroup directly races with the
+// Wait done by Close.
 func (s *KVStore) WaitGroup() *sync.WaitGroup {
 	return &s.wg
 }
@@ -776,7 +830,6 @@ func (s *KVStore) updateMinSafeTS(txnScope string, storeIDs []uint64) {
 }
 
 func (s *KVStore) safeTSUpdater() {
-	defer s.wg.Done()
 	t := time.NewTicker(safeTSUpdateInterval)
 	if _, e := util.EvalFailpoint("mockFastSafeTSUpdater"); e == nil {
 		t.Reset(time.Millisecond * 100)

tikv/pool.go

@@ -15,31 +15,65 @@
 package tikv
 
 import (
+	"sync"
 	"time"
 
+	"github.com/pingcap/errors"
 	"github.com/tiancaiamao/gp"
 )
 
+// ErrPoolClosed is returned by Pool.Run when the pool has been closed. The
+// submitted function is rejected and will never be executed.
+var ErrPoolClosed = errors.New("the goroutine pool is closed")
+
 // Pool is a simple interface for goroutine pool.
 type Pool interface {
+	// Run executes the function in a separate goroutine. A non-nil error must
+	// be returned if and only if the function is rejected, which means it will
+	// never be executed and the caller owns the cleanup of anything it has
+	// prepared for the function.
 	Run(func()) error
+	// Close releases the goroutines of the pool. In-flight functions are not
+	// interrupted. It must be safe to call Close more than once.
 	Close()
 }
 
 // Spool is a simple implementation of Pool.
 type Spool struct {
 	gp.Pool
+	mu struct {
+		sync.RWMutex
+		closed bool
+	}
 }
 
 // NewSpool creates a new Spool.
 func NewSpool(n int, dur time.Duration) *Spool {
-	return &Spool{
-		*gp.New(n, dur),
-	}
+	return &Spool{Pool: *gp.New(n, dur)}
 }
 
 // Run implements Pool.Run.
 func (p *Spool) Run(fn func()) error {
+	// The underlying pool silently drops the functions submitted after it is
+	// closed. Submitting under the read lock makes the outcome unambiguous: the
+	// function is either accepted by a pool that is still open, or rejected
+	// with ErrPoolClosed here.
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	if p.mu.closed {
+		return ErrPoolClosed
+	}
 	p.Go(fn)
 	return nil
 }
+
+// Close implements Pool.Close.
+func (p *Spool) Close() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.mu.closed {
+		return
+	}
+	p.mu.closed = true
+	p.Pool.Close()
+}

tikv/pool_test.go

@@ -0,0 +1,87 @@
+// Copyright 2025 TiKV Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tikv
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/internal/mockstore/mocktikv"
+	"github.com/tikv/client-go/v2/testutils"
+)
+
+func TestSpoolRejectsTasksAfterClose(t *testing.T) {
+	p := NewSpool(2, 10*time.Millisecond)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	require.NoError(t, p.Run(wg.Done))
+	wg.Wait()
+
+	p.Close()
+
+	var executed atomic.Bool
+	require.ErrorIs(t, p.Run(func() { executed.Store(true) }), ErrPoolClosed)
+	require.False(t, executed.Load())
+
+	// Close must be idempotent: closing the underlying pool twice panics.
+	p.Close()
+}
+
+func TestKVStoreGoTracksBackgroundJobs(t *testing.T) {
+	client, cluster, pdClient, err := testutils.NewMockTiKV("", nil)
+	require.NoError(t, err)
+	mocktikv.BootstrapWithMultiStores(cluster, 1)
+	store, err := NewTestTiKVStore(client, pdClient, nil, nil, 0)
+	require.NoError(t, err)
+
+	started := make(chan struct{})
+	release := make(chan struct{})
+	finished := make(chan struct{})
+	require.NoError(t, store.Go(func() {
+		close(started)
+		<-release
+		close(finished)
+	}))
+	<-started
+
+	var closeErr error
+	returned := make(chan struct{})
+	go func() {
+		defer close(returned)
+		closeErr = store.Close()
+	}()
+
+	// Close waits for the background job instead of racing with it.
+	select {
+	case <-returned:
+		require.FailNow(t, "Close returned before the background job finished")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	close(release)
+	<-finished
+	<-returned
+	require.NoError(t, closeErr)
+
+	// A closing store admits no new background job, and reports the rejection
+	// so that the caller can clean up after it.
+	var executed atomic.Bool
+	require.ErrorIs(t, store.Go(func() { executed.Store(true) }), ErrStoreClosed)
+	require.False(t, executed.Load())
+}

txnkv/transaction/2pc.go

@@ -116,13 +116,14 @@ type kvstore interface {
 	GetTiKVClient() (client client.Client)
 	GetLockResolver() *txnlock.LockResolver
 	Ctx() context.Context
-	WaitGroup() *sync.WaitGroup
 	// TxnLatches returns txnLatches.
 	TxnLatches() *latch.LatchesScheduler
 	GetClusterID() uint64
 	// IsClose checks whether the store is closed.
 	IsClose() bool
-	// Go run the function in a separate goroutine.
+	// Go runs the function in a background goroutine tracked by the store. It
+	// returns a non-nil error if and only if the function is rejected and will
+	// never be executed, which happens when the store is closing.
 	Go(f func()) error
 }
 
@@ -1096,7 +1097,7 @@ func (c *twoPhaseCommitter) doActionOnGroupMutations(bo *retry.Backoffer, action
 				zap.Uint64("sessionID", c.sessionID))
 			return nil
 		}
-		err = c.txn.spawnWithStorePool(func() {
+		err = c.txn.spawn(func() {
 			if c.sessionID > 0 {
 				if v, err := util.EvalFailpoint("beforeCommitSecondaries"); err == nil {
 					if s, ok := v.(string); !ok {
@@ -1478,7 +1479,7 @@ func broadcastToAllStores(
 			wg.Add(1)
 			target := s
 
-			err := txn.spawnWithStorePool(func() {
+			err := txn.spawn(func() {
 				defer wg.Done()
 				defer func() { <-rateLimit }()
 
@@ -1522,7 +1523,7 @@ func broadcastToAllStores(
 		wg.Wait()
 	}
 
-	if err := txn.spawnWithStorePool(broadcastFunc); err != nil {
+	if err := txn.spawn(broadcastFunc); err != nil {
 		logutil.Logger(store.Ctx()).Error("failed to spawn goroutine for broadcasting txn status",
 			zap.Error(err))
 	}
@@ -1675,7 +1676,7 @@ func (c *twoPhaseCommitter) cleanup(ctx context.Context) {
 		return
 	}
 	c.cleanWg.Add(1)
-	c.txn.spawn(func() {
+	err := c.txn.spawn(func() {
 		defer c.cleanWg.Done()
 
 		if _, err := util.EvalFailpoint("commitFailedSkipCleanup"); err == nil {
@@ -1735,6 +1736,14 @@ func (c *twoPhaseCommitter) cleanup(ctx context.Context) {
 			}
 		}
 	})
+	if err != nil {
+		// The cleanup will never run, release the waiter so that it does not
+		// block forever.
+		c.cleanWg.Done()
+		logutil.Logger(ctx).Warn("twoPhaseCommitter fail to spawn the cleanup goroutine",
+			zap.Uint64("txnStartTS", c.startTS), zap.Bool("isPessimistic", c.isPessimistic),
+			zap.Bool("isOnePC", c.isOnePC()), zap.Error(err))
+	}
 }
 
 // execute executes the two-phase commit protocol.
@@ -2056,7 +2065,7 @@ func (c *twoPhaseCommitter) execute(ctx context.Context) (err error) {
 				zap.Uint64("sessionID", c.sessionID))
 			return nil
 		}
-		c.txn.spawn(func() {
+		err := c.txn.spawn(func() {
 			if _, err := util.EvalFailpoint("asyncCommitDoNothing"); err == nil {
 				return
 			}
@@ -2067,6 +2076,11 @@ func (c *twoPhaseCommitter) execute(ctx context.Context) (err error) {
 					zap.Uint64("startTS", c.startTS), zap.Uint64("commitTS", c.commitTS), zap.Error(err))
 			}
 		})
+		if err != nil {
+			logutil.Logger(ctx).Warn("2PC fail to spawn the async commit goroutine",
+				zap.Uint64("startTS", c.startTS), zap.Uint64("commitTS", c.commitTS),
+				zap.Uint64("sessionID", c.sessionID), zap.Error(err))
+		}
 		return nil
 	}
 	return c.commitTxn(ctx, commitDetail)