feat: add leadership package

cmd/serve.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"github.com/formancehq/go-libs/v2/logging"
 	"github.com/formancehq/ledger/internal/api/common"
+	"github.com/formancehq/ledger/internal/leadership"
 	systemstore "github.com/formancehq/ledger/internal/storage/system"
 	"net/http"
 	"net/http/pprof"
@@ -108,6 +109,7 @@ func NewServeCommand() *cobra.Command {
 				}),
 				bus.NewFxModule(),
 				ballast.Module(serveConfiguration.ballastSize),
+				leadership.NewFXModule(),
 				api.Module(api.Config{
 					Version: Version,
 					Debug:   service.IsDebug(cmd),
@@ -122,15 +124,15 @@ func NewServeCommand() *cobra.Command {
 				}),
 				fx.Decorate(func(
 					params struct {
-					fx.In
+						fx.In
 
-					Handler          chi.Router
-					HealthController *health.HealthController
-					Logger           logging.Logger
+						Handler          chi.Router
+						HealthController *health.HealthController
+						Logger           logging.Logger
 
-					MeterProvider *metric.MeterProvider         `optional:"true"`
-					Exporter      *otlpmetrics.InMemoryExporter `optional:"true"`
-				},
+						MeterProvider *metric.MeterProvider         `optional:"true"`
+						Exporter      *otlpmetrics.InMemoryExporter `optional:"true"`
+					},
 				) chi.Router {
 					return assembleFinalRouter(
 						service.IsDebug(cmd),

internal/doc.go

@@ -1,2 +1,2 @@
-//go:generate gomarkdoc -o README.md --repository.default-branch main
+//go:generate gomarkdoc -o README.md --repository.default-branch main --repository.url https://github.com/formancehq/ledger
 package ledger

internal/leadership/broadcaster.go

@@ -0,0 +1,93 @@
+package leadership
+
+import (
+	"sync"
+)
+
+type listener struct {
+	channel chan Leadership
+}
+
+type Broadcaster struct {
+	mu sync.Mutex
+	t  *Leadership
+
+	inner []listener
+	outer chan Leadership
+}
+
+func (h *Broadcaster) Actual() Leadership {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	if h.t == nil {
+		return Leadership{}
+	}
+	return *h.t
+}
+
+func (h *Broadcaster) Subscribe() (<-chan Leadership, func()) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	newChannel := make(chan Leadership, 1)
+	l := listener{
+		channel: newChannel,
+	}
+	h.inner = append(h.inner, l)
+	if h.t != nil {
+		newChannel <- *h.t
+	}
+
+	return newChannel, func() {
+		h.mu.Lock()
+		defer h.mu.Unlock()
+
+		index := -1
+		for i, listener := range h.inner {
+			if listener == l {
+				index = i
+				break
+			}
+		}
+
+		if index < len(h.inner)-1 {
+			h.inner = append(h.inner[:index], h.inner[index+1:]...)
+		} else {
+			h.inner = h.inner[:index]
+		}
+	}
+}
+
+func (h *Broadcaster) Broadcast(t Leadership) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	h.t = &t
+
+	for _, inner := range h.inner {
+		inner.channel <- t
+	}
+}
+
+func (h *Broadcaster) Close() {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	for _, inner := range h.inner {
+		close(inner.channel)
+	}
+}
+
+func (h *Broadcaster) CountListeners() int {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	return len(h.inner)
+}
+
+func NewSignal() *Broadcaster {
+	return &Broadcaster{
+		outer: make(chan Leadership),
+	}
+}

internal/leadership/context.go

@@ -0,0 +1,43 @@
+package leadership
+
+import (
+	"context"
+	"sync"
+)
+
+type contextKey struct{}
+
+var holderContextKey contextKey = struct{}{}
+
+func ContextWithLeadershipInfo(ctx context.Context) context.Context {
+	return context.WithValue(ctx, holderContextKey, &holder{})
+}
+
+func IsLeader(ctx context.Context) bool {
+	h := ctx.Value(holderContextKey)
+	if h == nil {
+		return false
+	}
+	holder := h.(*holder)
+	holder.Lock()
+	defer holder.Unlock()
+
+	return holder.isLeader
+}
+
+func setIsLeader(ctx context.Context, isLeader bool) {
+	h := ctx.Value(holderContextKey)
+	if h == nil {
+		return
+	}
+	holder := h.(*holder)
+	holder.Lock()
+	defer holder.Unlock()
+
+	holder.isLeader = isLeader
+}
+
+type holder struct {
+	sync.Mutex
+	isLeader bool
+}

internal/leadership/leadership.go

@@ -0,0 +1,6 @@
+package leadership
+
+type Leadership struct {
+	Acquired bool
+	DB       *Mutex
+}

internal/leadership/locker.go

@@ -0,0 +1,54 @@
+package leadership
+
+import (
+	"context"
+	"fmt"
+	"github.com/uptrace/bun"
+)
+
+const leadershipAdvisoryLockKey = 123456789
+
+type DBHandle interface {
+	bun.IDB
+	Close() error
+}
+
+// Locker take a lock at process level
+// It returns a bun.IDB which MUST be invalidated when the lock is lost
+type Locker interface {
+	Take(ctx context.Context) (DBHandle, error)
+}
+
+type defaultLocker struct {
+	db *bun.DB
+}
+
+func (p *defaultLocker) Take(ctx context.Context) (DBHandle, error) {
+	conn, err := p.db.Conn(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("error opening new connection: %w", err)
+	}
+
+	ret := conn.QueryRowContext(ctx, "select pg_try_advisory_lock(?)", leadershipAdvisoryLockKey)
+	if ret.Err() != nil {
+		_ = conn.Close()
+		return nil, fmt.Errorf("error acquiring lock: %w", ret.Err())
+	}
+
+	var acquired bool
+	if err := ret.Scan(&acquired); err != nil {
+		_ = conn.Close()
+		return nil, err
+	}
+
+	if !acquired {
+		_ = conn.Close()
+		return nil, nil
+	}
+
+	return conn, nil
+}
+
+func NewDefaultLocker(db *bun.DB) Locker {
+	return &defaultLocker{db: db}
+}

internal/leadership/main_test.go

@@ -0,0 +1,24 @@
+//go:build it
+
+package leadership
+
+import (
+	. "github.com/formancehq/go-libs/v2/testing/utils"
+	"testing"
+
+	"github.com/formancehq/go-libs/v2/logging"
+	"github.com/formancehq/go-libs/v2/testing/docker"
+	"github.com/formancehq/go-libs/v2/testing/platform/pgtesting"
+)
+
+var (
+	srv           *pgtesting.PostgresServer
+)
+
+func TestMain(m *testing.M) {
+	WithTestMain(func(t *TestingTForMain) int {
+		srv = pgtesting.CreatePostgresServer(t, docker.NewPool(t, logging.Testing()), pgtesting.WithExtension("pgcrypto"))
+
+		return m.Run()
+	})
+}

internal/leadership/manager.go

@@ -0,0 +1,129 @@
+package leadership
+
+import (
+	"context"
+	"github.com/formancehq/go-libs/v2/logging"
+	"github.com/uptrace/bun"
+	"time"
+)
+
+type Manager struct {
+	locker      Locker
+	changes     *Broadcaster
+	logger      logging.Logger
+	retryPeriod time.Duration
+	stopChannel chan chan struct{}
+}
+
+func (m *Manager) Run(ctx context.Context) {
+	var (
+		dbMutex   *Mutex
+		nextRetry = time.After(time.Duration(0))
+		nextPing  <-chan time.Time
+	)
+	for {
+		select {
+		case ch := <-m.stopChannel:
+			if dbMutex != nil {
+				m.logger.Info("leadership lost")
+				dbMutex.Exec(func(_ bun.IDB) {
+					_ = dbMutex.db.Close()
+				})
+
+				setIsLeader(ctx, false)
+				m.changes.Broadcast(Leadership{})
+			}
+			close(ch)
+			close(m.stopChannel)
+			return
+		case <-nextRetry:
+			db, err := m.locker.Take(ctx)
+			if err != nil || db == nil {
+				if err != nil {
+					m.logger.Error("error acquiring lock", err)
+				}
+				nextRetry = time.After(m.retryPeriod)
+				continue
+			}
+
+			dbMutex = NewMutex(db)
+
+			m.changes.Broadcast(Leadership{
+				DB:       dbMutex,
+				Acquired: true,
+			})
+			m.logger.Info("leadership acquired")
+
+			setIsLeader(ctx, true)
+
+			nextPing = time.After(m.retryPeriod)
+
+		// Ping the database to check the connection status
+		// If the connection is lost, signal the listeners about the leadership loss
+		case <-nextPing:
+			dbMutex.Exec(func(db bun.IDB) {
+				_, err := db.
+					NewSelect().
+					ColumnExpr("1 as v").
+					Count(ctx)
+				if err != nil {
+					m.logger.Error("error pinging db", err)
+					_ = dbMutex.db.Close()
+					dbMutex = nil
+
+					setIsLeader(ctx, false)
+					m.changes.Broadcast(Leadership{})
+
+					nextRetry = time.After(m.retryPeriod)
+				} else {
+					nextPing = time.After(m.retryPeriod)
+				}
+			})
+		}
+	}
+}
+
+func (m *Manager) Stop(ctx context.Context) error {
+	select {
+	// if already closed
+	case <-m.stopChannel:
+		return nil
+	default:
+		ch := make(chan struct{})
+		m.stopChannel <- ch
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ch:
+			return nil
+		}
+	}
+}
+
+func (m *Manager) GetSignal() *Broadcaster {
+	return m.changes
+}
+
+func NewManager(locker Locker, logger logging.Logger, options ...Option) *Manager {
+	l := &Manager{
+		locker:      locker,
+		logger:      logger,
+		changes:     NewSignal(),
+		retryPeriod: 2 * time.Second,
+		stopChannel: make(chan chan struct{}),
+	}
+
+	for _, option := range options {
+		option(l)
+	}
+
+	return l
+}
+
+type Option func(leadership *Manager)
+
+func WithRetryPeriod(duration time.Duration) Option {
+	return func(leadership *Manager) {
+		leadership.retryPeriod = duration
+	}
+}