Merge accumulated changes related to message queue (#5098)

* balance partitions on brokers

* prepare topic partition first and then publish, move partition

* purge unused APIs

* clean up

* adjust logs

* add BalanceTopics() grpc API

* configure topic

* configure topic command

* refactor

* repair missing partitions

* sequence of operations to ensure ordering

* proto to close publishers and consumers

* rename file

* topic partition versioned by unixTimeNs

* create local topic partition

* close publishers

* randomize the client name

* wait until no publishers

* logs

* close stop publisher channel

* send last ack

* comments

* comment

* comments

* support list of brokers

* add cli options

* Update .gitignore

* logs

* return io.eof directly

* refactor

* optionally create topic

* refactoring

* detect consumer disconnection

* sub client wait for more messages

* subscribe by time stamp

* rename

* rename to sub_balancer

* rename

* adjust comments

* rename

* fix compilation

* rename

* rename

* SubscriberToSubCoordinator

* sticky rebalance

* go fmt

* add tests

* balance partitions on brokers

* prepare topic partition first and then publish, move partition

* purge unused APIs

* clean up

* adjust logs

* add BalanceTopics() grpc API

* configure topic

* configure topic command

* refactor

* repair missing partitions

* sequence of operations to ensure ordering

* proto to close publishers and consumers

* rename file

* topic partition versioned by unixTimeNs

* create local topic partition

* close publishers

* randomize the client name

* wait until no publishers

* logs

* close stop publisher channel

* send last ack

* comments

* comment

* comments

* support list of brokers

* add cli options

* Update .gitignore

* logs

* return io.eof directly

* refactor

* optionally create topic

* refactoring

* detect consumer disconnection

* sub client wait for more messages

* subscribe by time stamp

* rename

* rename to sub_balancer

* rename

* adjust comments

* rename

* fix compilation

* rename

* rename

* SubscriberToSubCoordinator

* sticky rebalance

* go fmt

* add tests

* tracking topic=>broker

* merge

* comment

7 files changed, 242 insertions, 70 deletions

diff --git a/weed/mq/topic/local_manager.go b/weed/mq/topic/local_manager.go
index 0c54f2bb1..173df090d 100644
--- a/weed/mq/topic/local_manager.go
+++ b/weed/mq/topic/local_manager.go
@@ -23,10 +23,7 @@ func NewLocalTopicManager() *LocalTopicManager {
 func (manager *LocalTopicManager) AddTopicPartition(topic Topic, localPartition *LocalPartition) {
 	localTopic, ok := manager.topics.Get(topic.String())
 	if !ok {
-		localTopic = &LocalTopic{
-			Topic:      topic,
-			Partitions: make([]*LocalPartition, 0),
-		}
+		localTopic = NewLocalTopic(topic)
 	}
 	if !manager.topics.SetIfAbsent(topic.String(), localTopic) {
 		localTopic, _ = manager.topics.Get(topic.String())
@@ -59,6 +56,22 @@ func (manager *LocalTopicManager) RemoveTopicPartition(topic Topic, partition Pa
 	return localTopic.removePartition(partition)
 }
 
+func (manager *LocalTopicManager) ClosePublishers(topic Topic, unixTsNs int64) (removed bool) {
+	localTopic, ok := manager.topics.Get(topic.String())
+	if !ok {
+		return false
+	}
+	return localTopic.closePartitionPublishers(unixTsNs)
+}
+
+func (manager *LocalTopicManager) CloseSubscribers(topic Topic, unixTsNs int64) (removed bool) {
+	localTopic, ok := manager.topics.Get(topic.String())
+	if !ok {
+		return false
+	}
+	return localTopic.closePartitionSubscribers(unixTsNs)
+}
+
 func (manager *LocalTopicManager) CollectStats(duration time.Duration) *mq_pb.BrokerStats {
 	stats := &mq_pb.BrokerStats{
 		Stats: make(map[string]*mq_pb.TopicPartitionStats),
@@ -101,3 +114,11 @@ func (manager *LocalTopicManager) CollectStats(duration time.Duration) *mq_pb.Br
 	return stats
 
 }
+
+func (manager *LocalTopicManager) WaitUntilNoPublishers(topic Topic) {
+	localTopic, ok := manager.topics.Get(topic.String())
+	if !ok {
+		return
+	}
+	localTopic.WaitUntilNoPublishers()
+}
diff --git a/weed/mq/topic/local_partition.go b/weed/mq/topic/local_partition.go
index 49b639dfa..aa1274ff5 100644
--- a/weed/mq/topic/local_partition.go
+++ b/weed/mq/topic/local_partition.go
@@ -11,19 +11,23 @@ import (
 
 type LocalPartition struct {
 	Partition
-	isLeader        bool
-	FollowerBrokers []pb.ServerAddress
-	logBuffer       *log_buffer.LogBuffer
-	ConsumerCount   int32
+	isLeader          bool
+	FollowerBrokers   []pb.ServerAddress
+	logBuffer         *log_buffer.LogBuffer
+	ConsumerCount     int32
+	StopPublishersCh  chan struct{}
+	Publishers        *LocalPartitionPublishers
+	StopSubscribersCh chan struct{}
+	Subscribers       *LocalPartitionSubscribers
 }
 
-func NewLocalPartition(topic Topic, partition Partition, isLeader bool, followerBrokers []pb.ServerAddress) *LocalPartition {
+func NewLocalPartition(partition Partition, isLeader bool, followerBrokers []pb.ServerAddress) *LocalPartition {
 	return &LocalPartition{
 		Partition:       partition,
 		isLeader:        isLeader,
 		FollowerBrokers: followerBrokers,
 		logBuffer: log_buffer.NewLogBuffer(
-			fmt.Sprintf("%s/%s/%4d-%4d", topic.Namespace, topic.Name, partition.RangeStart, partition.RangeStop),
+			fmt.Sprintf("%d/%4d-%4d", partition.UnixTimeNs, partition.RangeStart, partition.RangeStop),
 			2*time.Minute,
 			func(startTime, stopTime time.Time, buf []byte) {
 
@@ -32,34 +36,43 @@ func NewLocalPartition(topic Topic, partition Partition, isLeader bool, follower
 
 			},
 		),
+		Publishers:  NewLocalPartitionPublishers(),
+		Subscribers: NewLocalPartitionSubscribers(),
 	}
 }
 
 type OnEachMessageFn func(logEntry *filer_pb.LogEntry) error
 
-func (p LocalPartition) Publish(message *mq_pb.DataMessage) {
+func (p *LocalPartition) Publish(message *mq_pb.DataMessage) {
 	p.logBuffer.AddToBuffer(message.Key, message.Value, time.Now().UnixNano())
 }
 
-func (p LocalPartition) Subscribe(clientName string, startReadTime time.Time, eachMessageFn OnEachMessageFn) {
-	p.logBuffer.LoopProcessLogData(clientName, startReadTime, 0, func() bool {
-		return true
-	}, eachMessageFn)
+func (p *LocalPartition) Subscribe(clientName string, startReadTime time.Time, onNoMessageFn func() bool, eachMessageFn OnEachMessageFn) {
+	p.logBuffer.LoopProcessLogData(clientName, startReadTime, 0, onNoMessageFn, eachMessageFn)
 }
 
 func FromPbBrokerPartitionAssignment(self pb.ServerAddress, assignment *mq_pb.BrokerPartitionAssignment) *LocalPartition {
-	isLeaer := assignment.LeaderBroker == string(self)
-	localPartition := &LocalPartition{
-		Partition: FromPbPartition(assignment.Partition),
-		isLeader:  isLeaer,
-	}
-	if !isLeaer {
-		return localPartition
-	}
+	isLeader := assignment.LeaderBroker == string(self)
 	followers := make([]pb.ServerAddress, len(assignment.FollowerBrokers))
-	for i, follower := range assignment.FollowerBrokers {
-		followers[i] = pb.ServerAddress(follower)
+	for i, followerBroker := range assignment.FollowerBrokers {
+		followers[i] = pb.ServerAddress(followerBroker)
+	}
+	return NewLocalPartition(FromPbPartition(assignment.Partition), isLeader, followers)
+}
+
+func (p *LocalPartition) closePublishers() {
+	p.Publishers.SignalShutdown()
+	close(p.StopPublishersCh)
+}
+func (p *LocalPartition) closeSubscribers() {
+	p.Subscribers.SignalShutdown()
+}
+
+func (p *LocalPartition) WaitUntilNoPublishers() {
+	for {
+		if p.Publishers.IsEmpty() {
+			return
+		}
+		time.Sleep(113 * time.Millisecond)
 	}
-	localPartition.FollowerBrokers = followers
-	return localPartition
 }
diff --git a/weed/mq/topic/local_partition_publishers.go b/weed/mq/topic/local_partition_publishers.go
new file mode 100644
index 000000000..367ccce5f
--- /dev/null
+++ b/weed/mq/topic/local_partition_publishers.go
@@ -0,0 +1,52 @@
+package topic
+
+import "sync"
+
+type LocalPartitionPublishers struct {
+	publishers     map[string]*LocalPublisher
+	publishersLock sync.RWMutex
+}
+type LocalPublisher struct {
+}
+
+func NewLocalPublisher() *LocalPublisher {
+	return &LocalPublisher{}
+}
+func (p *LocalPublisher) SignalShutdown() {
+}
+
+func NewLocalPartitionPublishers() *LocalPartitionPublishers {
+	return &LocalPartitionPublishers{
+		publishers: make(map[string]*LocalPublisher),
+	}
+}
+
+func (p *LocalPartitionPublishers) AddPublisher(clientName string, publisher *LocalPublisher) {
+	p.publishersLock.Lock()
+	defer p.publishersLock.Unlock()
+
+	p.publishers[clientName] = publisher
+}
+
+func (p *LocalPartitionPublishers) RemovePublisher(clientName string) {
+	p.publishersLock.Lock()
+	defer p.publishersLock.Unlock()
+
+	delete(p.publishers, clientName)
+}
+
+func (p *LocalPartitionPublishers) SignalShutdown() {
+	p.publishersLock.RLock()
+	defer p.publishersLock.RUnlock()
+
+	for _, publisher := range p.publishers {
+		publisher.SignalShutdown()
+	}
+}
+
+func (p *LocalPartitionPublishers) IsEmpty() bool {
+	p.publishersLock.RLock()
+	defer p.publishersLock.RUnlock()
+
+	return len(p.publishers) == 0
+}
diff --git a/weed/mq/topic/local_partition_subscribers.go b/weed/mq/topic/local_partition_subscribers.go
new file mode 100644
index 000000000..e177ec7e8
--- /dev/null
+++ b/weed/mq/topic/local_partition_subscribers.go
@@ -0,0 +1,49 @@
+package topic
+
+import "sync"
+
+type LocalPartitionSubscribers struct {
+	Subscribers     map[string]*LocalSubscriber
+	SubscribersLock sync.RWMutex
+}
+type LocalSubscriber struct {
+	stopCh chan struct{}
+}
+
+func NewLocalSubscriber() *LocalSubscriber {
+	return &LocalSubscriber{
+		stopCh: make(chan struct{}, 1),
+	}
+}
+func (p *LocalSubscriber) SignalShutdown() {
+	close(p.stopCh)
+}
+
+func NewLocalPartitionSubscribers() *LocalPartitionSubscribers {
+	return &LocalPartitionSubscribers{
+		Subscribers: make(map[string]*LocalSubscriber),
+	}
+}
+
+func (p *LocalPartitionSubscribers) AddSubscriber(clientName string, Subscriber *LocalSubscriber) {
+	p.SubscribersLock.Lock()
+	defer p.SubscribersLock.Unlock()
+
+	p.Subscribers[clientName] = Subscriber
+}
+
+func (p *LocalPartitionSubscribers) RemoveSubscriber(clientName string) {
+	p.SubscribersLock.Lock()
+	defer p.SubscribersLock.Unlock()
+
+	delete(p.Subscribers, clientName)
+}
+
+func (p *LocalPartitionSubscribers) SignalShutdown() {
+	p.SubscribersLock.RLock()
+	defer p.SubscribersLock.RUnlock()
+
+	for _, Subscriber := range p.Subscribers {
+		Subscriber.SignalShutdown()
+	}
+}
diff --git a/weed/mq/topic/local_topic.go b/weed/mq/topic/local_topic.go
index ef3c0e65e..7825d2168 100644
--- a/weed/mq/topic/local_topic.go
+++ b/weed/mq/topic/local_topic.go
@@ -1,10 +1,19 @@
 package topic
 
+import "sync"
+
 type LocalTopic struct {
 	Topic
 	Partitions []*LocalPartition
 }
 
+func NewLocalTopic(topic Topic) *LocalTopic {
+	return &LocalTopic{
+		Topic:      topic,
+		Partitions: make([]*LocalPartition, 0),
+	}
+}
+
 func (localTopic *LocalTopic) findPartition(partition Partition) *LocalPartition {
 	for _, localPartition := range localTopic.Partitions {
 		if localPartition.Partition.Equals(partition) {
@@ -27,3 +36,52 @@ func (localTopic *LocalTopic) removePartition(partition Partition) bool {
 	localTopic.Partitions = append(localTopic.Partitions[:foundPartitionIndex], localTopic.Partitions[foundPartitionIndex+1:]...)
 	return true
 }
+
+func (localTopic *LocalTopic) closePartitionPublishers(unixTsNs int64) bool {
+	var wg sync.WaitGroup
+	for _, localPartition := range localTopic.Partitions {
+		if localPartition.UnixTimeNs != unixTsNs {
+			continue
+		}
+		wg.Add(1)
+		go func(localPartition *LocalPartition) {
+			defer wg.Done()
+			localPartition.closePublishers()
+		}(localPartition)
+	}
+	wg.Wait()
+	return true
+}
+
+func (localTopic *LocalTopic) closePartitionSubscribers(unixTsNs int64) bool {
+	var wg sync.WaitGroup
+	for _, localPartition := range localTopic.Partitions {
+		if localPartition.UnixTimeNs != unixTsNs {
+			continue
+		}
+		wg.Add(1)
+		go func(localPartition *LocalPartition) {
+			defer wg.Done()
+			localPartition.closeSubscribers()
+		}(localPartition)
+	}
+	wg.Wait()
+	return true
+}
+
+func (localTopic *LocalTopic) WaitUntilNoPublishers() {
+	for {
+		var wg sync.WaitGroup
+		for _, localPartition := range localTopic.Partitions {
+			wg.Add(1)
+			go func(localPartition *LocalPartition) {
+				defer wg.Done()
+				localPartition.WaitUntilNoPublishers()
+			}(localPartition)
+		}
+		wg.Wait()
+		if len(localTopic.Partitions) == 0 {
+			return
+		}
+	}
+}
diff --git a/weed/mq/topic/partition.go b/weed/mq/topic/partition.go
index 79c830f13..ca34c2390 100644
--- a/weed/mq/topic/partition.go
+++ b/weed/mq/topic/partition.go
@@ -8,6 +8,7 @@ type Partition struct {
 	RangeStart int32
 	RangeStop  int32 // exclusive
 	RingSize   int32
+	UnixTimeNs int64 // in nanoseconds
 }
 
 func (partition Partition) Equals(other Partition) bool {
@@ -20,6 +21,9 @@ func (partition Partition) Equals(other Partition) bool {
 	if partition.RingSize != other.RingSize {
 		return false
 	}
+	if partition.UnixTimeNs != other.UnixTimeNs {
+		return false
+	}
 	return true
 }
 
@@ -28,10 +32,11 @@ func FromPbPartition(partition *mq_pb.Partition) Partition {
 		RangeStart: partition.RangeStart,
 		RangeStop:  partition.RangeStop,
 		RingSize:   partition.RingSize,
+		UnixTimeNs: partition.UnixTimeNs,
 	}
 }
 
-func SplitPartitions(targetCount int32) []*Partition {
+func SplitPartitions(targetCount int32, ts int64) []*Partition {
 	partitions := make([]*Partition, 0, targetCount)
 	partitionSize := PartitionCount / targetCount
 	for i := int32(0); i < targetCount; i++ {
@@ -43,7 +48,17 @@ func SplitPartitions(targetCount int32) []*Partition {
 			RangeStart: i * partitionSize,
 			RangeStop:  partitionStop,
 			RingSize:   PartitionCount,
+			UnixTimeNs: ts,
 		})
 	}
 	return partitions
 }
+
+func (partition Partition) ToPbPartition() *mq_pb.Partition {
+	return &mq_pb.Partition{
+		RangeStart: partition.RangeStart,
+		RangeStop:  partition.RangeStop,
+		RingSize:   partition.RingSize,
+		UnixTimeNs: partition.UnixTimeNs,
+	}
+}
diff --git a/weed/mq/topic/topic.go b/weed/mq/topic/topic.go
index 3d457e6f1..6932fcb56 100644
--- a/weed/mq/topic/topic.go
+++ b/weed/mq/topic/topic.go
@@ -2,9 +2,7 @@ package topic
 
 import (
 	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/filer"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
-	"time"
 )
 
 type Topic struct {
@@ -25,47 +23,13 @@ func FromPbTopic(topic *mq_pb.Topic) Topic {
 	}
 }
 
-func (tp Topic) String() string {
-	return fmt.Sprintf("%s.%s", tp.Namespace, tp.Name)
-}
-
-type Segment struct {
-	Topic        Topic
-	Id           int32
-	Partition    Partition
-	LastModified time.Time
-}
-
-func FromPbSegment(segment *mq_pb.Segment) *Segment {
-	return &Segment{
-		Topic: Topic{
-			Namespace: segment.Namespace,
-			Name:      segment.Topic,
-		},
-		Id: segment.Id,
-		Partition: Partition{
-			RangeStart: segment.Partition.RangeStart,
-			RangeStop:  segment.Partition.RangeStop,
-			RingSize:   segment.Partition.RingSize,
-		},
+func (tp Topic) ToPbTopic() *mq_pb.Topic {
+	return &mq_pb.Topic{
+		Namespace: tp.Namespace,
+		Name:      tp.Name,
 	}
 }
 
-func (segment *Segment) ToPbSegment() *mq_pb.Segment {
-	return &mq_pb.Segment{
-		Namespace: string(segment.Topic.Namespace),
-		Topic:     segment.Topic.Name,
-		Id:        segment.Id,
-		Partition: &mq_pb.Partition{
-			RingSize:   segment.Partition.RingSize,
-			RangeStart: segment.Partition.RangeStart,
-			RangeStop:  segment.Partition.RangeStop,
-		},
-	}
-}
-
-func (segment *Segment) DirAndName() (dir string, name string) {
-	dir = fmt.Sprintf("%s/%s/%s", filer.TopicsDir, segment.Topic.Namespace, segment.Topic.Name)
-	name = fmt.Sprintf("%4d.segment", segment.Id)
-	return
+func (tp Topic) String() string {
+	return fmt.Sprintf("%s.%s", tp.Namespace, tp.Name)
 }