Watch for group member subscription changes

We used to set watch only on the group membership. To
notify group members about subscription changes members
had to deregister and then register again.

Author: horkhe

CHANGELOG.md

@@ -2,9 +2,14 @@
 
 #### Version 0.14.1 (TBD)
 
+Implemented:
+* Consumer.RebalanceTimeout was removed, so rebalancing is triggered as soon
+  as membership status of a consumer group or subscription of a consumer group
+  member changes.
+
 Fixed:
 * [#120](https://github.com/mailgun/kafka-pixy/pull/120) Consumption from a
-  topic stopped a group 
+  topic stopped for a group. 
 
 #### Version 0.14.0 (2017-09-11)
 

config/config.go

@@ -376,8 +376,6 @@ func (p *Proxy) validate() error {
 		return errors.New("consumer.offsets_commit_interval must be > 0")
 	case p.Consumer.OffsetsCommitTimeout <= 0:
 		return errors.New("consumer.offsets_commit_timeout must be > 0")
-	case p.Consumer.RebalanceDelay <= 0:
-		return errors.New("consumer.rebalance_delay must be > 0")
 	case p.Consumer.SubscriptionTimeout <= 0:
 		return errors.New("consumer.subscription_timeout must be > 0")
 	case p.Consumer.RetryBackoff <= 0:
@@ -428,7 +426,6 @@ func defaultProxyWithClientID(clientID string) *Proxy {
 	c.Consumer.MaxRetries = -1
 	c.Consumer.OffsetsCommitInterval = 500 * time.Millisecond
 	c.Consumer.OffsetsCommitTimeout = 1500 * time.Millisecond
-	c.Consumer.RebalanceDelay = 250 * time.Millisecond
 	c.Consumer.SubscriptionTimeout = 15 * time.Second
 	c.Consumer.RetryBackoff = 500 * time.Millisecond
 	return c

consumer/subscriber/subscriber.go

@@ -1,6 +1,7 @@
 package subscriber
 
 import (
+	"context"
 	"sort"
 	"sync"
 	"time"
@@ -23,13 +24,16 @@ const safeClaimRetriesCount = 10
 // leave and update their subscriptions, and generates notifications of such
 // changes. It also provides an API to for a partition consumer to claim and
 // release a group-topic-partition.
+//
+// FIXME: It is assumed that all members of the group are registered with the
+// FIXME: `static` pattern. If a member that pattern is either `white_list` or
+// FIXME: `black_list` joins the group the result will be unpredictable.
 type T struct {
 	actDesc          *actor.Descriptor
 	cfg              *config.Proxy
 	group            string
 	groupZNode       *kazoo.Consumergroup
 	groupMemberZNode *kazoo.ConsumergroupInstance
-	topics           []string
 	topicsCh         chan []string
 	subscriptionsCh  chan map[string][]string
 	stopCh           chan none.T
@@ -149,133 +153,126 @@ func (ss *T) run() {
 
 	var (
 		nilOrSubscriptionsCh     chan<- map[string][]string
-		nilOrGroupUpdatedCh      <-chan zk.Event
+		nilOrWatchCh             <-chan none.T
 		nilOrTimeoutCh           <-chan time.Time
-		pendingTopics            []string
-		pendingSubscriptions     map[string][]string
+		cancelWatch              context.CancelFunc
 		shouldSubmitTopics       = false
-		shouldFetchMembers       = false
 		shouldFetchSubscriptions = false
-		members                  []*kazoo.ConsumergroupInstance
+		topics                   []string
+		subscriptions            map[string][]string
 	)
 	for {
 		select {
-		case topics := <-ss.topicsCh:
-			pendingTopics = normalizeTopics(topics)
-			shouldSubmitTopics = !topicsEqual(pendingTopics, ss.topics)
-		case nilOrSubscriptionsCh <- pendingSubscriptions:
+		case topics = <-ss.topicsCh:
+			sort.Strings(topics)
+			shouldSubmitTopics = true
+		case nilOrSubscriptionsCh <- subscriptions:
 			nilOrSubscriptionsCh = nil
-		case <-nilOrGroupUpdatedCh:
-			nilOrGroupUpdatedCh = nil
-			shouldFetchMembers = true
+		case <-nilOrWatchCh:
+			nilOrWatchCh = nil
+			cancelWatch()
+			shouldFetchSubscriptions = true
 		case <-nilOrTimeoutCh:
+			nilOrTimeoutCh = nil
 		case <-ss.stopCh:
+			if cancelWatch != nil {
+				cancelWatch()
+			}
 			return
 		}
 
 		if shouldSubmitTopics {
-			if err = ss.submitTopics(pendingTopics); err != nil {
+			if err = ss.submitTopics(topics); err != nil {
 				ss.actDesc.Log().WithError(err).Error("Failed to submit topics")
 				nilOrTimeoutCh = time.After(ss.cfg.Consumer.RetryBackoff)
 				continue
 			}
-			ss.actDesc.Log().Infof("Submitted: topics=%v", pendingTopics)
+			ss.actDesc.Log().Infof("Submitted: topics=%v", topics)
 			shouldSubmitTopics = false
-			shouldFetchMembers = true
-		}
-
-		if shouldFetchMembers {
-			members, nilOrGroupUpdatedCh, err = ss.groupZNode.WatchInstances()
-			if err != nil {
-				ss.actDesc.Log().WithError(err).Error("Failed to watch members")
-				nilOrTimeoutCh = time.After(ss.cfg.Consumer.RetryBackoff)
-				continue
+			if cancelWatch != nil {
+				cancelWatch()
 			}
-			shouldFetchMembers = false
 			shouldFetchSubscriptions = true
-			// To avoid unnecessary rebalancing in case of a deregister/register
-			// sequences that happen when a member updates its topic subscriptions,
-			// we delay subscription fetching.
-			nilOrTimeoutCh = time.After(ss.cfg.Consumer.RebalanceDelay)
-			continue
 		}
 
 		if shouldFetchSubscriptions {
-			pendingSubscriptions, err = ss.fetchSubscriptions(members)
+			subscriptions, nilOrWatchCh, cancelWatch, err = ss.fetchSubscriptions()
 			if err != nil {
 				ss.actDesc.Log().WithError(err).Error("Failed to fetch subscriptions")
 				nilOrTimeoutCh = time.After(ss.cfg.Consumer.RetryBackoff)
 				continue
 			}
 			shouldFetchSubscriptions = false
-			ss.actDesc.Log().Infof("Fetched subscriptions: %v", pendingSubscriptions)
+			ss.actDesc.Log().Infof("Fetched subscriptions: %v", subscriptions)
 			nilOrSubscriptionsCh = ss.subscriptionsCh
 		}
 	}
 }
 
-// fetchSubscriptions retrieves registration records for the specified members
-// from ZooKeeper.
-//
-// FIXME: It is assumed that all members of the group are registered with the
-// FIXME: `static` pattern. If a member that pattern is either `white_list` or
-// FIXME: `black_list` joins the group the result will be unpredictable.
-func (ss *T) fetchSubscriptions(members []*kazoo.ConsumergroupInstance) (map[string][]string, error) {
+// fetchSubscriptions retrieves subscription topics for all group members and
+// returns a channel that will be closed
+func (ss *T) fetchSubscriptions() (map[string][]string, <-chan none.T, context.CancelFunc, error) {
+	members, groupUpdateWatchCh, err := ss.groupZNode.WatchInstances()
+	if err != nil {
+		return nil, nil, nil, errors.Wrapf(err, "failed to watch members")
+	}
+
+	memberUpdateWatchChs := make([]<-chan zk.Event, len(members))
 	subscriptions := make(map[string][]string, len(members))
-	for _, member := range members {
+	for i, member := range members {
 		var registration *kazoo.Registration
-		registration, err := member.Registration()
+		registration, memberUpdateWatchCh, err := member.WatchRegistration()
 		for err != nil {
-			return nil, errors.Wrapf(err, "Failed to fetch registration, member=%s", member.ID)
+			return nil, nil, nil, errors.Wrapf(err, "failed to watch registration, member=%s", member.ID)
 		}
-		// Sort topics to ensure deterministic output.
+		memberUpdateWatchChs[i] = memberUpdateWatchCh
+
 		topics := make([]string, 0, len(registration.Subscription))
 		for topic := range registration.Subscription {
 			topics = append(topics, topic)
 		}
-		subscriptions[member.ID] = normalizeTopics(topics)
+		// Sort topics to ensure deterministic output.
+		sort.Strings(topics)
+		subscriptions[member.ID] = topics
 	}
-	return subscriptions, nil
+
+	aggregateWatchCh := make(chan none.T)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go forwardWatch(ctx, groupUpdateWatchCh, aggregateWatchCh)
+	for _, memberUpdateWatchCh := range memberUpdateWatchChs {
+		go forwardWatch(ctx, memberUpdateWatchCh, aggregateWatchCh)
+	}
+
+	return subscriptions, aggregateWatchCh, cancel, nil
 }
 
 func (ss *T) submitTopics(topics []string) error {
-	if len(ss.topics) != 0 {
+	if len(topics) == 0 {
 		err := ss.groupMemberZNode.Deregister()
 		if err != nil && err != kazoo.ErrInstanceNotRegistered {
-			return errors.Wrap(err, "Failed to deregister")
+			return errors.Wrap(err, "failed to deregister")
 		}
+		return nil
 	}
-	ss.topics = nil
-	if len(topics) != 0 {
-		err := ss.groupMemberZNode.Register(topics)
-		for err != nil {
-			return errors.Wrap(err, "Failed to register")
-		}
+	err := ss.groupMemberZNode.Register(topics)
+	for err != nil {
+		return errors.Wrap(err, "failed to register")
 	}
-	ss.topics = topics
 	return nil
 }
 
-func normalizeTopics(s []string) []string {
-	if len(s) == 0 {
-		return nil
-	}
-	sort.Strings(s)
-	return s
+func millisSince(t time.Time) time.Duration {
+	return time.Now().Sub(t) / time.Millisecond * time.Millisecond
 }
 
-func topicsEqual(lhs, rhs []string) bool {
-	if len(lhs) != len(rhs) {
-		return false
-	}
-	for i := range lhs {
-		if lhs[i] != rhs[i] {
-			return false
+func forwardWatch(ctx context.Context, watchCh <-chan zk.Event, downstreamCh chan<- none.T) {
+	select {
+	case <-watchCh:
+		select {
+		case downstreamCh <- none.V:
+		case <-ctx.Done():
 		}
+	case <-ctx.Done():
 	}
-	return true
-}
-
-func millisSince(t time.Time) time.Duration {
-	return time.Now().Sub(t) / time.Millisecond * time.Millisecond
 }

consumer/subscriber/subscriber_test.go

@@ -1,6 +1,7 @@
 package subscriber
 
 import (
+	"reflect"
 	"sync"
 	"testing"
 	"time"
@@ -35,25 +36,6 @@ func (s *GroupMemberSuite) SetUpTest(c *C) {
 	s.ns = actor.Root().NewChild("T")
 }
 
-func (s *GroupMemberSuite) TestNormalizeTopics(c *C) {
-	c.Assert(normalizeTopics(nil), DeepEquals, []string(nil))
-	c.Assert(normalizeTopics([]string{}), DeepEquals, []string(nil))
-	c.Assert(normalizeTopics([]string{"c", "a", "b"}), DeepEquals, []string{"a", "b", "c"})
-
-	c.Assert(normalizeTopics([]string{"c", "a", "b"}), Not(DeepEquals), []string{"a", "b"})
-}
-
-func (s *GroupMemberSuite) TestTopicsEqual(c *C) {
-	c.Assert(topicsEqual([]string{}, nil), Equals, true)
-	c.Assert(topicsEqual(nil, []string{}), Equals, true)
-	c.Assert(topicsEqual([]string{}, []string{}), Equals, true)
-	c.Assert(topicsEqual([]string{"a"}, []string{"a"}), Equals, true)
-	c.Assert(topicsEqual([]string{"a", "b", "c"}, []string{"a", "b", "c"}), Equals, true)
-
-	c.Assert(topicsEqual([]string{"a", "b", "c"}, []string{"a", "b"}), Equals, false)
-	c.Assert(topicsEqual([]string{"a", "b"}, []string{"b", "a"}), Equals, false)
-}
-
 // When a list of topics is sent to the `topics()` channel, a membership change
 // is received with the same list of topics for the registrator name.
 func (s *GroupMemberSuite) TestSimpleSubscribe(c *C) {
@@ -86,8 +68,8 @@ func (s *GroupMemberSuite) TestSubscribeSequence(c *C) {
 		map[string][]string{"m1": {"bazz", "blah"}})
 }
 
-// If a group member resubscribes to the same list of topics, then nothing is
-// updated.
+// If a group member resubscribes to the same list of topics, then the same
+// member subscriptions are returned.
 func (s *GroupMemberSuite) TestReSubscribe(c *C) {
 	cfg1 := newConfig("m1")
 	ss1 := Spawn(s.ns.NewChild("m1"), "g1", cfg1, s.kazooClt)
@@ -103,20 +85,15 @@ func (s *GroupMemberSuite) TestReSubscribe(c *C) {
 		"m1": {"bar", "foo"},
 		"m2": {"bar", "bazz"},
 	}
-	c.Assert(<-ss1.Subscriptions(), DeepEquals, membership)
-	c.Assert(<-ss2.Subscriptions(), DeepEquals, membership)
+	wait4Subscription(c, ss1.Subscriptions(), membership, 3*time.Second)
+	wait4Subscription(c, ss2.Subscriptions(), membership, 3*time.Second)
 
 	// When
 	ss1.Topics() <- []string{"foo", "bar"}
 
 	// Then
-	select {
-	case update := <-ss1.Subscriptions():
-		c.Errorf("Unexpected update: %v", update)
-	case update := <-ss2.Subscriptions():
-		c.Errorf("Unexpected update: %v", update)
-	case <-time.After(300 * time.Millisecond):
-	}
+	wait4Subscription(c, ss1.Subscriptions(), membership, 3*time.Second)
+	wait4Subscription(c, ss2.Subscriptions(), membership, 3*time.Second)
 }
 
 // To unsubscribe from all topics an empty topic list can be sent.
@@ -177,11 +154,9 @@ func (s *GroupMemberSuite) TestSomethingAfterNothingBug(c *C) {
 	defer ss1.Stop()
 
 	ss1.Topics() <- []string{"foo"}
-	c.Assert(<-ss1.Subscriptions(), DeepEquals,
-		map[string][]string{"m1": {"foo"}})
+	c.Assert(<-ss1.Subscriptions(), DeepEquals, map[string][]string{"m1": {"foo"}})
 	ss1.Topics() <- []string{}
-	c.Assert(<-ss1.Subscriptions(), DeepEquals,
-		map[string][]string{})
+	c.Assert(<-ss1.Subscriptions(), DeepEquals, map[string][]string{})
 
 	// When
 	ss1.Topics() <- []string{"foo"}
@@ -237,20 +212,17 @@ func (s *GroupMemberSuite) TestRedundantUpdateBug(c *C) {
 	ss1.Topics() <- []string{"foo", "bar"}
 	ss2.Topics() <- []string{"foo", "bazz", "blah"}
 
-	c.Assert(<-ss1.Subscriptions(), DeepEquals,
-		map[string][]string{
-			"m1": {"bar", "foo"},
-			"m2": {"bazz", "blah", "foo"}})
+	membership := map[string][]string{
+		"m1": {"bar", "foo"},
+		"m2": {"bazz", "blah", "foo"}}
+	c.Assert(<-ss1.Subscriptions(), DeepEquals, membership)
 
 	// When
 	ss2.Topics() <- []string{"bar"}
 	ss2.Topics() <- []string{"foo", "bazz", "blah"}
 
 	// Then
-	c.Assert(<-ss1.Subscriptions(), DeepEquals,
-		map[string][]string{
-			"m1": {"bar", "foo"},
-			"m2": {"bazz", "blah", "foo"}})
+	wait4Subscription(c, ss1.Subscriptions(), membership, 3*time.Second)
 }
 
 // When a group registrator claims a topic partitions it becomes its owner.
@@ -426,6 +398,18 @@ func partitionOwner(gm *T, topic string, partition int32) (string, error) {
 func newConfig(clientId string) *config.Proxy {
 	cfg := config.DefaultProxy()
 	cfg.ClientID = clientId
-	cfg.Consumer.RebalanceDelay = 100 * time.Millisecond
 	return cfg
 }
+
+func wait4Subscription(c *C, ch <-chan map[string][]string, want map[string][]string, timeout time.Duration) {
+	for {
+		select {
+		case got := <-ch:
+			if reflect.DeepEqual(got, want) {
+				return
+			}
+		case <-time.After(timeout):
+			c.Error("Timeout waiting for %v")
+		}
+	}
+}

default.yaml

@@ -126,10 +126,6 @@ proxies:
       # How frequently to commit offsets to Kafka.
       offsets_commit_interval: 500ms
 
-      # Consumer should wait this long after it gets notification that a
-      # consumer joined/left its consumer group before starting rebalancing.
-      rebalance_delay: 250ms
-
       # If a request to a Kafka-Pixy fails for any reason, then it should wait this
       # long before retrying.
       retry_backoff: 500ms