refactor

ToDo
1. handle unassignment on client side
2. get notified of completed unassignment
3. in consumer_group.go, balance and send unassignment

weed/mq/client/sub_client/subscribe.go

@@ -30,11 +30,12 @@ func (sub *TopicSubscriber) startProcessors() {
 	var wg sync.WaitGroup
 	semaphore := make(chan struct{}, sub.SubscriberConfig.MaxPartitionCount)
 
-	for assigned := range sub.brokerPartitionAssignmentChan {
+	for message := range sub.brokerPartitionAssignmentChan {
+		if assigned := message.GetAssignment(); assigned != nil {
 			wg.Add(1)
 			semaphore <- struct{}{}
 
-		topicPartition := topic.FromPbPartition(assigned.Partition)
+			topicPartition := topic.FromPbPartition(assigned.PartitionAssignment.Partition)
 
 			// wait until no covering partition is still in progress
 			sub.waitUntilNoOverlappingPartitionInFlight(topicPartition)
@@ -60,7 +61,11 @@ func (sub *TopicSubscriber) startProcessors() {
 				} else {
 					glog.V(0).Infof("subscriber %s/%s partition %+v at %v completed", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup, assigned.Partition, assigned.LeaderBroker)
 				}
-		}(assigned, topicPartition)
+			}(assigned.PartitionAssignment, topicPartition)
+		}
+		if unAssignment := message.GetUnAssignment(); unAssignment != nil {
+
+		}
 	}
 
 	wg.Wait()

weed/mq/client/sub_client/subscriber.go

@@ -29,7 +29,7 @@ type OnCompletionFunc func()
 type TopicSubscriber struct {
 	SubscriberConfig              *SubscriberConfiguration
 	ContentConfig                 *ContentConfiguration
-	brokerPartitionAssignmentChan chan *mq_pb.BrokerPartitionAssignment
+	brokerPartitionAssignmentChan chan *mq_pb.SubscriberToSubCoordinatorResponse
 	OnEachMessageFunc             OnEachMessageFunc
 	OnCompletionFunc              OnCompletionFunc
 	bootstrapBrokers              []string
@@ -42,7 +42,7 @@ func NewTopicSubscriber(bootstrapBrokers []string, subscriber *SubscriberConfigu
 	return &TopicSubscriber{
 		SubscriberConfig:              subscriber,
 		ContentConfig:                 content,
-		brokerPartitionAssignmentChan: make(chan *mq_pb.BrokerPartitionAssignment, 1024),
+		brokerPartitionAssignmentChan: make(chan *mq_pb.SubscriberToSubCoordinatorResponse, 1024),
 		bootstrapBrokers:              bootstrapBrokers,
 		waitForMoreMessage:            true,
 		activeProcessors:              make(map[topic.Partition]*ProcessorState),

weed/mq/sub_coordinator/consumer_group.go

@@ -131,9 +131,11 @@ func (cg *ConsumerGroup) BalanceConsumerGroupInstances(knownPartitionSlotToBroke
 			partitionSlots = make([]*PartitionSlotToConsumerInstance, 0)
 		}
 		consumerGroupInstance.Partitions = ToPartitions(partitionSlotToBrokerList.RingSize, partitionSlots)
-		assignedPartitions := make([]*mq_pb.BrokerPartitionAssignment, len(partitionSlots))
-		for i, partitionSlot := range partitionSlots {
-			assignedPartitions[i] = &mq_pb.BrokerPartitionAssignment{
+		for _, partitionSlot := range partitionSlots {
+			consumerGroupInstance.ResponseChan <- &mq_pb.SubscriberToSubCoordinatorResponse{
+				Message: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment_{
+					Assignment: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment{
+						PartitionAssignment: &mq_pb.BrokerPartitionAssignment{
 							Partition: &mq_pb.Partition{
 								RangeStop:  partitionSlot.RangeStop,
 								RangeStart: partitionSlot.RangeStart,
@@ -142,16 +144,11 @@ func (cg *ConsumerGroup) BalanceConsumerGroupInstances(knownPartitionSlotToBroke
 							},
 							LeaderBroker:   partitionSlot.Broker,
 							FollowerBroker: partitionSlot.FollowerBroker,
-			}
-		}
-		response := &mq_pb.SubscriberToSubCoordinatorResponse{
-			Message: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment_{
-				Assignment: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment{
-					PartitionAssignments: assignedPartitions,
 						},
 					},
+				},
+			}
 		}
-		consumerGroupInstance.ResponseChan <- response
 	}
 
 }

weed/pb/mq.proto

@@ -183,11 +183,14 @@ message SubscriberToSubCoordinatorRequest {
 }
 message SubscriberToSubCoordinatorResponse {
     message Assignment {
-        int64 generation = 1;
-        repeated BrokerPartitionAssignment partition_assignments = 2;
+        BrokerPartitionAssignment partition_assignment = 1;
+    }
+    message UnAssignment {
+        Partition partition = 1;
     }
     oneof message {
         Assignment assignment = 1;
+        UnAssignment un_assignment = 2;
     }
 }