Пример #1
0
// 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
		}
	}
}
Пример #2
0
// 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
		}
	}
}