From 8eec8ba771db919ad638141103cafd0658548ef9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E8=B6=85?= <cclcwangchao@hotmail.com>
Date: Mon, 20 May 2024 20:13:15 +0800
Subject: [PATCH] ddl: remove `mock.Context` ref in
 `ddl/session/session_pool.go` (#53378)

close pingcap/tidb#53377
---
 pkg/ddl/export_test.go                   | 10 +++++++++-
 pkg/ddl/internal/session/BUILD.bazel     |  2 +-
 pkg/ddl/internal/session/session_pool.go | 16 ++++------------
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/pkg/ddl/export_test.go b/pkg/ddl/export_test.go
index 87d8920486c45e..6625aef25d8b49 100644
--- a/pkg/ddl/export_test.go
+++ b/pkg/ddl/export_test.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/ngaut/pools"
 	"github.com/pingcap/tidb/pkg/ddl/copr"
 	"github.com/pingcap/tidb/pkg/ddl/internal/session"
 	"github.com/pingcap/tidb/pkg/kv"
@@ -26,6 +27,7 @@ import (
 	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/mock"
 )
 
 type resultChanForTest struct {
@@ -47,7 +49,12 @@ func FetchChunk4Test(copCtx copr.CopContext, tbl table.PhysicalTable, startKey,
 	}
 	taskCh := make(chan *reorgBackfillTask, 5)
 	resultCh := make(chan IndexRecordChunk, 5)
-	sessPool := session.NewSessionPool(nil, store)
+	resPool := pools.NewResourcePool(func() (pools.Resource, error) {
+		ctx := mock.NewContext()
+		ctx.Store = store
+		return ctx, nil
+	}, 8, 8, 0)
+	sessPool := session.NewSessionPool(resPool, store)
 	pool := newCopReqSenderPool(context.Background(), copCtx, store, taskCh, sessPool, nil)
 	pool.chunkSender = &resultChanForTest{ch: resultCh}
 	pool.adjustSize(1)
@@ -55,6 +62,7 @@ func FetchChunk4Test(copCtx copr.CopContext, tbl table.PhysicalTable, startKey,
 	rs := <-resultCh
 	close(taskCh)
 	pool.close(false)
+	sessPool.Close()
 	return rs.Chunk
 }
 
diff --git a/pkg/ddl/internal/session/BUILD.bazel b/pkg/ddl/internal/session/BUILD.bazel
index 62c006a7b1e4cc..0221a13ca544d3 100644
--- a/pkg/ddl/internal/session/BUILD.bazel
+++ b/pkg/ddl/internal/session/BUILD.bazel
@@ -18,7 +18,7 @@ go_library(
         "//pkg/sessionctx",
         "//pkg/sessiontxn",
         "//pkg/util/chunk",
-        "//pkg/util/mock",
+        "//pkg/util/intest",
         "//pkg/util/sqlexec",
         "@com_github_ngaut_pools//:pools",
         "@com_github_pingcap_errors//:errors",
diff --git a/pkg/ddl/internal/session/session_pool.go b/pkg/ddl/internal/session/session_pool.go
index 66b8cef0372b8c..6a2feb91edc55b 100644
--- a/pkg/ddl/internal/session/session_pool.go
+++ b/pkg/ddl/internal/session/session_pool.go
@@ -25,7 +25,7 @@ import (
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/sessionctx"
-	"github.com/pingcap/tidb/pkg/util/mock"
+	"github.com/pingcap/tidb/pkg/util/intest"
 )
 
 // Pool is used to new Session.
@@ -40,18 +40,14 @@ type Pool struct {
 
 // NewSessionPool creates a new Session pool.
 func NewSessionPool(resPool *pools.ResourcePool, store kv.Storage) *Pool {
+	intest.AssertNotNil(resPool)
+	intest.AssertNotNil(store)
 	return &Pool{resPool: resPool, store: store}
 }
 
 // Get gets sessionCtx from context resource pool.
 // Please remember to call Put after you finished using sessionCtx.
 func (sg *Pool) Get() (sessionctx.Context, error) {
-	if sg.resPool == nil {
-		ctx := mock.NewContext()
-		ctx.Store = sg.store
-		return ctx, nil
-	}
-
 	sg.mu.Lock()
 	if sg.mu.closed {
 		sg.mu.Unlock()
@@ -78,10 +74,6 @@ func (sg *Pool) Get() (sessionctx.Context, error) {
 
 // Put returns sessionCtx to context resource pool.
 func (sg *Pool) Put(ctx sessionctx.Context) {
-	if sg.resPool == nil {
-		return
-	}
-
 	// no need to protect sg.resPool, even the sg.resPool is closed, the ctx still need to
 	// Put into resPool, because when resPool is closing, it will wait all the ctx returns, then resPool finish closing.
 	sg.resPool.Put(ctx.(pools.Resource))
@@ -93,7 +85,7 @@ func (sg *Pool) Close() {
 	sg.mu.Lock()
 	defer sg.mu.Unlock()
 	// prevent closing resPool twice.
-	if sg.mu.closed || sg.resPool == nil {
+	if sg.mu.closed {
 		return
 	}
 	logutil.DDLLogger().Info("closing session pool")