Fix: #1314 partition watcher doesn't reacts on partition number changing

consumergroup.go

@@ -497,7 +497,7 @@ func (g *Generation) heartbeatLoop(interval time.Duration) {
 // a bad spot and should rebalance. Commonly you will see an error here if there
 // is a problem with the connection to the coordinator and a rebalance will
 // establish a new connection to the coordinator.
-func (g *Generation) partitionWatcher(interval time.Duration, topic string) {
+func (g *Generation) partitionWatcher(interval time.Duration, topic string, startPartitions int) {
 	g.Start(func(ctx context.Context) {
 		g.log(func(l Logger) {
 			l.Printf("started partition watcher for group, %v, topic %v [%v]", g.GroupID, topic, interval)
@@ -509,14 +509,6 @@ func (g *Generation) partitionWatcher(interval time.Duration, topic string) {
 		ticker := time.NewTicker(interval)
 		defer ticker.Stop()
 
-		ops, err := g.conn.readPartitions(topic)
-		if err != nil {
-			g.logError(func(l Logger) {
-				l.Printf("Problem getting partitions during startup, %v\n, Returning and setting up nextGeneration", err)
-			})
-			return
-		}
-		oParts := len(ops)
 		for {
 			select {
 			case <-ctx.Done():
@@ -525,7 +517,7 @@ func (g *Generation) partitionWatcher(interval time.Duration, topic string) {
 				ops, err := g.conn.readPartitions(topic)
 				switch {
 				case err == nil, errors.Is(err, UnknownTopicOrPartition):
-					if len(ops) != oParts {
+					if len(ops) != startPartitions {
 						g.log(func(l Logger) {
 							l.Printf("Partition changes found, rebalancing group: %v.", g.GroupID)
 						})
@@ -651,11 +643,17 @@ func NewConsumerGroup(config ConsumerGroupConfig) (*ConsumerGroup, error) {
 	}
 
 	cg := &ConsumerGroup{
-		config: config,
-		next:   make(chan *Generation),
-		errs:   make(chan error),
-		done:   make(chan struct{}),
+		config:             config,
+		partitionsPerTopic: make(map[string]int, len(config.Topics)),
+		next:               make(chan *Generation),
+		errs:               make(chan error),
+		done:               make(chan struct{}),
+	}
+
+	for _, topic := range config.Topics {
+		cg.partitionsPerTopic[topic] = 0
 	}
+
 	cg.wg.Add(1)
 	go func() {
 		cg.run()
@@ -670,9 +668,10 @@ func NewConsumerGroup(config ConsumerGroupConfig) (*ConsumerGroup, error) {
 // Generation is where partition assignments and offset management occur.
 // Callers will use Next to get a handle to the Generation.
 type ConsumerGroup struct {
-	config ConsumerGroupConfig
-	next   chan *Generation
-	errs   chan error
+	config             ConsumerGroupConfig
+	partitionsPerTopic map[string]int
+	next               chan *Generation
+	errs               chan error
 
 	closeOnce sync.Once
 	wg        sync.WaitGroup
@@ -789,25 +788,21 @@ func (cg *ConsumerGroup) nextGeneration(memberID string) (string, error) {
 	}
 	defer conn.Close()
 
-	var generationID int32
-	var groupAssignments GroupMemberAssignments
-	var assignments map[string][]int32
-
 	// join group.  this will join the group and prepare assignments if our
 	// consumer is elected leader.  it may also change or assign the member ID.
-	memberID, generationID, groupAssignments, err = cg.joinGroup(conn, memberID)
+	memberID, generationID, groupAssignments, iAmLeader, err := cg.joinGroup(conn, memberID)
 	if err != nil {
 		cg.withErrorLogger(func(log Logger) {
 			log.Printf("Failed to join group %s: %v", cg.config.ID, err)
 		})
 		return memberID, err
 	}
 	cg.withLogger(func(log Logger) {
-		log.Printf("Joined group %s as member %s in generation %d", cg.config.ID, memberID, generationID)
+		log.Printf("Joined group %s as member %s (leader %t) in generation %d", cg.config.ID, memberID, iAmLeader, generationID)
 	})
 
 	// sync group
-	assignments, err = cg.syncGroup(conn, memberID, generationID, groupAssignments)
+	assignments, err := cg.syncGroup(conn, memberID, generationID, groupAssignments)
 	if err != nil {
 		cg.withErrorLogger(func(log Logger) {
 			log.Printf("Failed to sync group %s: %v", cg.config.ID, err)
@@ -843,9 +838,9 @@ func (cg *ConsumerGroup) nextGeneration(memberID string) (string, error) {
 	// any of these functions exit, then the generation is determined to be
 	// complete.
 	gen.heartbeatLoop(cg.config.HeartbeatInterval)
-	if cg.config.WatchPartitionChanges {
-		for _, topic := range cg.config.Topics {
-			gen.partitionWatcher(cg.config.PartitionWatchInterval, topic)
+	if cg.config.WatchPartitionChanges && iAmLeader {
+		for topic, startPartitions := range cg.partitionsPerTopic {
+			gen.partitionWatcher(cg.config.PartitionWatchInterval, topic, startPartitions)
 		}
 	}
 
@@ -925,24 +920,24 @@ func (cg *ConsumerGroup) coordinator() (coordinator, error) {
 // the leader.  Otherwise, GroupMemberAssignments will be nil.
 //
 // Possible kafka error codes returned:
-//  * GroupLoadInProgress:
-//  * GroupCoordinatorNotAvailable:
-//  * NotCoordinatorForGroup:
-//  * InconsistentGroupProtocol:
-//  * InvalidSessionTimeout:
-//  * GroupAuthorizationFailed:
-func (cg *ConsumerGroup) joinGroup(conn coordinator, memberID string) (string, int32, GroupMemberAssignments, error) {
+//   - GroupLoadInProgress:
+//   - GroupCoordinatorNotAvailable:
+//   - NotCoordinatorForGroup:
+//   - InconsistentGroupProtocol:
+//   - InvalidSessionTimeout:
+//   - GroupAuthorizationFailed:
+func (cg *ConsumerGroup) joinGroup(conn coordinator, memberID string) (string, int32, GroupMemberAssignments, bool, error) {
 	request, err := cg.makeJoinGroupRequestV1(memberID)
 	if err != nil {
-		return "", 0, nil, err
+		return "", 0, nil, false, err
 	}
 
 	response, err := conn.joinGroup(request)
 	if err == nil && response.ErrorCode != 0 {
 		err = Error(response.ErrorCode)
 	}
 	if err != nil {
-		return "", 0, nil, err
+		return "", 0, nil, false, err
 	}
 
 	memberID = response.MemberID
@@ -953,10 +948,12 @@ func (cg *ConsumerGroup) joinGroup(conn coordinator, memberID string) (string, i
 	})
 
 	var assignments GroupMemberAssignments
-	if iAmLeader := response.MemberID == response.LeaderID; iAmLeader {
+	iAmLeader := response.MemberID == response.LeaderID
+
+	if iAmLeader {
 		v, err := cg.assignTopicPartitions(conn, response)
 		if err != nil {
-			return memberID, 0, nil, err
+			return memberID, 0, nil, false, err
 		}
 		assignments = v
 
@@ -973,7 +970,7 @@ func (cg *ConsumerGroup) joinGroup(conn coordinator, memberID string) (string, i
 		l.Printf("joinGroup succeeded for response, %v.  generationID=%v, memberID=%v", cg.config.ID, response.GenerationID, response.MemberID)
 	})
 
-	return memberID, generationID, assignments, nil
+	return memberID, generationID, assignments, iAmLeader, nil
 }
 
 // makeJoinGroupRequestV1 handles the logic of constructing a joinGroup
@@ -1036,6 +1033,14 @@ func (cg *ConsumerGroup) assignTopicPartitions(conn coordinator, group joinGroup
 		return nil, err
 	}
 
+	// resetting old values of partitions per topic
+	cg.partitionsPerTopic = make(map[string]int, len(cg.config.Topics))
+
+	// setting new values of partitions per topic
+	for _, partition := range partitions {
+		cg.partitionsPerTopic[partition.Topic] += 1
+	}
+
 	cg.withLogger(func(l Logger) {
 		l.Printf("using '%v' balancer to assign group, %v", group.GroupProtocol, cg.config.ID)
 		for _, member := range members {
@@ -1073,11 +1078,11 @@ func (cg *ConsumerGroup) makeMemberProtocolMetadata(in []joinGroupResponseMember
 // Readers subscriptions topic => partitions
 //
 // Possible kafka error codes returned:
-//  * GroupCoordinatorNotAvailable:
-//  * NotCoordinatorForGroup:
-//  * IllegalGeneration:
-//  * RebalanceInProgress:
-//  * GroupAuthorizationFailed:
+//   - GroupCoordinatorNotAvailable:
+//   - NotCoordinatorForGroup:
+//   - IllegalGeneration:
+//   - RebalanceInProgress:
+//   - GroupAuthorizationFailed:
 func (cg *ConsumerGroup) syncGroup(conn coordinator, memberID string, generationID int32, memberAssignments GroupMemberAssignments) (map[string][]int32, error) {
 	request := cg.makeSyncGroupRequestV0(memberID, generationID, memberAssignments)
 	response, err := conn.syncGroup(request)

consumergroup_test.go

@@ -637,7 +637,7 @@ func TestGenerationExitsOnPartitionChange(t *testing.T) {
 
 	done := make(chan struct{})
 	go func() {
-		gen.partitionWatcher(watchTime, "topic-1")
+		gen.partitionWatcher(watchTime, "topic-1", 1)
 		close(done)
 	}()