package aether import ( "context" "encoding/json" "fmt" "log" "sync" "github.com/google/uuid" "github.com/nats-io/nats.go" ) // NATSEventBus is an EventBus that broadcasts events across all cluster nodes using NATS type NATSEventBus struct { *EventBus // Embed base EventBus for local subscriptions nc *nats.Conn // NATS connection subscriptions []*nats.Subscription namespaceSubscribers map[string]int // Track number of subscribers per namespace nodeID string // Unique ID for this node mutex sync.Mutex ctx context.Context cancel context.CancelFunc } // eventMessage is the wire format for events sent over NATS type eventMessage struct { NodeID string `json:"node_id"` NamespaceID string `json:"namespace_id"` Event *Event `json:"event"` } // NewNATSEventBus creates a new NATS-backed event bus func NewNATSEventBus(nc *nats.Conn) (*NATSEventBus, error) { ctx, cancel := context.WithCancel(context.Background()) neb := &NATSEventBus{ EventBus: NewEventBus(), nc: nc, nodeID: uuid.New().String(), subscriptions: make([]*nats.Subscription, 0), namespaceSubscribers: make(map[string]int), ctx: ctx, cancel: cancel, } return neb, nil } // Subscribe creates a local subscription and ensures NATS subscription exists for the namespace func (neb *NATSEventBus) Subscribe(namespaceID string) <-chan *Event { neb.mutex.Lock() defer neb.mutex.Unlock() // Create local subscription first ch := neb.EventBus.Subscribe(namespaceID) // Check if this is the first subscriber for this namespace count := neb.namespaceSubscribers[namespaceID] if count == 0 { // First subscriber - create NATS subscription subject := fmt.Sprintf("aether.events.%s", namespaceID) sub, err := neb.nc.Subscribe(subject, func(msg *nats.Msg) { neb.handleNATSEvent(msg) }) if err != nil { log.Printf("[NATSEventBus] Failed to subscribe to NATS subject %s: %v", subject, err) } else { neb.subscriptions = append(neb.subscriptions, sub) log.Printf("[NATSEventBus] Node %s subscribed to %s", neb.nodeID, subject) } } neb.namespaceSubscribers[namespaceID] = count + 1 return ch } // Unsubscribe removes a local subscription and cleans up NATS subscription if no more subscribers func (neb *NATSEventBus) Unsubscribe(namespaceID string, ch <-chan *Event) { neb.mutex.Lock() defer neb.mutex.Unlock() neb.EventBus.Unsubscribe(namespaceID, ch) count := neb.namespaceSubscribers[namespaceID] if count > 0 { count-- neb.namespaceSubscribers[namespaceID] = count if count == 0 { delete(neb.namespaceSubscribers, namespaceID) log.Printf("[NATSEventBus] No more subscribers for namespace %s on node %s", namespaceID, neb.nodeID) } } } // handleNATSEvent processes events received from NATS func (neb *NATSEventBus) handleNATSEvent(msg *nats.Msg) { var eventMsg eventMessage if err := json.Unmarshal(msg.Data, &eventMsg); err != nil { log.Printf("[NATSEventBus] Failed to unmarshal event: %v", err) return } // Skip events that originated from this node (already delivered locally) if eventMsg.NodeID == neb.nodeID { return } // Forward to local EventBus subscribers neb.EventBus.Publish(eventMsg.NamespaceID, eventMsg.Event) } // Publish publishes an event both locally and to NATS for cross-node broadcasting func (neb *NATSEventBus) Publish(namespaceID string, event *Event) { // First publish locally neb.EventBus.Publish(namespaceID, event) // Then publish to NATS for other nodes subject := fmt.Sprintf("aether.events.%s", namespaceID) eventMsg := eventMessage{ NodeID: neb.nodeID, NamespaceID: namespaceID, Event: event, } data, err := json.Marshal(eventMsg) if err != nil { log.Printf("[NATSEventBus] Failed to marshal event for NATS: %v", err) return } if err := neb.nc.Publish(subject, data); err != nil { log.Printf("[NATSEventBus] Failed to publish event to NATS: %v", err) return } } // Stop closes the NATS event bus and all subscriptions func (neb *NATSEventBus) Stop() { neb.mutex.Lock() defer neb.mutex.Unlock() neb.cancel() for _, sub := range neb.subscriptions { if err := sub.Unsubscribe(); err != nil { log.Printf("[NATSEventBus] Error unsubscribing: %v", err) } } neb.subscriptions = nil neb.EventBus.Stop() log.Printf("[NATSEventBus] Node %s stopped", neb.nodeID) }