From cc61c88605964c9a524622c5e3d6b25243a733e2 Mon Sep 17 00:00:00 2001 From: YangKeao Date: Sun, 4 Feb 2024 22:35:42 +0800 Subject: [PATCH] This is an automated cherry-pick of #50969 Signed-off-by: ti-chi-bot --- pkg/server/tests/cursor/cursor_test.go | 159 +++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 pkg/server/tests/cursor/cursor_test.go diff --git a/pkg/server/tests/cursor/cursor_test.go b/pkg/server/tests/cursor/cursor_test.go new file mode 100644 index 0000000000000..12991ad6ed343 --- /dev/null +++ b/pkg/server/tests/cursor/cursor_test.go @@ -0,0 +1,159 @@ +// 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 cursor + +import ( + "context" + "encoding/binary" + "fmt" + "math/rand" + "testing" + + "github.com/pingcap/failpoint" + "github.com/pingcap/tidb/pkg/config" + tmysql "github.com/pingcap/tidb/pkg/parser/mysql" + server2 "github.com/pingcap/tidb/pkg/server" + "github.com/pingcap/tidb/pkg/testkit" + "github.com/stretchr/testify/require" +) + +func TestCursorFetchErrorInFetch(t *testing.T) { + tmpStoragePath := t.TempDir() + restore := config.RestoreFunc() + defer restore() + config.UpdateGlobal(func(conf *config.Config) { + conf.TempStoragePath = tmpStoragePath + }) + + store, dom := testkit.CreateMockStoreAndDomain(t) + srv := server2.CreateMockServer(t, store) + srv.SetDomain(dom) + defer srv.Close() + + appendUint32 := binary.LittleEndian.AppendUint32 + ctx := context.Background() + c := server2.CreateMockConn(t, srv) + + tk := testkit.NewTestKitWithSession(t, store, c.Context().Session) + tk.MustExec("use test") + tk.MustExec("drop table if exists t") + tk.MustExec("create table t(id int, payload BLOB)") + payload := make([]byte, 512) + for i := 0; i < 2048; i++ { + rand.Read(payload) + tk.MustExec("insert into t values (?, ?)", i, payload) + } + + tk.MustExec("set global tidb_enable_tmp_storage_on_oom = ON") + tk.MustExec("set global tidb_mem_oom_action = 'CANCEL'") + defer tk.MustExec("set global tidb_mem_oom_action= DEFAULT") + require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 1) + + // execute a normal statement, it'll spill to disk + stmt, _, _, err := c.Context().Prepare("select * from t") + require.NoError(t, err) + + tk.MustExec(fmt.Sprintf("set tidb_mem_quota_query=%d", 1)) + + require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/util/chunk/get-chunk-error", "return(true)")) + defer func() { + require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/util/chunk/get-chunk-error")) + }() + + require.NoError(t, c.Dispatch(ctx, append( + appendUint32([]byte{tmysql.ComStmtExecute}, uint32(stmt.ID())), + tmysql.CursorTypeReadOnly, 0x1, 0x0, 0x0, 0x0, + ))) + + require.ErrorContains(t, c.Dispatch(ctx, appendUint32(appendUint32([]byte{tmysql.ComStmtFetch}, uint32(stmt.ID())), 1024)), "fail to get chunk for test") + // after getting a failed FETCH, the cursor should have been reseted + require.False(t, stmt.GetCursorActive()) + require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 0) + require.Len(t, tk.Session().GetSessionVars().DiskTracker.GetChildrenForTest(), 0) +} + +func TestCursorFetchShouldSpill(t *testing.T) { + restore := config.RestoreFunc() + defer restore() + config.UpdateGlobal(func(conf *config.Config) { + conf.TempStoragePath = t.TempDir() + }) + + store, dom := testkit.CreateMockStoreAndDomain(t) + srv := server2.CreateMockServer(t, store) + srv.SetDomain(dom) + defer srv.Close() + + require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/server/testCursorFetchSpill", "return(true)")) + defer func() { + require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/server/testCursorFetchSpill")) + }() + + appendUint32 := binary.LittleEndian.AppendUint32 + ctx := context.Background() + c := server2.CreateMockConn(t, srv) + + tk := testkit.NewTestKitWithSession(t, store, c.Context().Session) + tk.MustExec("use test") + tk.MustExec("drop table if exists t") + tk.MustExec("create table t(id_1 int, id_2 int)") + tk.MustExec("insert into t values (1, 1), (1, 2)") + tk.MustExec("set global tidb_enable_tmp_storage_on_oom = ON") + tk.MustExec("set global tidb_mem_oom_action = 'CANCEL'") + defer tk.MustExec("set global tidb_mem_oom_action= DEFAULT") + require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 1) + + // execute a normal statement, it'll spill to disk + stmt, _, _, err := c.Context().Prepare("select * from t") + require.NoError(t, err) + + tk.MustExec(fmt.Sprintf("set tidb_mem_quota_query=%d", 1)) + + require.NoError(t, c.Dispatch(ctx, append( + appendUint32([]byte{tmysql.ComStmtExecute}, uint32(stmt.ID())), + tmysql.CursorTypeReadOnly, 0x1, 0x0, 0x0, 0x0, + ))) +} + +func TestCursorFetchExecuteCheck(t *testing.T) { + store, dom := testkit.CreateMockStoreAndDomain(t) + srv := server2.CreateMockServer(t, store) + srv.SetDomain(dom) + defer srv.Close() + + appendUint32 := binary.LittleEndian.AppendUint32 + ctx := context.Background() + c := server2.CreateMockConn(t, srv) + + stmt, _, _, err := c.Context().Prepare("select 1") + require.NoError(t, err) + + // execute with wrong ID + require.Error(t, c.Dispatch(ctx, append( + appendUint32([]byte{tmysql.ComStmtExecute}, uint32(stmt.ID()+1)), + tmysql.CursorTypeReadOnly, 0x1, 0x0, 0x0, 0x0, + ))) + + // execute with wrong flag + require.Error(t, c.Dispatch(ctx, append( + appendUint32([]byte{tmysql.ComStmtExecute}, uint32(stmt.ID())), + tmysql.CursorTypeReadOnly|tmysql.CursorTypeForUpdate, 0x1, 0x0, 0x0, 0x0, + ))) + + require.Error(t, c.Dispatch(ctx, append( + appendUint32([]byte{tmysql.ComStmtExecute}, uint32(stmt.ID())), + tmysql.CursorTypeReadOnly|tmysql.CursorTypeScrollable, 0x1, 0x0, 0x0, 0x0, + ))) +}