Merge pull request #84 from viartemev/add-proper-connection-closing

Add proper connection closing

Author: viartemev

src/main/kotlin/com/viartemev/thewhiterabbit/channel/Channel.kt

@@ -1,6 +1,22 @@
 package com.viartemev.thewhiterabbit.channel
 
 import com.rabbitmq.client.Channel
+import com.rabbitmq.client.Connection
 import com.viartemev.thewhiterabbit.consumer.ConfirmConsumer
 
-fun Channel.consumer(queue: String) = ConfirmConsumer(this, queue)
+fun Channel.consumer(queue: String) = ConfirmConsumer(this, queue)
+
+suspend fun Connection.channel(block: suspend Channel.() -> Unit): Channel {
+    val channel = this.createChannel()
+    channel.use { block(it) }
+    return channel
+}
+
+suspend fun Channel.consume(queue: String, block: suspend ConfirmConsumer.() -> Unit) {
+    val consumer = this.consumer(queue)
+    try {
+        block(consumer)
+    } finally {
+        consumer.cancel()
+    }
+}

src/main/kotlin/com/viartemev/thewhiterabbit/channel/ConfirmChannel.kt

@@ -17,3 +17,14 @@ class ConfirmChannel internal constructor(private val channel: Channel) : Channe
  * @see com.rabbitmq.client.Channel.confirmSelect()
  */
 fun Connection.createConfirmChannel(): ConfirmChannel = ConfirmChannel(this.createChannel())
+
+suspend fun Connection.confirmChannel(block: suspend ConfirmChannel.() -> Unit): ConfirmChannel {
+    val confirmChannel = this.createConfirmChannel()
+    confirmChannel.use { block(it) }
+    return confirmChannel
+}
+
+suspend fun ConfirmChannel.publish(block: suspend ConfirmPublisher.() -> Unit) {
+    val publisher = this.publisher()
+    block(publisher)
+}

src/main/kotlin/com/viartemev/thewhiterabbit/consumer/ConfirmConsumer.kt

