fix: core - publisher confirmations - add support for confirmation of multiple messages (#201)

mi-char · web-flow · commit d9d1435721b5 · 2023-03-10T16:04:18.000+01:00
fix: core - publisher confirmations - add support for confirmation of multiple messages
diff --git a/core/src/main/scala/com/avast/clients/rabbitmq/publisher/PublishConfirmsRabbitMQProducer.scala b/core/src/main/scala/com/avast/clients/rabbitmq/publisher/PublishConfirmsRabbitMQProducer.scala
@@ -1,9 +1,9 @@
 package com.avast.clients.rabbitmq.publisher
 
+import cats.effect.Concurrent.ops.toAllConcurrentOps
 import cats.effect.concurrent.Deferred
 import cats.effect.{Blocker, ConcurrentEffect, ContextShift}
-import cats.syntax.flatMap._
-import cats.syntax.functor._
+import cats.implicits._
 import com.avast.bytes.Bytes
 import com.avast.clients.rabbitmq.api.{MaxAttemptsReached, MessageProperties, NotAcknowledgedPublish}
 import com.avast.clients.rabbitmq.logging.ImplicitContextLogger
@@ -62,7 +62,8 @@ class PublishConfirmsRabbitMQProducer[F[_], A: ProductConverter](name: String,
         _ <- result match {
           case Left(err) =>
             val sendResult = if (sendAttempts > 1) {
-              sendWithAck(routingKey, body, properties, attemptCount + 1)
+              logger.plainTrace(s"Republishing nacked message with sequenceNumber $sequenceNumber") >>
+                sendWithAck(routingKey, body, properties, attemptCount + 1)
             } else {
               F.raiseError(err)
             }
@@ -74,26 +75,40 @@ class PublishConfirmsRabbitMQProducer[F[_], A: ProductConverter](name: String,
     }
   }
 
-  private object DefaultConfirmListener extends ConfirmListener {
+  private[rabbitmq] object DefaultConfirmListener extends ConfirmListener {
 
     override def handleAck(deliveryTag: Long, multiple: Boolean): Unit = {
       startAndForget {
-        logger.plainTrace(s"Acked $deliveryTag") >> completeDefer(deliveryTag, Right(()))
+        logger.plainTrace(s"Acked $deliveryTag, multiple: $multiple") >> completeDefer(deliveryTag, multiple, Right(()))
       }
     }
 
     override def handleNack(deliveryTag: Long, multiple: Boolean): Unit = {
       startAndForget {
-        logger.plainTrace(s"Not acked $deliveryTag") >> completeDefer(
+        logger.plainTrace(s"Nacked $deliveryTag, multiple: $multiple") >> completeDefer(
           deliveryTag,
-          Left(NotAcknowledgedPublish(s"Message $deliveryTag not acknowledged by broker", messageId = deliveryTag)))
+          multiple,
+          Left(NotAcknowledgedPublish(s"Broker was unable to process the message", messageId = deliveryTag)))
       }
     }
 
-    private def completeDefer(deliveryTag: Long, result: Either[NotAcknowledgedPublish, Unit]): F[Unit] = {
-      confirmationCallbacks.get(deliveryTag) match {
-        case Some(callback) => callback.complete(result)
-        case None => logger.plainWarn("Received confirmation for unknown delivery tag. That is unexpected state.")
+    private def completeDefer(deliveryTag: Long, multiple: Boolean, result: Either[NotAcknowledgedPublish, Unit]): F[Unit] = {
+      if (multiple) {
+        confirmationCallbacks
+          .filter {
+            case (sequenceNumber, _) => sequenceNumber <= deliveryTag
+          }
+          .values
+          .toList
+          .traverse { callback =>
+            callback.complete(result).start
+          }
+          .void
+      } else {
+        confirmationCallbacks.get(deliveryTag) match {
+          case Some(callback) => callback.complete(result)
+          case None => logger.plainError(s"Received confirmation for unknown delivery tag $deliveryTag. That is unexpected state.")
+        }
       }
     }
   }
diff --git a/core/src/test/scala/com/avast/clients/rabbitmq/PublisherConfirmsRabbitMQProducerTest.scala b/core/src/test/scala/com/avast/clients/rabbitmq/PublisherConfirmsRabbitMQProducerTest.scala
@@ -1,6 +1,5 @@
 package com.avast.clients.rabbitmq
 
-import cats.implicits.catsSyntaxParallelAp
 import com.avast.bytes.Bytes
 import com.avast.clients.rabbitmq.api.{MessageProperties, NotAcknowledgedPublish}
 import com.avast.clients.rabbitmq.logging.ImplicitContextLogger
@@ -9,7 +8,7 @@ import com.avast.metrics.scalaeffectapi.Monitor
 import com.rabbitmq.client.impl.recovery.AutorecoveringChannel
 import monix.eval.Task
 import monix.execution.Scheduler.Implicits.global
-import org.junit.runner.manipulation.InvalidOrderingException
+import monix.execution.atomic.AtomicInt
 import org.mockito.Matchers
 import org.mockito.Matchers.any
 import org.mockito.Mockito.{times, verify, when}
@@ -23,11 +22,13 @@ class PublisherConfirmsRabbitMQProducerTest extends TestBase {
   test("Message is acked after one retry") {
     val exchangeName = Random.nextString(10)
     val routingKey = Random.nextString(10)
-    val seqNumber = 1L
-    val seqNumber2 = 2L
 
+    val nextSeqNumber = AtomicInt(0)
     val channel = mock[AutorecoveringChannel]
-    when(channel.getNextPublishSeqNo).thenReturn(seqNumber, seqNumber2)
+    when(channel.getNextPublishSeqNo).thenAnswer(_ => nextSeqNumber.get())
+    when(channel.basicPublish(any(), any(), any(), any())).thenAnswer(_ => {
+      nextSeqNumber.increment()
+    })
 
     val producer = new PublishConfirmsRabbitMQProducer[Task, Bytes](
       name = "test",
@@ -44,13 +45,15 @@ class PublisherConfirmsRabbitMQProducerTest extends TestBase {
 
     val body = Bytes.copyFrom(Array.fill(499)(32.toByte))
 
-    val publishTask = producer.send(routingKey, body).runToFuture
+    val publishFuture = producer.send(routingKey, body).runToFuture
+
+    while (nextSeqNumber.get() < 1) { Thread.sleep(5) }
+    producer.DefaultConfirmListener.handleNack(0, multiple = false)
 
-    updateMessageState(producer, seqNumber)(Left(NotAcknowledgedPublish("abcd", messageId = seqNumber))).parProduct {
-      updateMessageState(producer, seqNumber2)(Right())
-    }.await
+    while (nextSeqNumber.get() < 2) { Thread.sleep(5) }
+    producer.DefaultConfirmListener.handleAck(1, multiple = false)
 
-    Await.result(publishTask, 10.seconds)
+    Await.result(publishFuture, 10.seconds)
 
     verify(channel, times(2))
       .basicPublish(Matchers.eq(exchangeName), Matchers.eq(routingKey), any(), Matchers.eq(body.toByteArray))
@@ -59,10 +62,13 @@ class PublisherConfirmsRabbitMQProducerTest extends TestBase {
   test("Message not acked returned if number of attempts exhausted") {
     val exchangeName = Random.nextString(10)
     val routingKey = Random.nextString(10)
-    val seqNumber = 1L
 
+    val nextSeqNumber = AtomicInt(0)
     val channel = mock[AutorecoveringChannel]
-    when(channel.getNextPublishSeqNo).thenReturn(seqNumber)
+    when(channel.getNextPublishSeqNo).thenAnswer(_ => nextSeqNumber.get())
+    when(channel.basicPublish(any(), any(), any(), any())).thenAnswer(_ => {
+      nextSeqNumber.increment()
+    })
 
     val producer = new PublishConfirmsRabbitMQProducer[Task, Bytes](
       name = "test",
@@ -77,27 +83,35 @@ class PublisherConfirmsRabbitMQProducerTest extends TestBase {
       sendAttempts = 1
     )
 
-    val body = Bytes.copyFrom(Array.fill(499)(32.toByte))
+    val body = Bytes.copyFrom(Array.fill(499)(64.toByte))
 
     val publishTask = producer.send(routingKey, body).runToFuture
 
+    while (nextSeqNumber.get() < 1) {
+      Thread.sleep(5)
+    }
+
+    producer.DefaultConfirmListener.handleNack(0, multiple = false)
+
     assertThrows[NotAcknowledgedPublish] {
-      updateMessageState(producer, seqNumber)(Left(NotAcknowledgedPublish("abcd", messageId = seqNumber))).await
-      Await.result(publishTask, 10.seconds)
+      Await.result(publishTask, 1.seconds)
     }
 
     verify(channel).basicPublish(Matchers.eq(exchangeName), Matchers.eq(routingKey), any(), Matchers.eq(body.toByteArray))
   }
 
-  test("Multiple messages are fully acked") {
+  test("Multiple messages are fully acked one by one") {
     val exchangeName = Random.nextString(10)
     val routingKey = Random.nextString(10)
 
-    val channel = mock[AutorecoveringChannel]
+    val seqNumbers = (0 to 499).toList
+    val nextSeqNumber = AtomicInt(0)
 
-    val seqNumbers = 1 to 500
-    val iterator = seqNumbers.iterator
-    when(channel.getNextPublishSeqNo).thenAnswer(_ => { iterator.next() })
+    val channel = mock[AutorecoveringChannel]
+    when(channel.getNextPublishSeqNo).thenAnswer(_ => nextSeqNumber.get())
+    when(channel.basicPublish(any(), any(), any(), any())).thenAnswer(_ => {
+      nextSeqNumber.increment()
+    })
 
     val producer = new PublishConfirmsRabbitMQProducer[Task, Bytes](
       name = "test",
@@ -112,38 +126,64 @@ class PublisherConfirmsRabbitMQProducerTest extends TestBase {
       sendAttempts = 2
     )
 
-    val body = Bytes.copyFrom(Array.fill(499)(32.toByte))
+    val body = Bytes.copyFrom(Array.fill(499)(Random.nextInt(255).toByte))
 
-    val publishTasks = Task.parSequenceUnordered {
-      seqNumbers.map { _ =>
-        producer.send(routingKey, body)
-      }
+    val publishFuture = Task.parSequence {
+      seqNumbers.map(_ => producer.send(routingKey, body))
     }.runToFuture
 
-    Task
-      .parSequenceUnordered(seqNumbers.map { seqNumber =>
-        updateMessageState(producer, seqNumber)(Right())
-      })
-      .await(15.seconds)
+    seqNumbers.foreach { seqNumber =>
+      while (nextSeqNumber.get() <= seqNumber) { Thread.sleep(5) }
+      producer.DefaultConfirmListener.handleAck(seqNumber, multiple = false)
+    }
 
-    Await.result(publishTasks, 15.seconds)
+    Await.result(publishFuture, 15.seconds)
 
+    assertResult(seqNumbers.length)(nextSeqNumber.get())
     verify(channel, times(seqNumbers.length))
       .basicPublish(Matchers.eq(exchangeName), Matchers.eq(routingKey), any(), Matchers.eq(body.toByteArray))
   }
 
-  private def updateMessageState(producer: PublishConfirmsRabbitMQProducer[Task, Bytes], messageId: Long, attempt: Int = 1)(
-      result: Either[NotAcknowledgedPublish, Unit]): Task[Unit] = {
-    Task
-      .delay(producer.confirmationCallbacks.get(messageId))
-      .flatMap {
-        case Some(value) => value.complete(result)
-        case None =>
-          if (attempt < 90) {
-            Task.sleep(100.millis) >> updateMessageState(producer, messageId, attempt + 1)(result)
-          } else {
-            throw new InvalidOrderingException(s"The message ID $messageId is not present in the list of callbacks")
-          }
-      }
+  test("Multiple messages are fully acked at once") {
+    val exchangeName = Random.nextString(10)
+    val routingKey = Random.nextString(10)
+
+    val messageCount = 500
+    val seqNumbers = (0 until messageCount).toList
+    val nextSeqNumber = AtomicInt(0)
+
+    val channel = mock[AutorecoveringChannel]
+    when(channel.getNextPublishSeqNo).thenAnswer(_ => nextSeqNumber.get())
+    when(channel.basicPublish(any(), any(), any(), any())).thenAnswer(_ => {
+      nextSeqNumber.increment()
+    })
+
+    val producer = new PublishConfirmsRabbitMQProducer[Task, Bytes](
+      name = "test",
+      exchangeName = exchangeName,
+      channel = channel,
+      monitor = Monitor.noOp(),
+      defaultProperties = MessageProperties.empty,
+      reportUnroutable = false,
+      sizeLimitBytes = None,
+      blocker = TestBase.testBlocker,
+      logger = ImplicitContextLogger.createLogger,
+      sendAttempts = 2
+    )
+
+    val body = Bytes.copyFrom(Array.fill(499)(Random.nextInt(255).toByte))
+
+    val publishFuture = Task.parSequence {
+      seqNumbers.map(_ => producer.send(routingKey, body))
+    }.runToFuture
+
+    while (nextSeqNumber.get() < messageCount) { Thread.sleep(5) }
+    producer.DefaultConfirmListener.handleAck(messageCount, multiple = true)
+
+    Await.result(publishFuture, 15.seconds)
+
+    assertResult(seqNumbers.length)(nextSeqNumber.get())
+    verify(channel, times(seqNumbers.length))
+      .basicPublish(Matchers.eq(exchangeName), Matchers.eq(routingKey), any(), Matchers.eq(body.toByteArray))
   }
 }
