Add linearizable support to SQL VSchema management

go/cmd/vtcombo/cli/vschema_watcher.go

@@ -56,17 +56,20 @@ func loadKeyspacesFromDir(ctx context.Context, dir string, ts *topo.Server) {
 				log.Fatalf("Unable to read keyspace file %v: %v", ksFile, err)
 			}
 
-			keyspace := &vschemapb.Keyspace{}
-			err = json.Unmarshal(jsonData, keyspace)
+			ksvs := &topo.KeyspaceVSchemaInfo{
+				Name:     ks.Name,
+				Keyspace: &vschemapb.Keyspace{},
+			}
+			err = json.Unmarshal(jsonData, ksvs.Keyspace)
 			if err != nil {
 				log.Fatalf("Unable to parse keyspace file %v: %v", ksFile, err)
 			}
 
-			_, err = vindexes.BuildKeyspace(keyspace, env.Parser())
+			_, err = vindexes.BuildKeyspace(ksvs.Keyspace, env.Parser())
 			if err != nil {
 				log.Fatalf("Invalid keyspace definition: %v", err)
 			}
-			ts.SaveVSchema(ctx, ks.Name, keyspace)
+			ts.SaveVSchema(ctx, ksvs)
 			log.Infof("Loaded keyspace %v from %v\n", ks.Name, ksFile)
 		}
 	}

go/test/endtoend/vtgate/vschema/vschema_test.go

@@ -23,15 +23,20 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/rand"
 
 	"vitess.io/vitess/go/mysql"
 	"vitess.io/vitess/go/test/endtoend/cluster"
 	"vitess.io/vitess/go/test/endtoend/utils"
+	"vitess.io/vitess/go/vt/vtgate"
 )
 
 var (
@@ -164,14 +169,88 @@ func TestVSchema(t *testing.T) {
 	utils.AssertMatches(t, conn, "delete from vt_user", `[]`)
 
 	if utils.BinaryIsAtLeastAtVersion(22, "vtgate") {
+		// Don't allow any users to modify the vschema via the SQL API
+		// in order to test that behavior.
 		writeConfig(configFile, map[string]string{
 			"vschema_ddl_authorized_users": "",
 		})
-
+		// Allow anyone to modify the vschema via the SQL API again when
+		// the test completes.
+		defer func() {
+			writeConfig(configFile, map[string]string{
+				"vschema_ddl_authorized_users": "%",
+			})
+		}()
 		require.EventuallyWithT(t, func(t *assert.CollectT) {
 			_, err = conn.ExecuteFetch("ALTER VSCHEMA DROP TABLE main", 1000, false)
 			assert.Error(t, err)
 			assert.ErrorContains(t, err, "is not authorized to perform vschema operations")
 		}, 5*time.Second, 100*time.Millisecond)
 	}
 }
+
+// TestVSchemaSQLAPIConcurrency tests that we prevent lost writes when we have
+// concurrent vschema changes being made via the SQL API.
+func TestVSchemaSQLAPIConcurrency(t *testing.T) {
+	if !utils.BinaryIsAtLeastAtVersion(22, "vtgate") {
+		t.Skip("This test requires vtgate version 22 or higher")
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+	defer cancel()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	initialVSchema, err := conn.ExecuteFetch("SHOW VSCHEMA TABLES", -1, false)
+	require.NoError(t, err)
+	baseTableName := "t"
+	numTables := 1000
+	mysqlConns := make([]*mysql.Conn, numTables)
+	for i := 0; i < numTables; i++ {
+		c, err := mysql.Connect(ctx, &vtParams)
+		require.NoError(t, err)
+		mysqlConns[i] = c
+		defer c.Close()
+	}
+
+	isVersionMismatchErr := func(err error) bool {
+		// The error we get is an SQL error so we have to do string matching.
+		return err != nil && strings.Contains(err.Error(), vtgate.ErrStaleVSchema.Error())
+	}
+
+	wg := sync.WaitGroup{}
+	preventedLostWrites := atomic.Bool{}
+	for i := 0; i < numTables; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			time.Sleep(time.Duration(rand.Intn(100) * int(time.Nanosecond)))
+			tableName := fmt.Sprintf("%s%d", baseTableName, i)
+			_, err = mysqlConns[i].ExecuteFetch(fmt.Sprintf("ALTER VSCHEMA ADD TABLE %s", tableName), -1, false)
+			if isVersionMismatchErr(err) {
+				preventedLostWrites.Store(true)
+			} else {
+				require.NoError(t, err)
+				time.Sleep(time.Duration(rand.Intn(75) * int(time.Nanosecond)))
+				_, err = mysqlConns[i].ExecuteFetch(fmt.Sprintf("ALTER VSCHEMA DROP TABLE %s", tableName), -1, false)
+				if isVersionMismatchErr(err) {
+					preventedLostWrites.Store(true)
+				} else {
+					require.NoError(t, err)
+				}
+			}
+		}()
+	}
+	wg.Wait()
+	require.True(t, preventedLostWrites.Load())
+
+	// Cleanup any tables that were not dropped because the DROP query
+	// failed due to a bad node version.
+	for i := 0; i < numTables; i++ {
+		tableName := fmt.Sprintf("%s%d", baseTableName, i)
+		_, _ = mysqlConns[i].ExecuteFetch(fmt.Sprintf("ALTER VSCHEMA DROP TABLE %s", tableName), -1, false)
+	}
+	// Confirm that we're back to the initial state.
+	utils.AssertMatches(t, conn, "SHOW VSCHEMA TABLES", fmt.Sprintf("%v", initialVSchema.Rows))
+}

