Close RabbitMQ connection after publishing to prevent file descriptor leaks

Author: knmarcin

binder/websocket.py

@@ -32,20 +32,24 @@ def list_rooms_for_user(self, user):
 
 def trigger(data, rooms):
     if 'rabbitmq' in getattr(settings, 'HIGH_TEMPLAR', {}):
-        import pika
-        from pika import BlockingConnection
-
-        connection_credentials = pika.PlainCredentials(settings.HIGH_TEMPLAR['rabbitmq']['username'],
-                                                       settings.HIGH_TEMPLAR['rabbitmq']['password'])
-        connection_parameters = pika.ConnectionParameters(settings.HIGH_TEMPLAR['rabbitmq']['host'],
-                                                          credentials=connection_credentials)
-        connection = BlockingConnection(parameters=connection_parameters)
-        channel = connection.channel()
-
-        channel.basic_publish('hightemplar', routing_key='*', body=jsondumps({
-            'data': data,
-            'rooms': rooms,
-        }))
+        connection = None
+        try:
+            import pika
+            from pika import BlockingConnection
+            connection_credentials = pika.PlainCredentials(settings.HIGH_TEMPLAR['rabbitmq']['username'],
+                                                           settings.HIGH_TEMPLAR['rabbitmq']['password'])
+            connection_parameters = pika.ConnectionParameters(settings.HIGH_TEMPLAR['rabbitmq']['host'],
+                                                              credentials=connection_credentials)
+            connection = BlockingConnection(parameters=connection_parameters)
+            channel = connection.channel()
+
+            channel.basic_publish('hightemplar', routing_key='*', body=jsondumps({
+                'data': data,
+                'rooms': rooms,
+            }))
+        finally:
+            if connection and not connection.is_closed:
+                connection.close()
     if getattr(settings, 'HIGH_TEMPLAR_URL', None):
         url = getattr(settings, 'HIGH_TEMPLAR_URL')
         try:

ci-requirements.txt

@@ -8,3 +8,4 @@ codecov
 coverage
 django-hijack<3.0.0
 openpyxl
+pika

setup.py

@@ -40,6 +40,7 @@
 		'Pillow >= 3.2.0',
 		'django-request-id >= 1.0.0',
 		'requests >= 2.13.0',
+		'pika == 1.3.2',
 	],
 	tests_require=[
 		'django-hijack >= 2.1.10, < 3.0.0',

tests/test_websocket.py

@@ -2,6 +2,7 @@
 from django.contrib.auth.models import User
 from unittest import mock
 from binder.views import JsonResponse
+from binder.websocket import trigger
 from .testapp.urls import room_controller
 from .testapp.models import Animal, Costume
 import requests
@@ -68,3 +69,27 @@ def test_post_succeeds_when_trigger_fails(self):
 		costume.save()
 
 		self.assertIsNotNone(costume.pk)
+
+
+class TriggerConnectionCloseTest(TestCase):
+	@override_settings(
+		HIGH_TEMPLAR={
+			'rabbitmq': {
+				'host': 'localhost',
+				'username': 'guest',
+				'password': 'guest'
+			}
+		}
+	)
+	@mock.patch('pika.BlockingConnection')
+	def test_trigger_calls_connection_close(self, mock_connection_class):
+		mock_connection = mock_connection_class.return_value
+		mock_connection.is_closed = False
+
+		data = {'id': 123}
+		rooms = [{'costume': 123}]
+
+		trigger(data, rooms)
+
+		mock_connection.close.assert_called_once()
+