More tests. Refactor

Signed-off-by: Rohit Nayak <rohit@planetscale.com>

go/vt/topo/keyspace_routing_rules_lock.go

@@ -21,13 +21,18 @@ import (
  "fmt"
 )
 
+// KeyspaceRoutingRulesLock is a wrapper over TopoLock, to serialize updates to the keyspace routing rules.
 type KeyspaceRoutingRulesLock struct {
  *TopoLock
 
+ // sourceKeyspace is only used for logging, at the moment
  sourceKeyspace string
 }
 
 func NewKeyspaceRoutingRulesLock(ctx context.Context, ts *Server, sourceKeyspace string) (*KeyspaceRoutingRulesLock, error) {
+ if sourceKeyspace == "" {
+ return nil, fmt.Errorf("sourceKeyspace is not specified")
+ }
  return &KeyspaceRoutingRulesLock{
  TopoLock: &TopoLock{
  Root: "", // global

go/vt/topo/keyspace_routing_rules_lock_test.go

@@ -0,0 +1,61 @@
+/*
+Copyright 2024 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 topo_test
+
+import (
+ "context"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+
+ "vitess.io/vitess/go/vt/topo"
+ "vitess.io/vitess/go/vt/topo/memorytopo"
+)
+
+// TestKeyspaceRoutingRulesLock tests that the lock is acquired and released correctly.
+func TestKeyspaceRoutingRulesLock(t *testing.T) {
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ ts := memorytopo.NewServer(ctx, "zone1")
+ defer ts.Close()
+ conn := ts.GetGlobalConn()
+ require.NotNil(t, conn)
+ _, err := conn.Create(ctx, topo.KeyspaceRoutingRulesFile, []byte(""))
+ require.NoError(t, err)
+
+ currentTopoLockTimeout := topo.LockTimeout
+ topo.LockTimeout = testLockTimeout
+ defer func() {
+ topo.LockTimeout = currentTopoLockTimeout
+ }()
+
+ lock, err := topo.NewKeyspaceRoutingRulesLock(ctx, ts, "ks1")
+ require.NoError(t, err)
+ lockCtx, unlock, err := lock.Lock(ctx)
+ require.NoError(t, err)
+ require.NoError(t, lock.Check(lockCtx))
+
+ // re-acquiring the lock should fail
+ _, _, err = lock.Lock(ctx)
+ require.Error(t, err)
+
+ unlock(&err)
+
+ // re-acquiring the lock should succeed
+ _, _, err = lock.Lock(ctx)
+ require.NoError(t, err)
+}

go/vt/topo/server.go

@@ -352,6 +352,10 @@ func (ts *Server) Close() {
  ts.cellConns = make(map[string]cellConn)
 }
 
+func (ts *Server) GetGlobalConn() Conn {
+ return ts.globalCell
+}
+
 func (ts *Server) clearCellAliasesCache() {
  cellsAliases.mu.Lock()
  defer cellsAliases.mu.Unlock()

go/vt/topo/topo_lock.go

@@ -18,6 +18,7 @@ package topo
 
 import (
  "context"
+ "fmt"
  "path"
 
  "vitess.io/vitess/go/trace"
@@ -26,21 +27,42 @@ import (
  "vitess.io/vitess/go/vt/vterrors"
 )
 
+// ITopoLock is the interface for a lock that can be used to lock a key in the topology server.
+// The lock is associated with a context and can be unlocked by calling the returned function.
+// Note that we don't need an Unlock method on the interface, as the Lock() function
+// returns a function that can be used to unlock the lock.
 type ITopoLock interface {
  Lock(ctx context.Context) (context.Context, func(*error), error)
- Unlock(ctx context.Context) error
  Check(ctx context.Context) error
 }
 
 type TopoLock struct {
- Root, Key, Action, Name string
+ Root string // topo path
+ Key string // the topo file to lock, relative to Root
+ Action string // action, for logging purposes
+ Name string // name, for logging purposes
 
  ts *Server
 }
 
-func (l *Lock) lock(ctx context.Context, ts *Server, root, key string) (LockDescriptor, error) {
- log.Infof("Locking %s/%s for action %v", root, key, l.Action)
+var _ ITopoLock = (*TopoLock)(nil)
+
+func (ts *Server) NewTopoLock(root, key, action, name string) *TopoLock {
+ return &TopoLock{
+ ts: ts,
+ Root: root,
+ Key: key,
+ Action: action,
+ Name: name,
+ }
+}
+
+func (tl *TopoLock) String() string {
+ return fmt.Sprintf("TopoLock{Root: %v, Key: %v, Action: %v, Name: %v}", tl.Root, tl.Key, tl.Action, tl.Name)
+}
 
+// perform the topo lock operation
+func (l *Lock) lock(ctx context.Context, ts *Server, root, key string) (LockDescriptor, error) {
  ctx, cancel := context.WithTimeout(ctx, getLockTimeout())
  defer cancel()
  span, ctx := trace.NewSpan(ctx, "TopoServer.LockKeyForAction")
@@ -74,15 +96,15 @@ func (l *Lock) unlock(ctx context.Context, ts *Server, root, key string, lockDes
 
  // first update the actionNode
  if actionError != nil {
- log.Infof("Unlocking keyspace %v for action %v with error %v", key, l.Action, actionError)
  l.Status = "Error: " + actionError.Error()
  } else {
- log.Infof("Unlocking keyspace %v for successful action %v", key, l.Action)
  l.Status = "Done"
  }
  return lockDescriptor.Unlock(ctx)
 }
 
+// Lock adds lock information to the context, checks that the lock is not already held, and locks it.
+// It returns a new context with the lock information and a function to unlock the lock.
 func (tl TopoLock) Lock(ctx context.Context) (context.Context, func(*error), error) {
  i, ok := ctx.Value(locksKey).(*locksInfo)
  if !ok {
@@ -124,6 +146,7 @@ func (tl TopoLock) Lock(ctx context.Context) (context.Context, func(*error), err
  }
 
  err := l.unlock(ctx, tl.ts, tl.Root, tl.Key, lockDescriptor, *finalErr)
+ // if we have an error, we log it, but we still want to delete the lock
  if *finalErr != nil {
  if err != nil {
  // both error are set, just log the unlock error
@@ -136,28 +159,20 @@ func (tl TopoLock) Lock(ctx context.Context) (context.Context, func(*error), err
  }, nil
 }
 
-func (tl TopoLock) Unlock(ctx context.Context) error {
- // TODO implement me
- panic("implement me")
-}
-
+// Check checks that the lock is held in the context: it just validates that the lockInfo is present in the context.
 func (tl TopoLock) Check(ctx context.Context) error {
- // TODO implement me
- panic("implement me")
-}
-
-var _ ITopoLock = (*TopoLock)(nil)
-
-func (ts *Server) NewTopoLock(root, key, action, name string) TopoLock {
- return TopoLock{
- ts: ts,
- Root: root,
- Key: key,
- Action: action,
- Name: name,
+ // extract the locksInfo pointer
+ i, ok := ctx.Value(locksKey).(*locksInfo)
+ if !ok {
+ return vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "%s is not locked (no locksInfo)", tl.String())
  }
-}
+ i.mu.Lock()
+ defer i.mu.Unlock()
 
-func (ts *Server) GetGlobalConn() Conn {
- return ts.globalCell
+ // find the individual entry
+ _, ok = i.info[tl.Key]
+ if !ok {
+ return vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "%s is not locked (no lockInfo in map)", tl.String())
+ }
+ return nil
 }

go/vt/topo/topo_lock_test.go

@@ -19,15 +19,19 @@ package topo_test
 import (
  "context"
  "testing"
-
- "vitess.io/vitess/go/vt/topo"
+ "time"
 
  "github.com/stretchr/testify/require"
 
+ "vitess.io/vitess/go/vt/topo"
  "vitess.io/vitess/go/vt/topo/memorytopo"
 )
 
-func TestTopoLock(t *testing.T) {
+// lower the lock timeout for testing
+const testLockTimeout = 3 * time.Second
+
+// TestTopoLockTimeout tests that the lock times out after the specified duration.
+func TestTopoLockTimeout(t *testing.T) {
  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()
  ts := memorytopo.NewServer(ctx, "zone1")
@@ -38,18 +42,42 @@ func TestTopoLock(t *testing.T) {
  if err != nil {
  return
  }
- _, err = conn.Create(ctx, "root/key2", []byte("value"))
+
+ currentTopoLockTimeout := topo.LockTimeout
+ topo.LockTimeout = testLockTimeout
+ defer func() {
+ topo.LockTimeout = currentTopoLockTimeout
+ }()
+
+ // acquire the lock
+ origCtx := ctx
+ tl1 := ts.NewTopoLock("root", "key1", "action1", "name")
+ _, unlock, err := tl1.Lock(origCtx)
+ require.NoError(t, err)
+ defer unlock(&err)
+
+ // re-acquiring the lock should fail
+ _, _, err2 := tl1.Lock(origCtx)
+ require.Errorf(t, err2, "deadline exceeded")
+}
+
+// TestTopoLockBasic tests basic lock operations.
+func TestTopoLockBasic(t *testing.T) {
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ ts := memorytopo.NewServer(ctx, "zone1")
+ defer ts.Close()
+ conn := ts.GetGlobalConn()
+ require.NotNil(t, conn)
+ _, err := conn.Create(ctx, "root/key1", []byte("value"))
  if err != nil {
  return
  }
 
- var tl1, tl2 topo.ITopoLock
-
  origCtx := ctx
- tl1 = ts.NewTopoLock("root", "key1", "action1", "name")
+ tl1 := ts.NewTopoLock("root", "key1", "action1", "name")
  ctx, unlock, err := tl1.Lock(origCtx)
  require.NoError(t, err)
- require.NotNil(t, unlock)
 
  // locking the same key again, without unlocking, should return an error
  _, _, err2 := tl1.Lock(ctx)
@@ -59,14 +87,16 @@ func TestTopoLock(t *testing.T) {
  unlock(&err)
  ctx, unlock, err = tl1.Lock(origCtx)
  require.NoError(t, err)
- require.NotNil(t, unlock)
  defer unlock(&err)
 
  // locking another key should work
- tl2 = ts.NewTopoLock("root", "key2", "action2", "name")
+ _, err = conn.Create(ctx, "root/key2", []byte("value"))
+ if err != nil {
+ return
+ }
+ tl2 := ts.NewTopoLock("root", "key2", "action2", "name")
  _, unlock2, err := tl2.Lock(ctx)
  require.NoError(t, err)
- require.NotNil(t, unlock2)
  defer unlock2(&err)
 
 }

go/vt/vtctl/workflow/utils.go

@@ -882,6 +882,7 @@ func updateKeyspaceRoutingRule(ctx context.Context, ts *topo.Server, sourceKeysp
  return lockErr
  }
  defer unlock(&err)
+
  rules, err := topotools.GetKeyspaceRoutingRules(lockCtx, ts)
  if err != nil {
  return err