@@ -20,7 +20,13 @@ class ConfirmConsumer internal constructor(private val AMQPChannel: Channel, AMQ
     init {
         AMQPChannel.basicQos(prefetchSize, false)
         consTag = AMQPChannel.basicConsume(AMQPQueue, false,
-            { _, message -> continuations.sendBlocking(message) },
+            { consumerTag, message ->
+                try {
+                    continuations.sendBlocking(message)
+                } catch (e: Exception) {
+                    logger.info { "Consumer $consumerTag has been cancelled" }
+                }
+            },
             { consumerTag ->
                 logger.info { "Consumer $consumerTag has been cancelled" }
                 continuations.cancel()
@@ -44,4 +50,9 @@ class ConfirmConsumer internal constructor(private val AMQPChannel: Channel, AMQ
             throw AcknowledgeException(errorMessage)
         }
     }
+
+    fun cancel() {
+        AMQPChannel.basicCancel(consTag)
+        continuations.cancel()
+    }
 }

src/test/kotlin/com/viartemev/thewhiterabbit/consumer/ConfirmConsumerTest.kt

@@ -2,7 +2,8 @@ package com.viartemev.thewhiterabbit.consumer
 
 import com.rabbitmq.client.ConnectionFactory
 import com.rabbitmq.client.Delivery
-import com.viartemev.thewhiterabbit.channel.consumer
+import com.viartemev.thewhiterabbit.channel.channel
+import com.viartemev.thewhiterabbit.channel.consume
 import com.viartemev.thewhiterabbit.queue.QueueSpecification
 import com.viartemev.thewhiterabbit.queue.declareQueue
 import kotlinx.coroutines.delay
@@ -29,11 +30,12 @@ class ConfirmConsumerTest {
     @Test
     fun `test message consuming`() {
         factory.newConnection().use { connection ->
-            connection.createChannel().use { channel ->
-                runBlocking {
-                    channel.declareQueue(QueueSpecification(QUEUE_NAME))
-                    val consumer = channel.consumer(QUEUE_NAME)
-                    for (i in 1..3) consumer.consumeWithConfirm({ handleDelivery(it) })
+            runBlocking {
+                connection.channel {
+                    declareQueue(QueueSpecification(QUEUE_NAME))
+                    consume(QUEUE_NAME) {
+                        for (i in 1..3) consumeWithConfirm({ handleDelivery(it) })
+                    }
                 }
             }
         }

src/test/kotlin/com/viartemev/thewhiterabbit/publisher/PublisherTest.kt

@@ -2,7 +2,8 @@ package com.viartemev.thewhiterabbit.publisher
 
 import com.rabbitmq.client.ConnectionFactory
 import com.rabbitmq.client.MessageProperties
-import com.viartemev.thewhiterabbit.channel.createConfirmChannel
+import com.viartemev.thewhiterabbit.channel.confirmChannel
+import com.viartemev.thewhiterabbit.channel.publish
 import com.viartemev.thewhiterabbit.queue.QueueSpecification
 import com.viartemev.thewhiterabbit.queue.declareQueue
 import kotlinx.coroutines.async
@@ -33,14 +34,16 @@ class PublisherTest {
 
     @Test
     fun `test one message publishing`() {
-        factory.newConnection().use { connection ->
-            connection.createConfirmChannel().use { channel ->
-                val publisher = channel.publisher()
-                runBlocking {
-                    channel.declareQueue(QueueSpecification(QUEUE_NAME))
-                    val message = createMessage("Hello")
-                    val ack = publisher.publishWithConfirm(message)
-                    assertTrue { ack }
+        factory.newConnection().use {
+            val connection = it
+            runBlocking {
+                connection.confirmChannel {
+                    declareQueue(QueueSpecification(QUEUE_NAME))
+                    publish {
+                        val message = createMessage("Hello")
+                        val ack = publishWithConfirm(message)
+                        assertTrue { ack }
+                    }
                 }
             }
         }
@@ -50,19 +53,19 @@ class PublisherTest {
     fun `test n-messages publishing manually`() {
         val times = 10
         factory.newConnection().use { connection ->
-            connection.createConfirmChannel().use { channel ->
-                val publisher = channel.publisher()
-                runBlocking {
-                    channel.declareQueue(QueueSpecification(QUEUE_NAME))
-                    val acks = coroutineScope {
-
-                        (1..times).map {
-                            async {
-                                publisher.publishWithConfirm(createMessage("Hello #$it"))
-                            }
-                        }.awaitAll()
+            runBlocking {
+                connection.confirmChannel {
+                    declareQueue(QueueSpecification(QUEUE_NAME))
+                    publish {
+                        val acks = coroutineScope {
+                            (1..times).map {
+                                async {
+                                    publishWithConfirm(createMessage("Hello #$it"))
+                                }
+                            }.awaitAll()
+                        }
+                        assertTrue { acks.all { true } }
                     }
-                    assertTrue { acks.all { true } }
                 }
             }
         }
@@ -72,13 +75,14 @@ class PublisherTest {
     fun `test n-messages publishing`() {
         val times = 10
         factory.newConnection().use { connection ->
-            connection.createConfirmChannel().use { channel ->
-                val publisher = channel.publisher()
-                runBlocking {
-                    channel.declareQueue(QueueSpecification(QUEUE_NAME))
-                    val messages = (1..times).map { createMessage("Hello #$it") }
-                    val acks = publisher.publishWithConfirm(messages).awaitAll()
-                    assertTrue { acks.all { true } }
+            runBlocking {
+                connection.confirmChannel {
+                    declareQueue(QueueSpecification(QUEUE_NAME))
+                    publish {
+                        val messages = (1..times).map { createMessage("Hello #$it") }
+                        val acks = publishWithConfirm(messages).awaitAll()
+                        assertTrue { acks.all { true } }
+                    }
                 }
             }
         }