// ListenSubscriptions returns a stream of matching subscriptions for the current node func (lb *LogBroker) ListenSubscriptions(request *api.ListenSubscriptionsRequest, stream api.LogBroker_ListenSubscriptionsServer) error { remote, err := ca.RemoteNode(stream.Context()) if err != nil { return err } log := log.G(stream.Context()).WithFields( logrus.Fields{ "method": "(*LogBroker).ListenSubscriptions", "node": remote.NodeID, }, ) subscriptions, subscriptionCh, subscriptionCancel := lb.watchSubscriptions() defer subscriptionCancel() log.Debug("node registered") // Start by sending down all active subscriptions. for _, subscription := range subscriptions { select { case <-stream.Context().Done(): return stream.Context().Err() case <-lb.pctx.Done(): return nil default: } if err := stream.Send(subscription); err != nil { log.Error(err) return err } } // Send down new subscriptions. // TODO(aluzzardi): We should filter by relevant tasks for this node rather for { select { case v := <-subscriptionCh: subscription := v.(*api.SubscriptionMessage) if err := stream.Send(subscription); err != nil { log.Error(err) return err } case <-stream.Context().Done(): return stream.Context().Err() case <-lb.pctx.Done(): return nil } } }
// ListenSubscriptions returns a stream of matching subscriptions for the current node func (lb *LogBroker) ListenSubscriptions(request *api.ListenSubscriptionsRequest, stream api.LogBroker_ListenSubscriptionsServer) error { remote, err := ca.RemoteNode(stream.Context()) if err != nil { return err } lb.nodeConnected(remote.NodeID) defer lb.nodeDisconnected(remote.NodeID) log := log.G(stream.Context()).WithFields( logrus.Fields{ "method": "(*LogBroker).ListenSubscriptions", "node": remote.NodeID, }, ) subscriptions, subscriptionCh, subscriptionCancel := lb.watchSubscriptions(remote.NodeID) defer subscriptionCancel() log.Debug("node registered") activeSubscriptions := make(map[string]*subscription) defer func() { // If the worker quits, mark all active subscriptions as finished. for _, subscription := range activeSubscriptions { subscription.Done(remote.NodeID, fmt.Errorf("node %s disconnected unexpectedly", remote.NodeID)) } }() // Start by sending down all active subscriptions. for _, subscription := range subscriptions { select { case <-stream.Context().Done(): return stream.Context().Err() case <-lb.pctx.Done(): return nil default: } if err := stream.Send(subscription.message); err != nil { log.Error(err) return err } activeSubscriptions[subscription.message.ID] = subscription } // Send down new subscriptions. for { select { case v := <-subscriptionCh: subscription := v.(*subscription) if subscription.message.Close { log.WithField("subscription.id", subscription.message.ID).Debug("subscription closed") delete(activeSubscriptions, subscription.message.ID) } else { // Avoid sending down the same subscription multiple times if _, ok := activeSubscriptions[subscription.message.ID]; ok { continue } activeSubscriptions[subscription.message.ID] = subscription log.WithField("subscription.id", subscription.message.ID).Debug("subscription added") } if err := stream.Send(subscription.message); err != nil { log.Error(err) return err } case <-stream.Context().Done(): return stream.Context().Err() case <-lb.pctx.Done(): return nil } } }