Building a Distributed Locking System in Go with ZooKeeper: From Theory to Production

go get github.com/samuel/go-zookeeper/zk

package lock
import (
	"fmt"
	"sort"
	"time"
	"github.com/samuel/go-zookeeper/zk"
)
type DistributedLocker struct {
	conn      *zk.Conn
	basePath  string
	lockPath  string
	timeout   time.Duration
	sessionID string
}
// NewDistributedLocker creates a new distributed locker instance
func NewDistributedLocker(hosts []string, basePath string, 
	lockTimeout time.Duration) (*DistributedLocker, error) {
	// Connect to ZooKeeper ensemble
	conn, _, err := zk.Connect(hosts, 10*time.Second)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to ZooKeeper: %w", err)
	}
	// Ensure base path exists (create if needed)
	exists, _, err := conn.Exists(basePath)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to check path existence: %w", err)
	}
	if !exists {
		// Create path with empty data - this is just infrastructure
		_, err = conn.Create(basePath, []byte{}, 0, zk.WorldACL(zk.PermAll))
		if err != nil && err != zk.ErrNodeExists {
			conn.Close()
			return nil, fmt.Errorf("failed to create lock path: %w", err)
		}
	}
	return &DistributedLocker{
		conn:      conn,
		basePath:  basePath,
		timeout:   lockTimeout,
		sessionID: fmt.Sprintf("%d", time.Now().UnixNano()),
	}, nil
}
// Lock attempts to acquire the distributed lock
func (dl *DistributedLocker) Lock() error {
	// Create our ephemeral sequential node
	// The flags: 0 is the version, zk.FlagEphemeral | zk.FlagSequence creates
	// an ephemeral sequential node that ZK automatically deletes on disconnection
	lockNode := fmt.Sprintf("%s/lock-", dl.basePath)
	createdPath, err := dl.conn.Create(lockNode, []byte(dl.sessionID), 
		zk.FlagEphemeral|zk.FlagSequence, zk.WorldACL(zk.PermAll))
	if err != nil {
		return fmt.Errorf("failed to create lock node: %w", err)
	}
	dl.lockPath = createdPath
	// Now we need to check if we have the lock or wait for it
	for {
		// Get all lock nodes and sort them
		children, _, err := dl.conn.Children(dl.basePath)
		if err != nil {
			return fmt.Errorf("failed to get children: %w", err)
		}
		sort.Strings(children)
		// Extract just the node name from our full path
		ourNodeName := createdPath[len(dl.basePath)+1:]
		// Find our position
		ourIndex := -1
		for i, child := range children {
			if child == ourNodeName {
				ourIndex = i
				break
			}
		}
		if ourIndex == -1 {
			return fmt.Errorf("our lock node disappeared mysteriously (bad sign)")
		}
		// If we're first, we have the lock!
		if ourIndex == 0 {
			return nil
		}
		// Otherwise, watch the node ahead of us
		// If it disappears, we'll be notified and try again
		nodeBefore := children[ourIndex-1]
		watchPath := fmt.Sprintf("%s/%s", dl.basePath, nodeBefore)
		exists, _, eventChan, err := dl.conn.ExistsW(watchPath)
		if err != nil {
			return fmt.Errorf("failed to set watch: %w", err)
		}
		if !exists {
			// The node we were watching is already gone, retry immediately
			continue
		}
		// Wait for either a deletion event or timeout
		select {
		case event := <-eventChan:
			if event.Type == zk.EventNodeDeleted {
				// The node ahead of us was deleted, try to acquire the lock
				continue
			}
		case <-time.After(dl.timeout):
			// Clean up on timeout
			dl.conn.Delete(dl.lockPath, -1)
			return fmt.Errorf("lock acquisition timeout")
		}
	}
}
// Unlock releases the distributed lock
func (dl *DistributedLocker) Unlock() error {
	if dl.lockPath == "" {
		return fmt.Errorf("lock not acquired")
	}
	err := dl.conn.Delete(dl.lockPath, -1)
	if err != nil && err != zk.ErrNoNode {
		return fmt.Errorf("failed to delete lock node: %w", err)
	}
	dl.lockPath = ""
	return nil
}
// Close closes the ZooKeeper connection
func (dl *DistributedLocker) Close() error {
	if dl.lockPath != "" {
		dl.conn.Delete(dl.lockPath, -1)
	}
	dl.conn.Close()
	return nil
}

