diff --git a/go/vt/vtgate/engine/plan.go b/go/vt/vtgate/engine/plan.go index 18aee6fff6d..a0488e591ca 100644 --- a/go/vt/vtgate/engine/plan.go +++ b/go/vt/vtgate/engine/plan.go @@ -61,7 +61,7 @@ type PlanKey struct { } func (pk PlanKey) DebugString() string { - return fmt.Sprintf("CurrentKeyspace: %s, Query: %s, SetVarComment: %s, Collation: %d", pk.CurrentKeyspace, pk.Query, pk.SetVarComment, pk.Collation) + return fmt.Sprintf("CurrentKeyspace: %s, Destination: %s, Query: %s, SetVarComment: %s, Collation: %d", pk.CurrentKeyspace, pk.Destination, pk.Query, pk.SetVarComment, pk.Collation) } func (pk PlanKey) Hash() theine.HashKey256 { diff --git a/go/vt/vtgate/engine/plankey_test.go b/go/vt/vtgate/engine/plankey_test.go new file mode 100644 index 00000000000..bbab7b42d89 --- /dev/null +++ b/go/vt/vtgate/engine/plankey_test.go @@ -0,0 +1,17 @@ +/* +Copyright 2025 The Vitess Authors. + +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 engine diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go index 0b96c01324f..5c0c37b2436 100644 --- a/go/vt/vtgate/executor.go +++ b/go/vt/vtgate/executor.go @@ -1172,11 +1172,11 @@ func (e *Executor) getCachedOrBuild( planCachable := sqlparser.CachePlan(stmt) && vcursor.CachePlan() if planCachable { // build Plan key - planKey := e.hashPlan(ctx, vcursor, query, setVarComment) + planKey := createPlanKey(ctx, vcursor, query, setVarComment) var plan *engine.Plan var err error - plan, logStats.CachedPlan, err = e.plans.GetOrLoad(planKey, e.epoch.Load(), func() (*engine.Plan, error) { + plan, logStats.CachedPlan, err = e.plans.GetOrLoad(planKey.Hash(), e.epoch.Load(), func() (*engine.Plan, error) { return e.buildStatement(ctx, vcursor, query, stmt, reservedVars, bindVarNeeds, qh) }) return plan, err @@ -1184,36 +1184,43 @@ func (e *Executor) getCachedOrBuild( return e.buildStatement(ctx, vcursor, query, stmt, reservedVars, bindVarNeeds, qh) } -func (e *Executor) hashPlan(ctx context.Context, vcursor *econtext.VCursorImpl, query string, setVarComment string) theine.HashKey256 { - var allDest []string - currDest := vcursor.Destination() - if currDest != nil { - switch currDest.(type) { - case key.DestinationKeyspaceID, key.DestinationKeyspaceIDs: - resolved, _, err := vcursor.ResolveDestinations(ctx, vcursor.GetKeyspace(), nil, []key.Destination{currDest}) - if err == nil && len(resolved) > 0 { - shards := make([]string, len(resolved)) - for i := 0; i < len(shards); i++ { - shards[i] = resolved[i].Target.GetShard() - } - sort.Strings(shards) - allDest = shards - } - default: - allDest = []string{currDest.String()} - } - } +func createPlanKey(ctx context.Context, vcursor *econtext.VCursorImpl, query string, setVarComment string) engine.PlanKey { + allDest := getDestinations(ctx, vcursor) - pk := engine.PlanKey{ + return engine.PlanKey{ CurrentKeyspace: vcursor.GetKeyspace(), Destination: strings.Join(allDest, ","), Query: query, SetVarComment: setVarComment, Collation: vcursor.ConnCollation(), } +} + +func getDestinations(ctx context.Context, vcursor *econtext.VCursorImpl) []string { + currDest := vcursor.Destination() + if currDest == nil { + return nil + } + + switch currDest.(type) { + case key.DestinationKeyspaceID, key.DestinationKeyspaceIDs: + // these need to be resolved to shards + default: + return []string{currDest.String()} + } + + resolved, _, err := vcursor.ResolveDestinations(ctx, vcursor.GetKeyspace(), nil, []key.Destination{currDest}) + if err != nil || len(resolved) <= 0 { + return nil + } + + shards := make([]string, len(resolved)) + for i := 0; i < len(shards); i++ { + shards[i] = resolved[i].Target.GetShard() + } + sort.Strings(shards) - planKey := pk.Hash() - return planKey + return shards } func (e *Executor) buildStatement( diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index 6c499b5230b..b519a91a1e4 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -61,6 +61,79 @@ import ( "vitess.io/vitess/go/vt/vtgate/vtgateservice" ) +func TestPlanKey(t *testing.T) { + ks1 := &vindexes.Keyspace{Name: "ks1"} + ks1Schema := &vindexes.KeyspaceSchema{Keyspace: ks1} + vschemaWith1KS := &vindexes.VSchema{ + Keyspaces: map[string]*vindexes.KeyspaceSchema{ + ks1.Name: ks1Schema, + }, + } + + type testCase struct { + vschema *vindexes.VSchema + targetString string + expectedPlanPrefixKey string + } + + tests := []testCase{{ + vschema: vschemaWith1KS, + targetString: "", + expectedPlanPrefixKey: "CurrentKeyspace: ks1, Destination: , Query: SELECT 1, SetVarComment: , Collation: 255", + }, { + vschema: vschemaWith1KS, + targetString: "ks1@replica", + expectedPlanPrefixKey: "CurrentKeyspace: ks1, Destination: , Query: SELECT 1, SetVarComment: , Collation: 255", + }, { + vschema: vschemaWith1KS, + targetString: "ks1:-80", + expectedPlanPrefixKey: "CurrentKeyspace: ks1, Destination: , Query: SELECT 1, SetVarComment: , Collation: 255", + }, { + vschema: vschemaWith1KS, + targetString: "ks1[deadbeef]", + expectedPlanPrefixKey: "CurrentKeyspace: ks1, Destination: , Query: SELECT 1, SetVarComment: , Collation: 255", + }, { + vschema: vschemaWith1KS, + targetString: "", + expectedPlanPrefixKey: "CurrentKeyspace: ks1, Destination: , Query: SELECT 1, SetVarComment: , Collation: 255", + }, { + vschema: vschemaWith1KS, + targetString: "ks1@replica", + expectedPlanPrefixKey: "CurrentKeyspace: ks1, Destination: , Query: SELECT 1, SetVarComment: , Collation: 255", + }} + e, _, _, _, ctx := createExecutorEnv(t) + e.vschema = vschemaWith1KS + cfg := econtext.VCursorConfig{ + Collation: collations.CollationUtf8mb4ID, + DefaultTabletType: topodatapb.TabletType_PRIMARY, + } + + ss := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}) + vc, _ := econtext.NewVCursorImpl(ss, makeComments(""), e, nil, e.vm, e.VSchema(), e.resolver.resolver, nil, nullResultsObserver{}, cfg) + + for i, tc := range tests { + t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { + ss.SetTargetString(tc.targetString) + key := createPlanKey(ctx, vc, "SELECT 1", "") + require.Equal(t, tc.expectedPlanPrefixKey, key.DebugString(), "test case %d", i) + }) + + // t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { + // ss := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: false}) + // ss.SetTargetString(tc.targetString) + // // e := NewExecutor(context.Background(), nil, nil, "", nil, ExecutorConfig{}, false, nil, nil, plancontext.PlannerVersion(0), nil) + // // e.g + // // vc, err := econtext.NewVCursorImpl(ss, sqlparser.MarginComments{}, nil, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, srvtopo.NewResolver(&FakeTopoServer{}, nil, ""), nil, fakeObserver{}, cfg) + // // require.NoError(t, err) + // // vc.vschema = tc.vschema + // // + // // var buf strings.Builder + // // vc.KeyForPlan(context.Background(), "SELECT 1", &buf) + // require.Equal(t, tc.expectedPlanPrefixKey, buf.String()) + // }) + } +} + func TestExecutorResultsExceeded(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) @@ -1591,8 +1664,8 @@ func assertCacheContains(t *testing.T, e *Executor, vc *econtext.VCursorImpl, sq return true }) } else { - h := e.hashPlan(context.Background(), vc, sql, "") - plan, _ = e.plans.Get(h, e.epoch.Load()) + h := createPlanKey(context.Background(), vc, sql, "") + plan, _ = e.plans.Get(h.Hash(), e.epoch.Load()) } require.Truef(t, plan != nil, "plan not found for query: %s", sql) return plan diff --git a/go/vt/vtgate/executorcontext/vcursor_impl.go b/go/vt/vtgate/executorcontext/vcursor_impl.go index e8af00d7358..e899f700293 100644 --- a/go/vt/vtgate/executorcontext/vcursor_impl.go +++ b/go/vt/vtgate/executorcontext/vcursor_impl.go @@ -19,7 +19,6 @@ package executorcontext import ( "context" "fmt" - "io" "sort" "strings" "sync" @@ -215,6 +214,10 @@ func NewVCursorImpl( }, nil } +func (vc *VCursorImpl) GetSafeSession() *SafeSession { + return vc.SafeSession +} + func (vc *VCursorImpl) PrepareSetVarComment() string { var res []string vc.Session().GetSystemVariables(func(k, v string) { @@ -1368,40 +1371,6 @@ func (vc *VCursorImpl) FindMirrorRule(name sqlparser.TableName) (*vindexes.Mirro return mirrorRule, err } -func (vc *VCursorImpl) KeyForPlan(ctx context.Context, query string, buf io.StringWriter) { - _, _ = buf.WriteString(vc.keyspace) - _, _ = buf.WriteString(vindexes.TabletTypeSuffix[vc.tabletType]) - _, _ = buf.WriteString("+Collate:") - _, _ = buf.WriteString(vc.Environment().CollationEnv().LookupName(vc.config.Collation)) - - if vc.destination != nil { - switch vc.destination.(type) { - case key.DestinationKeyspaceID, key.DestinationKeyspaceIDs: - resolved, _, err := vc.ResolveDestinations(ctx, vc.keyspace, nil, []key.Destination{vc.destination}) - if err == nil && len(resolved) > 0 { - shards := make([]string, len(resolved)) - for i := 0; i < len(shards); i++ { - shards[i] = resolved[i].Target.GetShard() - } - sort.Strings(shards) - - _, _ = buf.WriteString("+KsIDsResolved:") - for i, s := range shards { - if i > 0 { - _, _ = buf.WriteString(",") - } - _, _ = buf.WriteString(s) - } - } - default: - _, _ = buf.WriteString("+") - _, _ = buf.WriteString(vc.destination.String()) - } - } - _, _ = buf.WriteString("+Query:") - _, _ = buf.WriteString(query) -} - func (vc *VCursorImpl) GetKeyspace() string { return vc.keyspace } diff --git a/go/vt/vtgate/executorcontext/vcursor_impl_test.go b/go/vt/vtgate/executorcontext/vcursor_impl_test.go index 54173bf63b0..748a77ffb9b 100644 --- a/go/vt/vtgate/executorcontext/vcursor_impl_test.go +++ b/go/vt/vtgate/executorcontext/vcursor_impl_test.go @@ -21,7 +21,6 @@ import ( "errors" "fmt" "strconv" - "strings" "testing" "time" @@ -29,7 +28,6 @@ import ( "vitess.io/vitess/go/streamlog" - "vitess.io/vitess/go/mysql/collations" "vitess.io/vitess/go/sqltypes" binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata" "vitess.io/vitess/go/vt/proto/vschema" @@ -193,15 +191,10 @@ func TestDestinationKeyspace(t *testing.T) { } var ( - ks1 = &vindexes.Keyspace{Name: "ks1"} - ks1Schema = &vindexes.KeyspaceSchema{Keyspace: ks1} - ks2 = &vindexes.Keyspace{Name: "ks2"} - ks2Schema = &vindexes.KeyspaceSchema{Keyspace: ks2} - vschemaWith1KS = &vindexes.VSchema{ - Keyspaces: map[string]*vindexes.KeyspaceSchema{ - ks1.Name: ks1Schema, - }, - } + ks1 = &vindexes.Keyspace{Name: "ks1"} + ks1Schema = &vindexes.KeyspaceSchema{Keyspace: ks1} + ks2 = &vindexes.Keyspace{Name: "ks2"} + ks2Schema = &vindexes.KeyspaceSchema{Keyspace: ks2} ) var vschemaWith2KS = &vindexes.VSchema{ @@ -253,58 +246,6 @@ func TestSetTarget(t *testing.T) { } } -func TestKeyForPlan(t *testing.T) { - type testCase struct { - vschema *vindexes.VSchema - targetString string - expectedPlanPrefixKey string - } - - tests := []testCase{{ - vschema: vschemaWith1KS, - targetString: "", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1@replica", - expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1:-80", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+DestinationShard(-80)+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1[deadbeef]", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+KsIDsResolved:80-+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1@replica", - expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }} - - for i, tc := range tests { - t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { - ss := NewSafeSession(&vtgatepb.Session{InTransaction: false}) - ss.SetTargetString(tc.targetString) - cfg := VCursorConfig{ - Collation: collations.CollationUtf8mb4ID, - DefaultTabletType: topodatapb.TabletType_PRIMARY, - } - vc, err := NewVCursorImpl(ss, sqlparser.MarginComments{}, &fakeExecutor{}, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, srvtopo.NewResolver(&FakeTopoServer{}, nil, ""), nil, fakeObserver{}, cfg) - require.NoError(t, err) - vc.vschema = tc.vschema - - var buf strings.Builder - vc.KeyForPlan(context.Background(), "SELECT 1", &buf) - require.Equal(t, tc.expectedPlanPrefixKey, buf.String()) - }) - } -} - func TestFirstSortedKeyspace(t *testing.T) { ks1Schema := &vindexes.KeyspaceSchema{Keyspace: &vindexes.Keyspace{Name: "xks1"}} ks2Schema := &vindexes.KeyspaceSchema{Keyspace: &vindexes.Keyspace{Name: "aks2"}}