server: add hooks to wait for background commit goroutines (#55608)

close #55607

pkg/kv/option.go

@@ -114,6 +114,8 @@ const (
 	SizeLimits
 	// SessionID marks the connection id, for logging and tracing.
 	SessionID
+	// BackgroundGoroutineLifecycleHooks is the hooks to track the start and end of background goroutine
+	BackgroundGoroutineLifecycleHooks
 )
 
 // TxnSizeLimits is the argument type for `SizeLimits` option

pkg/server/conn_test.go

@@ -819,6 +819,39 @@ func TestShutDown(t *testing.T) {
 	}
 }
 
+func TestCommitWaitGroup(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+
+	se, err := session.CreateSession4Test(store)
+	require.NoError(t, err)
+	tc := &TiDBContext{Session: se}
+
+	cfg := serverutil.NewTestConfig()
+	cfg.Port = 0
+	cfg.Status.StatusPort = 0
+	drv := NewTiDBDriver(store)
+	srv, err := NewServer(cfg, drv)
+	require.NoError(t, err)
+	srv.SetDomain(dom)
+
+	cc := &clientConn{server: srv}
+	cc.SetCtx(tc)
+	cc.CompareAndSwapStatus(cc.getStatus(), connStatusReading)
+	cc.getCtx().GetSessionVars().SetInTxn(false)
+	srv.clients[dom.NextConnID()] = cc
+
+	wg := cc.getCtx().GetCommitWaitGroup()
+	wg.Add(1)
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		wg.Done()
+	}()
+	begin := time.Now()
+	srv.DrainClients(time.Second, time.Second)
+	require.Greater(t, time.Since(begin), 100*time.Millisecond)
+	require.Less(t, time.Since(begin), time.Second)
+}
+
 type snapshotCache interface {
 	SnapCacheHitCount() int
 }

pkg/server/server.go

@@ -1021,13 +1021,29 @@ func (s *Server) DrainClients(drainWait time.Duration, cancelWait time.Duration)
 		for _, conn := range conns {
 			// Wait for the connections with explicit transaction or an executing auto-commit query.
 			if conn.getStatus() == connStatusReading && !conn.getCtx().GetSessionVars().InTxn() {
+				// The waitgroup is not protected by the `quitWaitingForConns`. However, the implementation
+				// of `client-go` will guarantee this `Wait` will return at least after killing the
+				// connections. We also wait for a similar `WaitGroup` on the store after killing the connections.
+				//
+				// Therefore, it'll not cause goroutine leak. Even if, it's not a big issue when the TiDB is
+				// going to shutdown.
+				//
+				// It should be waited for connections in all status, even if it's not in transactions and is reading
+				// from the client. Because it may run background commit goroutines at any time.
+				conn.getCtx().Session.GetCommitWaitGroup().Wait()
+
 				continue
 			}
 			select {
 			case <-conn.quit:
 			case <-quitWaitingForConns:
 				return
 			}
+
+			// Wait for the commit wait group after waiting for the `conn.quit` channel to make sure the foreground
+			// process has finished to avoid the situation that after waiting for the wait group, the transaction starts
+			// a new background goroutine and increase the wait group.
+			conn.getCtx().Session.GetCommitWaitGroup().Wait()
 		}
 	}()
 

pkg/session/session.go

@@ -223,6 +223,9 @@ type session struct {
 	sandBoxMode bool
 
 	cursorTracker cursor.Tracker
+
+	// Used to wait for all async commit background jobs to finish.
+	commitWaitGroup sync.WaitGroup
 }
 
 var parserPool = &sync.Pool{New: func() any { return parser.New() }}
