Use session advisory lock for leader election

Author: cursoragent

config/config.go

@@ -120,6 +120,12 @@ func load() Config {
 		}
 	}
 
+	if cfg.Elector.Enable {
+		if cfg.Elector.ID == "" || cfg.Elector.Salt == "" {
+			panic("elector configuration is required when enabled")
+		}
+	}
+
 	if cfg.Elector.Enable != cfg.Ingestors.CSFloat.Enable {
 		panic("both elector and ingestor require the same enable state")
 	}

domain/repository/factory.go

@@ -1,6 +1,7 @@
 package repository
 
 import (
+	"context"
 	"io"
 
 	"gorm.io/gorm"
@@ -18,6 +19,12 @@ type AdvisoryLock interface {
 	TryAdvisoryXactLock(id, salt string) (bool, error)
 }
 
+type AdvisoryLockSession interface {
+	io.Closer
+	TryAdvisoryLock(id, salt string) (bool, error)
+	AdvisoryUnlock(id, salt string) error
+}
+
 type PublicTransaction interface {
 	gorm.TxCommitter
 
@@ -38,4 +45,5 @@ type Factory interface {
 
 	NewPublicTransaction() PublicTransaction
 	RunInTransactionPublic(fn func(PublicTransaction) error) error
+	NewPublicAdvisoryLockSession(ctx context.Context) (AdvisoryLockSession, error)
 }

leader/leader.go

@@ -41,24 +41,34 @@ func (e *elector) Run(ctx context.Context, onLeader func()) {
 		case <-time.After(time.Until(nextMinute)):
 		}
 
-		tx := e.factory.NewPublicTransaction()
-		acquired, err := tx.TryAdvisoryXactLock(e.cfg.Elector.ID, e.cfg.Elector.Salt)
+		lockSession, err := e.factory.NewPublicAdvisoryLockSession(ctx)
+		if err != nil {
+			e.log.Warnf("failed to create advisory lock session: %v", err)
+			continue
+		}
+
+		acquired, err := lockSession.TryAdvisoryLock(e.cfg.Elector.ID, e.cfg.Elector.Salt)
 		if err != nil {
 			e.log.Warnf("failed to acquire leader lock: %v", err)
-			tx.Rollback()
+			lockSession.Close()
 			continue
 		}
 
 		if !acquired {
 			e.log.Warn("failed to acquire leader lock")
-			tx.Rollback()
+			lockSession.Close()
 			continue
 		}
 
 		e.log.Info("leader lock acquired")
 		onLeader()
 
-		tx.Rollback()
+		if err := lockSession.AdvisoryUnlock(e.cfg.Elector.ID, e.cfg.Elector.Salt); err != nil {
+			e.log.Warnf("failed to release leader lock: %v", err)
+		}
+		if err := lockSession.Close(); err != nil {
+			e.log.Warnf("failed to close advisory lock session: %v", err)
+		}
 		return
 	}
 }

repository/factory/factory.go

@@ -1,6 +1,7 @@
 package factory
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -210,6 +211,20 @@ func (f *factory) NewPublicTransaction() repository.PublicTransaction {
 	return newPublicTransaction(f.public.Begin())
 }
 
+func (f *factory) NewPublicAdvisoryLockSession(ctx context.Context) (repository.AdvisoryLockSession, error) {
+	sqlDB, err := f.public.DB()
+	if err != nil {
+		return nil, err
+	}
+
+	conn, err := sqlDB.Conn(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return newPublicAdvisoryLockSession(conn), nil
+}
+
 func (f *factory) RunInTransactionPublic(fn func(repository.PublicTransaction) error) error {
 	return f.public.Transaction(func(gormTx *gorm.DB) error {
 		tx := newPublicTransaction(gormTx)

repository/factory/transaction.go

@@ -1,6 +1,8 @@
 package factory
 
 import (
+	"context"
+	"database/sql"
 	"hash/fnv"
 
 	"reverse-watch/domain/repository"
@@ -72,13 +74,9 @@ func (t *publicTransaction) Reversal() repository.ReversalRepository {
 }
 
 func (t *publicTransaction) TryAdvisoryXactLock(id, salt string) (bool, error) {
-	h := fnv.New32a()
-	h.Write([]byte(salt))
-	h.Write([]byte(id))
-	lockKey := h.Sum32()
-
+	lockKey := advisoryLockKey(id, salt)
 	var hasLock bool
-	if err := t.tx.Raw("SELECT pg_try_advisory_xact_lock(?)", lockKey).Scan(&hasLock).Error; err != nil {
+	if err := t.tx.Raw("SELECT pg_try_advisory_lock(?)", lockKey).Scan(&hasLock).Error; err != nil {
 		return false, err
 	}
 
@@ -87,3 +85,40 @@ func (t *publicTransaction) TryAdvisoryXactLock(id, salt string) (bool, error) {
 	}
 	return true, nil
 }
+
+type publicAdvisoryLockSession struct {
+	conn *sql.Conn
+}
+
+func newPublicAdvisoryLockSession(conn *sql.Conn) *publicAdvisoryLockSession {
+	return &publicAdvisoryLockSession{conn: conn}
+}
+
+func (s *publicAdvisoryLockSession) Close() error {
+	return s.conn.Close()
+}
+
+func (s *publicAdvisoryLockSession) TryAdvisoryLock(id, salt string) (bool, error) {
+	lockKey := advisoryLockKey(id, salt)
+	var hasLock bool
+	if err := s.conn.QueryRowContext(context.Background(), "SELECT pg_try_advisory_lock($1)", lockKey).Scan(&hasLock); err != nil {
+		return false, err
+	}
+	return hasLock, nil
+}
+
+func (s *publicAdvisoryLockSession) AdvisoryUnlock(id, salt string) error {
+	lockKey := advisoryLockKey(id, salt)
+	var unlocked bool
+	if err := s.conn.QueryRowContext(context.Background(), "SELECT pg_advisory_unlock($1)", lockKey).Scan(&unlocked); err != nil {
+		return err
+	}
+	return nil
+}
+
+func advisoryLockKey(id, salt string) uint32 {
+	h := fnv.New32a()
+	h.Write([]byte(salt))
+	h.Write([]byte(id))
+	return h.Sum32()
+}