Wait For Inflight Messages (#104)

Gsantomaggio · web-flow · commit fe231a780cba · 2021-11-09T19:27:39.000+01:00
Add waitforinflightmessages. Fixes #103
diff --git a/pkg/stream/coordinator.go b/pkg/stream/coordinator.go
@@ -55,6 +55,7 @@ func (coordinator *Coordinator) NewProducer(
 	var producer = &Producer{id: lastId,
 		options:             parameters,
 		mutex:               &sync.Mutex{},
+		mutexPending:        &sync.Mutex{},
 		unConfirmedMessages: map[int64]*ConfirmationStatus{},
 		status:              open,
 		messageSequenceCh:   make(chan messageSequence, size),
@@ -91,7 +92,7 @@ func (coordinator *Coordinator) RemoveProducerById(id uint8, reason Event) error
 	reason.Name = producer.GetName()
 	tentatives := 0
 	for producer.lenUnConfirmed() > 0 && tentatives < 3 {
-		time.Sleep(500 * time.Millisecond)
+		time.Sleep(200 * time.Millisecond)
 		tentatives++
 	}
 	producer.FlushUnConfirmedMessages()
diff --git a/pkg/stream/producer.go b/pkg/stream/producer.go
@@ -67,6 +67,7 @@ type Producer struct {
 	unConfirmedMessages map[int64]*ConfirmationStatus
 	sequence            int64
 	mutex               *sync.Mutex
+	mutexPending        *sync.Mutex
 	publishConfirm      chan []*ConfirmationStatus
 	closeHandler        chan Event
 	status              int
@@ -160,6 +161,12 @@ func (producer *Producer) lenUnConfirmed() int {
 	return len(producer.unConfirmedMessages)
 }
 
+func (producer *Producer) lenPendingMessages() int {
+	producer.mutexPending.Lock()
+	defer producer.mutexPending.Unlock()
+	return len(producer.pendingMessages.messages)
+}
+
 func (producer *Producer) getUnConfirmed(sequence int64) *ConfirmationStatus {
 	producer.mutex.Lock()
 	defer producer.mutex.Unlock()
@@ -229,6 +236,7 @@ func (producer *Producer) startPublishTask() {
 						}
 						return
 					}
+					producer.mutexPending.Lock()
 					if producer.pendingMessages.size+msg.unCompressedSize >= producer.options.client.getTuneState().
 						requestedMaxFrameSize {
 						producer.sendBufferedMessages()
@@ -239,10 +247,13 @@ func (producer *Producer) startPublishTask() {
 					if len(producer.pendingMessages.messages) >= (producer.options.BatchSize) {
 						producer.sendBufferedMessages()
 					}
+					producer.mutexPending.Unlock()
 				}
 
 			case <-ticker.C:
+				producer.mutexPending.Lock()
 				producer.sendBufferedMessages()
+				producer.mutexPending.Unlock()
 			}
 
 		}
@@ -401,6 +412,7 @@ func (producer *Producer) aggregateEntities(msgs []messageSequence, size int, co
 
 /// the producer id is always the producer.GetID(). This function is needed only for testing
 // some condition, like simulate publish error, see
+
 func (producer *Producer) internalBatchSendProdId(messagesSequence []messageSequence, producerID uint8) error {
 	producer.options.client.socket.mutex.Lock()
 	defer producer.options.client.socket.mutex.Unlock()
@@ -456,7 +468,6 @@ func (producer *Producer) internalBatchSendProdId(messagesSequence []messageSequ
 
 		producer.setStatus(closed)
 		producer.FlushUnConfirmedMessages()
-		//time.Sleep(800 * time.Millisecond)
 		return err
 	}
 	return nil
@@ -480,7 +491,10 @@ func (producer *Producer) Close() error {
 	if producer.getStatus() == closed {
 		return AlreadyClosed
 	}
+
+	producer.waitForInflightMessages()
 	producer.setStatus(closed)
+
 	if !producer.options.client.socket.isOpen() {
 		return fmt.Errorf("tcp connection is closed")
 	}
@@ -508,6 +522,27 @@ func (producer *Producer) Close() error {
 	return nil
 }
 
+func (producer *Producer) waitForInflightMessages() {
+	// during the close there cloud be pending messages
+	// it waits for producer.options.BatchPublishingDelay
+	// to flush the last messages
+	// see issues/103
+
+	channelLength := len(producer.messageSequenceCh)
+	pendingMessagesLen := producer.lenPendingMessages()
+	tentatives := 0
+
+	for (channelLength > 0 || pendingMessagesLen > 0 || producer.lenUnConfirmed() > 0) && tentatives < 3 {
+		logs.LogDebug("waitForInflightMessages, channel: %d - pending messages len: %d - unconfirmed len: %d - retry: %d",
+			channelLength, pendingMessagesLen,
+			producer.lenUnConfirmed(), tentatives)
+		time.Sleep(time.Duration(2*producer.options.BatchPublishingDelay) * time.Millisecond)
+		channelLength = len(producer.messageSequenceCh)
+		pendingMessagesLen = producer.lenPendingMessages()
+		tentatives++
+	}
+}
+
 func (producer *Producer) GetStreamName() string {
 	if producer.options == nil {
 		return ""
diff --git a/pkg/stream/producer_test.go b/pkg/stream/producer_test.go
@@ -40,8 +40,6 @@ var _ = Describe("Streaming Producers", func() {
 		producer, err := testEnvironment.NewProducer(testProducerStream, nil)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(producer.BatchSend(CreateArrayMessagesForTesting(5))).NotTo(HaveOccurred())
-		// we can't close the subscribe until the publish is finished
-		time.Sleep(200 * time.Millisecond)
 		Expect(producer.Close()).NotTo(HaveOccurred())
 	})
 
@@ -55,8 +53,6 @@ var _ = Describe("Streaming Producers", func() {
 				Expect(err).NotTo(HaveOccurred())
 
 				Expect(producer.BatchSend(CreateArrayMessagesForTesting(5))).NotTo(HaveOccurred())
-				// we can't close the subscribe until the publish is finished
-				time.Sleep(200 * time.Millisecond)
 				err = producer.Close()
 				Expect(err).NotTo(HaveOccurred())
 			}(&wg)
@@ -217,6 +213,22 @@ var _ = Describe("Streaming Producers", func() {
 		Expect(producer.Close()).NotTo(HaveOccurred())
 	})
 
+	It("Wait for inflight messages", func() {
+		// https://github.com/rabbitmq/rabbitmq-stream-go-client/issues/103
+
+		producer, err := testEnvironment.NewProducer(testProducerStream, nil)
+		Expect(err).NotTo(HaveOccurred())
+
+		for i := 0; i < 65672; i++ {
+			Expect(producer.Send(amqp.NewMessage([]byte("h")))).NotTo(HaveOccurred())
+		}
+
+		Expect(producer.Close()).NotTo(HaveOccurred())
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
+		Expect(producer.lenPendingMessages()).To(Equal(0))
+		Expect(len(producer.messageSequenceCh)).To(Equal(0))
+	})
+
 	It("Handle close", func() {
 		var commandIdRecv int32
 
@@ -292,7 +304,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(101)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 		Expect(producer.Close()).NotTo(HaveOccurred())
 		// in this case must raise an error since the producer is closed
 		Expect(producer.Close()).To(HaveOccurred())
@@ -323,7 +335,7 @@ var _ = Describe("Streaming Producers", func() {
 		By("Max frame Error")
 		s := make([]byte, 1148576)
 		Expect(producer.Send(amqp.NewMessage(s))).To(HaveOccurred())
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 		Expect(producer.Close()).NotTo(HaveOccurred())
 
 		producer, err = testEnvironment.NewProducer(testProducerStream,
@@ -348,7 +360,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(101)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 		Expect(producer.Close()).NotTo(HaveOccurred())
 
 	})
@@ -381,7 +393,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(10)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 		Expect(producer.Close()).NotTo(HaveOccurred())
 	})
 
@@ -418,7 +430,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(10)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 		err = producer.Close()
 		Expect(err).NotTo(HaveOccurred())
 	})
@@ -561,7 +573,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(232)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 
 		// same test above but using batch send
 		var arr []message.StreamMessage
@@ -578,7 +590,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(12*20)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 
 		Expect(producer.Close()).NotTo(HaveOccurred())
 
@@ -651,7 +663,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(501)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 		atomic.StoreInt32(&messagesConfirmed, 0)
 
 		for z := 0; z < 501; z++ {
@@ -664,7 +676,7 @@ var _ = Describe("Streaming Producers", func() {
 		}, 5*time.Second).Should(Equal(int32(501*5)),
 			"confirm should receive same messages send by producer")
 
-		Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+		Expect(producer.lenUnConfirmed()).To(Equal(0))
 		Expect(producer.Close()).NotTo(HaveOccurred())
 	})
 
@@ -724,7 +736,7 @@ func testCompress(producer *Producer) {
 	}, 5*time.Second).Should(Equal(int32(457)),
 		"confirm should receive same messages send by producer")
 
-	Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+	Expect(producer.lenUnConfirmed()).To(Equal(0))
 	atomic.StoreInt32(&messagesConfirmed, 0)
 
 	for z := 0; z < 457; z++ {
@@ -789,7 +801,7 @@ func verifyProducerSent(producer *Producer, confirmationReceived *int32, message
 	}, 10*time.Second, 1*time.Second).Should(Equal(int32(messageSent)),
 		"confirm should receive same messages send by producer")
 
-	Expect(len(producer.unConfirmedMessages)).To(Equal(0))
+	Expect(producer.lenUnConfirmed()).To(Equal(0))
 }
 
 func runConcurrentlyAndWaitTillAllDone(threadCount int, wg *sync.WaitGroup, runner func(int)) {
