Using Tester for Lock (#7)

concurrency/lock_helper_test.go

@@ -0,0 +1,83 @@
+package concurrency
+
+import (
+	"errors"
+	"strconv"
+
+	"github.com/QuangTung97/zk"
+	"github.com/QuangTung97/zk/curator"
+)
+
+type simpleCounter struct {
+	client curator.FakeClientID
+}
+
+func newSimpleCounter(client curator.FakeClientID) *simpleCounter {
+	return &simpleCounter{
+		client: client,
+	}
+}
+
+func (l *simpleCounter) isLeader(sess *curator.Session) {
+	sess.Run(func(client curator.Client) {
+		client.Get("/counter", func(resp zk.GetResponse, err error) {
+			if err != nil {
+				if errors.Is(err, zk.ErrNoNode) {
+					l.increase(sess, 1, 0)
+					return
+				}
+				if errors.Is(err, zk.ErrConnectionClosed) {
+					sess.AddRetry(l.isLeader)
+					return
+				}
+				panic(err)
+			}
+
+			num, err := strconv.ParseInt(string(resp.Data), 10, 64)
+			if err != nil {
+				panic(err)
+			}
+			l.increase(sess, int(num)+1, resp.Stat.Version)
+		})
+	})
+}
+
+func numToBytes(val int) []byte {
+	return []byte(strconv.FormatInt(int64(val), 10))
+}
+
+func (l *simpleCounter) setCounterResp(sess *curator.Session) func(_ zk.SetResponse, err error) {
+	return func(_ zk.SetResponse, err error) {
+		if err != nil {
+			if errors.Is(err, zk.ErrConnectionClosed) {
+				sess.AddRetry(l.isLeader)
+				return
+			}
+			panic(err)
+		}
+		l.isLeader(sess)
+	}
+}
+
+func (l *simpleCounter) createCounterResp(sess *curator.Session) func(_ zk.CreateResponse, err error) {
+	return func(_ zk.CreateResponse, err error) {
+		if err != nil {
+			if errors.Is(err, zk.ErrConnectionClosed) {
+				sess.AddRetry(l.isLeader)
+				return
+			}
+			panic(err)
+		}
+		l.isLeader(sess)
+	}
+}
+
+func (l *simpleCounter) increase(sess *curator.Session, nextVal int, version int32) {
+	sess.Run(func(client curator.Client) {
+		if nextVal > 1 {
+			client.Set("/counter", numToBytes(nextVal), version, l.setCounterResp(sess))
+		} else {
+			client.Create("/counter", numToBytes(nextVal), 0, l.createCounterResp(sess))
+		}
+	})
+}

concurrency/lock_test.go

@@ -1,8 +1,10 @@
 package concurrency
 
 import (
+	"fmt"
 	"slices"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 
@@ -12,6 +14,7 @@ import (
 
 const client1 curator.FakeClientID = "client1"
 const client2 curator.FakeClientID = "client2"
+const client3 curator.FakeClientID = "client3"
 const initClient curator.FakeClientID = "init"
 
 func initStore(parent string) *curator.FakeZookeeper {
@@ -336,3 +339,62 @@ func TestSortString(t *testing.T) {
 		"A", "B", "D", "EE", "IA", "IY", "M", "Z",
 	}, s)
 }
+
+func TestLock_With_Tester(t *testing.T) {
+	l1 := NewLock("/workers", "node01")
+	l2 := NewLock("/workers", "node02")
+	l3 := NewLock("/workers", "node03")
+
+	store := initStore("/workers")
+
+	tester := curator.NewFakeZookeeperTester(store,
+		[]curator.FakeClientID{client1, client2, client3},
+		123,
+	)
+
+	startLock(l1, store, client1, newSimpleCounter(client1).isLeader)
+	startLock(l2, store, client2, newSimpleCounter(client2).isLeader)
+	startLock(l3, store, client3, newSimpleCounter(client3).isLeader)
+
+	tester.Begin()
+
+	steps := tester.RunSessionExpiredAndConnectionError(
+		10,
+		10,
+		2000,
+	)
+	assert.Equal(t, 2000, steps)
+
+	store.PrintData()
+}
+
+func TestLock_With_Tester__Multi_Times(t *testing.T) {
+	for k := 0; k < 1000; k++ {
+		seed := time.Now().UnixNano()
+		fmt.Println("SEED:", seed)
+
+		l1 := NewLock("/workers", "node01")
+		l2 := NewLock("/workers", "node02")
+		l3 := NewLock("/workers", "node03")
+
+		store := initStore("/workers")
+
+		tester := curator.NewFakeZookeeperTester(store,
+			[]curator.FakeClientID{client1, client2, client3},
+			seed,
+		)
+
+		startLock(l1, store, client1, newSimpleCounter(client1).isLeader)
+		startLock(l2, store, client2, newSimpleCounter(client2).isLeader)
+		startLock(l3, store, client3, newSimpleCounter(client3).isLeader)
+
+		tester.Begin()
+
+		steps := tester.RunSessionExpiredAndConnectionError(
+			20,
+			20,
+			2000,
+		)
+		assert.Equal(t, 2000, steps)
+	}
+}

curator/fake_client.go

@@ -384,7 +384,11 @@ func (s *FakeZookeeper) notifyChildrenWatches(parent *ZNode, path string) {
 
 // ConnError ...
 func (s *FakeZookeeper) ConnError(clientID FakeClientID) {
-	s.States[clientID].ConnErr = true
+	state := s.States[clientID]
+	if state.ConnErr {
+		panic("Can NOT call ConnError multiple times in a row")
+	}
+	state.ConnErr = true
 	s.runAllCallbacksWithConnectionError(clientID, true)
 }
 
@@ -513,6 +517,7 @@ func (s *FakeZookeeper) Retry(clientID FakeClientID) {
 	getActionWithType[RetryInput](s, clientID, "Retry")
 
 	state := s.States[clientID]
+	state.ConnErr = false
 	for _, fn := range state.PendingEvents {
 		fn()
 	}

curator/fake_property.go

@@ -41,6 +41,7 @@ RetrySelect:
 		for _, c := range f.clients {
 			if !f.store.States[c].HasSession {
 				sessionExpiredClients = append(sessionExpiredClients, c)
+				continue
 			}
 
 			if len(f.store.Pending[c]) == 0 {
@@ -74,6 +75,9 @@ func (f *FakeZookeeperTester) doSessionExpired(client FakeClientID) {
 }
 
 func (f *FakeZookeeperTester) doConnectionError(client FakeClientID) {
+	if f.store.States[client].ConnErr {
+		return
+	}
 	cmds := f.store.Pending[client]
 	if len(cmds) == 0 {
 		f.store.ConnError(client)
@@ -91,7 +95,7 @@ func (f *FakeZookeeperTester) RunSessionExpiredAndConnectionError(
 	sessionExpiredPercentage float64,
 	connectionErrorPercentage float64,
 	numSteps int,
-) {
+) int {
 	sessionExpiredEnd := int(sessionExpiredPercentage / 100.0 * randMax)
 	connectionErrorEnd := int(connectionErrorPercentage / 100.0 * randMax)
 
@@ -114,14 +118,18 @@ func (f *FakeZookeeperTester) RunSessionExpiredAndConnectionError(
 
 		client, ok := f.getActionableRandomClient()
 		if !ok {
-			return
+			return i + 1
 		}
 		genericCmd := f.store.Pending[client][0]
 		switch genericCmd.(type) {
 		case CreateInput:
 			f.store.CreateApply(client)
 		case GetInput:
 			f.store.GetApply(client)
+		case ChildrenInput:
+			f.store.ChildrenApply(client)
+		case DeleteInput:
+			f.store.DeleteApply(client)
 		case SetInput:
 			f.store.SetApply(client)
 		case RetryInput:
@@ -130,4 +138,6 @@ func (f *FakeZookeeperTester) RunSessionExpiredAndConnectionError(
 			panic(fmt.Sprintf("unknown command: %s%+v", reflect.TypeOf(genericCmd).String(), genericCmd))
 		}
 	}
+
+	return numSteps
 }

curator/fake_property_test.go

@@ -143,11 +143,12 @@ func TestFakeZookeeperTester_Master_Lock(t *testing.T) {
 
 	tester.Begin()
 
-	tester.RunSessionExpiredAndConnectionError(
+	steps := tester.RunSessionExpiredAndConnectionError(
 		10,
 		10,
 		1000,
 	)
+	assert.Equal(t, 1000, steps)
 
 	store.PrintData()
 	store.PrintPendingCalls()
@@ -159,7 +160,7 @@ func TestFakeZookeeperTester_Master_Lock(t *testing.T) {
 			break
 		}
 	}
-	assert.Equal(t, "283", string(node.Data))
+	assert.Equal(t, "274", string(node.Data))
 }
 
 func TestFakeZookeeperTester_Master_Lock__Multi_Times(t *testing.T) {