feat: monitor connection

Author: gfyrag

internal/leadership/leadership.go

@@ -2,5 +2,5 @@ package leadership
 
 type Leadership struct {
 	Acquired bool
-	DB       DBHandle
+	DB       *Mutex
 }

internal/leadership/manager.go

@@ -3,6 +3,7 @@ package leadership
 import (
 	"context"
 	"github.com/formancehq/go-libs/v2/logging"
+	"github.com/uptrace/bun"
 	"time"
 )
 
@@ -16,24 +17,27 @@ type Manager struct {
 
 func (m *Manager) Run(ctx context.Context) {
 	var (
-		db        DBHandle
+		dbMutex   *Mutex
 		nextRetry = time.After(time.Duration(0))
-		err       error
+		nextPing  <-chan time.Time
 	)
 	for {
 		select {
 		case ch := <-m.stopChannel:
-			if db != nil {
+			if dbMutex != nil {
 				m.logger.Info("leadership lost")
-				_ = db.Close()
+				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)
+			db, err := m.locker.Take(ctx)
 			if err != nil || db == nil {
 				if err != nil {
 					m.logger.Error("error acquiring lock", err)
@@ -42,13 +46,39 @@ func (m *Manager) Run(ctx context.Context) {
 				continue
 			}
 
+			dbMutex = NewMutex(db)
+
 			m.changes.Broadcast(Leadership{
-				DB:       db,
+				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)
+				}
+			})
 		}
 	}
 }

internal/leadership/manager_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/formancehq/go-libs/v2/bun/bunconnect"
 	"github.com/formancehq/go-libs/v2/logging"
 	"github.com/stretchr/testify/require"
+	"github.com/uptrace/bun"
 	"testing"
 	"time"
 )
@@ -52,15 +53,18 @@ func TestLeaderShip(t *testing.T) {
 	require.Equal(t, 1, leaderCount)
 	require.GreaterOrEqual(t, selectedLeader, 0)
 
-	// ensure the provided db connection is still functionnal
-	require.NoError(t, instances[selectedLeader].
+	// ensure the provided db connection is still functional
+	instances[selectedLeader].
 		GetSignal().
 		Actual().DB.
-		NewSelect().
-		Model(&map[string]any{}).
-		ColumnExpr("1 as v").
-		Scan(ctx),
-	)
+		Exec(func(db bun.IDB) {
+			require.NoError(t, db.
+				NewSelect().
+				Model(&map[string]any{}).
+				ColumnExpr("1 as v").
+				Scan(ctx),
+			)
+		})
 
 	require.NoError(t, instances[selectedLeader].Stop(ctx))
 

internal/leadership/mutex.go

@@ -1,20 +1,25 @@
 package leadership
 
-import "sync"
+import (
+	"github.com/uptrace/bun"
+	"sync"
+)
 
-type Mutex[T any] struct {
+type Mutex struct {
 	*sync.Mutex
-	t T
+	db DBHandle
 }
 
-func (m *Mutex[T]) Lock() (T, func()) {
+func (m *Mutex) Exec(fn func(db bun.IDB)) {
 	m.Mutex.Lock()
-	return m.t, m.Unlock
+	defer m.Mutex.Unlock()
+
+	fn(m.db)
 }
 
-func NewMutex[T any](t T) *Mutex[T] {
-	return &Mutex[T]{
+func NewMutex(db DBHandle) *Mutex {
+	return &Mutex{
 		Mutex: &sync.Mutex{},
-		t:     t,
+		db:    db,
 	}
 }