Use a truly random Destination for DestinationAnyShard

Signed-off-by: Aaron Young <[email protected]>

go/vt/key/destination.go

@@ -19,6 +19,7 @@ package key
 import (
 	"bytes"
 	"encoding/hex"
+	"math/rand"
 	"strings"
 
 	"vitess.io/vitess/go/vt/vterrors"
@@ -27,6 +28,9 @@ import (
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 )
 
+// AnyShardPicker makes a choice on what shard to use when any shard will do. Used for testing.
+var AnyShardPicker DestinationAnyShardPicker = DestinationAnyShardPickerRandomShard{}
+
 // Destination is an interface definition for a query destination,
 // within a given Keyspace / Tablet Type. It is meant to be an internal
 // data structure, with multiple possible implementations.
@@ -369,12 +373,25 @@ func (d DestinationKeyspaceIDs) String() string {
 	return buffer.String()
 }
 
+// DestinationAnyShardPicker exposes an interface that will pick an index given a number of available shards.
+type DestinationAnyShardPicker interface {
+	// PickShard picks a shard given a number of shards
+	PickShard(shardCount int) int
+}
+
+// DestinationAnyShardPickerRandomShard picks a random shard.
+type DestinationAnyShardPickerRandomShard struct{}
+
+// PickShard is DestinationAnyShardPickerRandomShard's implmentation.
+func (dp DestinationAnyShardPickerRandomShard) PickShard(shardCount int) int {
+	return rand.Intn(shardCount)
+}
+
 //
 // DestinationAnyShard
 //
 
-// DestinationAnyShard is the destination for any one shard in the
-// keyspace. This usually maps to the first one in the list.
+// DestinationAnyShard is the destination for any one shard in the keyspace.
 // It implements the Destination interface.
 type DestinationAnyShard struct{}
 
@@ -388,7 +405,7 @@ func (d DestinationAnyShard) Resolve(allShards []*topodatapb.ShardReference, add
 	if len(allShards) == 0 {
 		return vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "no shard in keyspace")
 	}
-	return addShard(allShards[0].Name)
+	return addShard(allShards[AnyShardPicker.PickShard(len(allShards))].Name)
 }
 
 // String is part of the Destination interface.

go/vt/vtgate/engine/insert.go

@@ -78,6 +78,37 @@ type Insert struct {
 	MultiShardAutocommit bool
 }
 
+// NewQueryInsert creates an Insert with a query string.
+func NewQueryInsert(opcode InsertOpcode, keyspace *vindexes.Keyspace, query string) *Insert {
+	return &Insert{
+		Opcode:   opcode,
+		Keyspace: keyspace,
+		Query:    query,
+	}
+}
+
+// NewSimpleInsert creates an Insert for a Table.
+func NewSimpleInsert(opcode InsertOpcode, table *vindexes.Table, keyspace *vindexes.Keyspace) *Insert {
+	return &Insert{
+		Opcode:   opcode,
+		Table:    table,
+		Keyspace: keyspace,
+	}
+}
+
+// NewInsert creates a new Insert.
+func NewInsert(opcode InsertOpcode, keyspace *vindexes.Keyspace, vindexValues []sqltypes.PlanValue, table *vindexes.Table, prefix string, mid []string, suffix string) *Insert {
+	return &Insert{
+		Opcode:       opcode,
+		Keyspace:     keyspace,
+		VindexValues: vindexValues,
+		Table:        table,
+		Prefix:       prefix,
+		Mid:          mid,
+		Suffix:       suffix,
+	}
+}
+
 // MarshalJSON serializes the Insert into a JSON representation.
 // It's used for testing and diagnostics.
 func (ins *Insert) MarshalJSON() ([]byte, error) {