synapsecns/sanguine

package rest

import (
    "fmt"

    "github.com/puzpuzpuz/xsync"
    "github.com/synapsecns/sanguine/services/rfq/api/model"
)

// PubSubManager is a manager for a pubsub system.
type PubSubManager interface {
    AddSubscription(relayerAddr string, params model.SubscriptionParams) error
    RemoveSubscription(relayerAddr string, params model.SubscriptionParams) error
    IsSubscribed(relayerAddr string, origin, dest int) bool
}

type pubSubManagerImpl struct {
    subscriptions *xsync.MapOf[string, map[int]struct{}]
}

// NewPubSubManager creates a new pubsub manager.
func NewPubSubManager() PubSubManager {
    return &pubSubManagerImpl{
        subscriptions: xsync.NewMapOf[map[int]struct{}](),
    }
}

func (p *pubSubManagerImpl) AddSubscription(relayerAddr string, params model.SubscriptionParams) error {
    if params.Chains == nil {
        return fmt.Errorf("chains is nil")
    }

    sub, ok := p.subscriptions.Load(relayerAddr)
    if !ok {
        sub = make(map[int]struct{})
        for _, c := range params.Chains {
            sub[c] = struct{}{}
        }
        p.subscriptions.Store(relayerAddr, sub)
        return nil
    }
    for _, c := range params.Chains {
        sub[c] = struct{}{}
    }
    return nil
}

func (p *pubSubManagerImpl) RemoveSubscription(relayerAddr string, params model.SubscriptionParams) error {
    if params.Chains == nil {
        return fmt.Errorf("chains is nil")
    }

    sub, ok := p.subscriptions.Load(relayerAddr)
    if !ok {
        return fmt.Errorf("relayer %s has no subscriptions", relayerAddr)
    }

    for _, c := range params.Chains {
        _, ok := sub[c]
        if !ok {
            return fmt.Errorf("relayer %s is not subscribed to chain %d", relayerAddr, c)
        }
        delete(sub, c)
    }

    return nil
}

func (p *pubSubManagerImpl) IsSubscribed(relayerAddr string, origin, dest int) bool {
    sub, ok := p.subscriptions.Load(relayerAddr)
    if !ok {
        return false
    }
    _, ok = sub[origin]
    if !ok {
        return false
    }
    _, ok = sub[dest]
    return ok
}