ddl: remove mock.Context ref in ddl/session/session_pool.go

pingcap · May 20, 2024 · 8889dd0 · 8889dd0
1 parent f5ac93e
commit 8889dd0
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 9 deletions.
diff --git 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,14 +49,20 @@ 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)
 	pool.tasksCh <- task
 	rs := <-resultCh
 	close(taskCh)
 	pool.close(false)
+	resPool.Close()
 	return rs.Chunk
 }
 

diff --git 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
@@ -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()