refactoring and consolidation

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

core/services/advisorylock/advisorylock.go

@@ -84,3 +84,32 @@ func KeyFromUUID(b [16]byte) int64 {
 	h.Write(b[:])
 	return int64(h.Sum64()>>1) | 0x100000000
 }
+
+// KeyFromString converts an arbitrary string to an advisory lock key via FNV-1a hashing.
+func KeyFromString(s string) int64 {
+	h := fnv.New64a()
+	h.Write([]byte(s))
+	return int64(h.Sum64()>>1) | 0x100000000
+}
+
+// WithLockCtx is like WithLock but respects context cancellation.
+// If ctx is cancelled while waiting for the lock, the function returns ctx.Err().
+func WithLockCtx(ctx context.Context, db *gorm.DB, key int64, fn func() error) error {
+	sqlDB, err := db.DB()
+	if err != nil {
+		return fmt.Errorf("advisorylock: getting sql.DB: %w", err)
+	}
+	conn, err := sqlDB.Conn(ctx)
+	if err != nil {
+		return fmt.Errorf("advisorylock: getting connection: %w", err)
+	}
+	defer conn.Close()
+
+	if _, err := conn.ExecContext(ctx, "SELECT pg_advisory_lock($1)", key); err != nil {
+		return fmt.Errorf("advisorylock: acquiring lock %d: %w", key, err)
+	}
+	// Always release the lock, even if ctx is cancelled
+	defer conn.ExecContext(context.Background(), "SELECT pg_advisory_unlock($1)", key)
+
+	return fn()
+}

core/services/advisorylock/advisorylock_test.go

@@ -1,6 +1,7 @@
 package advisorylock
 
 import (
+	"context"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -108,5 +109,95 @@ var _ = Describe("AdvisoryLock", func() {
 
 			Expect(keys).To(HaveLen(6), "some advisory lock keys have the same value")
 		})
+
+		It("KeyFromString is deterministic", func() {
+			k1 := KeyFromString("foo")
+			k2 := KeyFromString("foo")
+			Expect(k1).To(Equal(k2), "KeyFromString should return the same value for the same input")
+		})
+
+		It("KeyFromString returns different keys for different inputs", func() {
+			kFoo := KeyFromString("foo")
+			kBar := KeyFromString("bar")
+			Expect(kFoo).ToNot(Equal(kBar), "KeyFromString should return different keys for different inputs")
+		})
+	})
+
+	Context("WithLockCtx (PostgreSQL)", func() {
+		var db *gorm.DB
+
+		BeforeEach(func() {
+			if runtime.GOOS == "darwin" {
+				Skip("testcontainers requires Docker, not available on macOS CI")
+			}
+			db = testutil.SetupTestDB()
+		})
+
+		It("acquires lock and executes the function", func() {
+			const lockKey int64 = 700
+			executed := false
+
+			err := WithLockCtx(context.Background(), db, lockKey, func() error {
+				executed = true
+				return nil
+			})
+			Expect(err).ToNot(HaveOccurred())
+			Expect(executed).To(BeTrue(), "function should have been executed under lock")
+		})
+
+		It("returns error when context is already cancelled", func() {
+			const lockKey int64 = 701
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel() // cancel immediately
+
+			err := WithLockCtx(ctx, db, lockKey, func() error {
+				Fail("function should not run with cancelled context")
+				return nil
+			})
+			Expect(err).To(HaveOccurred())
+		})
+
+		It("serializes concurrent WithLockCtx on same key", func() {
+			const lockKey int64 = 702
+
+			var (
+				mu          sync.Mutex
+				maxRunning  int32
+				running     int32
+				concurrency int32
+			)
+
+			var wg sync.WaitGroup
+			wg.Add(2)
+
+			for i := 0; i < 2; i++ {
+				go func() {
+					defer GinkgoRecover()
+					defer wg.Done()
+					err := WithLockCtx(context.Background(), db, lockKey, func() error {
+						cur := atomic.AddInt32(&running, 1)
+						mu.Lock()
+						if cur > maxRunning {
+							maxRunning = cur
+						}
+						if cur > 1 {
+							atomic.AddInt32(&concurrency, 1)
+						}
+						mu.Unlock()
+
+						time.Sleep(50 * time.Millisecond)
+
+						atomic.AddInt32(&running, -1)
+						return nil
+					})
+					Expect(err).ToNot(HaveOccurred())
+				}()
+			}
+
+			wg.Wait()
+
+			Expect(maxRunning).To(BeNumerically("<=", 1), "expected max 1 goroutine inside lock at a time")
+			Expect(concurrency).To(BeZero(), "detected concurrent execution inside advisory lock")
+		})
 	})
 })

core/services/advisorylock/leader_loop.go

@@ -0,0 +1,33 @@
+package advisorylock
+
+import (
+	"context"
+	"time"
+
+	"github.com/mudler/xlog"
+	"gorm.io/gorm"
+)
+
+// RunLeaderLoop runs fn on a fixed interval, guarded by a PostgreSQL advisory lock.
+// Only one instance across the cluster executes fn at a time. If the lock is not
+// acquired (another instance holds it), the tick is skipped.
+// The loop stops when ctx is cancelled.
+func RunLeaderLoop(ctx context.Context, db *gorm.DB, lockKey int64, interval time.Duration, fn func()) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			_, err := TryWithLock(db, lockKey, func() error {
+				fn()
+				return nil
+			})
+			if err != nil {
+				xlog.Error("Leader loop advisory lock error", "key", lockKey, "error", err)
+			}
+		}
+	}
+}

