Skip to content

Commit

Permalink
fix: core - publisher confirmations - add support for confirmation of…
Browse files Browse the repository at this point in the history 
… multiple messages (#201)

fix: core - publisher confirmations - add support for confirmation of multiple messages
  • Loading branch information
@mi-char
mi-char committed Mar 10, 2023
1 parent 3ab1cdf commit d9d1435
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 55 deletions.
37 changes: 26 additions & 11 deletions ...src/main/scala/com/avast/clients/rabbitmq/publisher/PublishConfirmsRabbitMQProducer.scala
View file
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
}
Expand All @@ -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.")
}
}
}
}
Expand Down
128 changes: 84 additions & 44 deletions core/src/test/scala/com/avast/clients/rabbitmq/PublisherConfirmsRabbitMQProducerTest.scala
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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}
Expand All @@ -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",
Expand All @@ -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))
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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))
}
}

0 comments on commit d9d1435

Please sign in to comment.