@@ -4604,3 +4607,8 @@ func GetDBNames(seVar *variable.SessionVars) []string {
 func (s *session) GetCursorTracker() cursor.Tracker {
 	return s.cursorTracker
 }
+
+// GetCommitWaitGroup returns the internal `sync.WaitGroup` for async commit and secondary key lock cleanup
+func (s *session) GetCommitWaitGroup() *sync.WaitGroup {
+	return &s.commitWaitGroup
+}

pkg/sessionctx/context.go

@@ -16,6 +16,7 @@ package sessionctx
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -207,6 +208,8 @@ type Context interface {
 	NewStmtIndexUsageCollector() *indexusage.StmtIndexUsageCollector
 	// GetCursorTracker returns the cursor tracker of the session
 	GetCursorTracker() cursor.Tracker
+	// GetCommitWaitGroup returns the wait group for async commit and secondary lock cleanup background goroutines
+	GetCommitWaitGroup() *sync.WaitGroup
 }
 
 // TxnFuture is an interface where implementations have a kv.Transaction field and after

pkg/sessiontxn/isolation/BUILD.bazel

@@ -39,6 +39,7 @@ go_library(
         "@com_github_tikv_client_go_v2//error",
         "@com_github_tikv_client_go_v2//kv",
         "@com_github_tikv_client_go_v2//oracle",
+        "@com_github_tikv_client_go_v2//txnkv/transaction",
         "@org_uber_go_zap//:zap",
     ],
 )

pkg/sessiontxn/isolation/base.go

@@ -38,6 +38,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/tracing"
 	tikvstore "github.com/tikv/client-go/v2/kv"
 	"github.com/tikv/client-go/v2/oracle"
+	"github.com/tikv/client-go/v2/txnkv/transaction"
 )
 
 // baseTxnContextProvider is a base class for the transaction context providers that implement `TxnContextProvider` in different isolation.
@@ -515,6 +516,19 @@ func (p *baseTxnContextProvider) SetOptionsOnTxnActive(txn kv.Transaction) {
 	}
 
 	txn.SetOption(kv.SessionID, p.sctx.GetSessionVars().ConnectionID)
+
+	// backgroundGoroutineWaitGroup is pre-initialized before the closure to avoid accessing `p.sctx` in the closure,
+	// which may cause unexpected race condition.
+	backgroundGoroutineWaitGroup := p.sctx.GetCommitWaitGroup()
+	lifecycleHooks := transaction.LifecycleHooks{
+		Pre: func() {
+			backgroundGoroutineWaitGroup.Add(1)
+		},
+		Post: func() {
+			backgroundGoroutineWaitGroup.Done()
+		},
+	}
+	txn.SetOption(kv.BackgroundGoroutineLifecycleHooks, lifecycleHooks)
 }
 
 func (p *baseTxnContextProvider) SetOptionsBeforeCommit(

pkg/store/driver/txn/txn_driver.go

@@ -38,6 +38,7 @@ import (
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"github.com/tikv/client-go/v2/tikvrpc/interceptor"
 	"github.com/tikv/client-go/v2/txnkv"
+	"github.com/tikv/client-go/v2/txnkv/transaction"
 	"github.com/tikv/client-go/v2/txnkv/txnsnapshot"
 	"go.uber.org/zap"
 )
@@ -308,6 +309,8 @@ func (txn *tikvTxn) SetOption(opt int, val any) {
 		txn.KVTxn.GetUnionStore().SetEntrySizeLimit(limits.Entry, limits.Total)
 	case kv.SessionID:
 		txn.KVTxn.SetSessionID(val.(uint64))
+	case kv.BackgroundGoroutineLifecycleHooks:
+		txn.KVTxn.SetBackgroundGoroutineLifecycleHooks(val.(transaction.LifecycleHooks))
 	}
 }
 

pkg/util/mock/context.go

@@ -18,6 +18,7 @@ package mock
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -611,6 +612,11 @@ func (*Context) GetCursorTracker() cursor.Tracker {
 	return nil
 }
 
+// GetCommitWaitGroup implements the sessionctx.Context interface
+func (*Context) GetCommitWaitGroup() *sync.WaitGroup {
+	return nil
+}
+
 // NewContext creates a new mocked sessionctx.Context.
 func NewContext() *Context {
 	ctx, cancel := context.WithCancel(context.Background())