Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@systay
systay committed Feb 18, 2025
1 parent 035eac2 commit 5bdb210
Show file tree
Hide file tree
Showing 6 changed files with 132 additions and 125 deletions.
2 changes: 1 addition & 1 deletion go/vt/vtgate/engine/plan.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
17 changes: 17 additions & 0 deletions go/vt/vtgate/engine/plankey_test.go
View file
Original file line number Diff line number Diff line change
@@ -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
55 changes: 31 additions & 24 deletions go/vt/vtgate/executor.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1172,48 +1172,55 @@ 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
}
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(
Expand Down
77 changes: 75 additions & 2 deletions go/vt/vtgate/executor_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand Down
39 changes: 4 additions & 35 deletions go/vt/vtgate/executorcontext/vcursor_impl.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ package executorcontext
import (
"context"
"fmt"
"io"
"sort"
"strings"
"sync"
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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
}
Expand Down
67 changes: 4 additions & 63 deletions go/vt/vtgate/executorcontext/vcursor_impl_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,15 +21,13 @@ import (
"errors"
"fmt"
"strconv"
"strings"
"testing"
"time"

"github.com/stretchr/testify/require"

"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"
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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"}}
Expand Down

0 comments on commit 5bdb210

Please sign in to comment.