utils: add subscription manager

Author: sputn1ck

utils/subscription.go

@@ -0,0 +1,115 @@
+package utils
+
+import (
+	"context"
+	"sync"
+	"time"
+)
+
+type Subscription[E any] interface {
+	Subscribe(ctx context.Context) (eventChan <-chan E, errChan <-chan error, err error)
+	HandleEvent(event E) error
+	HandleError(err error)
+}
+
+type SubscriptionManager[E any] struct {
+	subscription Subscription[E]
+	isSubscribed bool
+	mutex        sync.Mutex
+	backoff      time.Duration
+}
+
+func NewSubscriptionManager[E any](subscription Subscription[E]) *SubscriptionManager[E] {
+	return &SubscriptionManager[E]{
+		subscription: subscription,
+		backoff:      time.Second,
+	}
+}
+
+func (sm *SubscriptionManager[E]) IsSubscribed() bool {
+	sm.mutex.Lock()
+	defer sm.mutex.Unlock()
+	return sm.isSubscribed
+}
+
+func (sm *SubscriptionManager[E]) Start(ctx context.Context) {
+	sm.mutex.Lock()
+	if sm.isSubscribed {
+		sm.mutex.Unlock()
+		return
+	}
+	sm.mutex.Unlock()
+
+	go sm.manageSubscription(ctx)
+}
+
+func (sm *SubscriptionManager[E]) manageSubscription(ctx context.Context) {
+	defer func() {
+		sm.mutex.Lock()
+		sm.isSubscribed = false
+		sm.mutex.Unlock()
+	}()
+
+	for {
+		eventChan, errChan, err := sm.subscription.Subscribe(ctx)
+		if err != nil {
+			sm.subscription.HandleError(err)
+			if !sm.shouldRetry(ctx) {
+				return
+			}
+			continue
+		}
+
+		sm.mutex.Lock()
+		sm.isSubscribed = true
+		sm.mutex.Unlock()
+
+	handleLoop:
+		for {
+			select {
+			case event, ok := <-eventChan:
+				if !ok {
+					if !sm.shouldRetry(ctx) {
+						return
+					}
+					break
+				}
+				if err := sm.subscription.HandleEvent(event); err != nil {
+					sm.subscription.HandleError(err)
+				}
+			case err, ok := <-errChan:
+				if !ok || err == nil {
+					if !sm.shouldRetry(ctx) {
+						return
+					}
+					break
+				}
+				sm.subscription.HandleError(err)
+				if !sm.shouldRetry(ctx) {
+					return
+				}
+				break handleLoop
+			case <-ctx.Done():
+				return
+			}
+		}
+	}
+}
+
+func (sm *SubscriptionManager[E]) shouldRetry(ctx context.Context) bool {
+	sm.mutex.Lock()
+	sm.isSubscribed = false
+	sm.mutex.Unlock()
+
+	select {
+	case <-ctx.Done():
+		return false
+	case <-time.After(sm.backoff):
+		// Exponential backoff with cap
+		if sm.backoff < time.Minute {
+			sm.backoff *= 2
+		}
+		sm.Start(ctx)
+		return true
+	}
+}