chore: leader race fixes (#28)

Author: thevilledev

Makefile

@@ -5,9 +5,14 @@ all: fmt lint test
 fmt:
 	go fmt $$(go list ./...)
 
-lint:
+lint: vet
 	golangci-lint run
 
+lit: lint
+
+vet:
+	go vet $$(go list ./...)
+
 test:
 	go test -v -race -run ^Test -parallel=8 ./...
 

leader.go

@@ -45,28 +45,35 @@ func (s *leaderState) tryBecomeLeader(ctx context.Context) error {
 }
 
 func (s *leaderState) handleElection(ctx context.Context) error {
-	if s.manager.onElected == nil {
+	// Get callback reference under read lock
+	s.mu.RLock()
+	callback := s.manager.onElected
+	s.mu.RUnlock()
+
+	if callback == nil {
 		return nil
 	}
 
-	if err := s.manager.onElected(ctx); err != nil {
+	// Release lock during potentially long-running callback
+	if err := callback(ctx); err != nil {
 		log.Printf("Error in leader callback: %v\n", err)
-
 		s.setLeader(false)
 	}
 
 	return nil
 }
 
 func (s *leaderState) handleDemotion(ctx context.Context) {
-	// Get current state before demotion
-	wasLeader := s.getIsLeader()
+	s.mu.Lock()
+	defer s.mu.Unlock()
 
-	if wasLeader && s.manager.onDemoted != nil {
+	// Need atomic operation because we need to check state and potentially
+	// call callback within the same lock
+	if s.isLeader && s.manager.onDemoted != nil {
 		s.manager.onDemoted(ctx)
 	}
 
-	s.setLeader(false)
+	s.isLeader = false
 }
 
 func (s *leaderState) runLeaderMaintenance(ctx context.Context) error {

leader_test.go

@@ -264,6 +264,7 @@ func testConcurrentMaintenanceAndDemotion(t *testing.T, state *leaderState) {
 	// Start maintenance loop
 	go func() {
 		defer wGroup.Done()
+
 		state.setLeader(true)
 		_ = state.runLeaderMaintenance(ctx)
 	}()