send is_close message to broker

2024-03-31 01:28:40 -07:00
parent c9df613b6b
commit 546ae87c39
5 changed files with 261 additions and 232 deletions
--- a/weed/mq/client/pub_client/publish.go
+++ b/weed/mq/client/pub_client/publish.go
@@ -24,3 +24,16 @@ func (p *TopicPublisher) Publish(key, value []byte) error {
 		TsNs:  time.Now().UnixNano(),
 	})
 }
+
+func (p *TopicPublisher) FinishPublish() error {
+	if inputBuffers, found := p.partition2Buffer.AllIntersections(0, pub_balancer.MaxPartitionCount); found {
+		for _, inputBuffer := range inputBuffers {
+			inputBuffer.Enqueue(&mq_pb.DataMessage{
+				IsClose: true,
+				TsNs:  time.Now().UnixNano(),
+			})
+		}
+	}
+
+	return nil
+}