Count how many shards were hit by routes when using vexplain (#16802)

Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -132,7 +132,7 @@ func (t *noopVCursor) UnresolvedTransactions(ctx context.Context, keyspace strin
  panic("implement me")
 }
 
-func (t *noopVCursor) StartPrimitiveTrace() func() map[Primitive]RowsReceived {
+func (t *noopVCursor) StartPrimitiveTrace() func() Stats {
  panic("implement me")
 }
 

go/vt/vtgate/engine/plan_description.go

@@ -48,7 +48,8 @@ type PrimitiveDescription struct {
  InputName string
  Inputs []PrimitiveDescription
 
- Stats RowsReceived
+ RowsReceived RowsReceived
+ ShardsQueried *ShardsQueried
 }
 
 // MarshalJSON serializes the PlanDescription into a JSON representation.
@@ -92,15 +93,20 @@ func (pd PrimitiveDescription) MarshalJSON() ([]byte, error) {
  return nil, err
  }
  }
- if len(pd.Stats) > 0 {
- if err := marshalAdd(prepend, buf, "NoOfCalls", len(pd.Stats)); err != nil {
+ if len(pd.RowsReceived) > 0 {
+ if err := marshalAdd(prepend, buf, "NoOfCalls", len(pd.RowsReceived)); err != nil {
  return nil, err
  }
 
- if err := marshalAdd(prepend, buf, "AvgNumberOfRows", average(pd.Stats)); err != nil {
+ if err := marshalAdd(prepend, buf, "AvgNumberOfRows", average(pd.RowsReceived)); err != nil {
  return nil, err
  }
- if err := marshalAdd(prepend, buf, "MedianNumberOfRows", median(pd.Stats)); err != nil {
+ if err := marshalAdd(prepend, buf, "MedianNumberOfRows", median(pd.RowsReceived)); err != nil {
+ return nil, err
+ }
+ }
+ if pd.ShardsQueried != nil {
+ if err := marshalAdd(prepend, buf, "ShardsQueried", pd.ShardsQueried); err != nil {
  return nil, err
  }
  }
@@ -219,10 +225,16 @@ func marshalAdd(prepend string, buf *bytes.Buffer, name string, obj any) error {
 
 // PrimitiveToPlanDescription transforms a primitive tree into a corresponding PlanDescription tree
 // If stats is not nil, it will be used to populate the stats field of the PlanDescription
-func PrimitiveToPlanDescription(in Primitive, stats map[Primitive]RowsReceived) PrimitiveDescription {
+func PrimitiveToPlanDescription(in Primitive, stats *Stats) PrimitiveDescription {
  this := in.description()
  if stats != nil {
- this.Stats = stats[in]
+ this.RowsReceived = stats.InterOpStats[in]
+
+ // Only applies to Route primitive
+ v, ok := stats.ShardsStats[in]
+ if ok {
+ this.ShardsQueried = &v
+ }
  }
 
  inputs, infos := in.Inputs()

go/vt/vtgate/engine/primitive.go

@@ -143,7 +143,7 @@ type (
 
  // StartPrimitiveTrace starts a trace for the given primitive,
  // and returns a function to get the trace logs after the primitive execution.
- StartPrimitiveTrace() func() map[Primitive]RowsReceived
+ StartPrimitiveTrace() func() Stats
  }
 
  // SessionActions gives primitives ability to interact with the session state

go/vt/vtgate/engine/vexplain.go

@@ -44,7 +44,13 @@ type (
  Type sqlparser.VExplainType
  }
 
- RowsReceived []int
+ ShardsQueried int
+ RowsReceived []int
+
+ Stats struct {
+ InterOpStats map[Primitive]RowsReceived
+ ShardsStats map[Primitive]ShardsQueried
+ }
 )
 
 var _ Primitive = (*VExplain)(nil)
@@ -111,7 +117,7 @@ func (v *VExplain) NeedsTransaction() bool {
 
 // TryExecute implements the Primitive interface
 func (v *VExplain) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
- var stats func() map[Primitive]RowsReceived
+ var stats func() Stats
  if v.Type == sqlparser.TraceVExplainType {
  stats = vcursor.StartPrimitiveTrace()
  } else {
@@ -130,7 +136,7 @@ func noOpCallback(*sqltypes.Result) error {
 
 // TryStreamExecute implements the Primitive interface
 func (v *VExplain) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
- var stats func() map[Primitive]RowsReceived
+ var stats func() Stats
  if v.Type == sqlparser.TraceVExplainType {
  stats = vcursor.StartPrimitiveTrace()
  } else {
@@ -148,7 +154,7 @@ func (v *VExplain) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVa
  return callback(result)
 }
 
-func (v *VExplain) convertToResult(ctx context.Context, vcursor VCursor, stats func() map[Primitive]RowsReceived) (*sqltypes.Result, error) {
+func (v *VExplain) convertToResult(ctx context.Context, vcursor VCursor, stats func() Stats) (*sqltypes.Result, error) {
  switch v.Type {
  case sqlparser.QueriesVExplainType:
  result := convertToVExplainQueriesResult(vcursor.Session().GetVExplainLogs())
@@ -163,8 +169,9 @@ func (v *VExplain) convertToResult(ctx context.Context, vcursor VCursor, stats f
  }
 }
 
-func (v *VExplain) getExplainTraceOutput(getOpStats func() map[Primitive]RowsReceived) (*sqltypes.Result, error) {
- description := PrimitiveToPlanDescription(v.Input, getOpStats())
+func (v *VExplain) getExplainTraceOutput(getOpStats func() Stats) (*sqltypes.Result, error) {
+ stats := getOpStats()
+ description := PrimitiveToPlanDescription(v.Input, &stats)
 
  output, err := json.MarshalIndent(description, "", "\t")
  if err != nil {

go/vt/vtgate/executor_vexplain_test.go

@@ -94,9 +94,10 @@ func TestSimpleVexplainTrace(t *testing.T) {
  "Name": "TestExecutor",
  "Sharded": true
  },
- "NoOfCalls": 2,
+ "NoOfCalls": 1,
  "AvgNumberOfRows": 16,
  "MedianNumberOfRows": 16,
+ "ShardsQueried": 8,
  "FieldQuery": "select count(*), col2, weight_string(col2) from music where 1 != 1 group by col2, weight_string(col2)",
  "OrderBy": "(1|2) ASC",
  "Query": "select count(*), col2, weight_string(col2) from music group by col2, weight_string(col2) order by col2 asc",

go/vt/vtgate/vcursor_impl.go

@@ -135,7 +135,8 @@ type (
 
  // this is a map of the number of rows that every primitive has returned
  // if this field is nil, it means that we are not logging operator traffic
- primitiveStats map[engine.Primitive]engine.RowsReceived
+ interOpStats map[engine.Primitive]engine.RowsReceived
+ shardsStats map[engine.Primitive]engine.ShardsQueried
  }
 )
 
@@ -285,10 +286,14 @@ func (vc *vcursorImpl) UnresolvedTransactions(ctx context.Context, keyspace stri
  return vc.executor.UnresolvedTransactions(ctx, targets)
 }
 
-func (vc *vcursorImpl) StartPrimitiveTrace() func() map[engine.Primitive]engine.RowsReceived {
- vc.primitiveStats = make(map[engine.Primitive]engine.RowsReceived)
- return func() map[engine.Primitive]engine.RowsReceived {
- return vc.primitiveStats
+func (vc *vcursorImpl) StartPrimitiveTrace() func() engine.Stats {
+ vc.interOpStats = make(map[engine.Primitive]engine.RowsReceived)
+ vc.shardsStats = make(map[engine.Primitive]engine.ShardsQueried)
+ return func() engine.Stats {
+ return engine.Stats{
+ InterOpStats: vc.interOpStats,
+ ShardsStats: vc.shardsStats,
+ }
  }
 }
 
@@ -532,14 +537,20 @@ func (vc *vcursorImpl) ExecutePrimitive(ctx context.Context, primitive engine.Pr
 }
 
 func (vc *vcursorImpl) logOpTraffic(primitive engine.Primitive, res *sqltypes.Result) {
- if vc.primitiveStats != nil {
- rows := vc.primitiveStats[primitive]
+ if vc.interOpStats != nil {
+ rows := vc.interOpStats[primitive]
  if res == nil {
  rows = append(rows, 0)
  } else {
  rows = append(rows, len(res.Rows))
  }
- vc.primitiveStats[primitive] = rows
+ vc.interOpStats[primitive] = rows
+ }
+}
+
+func (vc *vcursorImpl) logShardsQueried(primitive engine.Primitive, shardsNb int) {
+ if vc.shardsStats != nil {
+ vc.shardsStats[primitive] += engine.ShardsQueried(shardsNb)
  }
 }
 
@@ -558,7 +569,7 @@ func (vc *vcursorImpl) ExecutePrimitiveStandalone(ctx context.Context, primitive
 }
 
 func (vc *vcursorImpl) wrapCallback(callback func(*sqltypes.Result) error, primitive engine.Primitive) func(*sqltypes.Result) error {
- if vc.primitiveStats == nil {
+ if vc.interOpStats == nil {
  return callback
  }
 
@@ -650,7 +661,7 @@ func (vc *vcursorImpl) ExecuteMultiShard(ctx context.Context, primitive engine.P
 
  qr, errs := vc.executor.ExecuteMultiShard(ctx, primitive, rss, commentedShardQueries(queries, vc.marginComments), vc.safeSession, canAutocommit, vc.ignoreMaxMemoryRows, vc.resultsObserver)
  vc.setRollbackOnPartialExecIfRequired(len(errs) != len(rss), rollbackOnError)
- vc.logOpTraffic(primitive, qr)
+ vc.logShardsQueried(primitive, len(rss))
  return qr, errs
 }
 
@@ -689,7 +700,7 @@ func (vc *vcursorImpl) ExecuteStandalone(ctx context.Context, primitive engine.P
  // The autocommit flag is always set to false because we currently don't
  // execute DMLs through ExecuteStandalone.
  qr, errs := vc.executor.ExecuteMultiShard(ctx, primitive, rss, bqs, NewAutocommitSession(vc.safeSession.Session), false /* autocommit */, vc.ignoreMaxMemoryRows, vc.resultsObserver)
- vc.logOpTraffic(primitive, qr)
+ vc.logShardsQueried(primitive, len(rss))
  return qr, vterrors.Aggregate(errs)
 }