-
Notifications
You must be signed in to change notification settings - Fork 5.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add a function to detach the TableReaderExecutor
Signed-off-by: Yang Keao <yangkeao@chunibyo.icu>
- Loading branch information
Showing
19 changed files
with
512 additions
and
128 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| // 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 executor | ||
|
|
||
| import ( | ||
| "github.com/pingcap/tidb/pkg/executor/internal/exec" | ||
| "github.com/pingcap/tidb/pkg/expression/contextsession" | ||
| ) | ||
|
|
||
| // Detach detaches the current executor from the session context. After detaching, the session context | ||
| // can be used to execute another statement while this executor is still running. The returning value | ||
| // shows whether this executor is able to be detached. | ||
| // | ||
| // NOTE: the implementation of `Detach` should guarantee that no matter whether it returns true or false, | ||
| // both the original executor and the returning executor should be able to be used correctly. This restriction | ||
| // is to make sure that if `Detach(a)` returns `true`, while other children of `a`'s parent returns `false`, | ||
| // the caller can still use the original one. | ||
| func Detach(originalExecutor exec.Executor) (exec.Executor, bool) { | ||
| newExecutor, ok := originalExecutor.Detach() | ||
| if !ok { | ||
| return nil, false | ||
| } | ||
|
|
||
| children := originalExecutor.AllChildren() | ||
| newChildren := make([]exec.Executor, len(children)) | ||
| for i, child := range children { | ||
| detached, ok := Detach(child) | ||
| if !ok { | ||
| return nil, false | ||
| } | ||
| newChildren[i] = detached | ||
| } | ||
| copy(newExecutor.AllChildren(), newChildren) | ||
|
|
||
| return newExecutor, true | ||
| } | ||
|
|
||
| func (treCtx tableReaderExecutorContext) Detach() tableReaderExecutorContext { | ||
| newCtx := treCtx | ||
|
|
||
| if ctx, ok := treCtx.ectx.(*contextsession.SessionExprContext); ok { | ||
| staticExprCtx := ctx.IntoStatic() | ||
|
|
||
| newCtx.dctx = newCtx.dctx.Detach() | ||
| newCtx.rctx = newCtx.rctx.Detach(staticExprCtx) | ||
| newCtx.buildPBCtx = newCtx.buildPBCtx.Detach(staticExprCtx) | ||
| newCtx.ectx = staticExprCtx | ||
| return newCtx | ||
| } | ||
|
|
||
| return treCtx | ||
| } | ||
|
|
||
| // Detach detaches the current executor from the session context. | ||
| func (e *TableReaderExecutor) Detach() (exec.Executor, bool) { | ||
| newExec := new(TableReaderExecutor) | ||
| *newExec = *e | ||
|
|
||
| newExec.tableReaderExecutorContext = newExec.tableReaderExecutorContext.Detach() | ||
|
|
||
| return newExec, true | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,174 @@ | ||
| // 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 executor_test | ||
|
|
||
| import ( | ||
| "context" | ||
| "strconv" | ||
| "sync" | ||
| "sync/atomic" | ||
| "testing" | ||
|
|
||
| "github.com/pingcap/tidb/pkg/executor/internal/exec" | ||
| "github.com/pingcap/tidb/pkg/parser/mysql" | ||
| "github.com/pingcap/tidb/pkg/testkit" | ||
| "github.com/pingcap/tidb/pkg/util/sqlexec" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| type exportExecutor interface { | ||
| GetExecutor4Test() any | ||
| } | ||
|
|
||
| func TestDetachAllContexts(t *testing.T) { | ||
| store := testkit.CreateMockStore(t) | ||
| tk := testkit.NewTestKit(t, store) | ||
|
|
||
| tk.MustExec("use test") | ||
| tk.Session().GetSessionVars().SetStatusFlag(mysql.ServerStatusCursorExists, true) | ||
| tk.MustExec("create table t (a int)") | ||
| tk.MustExec("insert into t values (1), (2), (3)") | ||
|
|
||
| rs, err := tk.Exec("select * from t") | ||
| require.NoError(t, err) | ||
| oldExecutor := rs.(exportExecutor).GetExecutor4Test().(exec.Executor) | ||
|
|
||
| drs := rs.(sqlexec.DetachableRecordSet) | ||
| srs, ok, err := drs.TryDetach() | ||
| require.True(t, ok) | ||
| require.NoError(t, err) | ||
|
|
||
| require.NotEqual(t, rs, srs) | ||
| newExecutor := srs.(exportExecutor).GetExecutor4Test().(exec.Executor) | ||
|
|
||
| require.NotEqual(t, oldExecutor, newExecutor) | ||
| // Children should be different | ||
| for i, child := range oldExecutor.AllChildren() { | ||
| require.NotEqual(t, child, newExecutor.AllChildren()[i]) | ||
| } | ||
|
|
||
| // Then execute another statement | ||
| tk.MustQuery("select * from t limit 1").Check(testkit.Rows("1")) | ||
| // The previous detached record set can still be used | ||
| // check data | ||
| chk := srs.NewChunk(nil) | ||
| err = srs.Next(context.Background(), chk) | ||
| require.NoError(t, err) | ||
| require.Equal(t, 3, chk.NumRows()) | ||
| require.Equal(t, int64(1), chk.GetRow(0).GetInt64(0)) | ||
| require.Equal(t, int64(2), chk.GetRow(1).GetInt64(0)) | ||
| require.Equal(t, int64(3), chk.GetRow(2).GetInt64(0)) | ||
| } | ||
|
|
||
| func TestAfterDetachSessionCanExecute(t *testing.T) { | ||
| // This test shows that the session can be safely used to execute another statement after detaching. | ||
| store := testkit.CreateMockStore(t) | ||
| tk := testkit.NewTestKit(t, store) | ||
|
|
||
| tk.MustExec("use test") | ||
| tk.Session().GetSessionVars().SetStatusFlag(mysql.ServerStatusCursorExists, true) | ||
| tk.MustExec("create table t (a int)") | ||
| for i := 0; i < 10000; i++ { | ||
| tk.MustExec("insert into t values (?)", i) | ||
| } | ||
|
|
||
| rs, err := tk.Exec("select * from t") | ||
| require.NoError(t, err) | ||
| drs, ok, err := rs.(sqlexec.DetachableRecordSet).TryDetach() | ||
| require.NoError(t, err) | ||
| require.True(t, ok) | ||
|
|
||
| // Now, the `drs` can be used concurrently with the session. | ||
| var wg sync.WaitGroup | ||
| var stop atomic.Bool | ||
| wg.Add(1) | ||
| go func() { | ||
| defer wg.Done() | ||
|
|
||
| for i := 0; i < 10000; i++ { | ||
| if stop.Load() { | ||
| return | ||
| } | ||
| tk.MustQuery("select * from t where a = ?", i).Check(testkit.Rows(strconv.Itoa(i))) | ||
| } | ||
| }() | ||
|
|
||
| chk := drs.NewChunk(nil) | ||
| expectedSelect := 0 | ||
| for { | ||
| err = drs.Next(context.Background(), chk) | ||
| require.NoError(t, err) | ||
|
|
||
| if chk.NumRows() == 0 { | ||
| break | ||
| } | ||
| for i := 0; i < chk.NumRows(); i++ { | ||
| require.Equal(t, int64(expectedSelect), chk.GetRow(i).GetInt64(0)) | ||
| expectedSelect++ | ||
| } | ||
| } | ||
| stop.Store(true) | ||
| wg.Wait() | ||
| } | ||
|
|
||
| func TestDetachWithParam(t *testing.T) { | ||
| store := testkit.CreateMockStore(t) | ||
| tk := testkit.NewTestKit(t, store) | ||
|
|
||
| tk.MustExec("use test") | ||
| tk.Session().GetSessionVars().SetStatusFlag(mysql.ServerStatusCursorExists, true) | ||
| tk.MustExec("create table t (a int primary key)") | ||
| for i := 0; i < 10000; i++ { | ||
| tk.MustExec("insert into t values (?)", i) | ||
| } | ||
|
|
||
| rs, err := tk.Exec("select * from t where a > ? and a < ?", 100, 200) | ||
| require.NoError(t, err) | ||
| drs, ok, err := rs.(sqlexec.DetachableRecordSet).TryDetach() | ||
| require.NoError(t, err) | ||
| require.True(t, ok) | ||
|
|
||
| // Now, execute another statement with different size of param. It'll not affect the execution of detached executor. | ||
| var wg sync.WaitGroup | ||
| var stop atomic.Bool | ||
| wg.Add(1) | ||
| go func() { | ||
| defer wg.Done() | ||
|
|
||
| for i := 0; i < 10000; i++ { | ||
| if stop.Load() { | ||
| return | ||
| } | ||
| tk.MustQuery("select * from t where a = ?", i).Check(testkit.Rows(strconv.Itoa(i))) | ||
| } | ||
| }() | ||
|
|
||
| chk := drs.NewChunk(nil) | ||
| expectedSelect := 101 | ||
| for { | ||
| err = drs.Next(context.Background(), chk) | ||
| require.NoError(t, err) | ||
|
|
||
| if chk.NumRows() == 0 { | ||
| break | ||
| } | ||
| for i := 0; i < chk.NumRows(); i++ { | ||
| require.Equal(t, int64(expectedSelect), chk.GetRow(i).GetInt64(0)) | ||
| expectedSelect++ | ||
| } | ||
| } | ||
| stop.Store(true) | ||
| wg.Wait() | ||
| } |
Oops, something went wrong.