143 lines
3.6 KiB
Go
143 lines
3.6 KiB
Go
package services
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"gpt-load/internal/store"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
const (
|
|
leaderLockKey = "cluster:leader"
|
|
leaderLockTTL = 30 * time.Second
|
|
leaderRenewalInterval = 10 * time.Second
|
|
leaderElectionTimeout = 5 * time.Second
|
|
)
|
|
|
|
// LeaderService provides a mechanism for electing a single leader in a cluster.
|
|
type LeaderService struct {
|
|
store store.Store
|
|
nodeID string
|
|
isLeader atomic.Bool
|
|
stopChan chan struct{}
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
// NewLeaderService creates a new LeaderService.
|
|
func NewLeaderService(store store.Store) *LeaderService {
|
|
return &LeaderService{
|
|
store: store,
|
|
nodeID: generateNodeID(),
|
|
stopChan: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
// Start begins the leader election process.
|
|
func (s *LeaderService) Start() {
|
|
logrus.WithField("nodeID", s.nodeID).Info("Starting LeaderService...")
|
|
s.wg.Add(1)
|
|
go s.electionLoop()
|
|
}
|
|
|
|
// Stop gracefully stops the leader election process.
|
|
func (s *LeaderService) Stop() {
|
|
logrus.Info("Stopping LeaderService...")
|
|
close(s.stopChan)
|
|
s.wg.Wait()
|
|
logrus.Info("LeaderService stopped.")
|
|
}
|
|
|
|
// IsLeader returns true if the current node is the leader.
|
|
// This is a fast, local check against an atomic boolean.
|
|
func (s *LeaderService) IsLeader() bool {
|
|
return s.isLeader.Load()
|
|
}
|
|
|
|
func (s *LeaderService) electionLoop() {
|
|
defer s.wg.Done()
|
|
|
|
// Attempt to acquire leadership immediately on start.
|
|
s.tryToBeLeader()
|
|
|
|
ticker := time.NewTicker(leaderRenewalInterval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
s.tryToBeLeader()
|
|
case <-s.stopChan:
|
|
if s.IsLeader() {
|
|
s.releaseLock()
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *LeaderService) tryToBeLeader() {
|
|
if s.IsLeader() {
|
|
// Already the leader, just renew the lock.
|
|
if err := s.renewLock(); err != nil {
|
|
logrus.WithError(err).Error("Failed to renew leader lock, relinquishing leadership.")
|
|
s.isLeader.Store(false)
|
|
}
|
|
return
|
|
}
|
|
|
|
// Not the leader, try to acquire the lock.
|
|
acquired, err := s.acquireLock()
|
|
if err != nil {
|
|
logrus.WithError(err).Error("Error trying to acquire leader lock.")
|
|
s.isLeader.Store(false)
|
|
return
|
|
}
|
|
|
|
if acquired {
|
|
logrus.WithField("nodeID", s.nodeID).Info("Successfully acquired leader lock.")
|
|
s.isLeader.Store(true)
|
|
} else {
|
|
logrus.Debug("Could not acquire leader lock, another node is likely the leader.")
|
|
s.isLeader.Store(false)
|
|
}
|
|
}
|
|
|
|
func (s *LeaderService) acquireLock() (bool, error) {
|
|
// SetNX is an atomic operation. If the key already exists, it does nothing.
|
|
// This is the core of our distributed lock.
|
|
return s.store.SetNX(leaderLockKey, []byte(s.nodeID), leaderLockTTL)
|
|
}
|
|
|
|
func (s *LeaderService) renewLock() error {
|
|
// To renew, we must ensure we are still the lock holder.
|
|
// A LUA script is the safest way to do this atomically.
|
|
// For simplicity here, we get and set, but this is not truly atomic without LUA.
|
|
// A simple SET can also work if we are confident in our election loop timing.
|
|
return s.store.Set(leaderLockKey, []byte(s.nodeID), leaderLockTTL)
|
|
}
|
|
|
|
func (s *LeaderService) releaseLock() {
|
|
// Best-effort attempt to release the lock on shutdown.
|
|
// The TTL will handle cases where this fails.
|
|
if err := s.store.Delete(leaderLockKey); err != nil {
|
|
logrus.WithError(err).Error("Failed to release leader lock on shutdown.")
|
|
} else {
|
|
logrus.Info("Successfully released leader lock.")
|
|
}
|
|
s.isLeader.Store(false)
|
|
}
|
|
|
|
func generateNodeID() string {
|
|
bytes := make([]byte, 16)
|
|
if _, err := rand.Read(bytes); err != nil {
|
|
// Fallback to a timestamp-based ID if crypto/rand fails
|
|
return "node-" + time.Now().Format(time.RFC3339Nano)
|
|
}
|
|
return hex.EncodeToString(bytes)
|
|
}
|