go/test/utils/diff.go

@@ -68,7 +68,7 @@ func MustMatchFn(ignoredFields ...string) func(t *testing.T, want, got any, errM
 		t.Helper()
 		diff := cmp.Diff(want, got, diffOpts...)
 		if diff != "" {
-			t.Fatalf("%v: (-want +got)\n%v", errMsg, diff)
+			require.FailNow(t, "%v: (-want +got)\n%v", errMsg, diff)
 		}
 	}
 }

go/vt/topo/helpers/copy.go

@@ -41,7 +41,6 @@ func CopyKeyspaces(ctx context.Context, fromTS, toTS *topo.Server, parser *sqlpa
 	}
 
 	for _, keyspace := range keyspaces {
-
 		ki, err := fromTS.GetKeyspace(ctx, keyspace)
 		if err != nil {
 			return fmt.Errorf("GetKeyspace(%v): %w", keyspace, err)
@@ -55,15 +54,15 @@ func CopyKeyspaces(ctx context.Context, fromTS, toTS *topo.Server, parser *sqlpa
 			}
 		}
 
-		vs, err := fromTS.GetVSchema(ctx, keyspace)
+		ksvs, err := fromTS.GetVSchema(ctx, keyspace)
 		switch {
 		case err == nil:
-			_, err = vindexes.BuildKeyspace(vs, parser)
+			_, err = vindexes.BuildKeyspace(ksvs.Keyspace, parser)
 			if err != nil {
 				log.Errorf("BuildKeyspace(%v): %v", keyspace, err)
 				break
 			}
-			if err := toTS.SaveVSchema(ctx, keyspace, vs); err != nil {
+			if err := toTS.SaveVSchema(ctx, ksvs); err != nil {
 				log.Errorf("SaveVSchema(%v): %v", keyspace, err)
 			}
 		case topo.IsErrType(err, topo.NoNode):

go/vt/topo/srv_vschema.go

@@ -171,10 +171,13 @@ func (ts *Server) RebuildSrvVSchema(ctx context.Context, cells []string) error {
 		go func(keyspace string) {
 			defer wg.Done()
 
-			k, err := ts.GetVSchema(ctx, keyspace)
+			ksvs, err := ts.GetVSchema(ctx, keyspace)
 			if IsErrType(err, NoNode) {
 				err = nil
-				k = &vschemapb.Keyspace{}
+				ksvs = &KeyspaceVSchemaInfo{
+					Name:     keyspace,
+					Keyspace: &vschemapb.Keyspace{},
+				}
 			}
 
 			mu.Lock()
@@ -184,7 +187,7 @@ func (ts *Server) RebuildSrvVSchema(ctx context.Context, cells []string) error {
 				finalErr = err
 				return
 			}
-			srvVSchema.Keyspaces[keyspace] = k
+			srvVSchema.Keyspaces[keyspace] = ksvs.Keyspace
 		}(keyspace)
 	}
 	wg.Wait()

go/vt/topo/test/vschema.go

@@ -44,9 +44,12 @@ func checkVSchema(t *testing.T, ctx context.Context, ts *topo.Server) {
 		t.Error(err)
 	}
 
-	err = ts.SaveVSchema(ctx, "test_keyspace", &vschemapb.Keyspace{
-		Tables: map[string]*vschemapb.Table{
-			"unsharded": {},
+	err = ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name: "test_keyspace",
+		Keyspace: &vschemapb.Keyspace{
+			Tables: map[string]*vschemapb.Table{
+				"unsharded": {},
+			},
 		},
 	})
 	if err != nil {
@@ -64,8 +67,11 @@ func checkVSchema(t *testing.T, ctx context.Context, ts *topo.Server) {
 		t.Errorf("GetVSchema: %s, want %s", got, want)
 	}
 
-	err = ts.SaveVSchema(ctx, "test_keyspace", &vschemapb.Keyspace{
-		Sharded: true,
+	err = ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name: "test_keyspace",
+		Keyspace: &vschemapb.Keyspace{
+			Sharded: true,
+		},
 	})
 	require.NoError(t, err)
 

go/vt/topo/topotests/srv_vschema_test.go

@@ -21,6 +21,7 @@ import (
 
 	"google.golang.org/protobuf/proto"
 
+	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -76,7 +77,10 @@ func TestRebuildVSchema(t *testing.T) {
 	keyspace1 := &vschemapb.Keyspace{
 		Sharded: true,
 	}
-	if err := ts.SaveVSchema(ctx, "ks1", keyspace1); err != nil {
+	if err := ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name:     "ks1",
+		Keyspace: keyspace1,
+	}); err != nil {
 		t.Fatalf("SaveVSchema(ks1) failed: %v", err)
 	}
 	if err := ts.RebuildSrvVSchema(ctx, cells); err != nil {
@@ -118,7 +122,10 @@ func TestRebuildVSchema(t *testing.T) {
 			},
 		},
 	}
-	if err := ts.SaveVSchema(ctx, "ks2", keyspace2); err != nil {
+	if err := ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name:     "ks2",
+		Keyspace: keyspace2,
+	}); err != nil {
 		t.Fatalf("SaveVSchema(ks1) failed: %v", err)
 	}
 	if err := ts.RebuildSrvVSchema(ctx, []string{"cell1"}); err != nil {
@@ -182,7 +189,10 @@ func TestRebuildVSchema(t *testing.T) {
 	wanted4.RoutingRules = rr
 
 	// Delete a keyspace, checks vschema entry in map goes away.
-	if err := ts.SaveVSchema(ctx, "ks2", &vschemapb.Keyspace{}); err != nil {
+	if err := ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name:     "ks2",
+		Keyspace: &vschemapb.Keyspace{},
+	}); err != nil {
 		t.Fatalf("SaveVSchema(ks1) failed: %v", err)
 	}
 	if err := ts.DeleteKeyspace(ctx, "ks2"); err != nil {

go/vt/topo/vschema.go

@@ -20,34 +20,58 @@ import (
 	"context"
 	"path"
 
-	"google.golang.org/protobuf/proto"
-
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/vterrors"
 
 	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
 )
 
-// SaveVSchema saves a Vschema. A valid Vschema should be passed in. It does not verify its correctness.
-// If the VSchema is empty, just remove it.
-func (ts *Server) SaveVSchema(ctx context.Context, keyspace string, vschema *vschemapb.Keyspace) error {
+// KeyspaceVSchemaInfo wraps a vschemapb.Keyspace and is a meta
+// struct that contains metadata to give the data more context
+// and convenience. This is the main way we interact with a
+// keyspace's vschema.
+type KeyspaceVSchemaInfo struct {
+	Name string
+	*vschemapb.Keyspace
+	version Version
+}
+
+func (k *KeyspaceVSchemaInfo) CloneVT() *KeyspaceVSchemaInfo {
+	if k == nil {
+		return (*KeyspaceVSchemaInfo)(nil)
+	}
+	kc := &KeyspaceVSchemaInfo{
+		Name:    k.Name,
+		version: Version(k.version),
+	}
+	if k.Keyspace != nil {
+		kc.Keyspace = k.Keyspace.CloneVT()
+	}
+	return kc
+}
+
+// SaveVSchema saves a Vschema. A valid Vschema should be passed in.
+// It does not verify its correctness beyond marshaling it.
+func (ts *Server) SaveVSchema(ctx context.Context, ksvs *KeyspaceVSchemaInfo) error {
 	if err := ctx.Err(); err != nil {
 		return err
 	}
 
-	nodePath := path.Join(KeyspacesPath, keyspace, VSchemaFile)
-	data, err := vschema.MarshalVT()
+	nodePath := path.Join(KeyspacesPath, ksvs.Name, VSchemaFile)
+	data, err := ksvs.MarshalVT()
 	if err != nil {
 		return err
 	}
 
-	_, err = ts.globalCell.Update(ctx, nodePath, data, nil)
+	version, err := ts.globalCell.Update(ctx, nodePath, data, ksvs.version)
 	if err != nil {
-		log.Errorf("failed to update vschema for keyspace %s: %v", keyspace, err)
-	} else {
-		log.Infof("successfully updated vschema for keyspace %s: %+v", keyspace, vschema)
+		log.Errorf("failed to update vschema for keyspace %s: %v", ksvs.Name, err)
+		return err
 	}
-	return err
+	ksvs.version = version
+	log.Infof("successfully updated vschema for keyspace %s: %+v", ksvs.Name, ksvs.Keyspace)
+
+	return nil
 }
 
 // DeleteVSchema delete the keyspace if it exists
@@ -61,35 +85,43 @@ func (ts *Server) DeleteVSchema(ctx context.Context, keyspace string) error {
 }
 
 // GetVSchema fetches the vschema from the topo.
-func (ts *Server) GetVSchema(ctx context.Context, keyspace string) (*vschemapb.Keyspace, error) {
+func (ts *Server) GetVSchema(ctx context.Context, keyspace string) (*KeyspaceVSchemaInfo, error) {
 	if err := ctx.Err(); err != nil {
 		return nil, err
 	}
 
 	nodePath := path.Join(KeyspacesPath, keyspace, VSchemaFile)
-	data, _, err := ts.globalCell.Get(ctx, nodePath)
+	data, version, err := ts.globalCell.Get(ctx, nodePath)
 	if err != nil {
 		return nil, err
 	}
-	var vs vschemapb.Keyspace
-	err = proto.Unmarshal(data, &vs)
+
+	vs := &vschemapb.Keyspace{}
+	err = vs.UnmarshalVT(data)
 	if err != nil {
 		return nil, vterrors.Wrapf(err, "bad vschema data: %q", data)
 	}
-	return &vs, nil
+	return &KeyspaceVSchemaInfo{
+		Name:     keyspace,
+		Keyspace: vs,
+		version:  version,
+	}, nil
 }
 
 // EnsureVSchema makes sure that a vschema is present for this keyspace or creates a blank one if it is missing
 func (ts *Server) EnsureVSchema(ctx context.Context, keyspace string) error {
-	vschema, err := ts.GetVSchema(ctx, keyspace)
+	ksvs, err := ts.GetVSchema(ctx, keyspace)
 	if err != nil && !IsErrType(err, NoNode) {
 		log.Infof("error in getting vschema for keyspace %s: %v", keyspace, err)
 	}
-	if vschema == nil || IsErrType(err, NoNode) {
-		err = ts.SaveVSchema(ctx, keyspace, &vschemapb.Keyspace{
-			Sharded:  false,
-			Vindexes: make(map[string]*vschemapb.Vindex),
-			Tables:   make(map[string]*vschemapb.Table),
+	if ksvs == nil || ksvs.Keyspace == nil || IsErrType(err, NoNode) {
+		err = ts.SaveVSchema(ctx, &KeyspaceVSchemaInfo{
+			Name: keyspace,
+			Keyspace: &vschemapb.Keyspace{
+				Sharded:  false,
+				Vindexes: make(map[string]*vschemapb.Vindex),
+				Tables:   make(map[string]*vschemapb.Table),
+			},
 		})
 		if err != nil {
 			log.Errorf("could not create blank vschema: %v", err)