master: make admin lock lease and release atomic (#12532)

Co-authored-by: Mansur Akhmadi <10739460+mansur-akhmadi@users.noreply.github.com>

weed/server/master_grpc_server_admin.go

@@ -87,6 +87,12 @@ func NewAdminLocks() *AdminLocks {
 func (locks *AdminLocks) isLocked(lockName string) (clientName string, message string, isLocked bool) {
 	locks.RLock()
 	defer locks.RUnlock()
+	return locks.lockStatusLocked(lockName)
+}
+
+// The *Locked helpers below expect the caller to hold locks, so that a lease
+// or release checks and changes the lock state in one critical section.
+func (locks *AdminLocks) lockStatusLocked(lockName string) (clientName string, message string, isLocked bool) {
 	adminLock, found := locks.locks[lockName]
 	if !found {
 		return "", "", false
@@ -95,9 +101,7 @@ func (locks *AdminLocks) isLocked(lockName string) (clientName string, message s
 	return adminLock.lastClient, adminLock.lastMessage, adminLock.accessLockTime.Add(LockDuration).After(time.Now())
 }
 
-func (locks *AdminLocks) isValidToken(lockName string, ts time.Time, token int64) bool {
-	locks.RLock()
-	defer locks.RUnlock()
+func (locks *AdminLocks) isValidTokenLocked(lockName string, ts time.Time, token int64) bool {
 	adminLock, found := locks.locks[lockName]
 	if !found {
 		return false
@@ -105,9 +109,7 @@ func (locks *AdminLocks) isValidToken(lockName string, ts time.Time, token int64
 	return adminLock.accessLockTime.Equal(ts) && adminLock.accessSecret == token
 }
 
-func (locks *AdminLocks) generateToken(lockName string, clientName string) (ts time.Time, token int64) {
-	locks.Lock()
-	defer locks.Unlock()
+func (locks *AdminLocks) generateTokenLocked(lockName string, clientName string) (ts time.Time, token int64) {
 	if existing, found := locks.locks[lockName]; found && existing.lastClient != clientName {
 		stats.MasterAdminLock.DeleteLabelValues(existing.lastClient)
 	}
@@ -121,11 +123,11 @@ func (locks *AdminLocks) generateToken(lockName string, clientName string) (ts t
 	return lock.accessLockTime, lock.accessSecret
 }
 
-func (locks *AdminLocks) deleteLock(lockName string) {
-	locks.Lock()
-	stats.MasterAdminLock.DeleteLabelValues(locks.locks[lockName].lastClient)
-	defer locks.Unlock()
-	delete(locks.locks, lockName)
+func (locks *AdminLocks) deleteLockLocked(lockName string) {
+	if adminLock, found := locks.locks[lockName]; found {
+		stats.MasterAdminLock.DeleteLabelValues(adminLock.lastClient)
+		delete(locks.locks, lockName)
+	}
 }
 
 func (ms *MasterServer) LeaseAdminToken(ctx context.Context, req *master_pb.LeaseAdminTokenRequest) (*master_pb.LeaseAdminTokenResponse, error) {
@@ -135,11 +137,13 @@ func (ms *MasterServer) LeaseAdminToken(ctx context.Context, req *master_pb.Leas
 		return resp, raft.NotLeaderError
 	}
 
-	if lastClient, lastMessage, isLocked := ms.adminLocks.isLocked(req.LockName); isLocked {
+	ms.adminLocks.Lock()
+	defer ms.adminLocks.Unlock()
+	if lastClient, lastMessage, isLocked := ms.adminLocks.lockStatusLocked(req.LockName); isLocked {
 		glog.V(4).Infof("LeaseAdminToken %v", lastClient)
-		if req.PreviousToken != 0 && ms.adminLocks.isValidToken(req.LockName, time.Unix(0, req.PreviousLockTime), req.PreviousToken) {
+		if req.PreviousToken != 0 && ms.adminLocks.isValidTokenLocked(req.LockName, time.Unix(0, req.PreviousLockTime), req.PreviousToken) {
 			// for renew
-			ts, token := ms.adminLocks.generateToken(req.LockName, req.ClientName)
+			ts, token := ms.adminLocks.generateTokenLocked(req.LockName, req.ClientName)
 			resp.Token, resp.LockTsNs = token, ts.UnixNano()
 			return resp, nil
 		}
@@ -147,7 +151,7 @@ func (ms *MasterServer) LeaseAdminToken(ctx context.Context, req *master_pb.Leas
 		return resp, fmt.Errorf("already locked by %v: %v", lastClient, lastMessage)
 	}
 	// for fresh lease request
-	ts, token := ms.adminLocks.generateToken(req.LockName, req.ClientName)
+	ts, token := ms.adminLocks.generateTokenLocked(req.LockName, req.ClientName)
 	resp.Token, resp.LockTsNs = token, ts.UnixNano()
 	return resp, nil
 }
@@ -159,8 +163,10 @@ func (ms *MasterServer) ReleaseAdminToken(ctx context.Context, req *master_pb.Re
 	if !ms.Topo.IsLeader() {
 		return resp, raft.NotLeaderError
 	}
-	if ms.adminLocks.isValidToken(req.LockName, time.Unix(0, req.PreviousLockTime), req.PreviousToken) {
-		ms.adminLocks.deleteLock(req.LockName)
+	ms.adminLocks.Lock()
+	defer ms.adminLocks.Unlock()
+	if ms.adminLocks.isValidTokenLocked(req.LockName, time.Unix(0, req.PreviousLockTime), req.PreviousToken) {
+		ms.adminLocks.deleteLockLocked(req.LockName)
 	}
 	return resp, nil
 }

weed/server/master_grpc_server_admin_lock_test.go

@@ -0,0 +1,149 @@
+package weed_server
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
+)
+
+func newAdminLockMaster() *MasterServer {
+	ms := newLeaderMaster()
+	ms.adminLocks = NewAdminLocks()
+	return ms
+}
+
+func leaseAdminLock(ms *MasterServer, client string, prev *master_pb.LeaseAdminTokenResponse) (*master_pb.LeaseAdminTokenResponse, error) {
+	req := &master_pb.LeaseAdminTokenRequest{LockName: "admin", ClientName: client}
+	if prev != nil {
+		req.PreviousToken, req.PreviousLockTime = prev.Token, prev.LockTsNs
+	}
+	return ms.LeaseAdminToken(context.Background(), req)
+}
+
+func releaseAdminLock(ms *MasterServer, lease *master_pb.LeaseAdminTokenResponse) error {
+	_, err := ms.ReleaseAdminToken(context.Background(), &master_pb.ReleaseAdminTokenRequest{
+		LockName: "admin", PreviousToken: lease.Token, PreviousLockTime: lease.LockTsNs,
+	})
+	return err
+}
+
+// runTogether starts fns at the same moment and waits for all of them. A panic
+// in any of them is returned as an error instead of crashing the test binary,
+// the same way it would take down the master (gRPC has no recovery there).
+func runTogether(fns ...func()) (panicErr error) {
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	start := make(chan struct{})
+	for _, fn := range fns {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					mu.Lock()
+					panicErr = fmt.Errorf("panic: %v", r)
+					mu.Unlock()
+				}
+			}()
+			<-start
+			fn()
+		}()
+	}
+	close(start)
+	wg.Wait()
+	return panicErr
+}
+
+func TestAdminLockLeaseRenewRelease(t *testing.T) {
+	ms := newAdminLockMaster()
+
+	a, err := leaseAdminLock(ms, "shellA", nil)
+	require.NoError(t, err)
+
+	_, err = leaseAdminLock(ms, "shellB", nil)
+	require.Error(t, err, "second client must be refused while the lock is held")
+
+	renewed, err := leaseAdminLock(ms, "shellA", a)
+	require.NoError(t, err)
+
+	_, err = leaseAdminLock(ms, "shellB", a)
+	require.Error(t, err, "a token replaced by a renewal must not renew again")
+
+	require.NoError(t, releaseAdminLock(ms, a))
+	_, err = leaseAdminLock(ms, "shellB", nil)
+	require.Error(t, err, "releasing with a stale token must keep the current holder")
+
+	require.NoError(t, releaseAdminLock(ms, renewed))
+	require.NoError(t, releaseAdminLock(ms, renewed), "releasing twice is a no-op")
+
+	_, err = leaseAdminLock(ms, "shellB", nil)
+	require.NoError(t, err, "lock is free after release")
+}
+
+func TestAdminLockReleaseUnknownLock(t *testing.T) {
+	ms := newAdminLockMaster()
+	require.NoError(t, releaseAdminLock(ms, &master_pb.LeaseAdminTokenResponse{Token: 1, LockTsNs: 1}))
+}
+
+// Two releases of the same token (a shell retrying on timeout) both passed the
+// token check before either deleted the lock; the second one then dereferenced
+// the missing entry and panicked.
+func TestAdminLockConcurrentReleaseOfSameToken(t *testing.T) {
+	ms := newAdminLockMaster()
+	for i := 0; i < 10000; i++ {
+		lease, err := leaseAdminLock(ms, "shell", nil)
+		require.NoError(t, err)
+		release := func() { _ = releaseAdminLock(ms, lease) }
+		require.NoError(t, runTogether(release, release), "iteration %d", i)
+		_, _, locked := ms.adminLocks.isLocked("admin")
+		require.False(t, locked, "iteration %d", i)
+	}
+}
+
+// Two clients leasing a free lock at the same time must not both get it.
+func TestAdminLockConcurrentLeaseGrantsOnce(t *testing.T) {
+	ms := newAdminLockMaster()
+	for i := 0; i < 2000; i++ {
+		ms.adminLocks = NewAdminLocks()
+		var mu sync.Mutex
+		granted := 0
+		lease := func(client string) func() {
+			return func() {
+				if _, err := leaseAdminLock(ms, client, nil); err == nil {
+					mu.Lock()
+					granted++
+					mu.Unlock()
+				}
+			}
+		}
+		require.NoError(t, runTogether(lease("shellA"), lease("shellB"), lease("shellC")))
+		require.Equal(t, 1, granted, "iteration %d: exclusive lock granted %d times", i, granted)
+	}
+}
+
+// Concurrent renewals with the same token must not all succeed: only one new
+// token can be the valid one.
+func TestAdminLockConcurrentRenewalOfSameToken(t *testing.T) {
+	ms := newAdminLockMaster()
+	for i := 0; i < 2000; i++ {
+		ms.adminLocks = NewAdminLocks()
+		held, err := leaseAdminLock(ms, "shell", nil)
+		require.NoError(t, err)
+		var mu sync.Mutex
+		renewed := 0
+		renew := func() {
+			if _, err := leaseAdminLock(ms, "shell", held); err == nil {
+				mu.Lock()
+				renewed++
+				mu.Unlock()
+			}
+		}
+		require.NoError(t, runTogether(renew, renew, renew, renew))
+		require.Equal(t, 1, renewed, "iteration %d: one token renewed %d times", i, renewed)
+	}
+}