core/services/advisorylock/leader_loop_test.go

@@ -0,0 +1,115 @@
+package advisorylock
+
+import (
+	"context"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/services/testutil"
+)
+
+var _ = Describe("RunLeaderLoop", func() {
+	Context("PostgreSQL leader election", func() {
+		BeforeEach(func() {
+			if runtime.GOOS == "darwin" {
+				Skip("testcontainers requires Docker, not available on macOS CI")
+			}
+		})
+
+		It("executes function on tick", func() {
+			db := testutil.SetupTestDB()
+			const lockKey int64 = 5000
+
+			var callCount int32
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			go RunLeaderLoop(ctx, db, lockKey, 50*time.Millisecond, func() {
+				atomic.AddInt32(&callCount, 1)
+			})
+
+			Eventually(func() int32 {
+				return atomic.LoadInt32(&callCount)
+			}, 500*time.Millisecond, 10*time.Millisecond).Should(BeNumerically(">=", 1),
+				"expected function to be called at least once")
+		})
+
+		It("stops when context is cancelled", func() {
+			db := testutil.SetupTestDB()
+			const lockKey int64 = 5001
+
+			var callCount int32
+			ctx, cancel := context.WithCancel(context.Background())
+
+			done := make(chan struct{})
+			go func() {
+				RunLeaderLoop(ctx, db, lockKey, 50*time.Millisecond, func() {
+					atomic.AddInt32(&callCount, 1)
+				})
+				close(done)
+			}()
+
+			// Let it run a bit then cancel
+			time.Sleep(150 * time.Millisecond)
+			cancel()
+
+			// RunLeaderLoop should return
+			Eventually(done, 500*time.Millisecond).Should(BeClosed())
+
+			// Record count after cancellation
+			countAfterCancel := atomic.LoadInt32(&callCount)
+			time.Sleep(150 * time.Millisecond)
+			countLater := atomic.LoadInt32(&callCount)
+
+			Expect(countLater).To(Equal(countAfterCancel),
+				"function should stop being called after context cancellation")
+		})
+
+		It("only one leader executes at a time (two concurrent loops)", func() {
+			db := testutil.SetupTestDB()
+			const lockKey int64 = 5002
+
+			var (
+				mu         sync.Mutex
+				maxRunning int32
+				running    int32
+			)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			fn := func() {
+				cur := atomic.AddInt32(&running, 1)
+				mu.Lock()
+				if cur > maxRunning {
+					maxRunning = cur
+				}
+				mu.Unlock()
+
+				time.Sleep(30 * time.Millisecond)
+
+				atomic.AddInt32(&running, -1)
+			}
+
+			// Start two competing leader loops with the same lock key
+			go RunLeaderLoop(ctx, db, lockKey, 50*time.Millisecond, fn)
+			go RunLeaderLoop(ctx, db, lockKey, 50*time.Millisecond, fn)
+
+			// Let them run for a while
+			time.Sleep(400 * time.Millisecond)
+			cancel()
+
+			mu.Lock()
+			observed := maxRunning
+			mu.Unlock()
+
+			Expect(observed).To(BeNumerically("<=", 1),
+				"expected at most 1 goroutine running the leader function at a time")
+		})
+	})
+})

core/services/agents/dispatcher.go

@@ -10,6 +10,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/mudler/LocalAI/core/services/messaging"
+
 	coreTypes "github.com/mudler/LocalAGI/core/types"
 	"github.com/mudler/cogito"
 	"github.com/mudler/xlog"
@@ -189,15 +191,13 @@ func (d *LocalDispatcher) buildLocalCallbacks(writer SSEWriter, messageID string
 
 // --- NATS Dispatcher (distributed) ---
 
-// NATSSub is a subscription that can be unsubscribed.
-type NATSSub interface {
-	Unsubscribe() error
-}
-
 // NATSClient combines publishing and subscribing for NATS.
+// It uses messaging.Subscription so that the concrete *messaging.Client
+// (whose subscribe methods return *nats.Subscription) satisfies it via
+// a thin adapter that widens the return type.
 type NATSClient interface {
 	Publish(subject string, data any) error
-	QueueSubscribe(subject, queue string, handler func(data []byte)) (NATSSub, error)
+	QueueSubscribe(subject, queue string, handler func(data []byte)) (messaging.Subscription, error)
 }
 
 // NATSDispatcher dispatches agent chats via NATS queue group.
@@ -209,7 +209,7 @@ type NATSDispatcher struct {
 	apiKey      string
 	subject     string
 	queue       string
-	sub         NATSSub // stored subscription for cleanup
+	sub         messaging.Subscription // stored subscription for cleanup
 }
 
 // NewNATSDispatcher creates a dispatcher that uses NATS for distribution.