bindinfo: add timeout for loading binding from storage

pkg/bindinfo/BUILD.bazel

@@ -31,13 +31,15 @@ go_library(
         "//pkg/util/chunk",
         "//pkg/util/hack",
         "//pkg/util/hint",
+        "//pkg/util/intest",
         "//pkg/util/kvcache",
         "//pkg/util/logutil",
         "//pkg/util/mathutil",
         "//pkg/util/memory",
         "//pkg/util/parser",
         "//pkg/util/sqlexec",
         "//pkg/util/stmtsummary/v2:stmtsummary",
+        "//pkg/util/stringutil",
         "//pkg/util/table-filter",
         "@com_github_ngaut_pools//:pools",
         "@com_github_pingcap_errors//:errors",

pkg/bindinfo/binding_cache.go

@@ -17,21 +17,36 @@ package bindinfo
 import (
 	"errors"
 	"sync"
+	"sync/atomic"
+	"time"
 
-	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/bindinfo/norm"
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/util/hack"
+	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/pingcap/tidb/pkg/util/kvcache"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/mathutil"
 	"github.com/pingcap/tidb/pkg/util/memory"
+	"github.com/pingcap/tidb/pkg/util/stringutil"
 	"go.uber.org/zap"
 )
 
+// GetBindingReturnNil is only for test
+var GetBindingReturnNil = stringutil.StringerStr("GetBindingReturnNil")
+
+// GetBindingReturnNilBool is only for test
+var GetBindingReturnNilBool atomic.Bool
+
+// GetBindingReturnNilAlways is only for test
+var GetBindingReturnNilAlways = stringutil.StringerStr("getBindingReturnNilAlways")
+
+// LoadBindingNothing is only for test
+var LoadBindingNothing = stringutil.StringerStr("LoadBindingNothing")
+
 // FuzzyBindingCache is based on BindingCache, and provide some more advanced features, like
 // fuzzy matching, loading binding if cache miss automatically (TODO).
 type FuzzyBindingCache interface {
@@ -57,10 +72,10 @@ type fuzzyBindingCache struct {
 	sql2FuzzyDigest map[string]string // sqlDigest --> fuzzyDigest
 
 	// loadBindingFromStorageFunc is used to load binding from storage if cache miss.
-	loadBindingFromStorageFunc func(sqlDigest string) (Bindings, error)
+	loadBindingFromStorageFunc func(sctx sessionctx.Context, sqlDigest string) (Bindings, error)
 }
 
-func newFuzzyBindingCache(loadBindingFromStorageFunc func(string) (Bindings, error)) FuzzyBindingCache {
+func newFuzzyBindingCache(loadBindingFromStorageFunc func(sessionctx.Context, string) (Bindings, error)) FuzzyBindingCache {
 	return &fuzzyBindingCache{
 		BindingCache:               newBindCache(),
 		fuzzy2SQLDigests:           make(map[string][]string),
@@ -77,7 +92,7 @@ func (fbc *fuzzyBindingCache) FuzzyMatchingBinding(sctx sessionctx.Context, fuzz
 	if fbc.loadBindingFromStorageFunc == nil {
 		return
 	}
-	fbc.loadFromStore(missingSQLDigest) // loadFromStore's SetBinding has a Mutex inside, so it's safe to call it without lock
+	fbc.loadFromStore(sctx, missingSQLDigest) // loadFromStore's SetBinding has a Mutex inside, so it's safe to call it without lock
 	matchedBinding, isMatched, _ = fbc.getFromMemory(sctx, fuzzyDigest, tableNames)
 	return
 }
@@ -92,7 +107,18 @@ func (fbc *fuzzyBindingCache) getFromMemory(sctx sessionctx.Context, fuzzyDigest
 	leastWildcards := len(tableNames) + 1
 	enableFuzzyBinding := sctx.GetSessionVars().EnableFuzzyBinding
 	for _, sqlDigest := range fbc.fuzzy2SQLDigests[fuzzyDigest] {
-		if bindings := bindingCache.GetBinding(sqlDigest); bindings != nil {
+		bindings := bindingCache.GetBinding(sqlDigest)
+		if intest.InTest {
+			if sctx.Value(GetBindingReturnNil) != nil {
+				if GetBindingReturnNilBool.CompareAndSwap(false, true) {
+					bindings = nil
+				}
+			}
+			if sctx.Value(GetBindingReturnNilAlways) != nil {
+				bindings = nil
+			}
+		}
+		if bindings != nil {
 			for _, binding := range bindings {
 				numWildcards, matched := fuzzyMatchBindingTableName(sctx.GetSessionVars().CurrentDB, tableNames, binding.TableNames)
 				if matched && numWildcards > 0 && sctx != nil && !enableFuzzyBinding {
@@ -112,11 +138,19 @@ func (fbc *fuzzyBindingCache) getFromMemory(sctx sessionctx.Context, fuzzyDigest
 	return matchedBinding, isMatched, missingSQLDigest
 }
 
-func (fbc *fuzzyBindingCache) loadFromStore(missingSQLDigest []string) {
+func (fbc *fuzzyBindingCache) loadFromStore(sctx sessionctx.Context, missingSQLDigest []string) {
+	if intest.InTest && sctx.Value(LoadBindingNothing) != nil {
+		return
+	}
 	for _, sqlDigest := range missingSQLDigest {
-		bindings, err := fbc.loadBindingFromStorageFunc(sqlDigest)
+		start := time.Now()
+		bindings, err := fbc.loadBindingFromStorageFunc(sctx, sqlDigest)
 		if err != nil {
-			logutil.BgLogger().Warn("loadBindingFromStorageFunc binding failed", zap.String("sqlDigest", sqlDigest), zap.Error(err))
+			logutil.BgLogger().Warn("failed to load binding from storage",
+				zap.String("sqlDigest", sqlDigest),
+				zap.Error(err),
+				zap.Duration("duration", time.Since(start)),
+			)
 			continue
 		}
 		// put binding into the cache
@@ -317,9 +351,6 @@ func (c *bindingCache) delete(key bindingCacheKey) bool {
 // The return value is not read-only, but it shouldn't be changed in the caller functions.
 // The function is thread-safe.
 func (c *bindingCache) GetBinding(sqlDigest string) Bindings {
-	failpoint.Inject("get_binding_return_nil", func(_ failpoint.Value) {
-		failpoint.Return(nil)
-	})
 	c.lock.Lock()
 	defer c.lock.Unlock()
 	return c.get(bindingCacheKey(sqlDigest))

pkg/bindinfo/binding_match.go

@@ -86,9 +86,11 @@ func matchSQLBinding(sctx sessionctx.Context, stmtNode ast.StmtNode, info *Bindi
 	if globalHandle == nil {
 		return
 	}
-	if binding, matched := globalHandle.MatchGlobalBinding(sctx, fuzzyDigest, tableNames); matched {
+	binding, matched = globalHandle.MatchGlobalBinding(sctx, fuzzyDigest, tableNames)
+	if matched {
 		return binding, matched, metrics.ScopeGlobal
 	}
+
 	return
 }
 

pkg/bindinfo/global_handle.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/format"
@@ -685,27 +686,38 @@ func (h *globalBindingHandle) Stats(_ *variable.SessionVars) (map[string]any, er
 }
 
 // LoadBindingsFromStorageToCache loads global bindings from storage to the memory cache.
-func (h *globalBindingHandle) LoadBindingsFromStorage(sqlDigest string) (Bindings, error) {
+func (h *globalBindingHandle) LoadBindingsFromStorage(sctx sessionctx.Context, sqlDigest string) (Bindings, error) {
 	if sqlDigest == "" {
 		return nil, nil
 	}
-	bindings, err, _ := h.syncBindingSingleflight.Do(sqlDigest, func() (any, error) {
+	timeout := time.Duration(sctx.GetSessionVars().LoadBindingTimeout) * time.Millisecond
+	resultChan := h.syncBindingSingleflight.DoChan(sqlDigest, func() (any, error) {
 		return h.loadBindingsFromStorageInternal(sqlDigest)
 	})
-	if err != nil {
-		logutil.BgLogger().Warn("fail to LoadBindingsFromStorageToCache", zap.Error(err))
-		return nil, err
-	}
-	if bindings == nil {
-		return nil, nil
+	select {
+	case result := <-resultChan:
+		if result.Err != nil {
+			return nil, result.Err
+		}
+		bindings := result.Val
+		if bindings == nil {
+			return nil, nil
+		}
+		return bindings.(Bindings), nil
+	case <-time.After(timeout):
+		return nil, errors.New("load bindings from storage timeout")
 	}
-	return bindings.(Bindings), err
 }
 
 func (h *globalBindingHandle) loadBindingsFromStorageInternal(sqlDigest string) (any, error) {
+	failpoint.Inject("load_bindings_from_storage_internal_timeout", func() {
+		time.Sleep(time.Second)
+	})
+
 	var bindings Bindings
+	selectStmt := fmt.Sprintf("SELECT original_sql, bind_sql, default_db, status, create_time, update_time, charset, collation, source, sql_digest, plan_digest FROM mysql.bind_info where sql_digest = '%s'", sqlDigest)
 	err := h.callWithSCtx(false, func(sctx sessionctx.Context) error {
-		rows, _, err := execRows(sctx, "SELECT original_sql, bind_sql, default_db, status, create_time, update_time, charset, collation, source, sql_digest, plan_digest FROM mysql.bind_info where sql_digest = ?", sqlDigest)
+		rows, _, err := execRows(sctx, selectStmt)
 		if err != nil {
 			return err
 		}

pkg/bindinfo/global_handle_test.go

@@ -485,7 +485,6 @@ func TestGlobalBinding(t *testing.T) {
 		_, fuzzyDigest = norm.NormalizeStmtForBinding(stmt, norm.WithFuzz(true))
 		_, matched = dom.BindHandle().MatchGlobalBinding(tk.Session(), fuzzyDigest, bindinfo.CollectTableNames(stmt))
 		require.False(t, matched) // dropped
-
 		bindHandle = bindinfo.NewGlobalBindingHandle(&mockSessionPool{tk.Session()})
 		err = bindHandle.LoadFromStorageToCache(true)
 		require.NoError(t, err)

pkg/bindinfo/session_handle.go

@@ -100,7 +100,8 @@ func (h *sessionBindingHandle) DropSessionBinding(sqlDigest string) error {
 
 // MatchSessionBinding returns the matched binding for this statement.
 func (h *sessionBindingHandle) MatchSessionBinding(sctx sessionctx.Context, fuzzyDigest string, tableNames []*ast.TableName) (matchedBinding Binding, isMatched bool) {
-	return h.ch.FuzzyMatchingBinding(sctx, fuzzyDigest, tableNames)
+	matchedBinding, isMatched = h.ch.FuzzyMatchingBinding(sctx, fuzzyDigest, tableNames)
+	return
 }
 
 // GetAllSessionBindings return all session bind info.

pkg/bindinfo/tests/BUILD.bazel

@@ -23,7 +23,6 @@ go_test(
         "//pkg/util",
         "//pkg/util/parser",
         "//pkg/util/stmtsummary",
-        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
         "@org_uber_go_goleak//:goleak",
     ],

pkg/bindinfo/tests/bind_test.go

@@ -20,7 +20,6 @@ import (
 	"strconv"
 	"testing"
 
-	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/bindinfo"
 	"github.com/pingcap/tidb/pkg/bindinfo/internal"
 	"github.com/pingcap/tidb/pkg/bindinfo/norm"
@@ -1099,13 +1098,15 @@ func TestFuzzyBindingHintsWithSourceReturning(t *testing.T) {
 		for _, currentDB := range []string{"db1", "db2", "db3"} {
 			tk.MustExec(`use ` + currentDB)
 			for _, db := range []string{"db1.", "db2.", "db3.", ""} {
-				require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/bindinfo/get_binding_return_nil", `return()`))
 				query := fmt.Sprintf(c.qTemplate, db)
 				tk.MustExec(query)
-				require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/bindinfo/get_binding_return_nil"))
 				tk.MustQuery(`show warnings`).Check(testkit.Rows()) // no warning
+				sctx := tk.Session()
+				sctx.SetValue(bindinfo.GetBindingReturnNil, true)
 				tk.MustExec(query)
+				sctx.ClearValue(bindinfo.GetBindingReturnNil)
 				tk.MustQuery(`select @@last_plan_from_binding`).Check(testkit.Rows("1"))
+				bindinfo.GetBindingReturnNilBool.Store(false)
 			}
 		}
 	}

pkg/bindinfo/tests/timeout/BUILD.bazel

@@ -0,0 +1,17 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_test")
+
+go_test(
+    name = "timeout_test",
+    timeout = "short",
+    srcs = [
+        "main_test.go",
+        "timeout_test.go",
+    ],
+    flaky = True,
+    deps = [
+        "//pkg/bindinfo",
+        "//pkg/testkit",
+        "//pkg/testkit/testsetup",
+        "@org_uber_go_goleak//:goleak",
+    ],
+)

pkg/bindinfo/tests/timeout/main_test.go

@@ -0,0 +1,35 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// 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 timeout
+
+import (
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/testkit/testsetup"
+	"go.uber.org/goleak"
+)
+
+func TestMain(m *testing.M) {
+	testsetup.SetupForCommonTest()
+	opts := []goleak.Option{
+		goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"),
+		goleak.IgnoreTopFunction("github.com/bazelbuild/rules_go/go/tools/bzltestutil.RegisterTimeoutHandler.func1"),
+		goleak.IgnoreTopFunction("github.com/lestrrat-go/httprc.runFetchWorker"),
+		goleak.IgnoreTopFunction("go.etcd.io/etcd/client/pkg/v3/logutil.(*MergeLogger).outputLoop"),
+		goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
+		goleak.IgnoreTopFunction("time.Sleep"),
+	}
+	goleak.VerifyTestMain(m, opts...)
+}

pkg/bindinfo/tests/timeout/timeout_test.go

@@ -0,0 +1,67 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// 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 timeout
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/bindinfo"
+	"github.com/pingcap/tidb/pkg/testkit"
+)
+
+func TestFuzzyBindingHintsWithSourceReturningTimeout(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec(`use test`)
+
+	for _, db := range []string{"db1", "db2", "db3"} {
+		tk.MustExec(`create database ` + db)
+		tk.MustExec(`use ` + db)
+		tk.MustExec(`create table t1 (a int, b int, c int, d int, key(a), key(b), key(c), key(d))`)
+		tk.MustExec(`create table t2 (a int, b int, c int, d int, key(a), key(b), key(c), key(d))`)
+		tk.MustExec(`create table t3 (a int, b int, c int, d int, key(a), key(b), key(c), key(d))`)
+	}
+	tk.MustExec(`set @@tidb_opt_enable_fuzzy_binding=1`)
+
+	for _, c := range []struct {
+		binding   string
+		qTemplate string
+	}{
+		// use index
+		{`create global binding using select /*+ use_index(t1, c) */ * from *.t1 where a=1`,
+			`select * from %st1 where a=1000`},
+	} {
+		tk.MustExec(c.binding)
+		for _, currentDB := range []string{"db1", "db2", "db3"} {
+			tk.MustExec(`use ` + currentDB)
+			for _, db := range []string{"db1.", "db2.", "db3.", ""} {
+				query := fmt.Sprintf(c.qTemplate, db)
+				tk.MustExec(query)
+				tk.MustQuery(`show warnings`).Check(testkit.Rows()) // no warning
+				sctx := tk.Session()
+				sctx.SetValue(bindinfo.GetBindingReturnNil, true)
+				sctx.SetValue(bindinfo.GetBindingReturnNilAlways, true)
+				sctx.SetValue(bindinfo.LoadBindingNothing, true)
+				tk.MustExec(query)
+				sctx.ClearValue(bindinfo.GetBindingReturnNil)
+				sctx.ClearValue(bindinfo.GetBindingReturnNilAlways)
+				sctx.ClearValue(bindinfo.LoadBindingNothing)
+				tk.MustQuery(`select @@last_plan_from_binding`).Check(testkit.Rows("0"))
+				bindinfo.GetBindingReturnNilBool.Store(false)
+			}
+		}
+	}
+}

pkg/sessionctx/variable/session.go

@@ -1067,6 +1067,9 @@ type SessionVars struct {
 	// See https://dev.mysql.com/doc/refman/5.7/en/server-system-variables.html#sysvar_max_execution_time
 	MaxExecutionTime uint64
 
+	// LoadBindingTimeout is the timeout for loading the bind info.
+	LoadBindingTimeout uint64
+
 	// TiKVClientReadTimeout is the timeout for readonly kv request in milliseconds, 0 means using default value
 	// See https://github.com/pingcap/tidb/blob/7105505a78fc886c33258caa5813baf197b15247/docs/design/2023-06-30-configurable-kv-timeout.md?plain=1#L14-L15
 	TiKVClientReadTimeout uint64