package main
import (
	"fmt"
	"log"
	"time"
	"yourmodule/lock"
)
func main() {
	// Configuration
	zkHosts := []string{"localhost:2181"}
	basePath := "/my-app/locks"
	lockTimeout := 30 * time.Second
	// Create locker
	locker, err := lock.NewDistributedLocker(zkHosts, basePath, lockTimeout)
	if err != nil {
		log.Fatalf("Failed to create locker: %v", err)
	}
	defer locker.Close()
	// Use it like a regular mutex
	if err := locker.Lock(); err != nil {
		log.Fatalf("Failed to acquire lock: %v", err)
	}
	defer locker.Unlock()
	// Critical section - only one service in here at a time
	fmt.Println("Lock acquired! Doing critical work...")
	criticalOperation()
	fmt.Println("Work complete, releasing lock")
}
func criticalOperation() {
	// Simulate some work that needs to be exclusively accessed
	time.Sleep(2 * time.Second)
}

// Monitor the connection state
go func() {
	for {
		select {
		case event := <-locker.conn.SessionEvent():
			if event == zk.StateDisconnected {
				log.Println("Lost connection to ZooKeeper!")
				// Handle reconnection or graceful shutdown
			}
		}
	}
}()

// Good pattern: defer ensures release even on panic
if err := locker.Lock(); err != nil {
	return err
}
defer locker.Unlock()
// Your critical section here

// For read locks, create: /locks/read-0000000001
// For write locks, create: /locks/write-0000000001
// In the ordering, all locks compete, but:
// - A read lock is satisfied when no write locks before it exist
// - A write lock is satisfied when no other locks before it exist
type LockType string
const (
	ReadLock  LockType = "read-"
	WriteLock LockType = "write-"
)
// Modify Lock() to accept lock type
func (dl *DistributedLocker) LockWithType(lockType LockType) error {
	lockNode := fmt.Sprintf("%s/%s", dl.basePath, lockType)
	// ... rest of the lock acquisition logic
	// When checking if we can acquire, also check for incompatible lock types
}

package lock
import (
	"testing"
	"time"
)
func TestDistributedLockSequence(t *testing.T) {
	zkHosts := []string{"localhost:2181"}
	basePath := "/test/locks"
	// Create two lockers
	locker1, err := NewDistributedLocker(zkHosts, basePath, 30*time.Second)
	if err != nil {
		t.Fatalf("Failed to create locker1: %v", err)
	}
	defer locker1.Close()
	locker2, err := NewDistributedLocker(zkHosts, basePath, 30*time.Second)
	if err != nil {
		t.Fatalf("Failed to create locker2: %v", err)
	}
	defer locker2.Close()
	// Acquire first lock
	if err := locker1.Lock(); err != nil {
		t.Fatalf("Failed to acquire lock: %v", err)
	}
	// Try to acquire second lock (should not block in this test)
	acquired := make(chan bool)
	go func() {
		err := locker2.Lock()
		acquired <- err == nil
	}()
	// Second locker should not acquire the lock immediately
	select {
	case <-acquired:
		t.Fatal("Second locker acquired lock while first still holds it")
	case <-time.After(100 * time.Millisecond):
		// Good, it's waiting
	}
	// Release first lock
	locker1.Unlock()
	// Now second locker should acquire it
	select {
	case ok := <-acquired:
		if !ok {
			t.Fatal("Second locker failed to acquire lock")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("Second locker never acquired lock")
	}
	locker2.Unlock()
}

start := time.Now()
if err := locker.Lock(); err != nil {
	log.Printf("Lock acquisition failed after %v", time.Since(start))
}
// Log acquisition time for monitoring
lockAcquisitionTime.Observe(time.Since(start).Seconds())

sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
go func() {
	<-sigChan
	locker.Unlock()
	locker.Close()
	os.Exit(0)
}()