Fix race condition causing many connections to RabbitMQ

Author: balis

functions/kubernetes/amqpConnector.js

@@ -1,17 +1,42 @@
+
 const amqplib = require('amqplib'),
     createJobMessage = require('../../common/jobMessage').createJobMessage;
-let channels = {};
+
 let conn = null;
+let connPromise = null;
+let channels = {};
+let channelPromises = {};
+
+async function getConnection() {
+    if (conn) return conn;
+    if (!connPromise) {
+        console.log("[AMQP] Creating new connection...");
+        connPromise = amqplib.connect(`amqp://${process.env.RABBIT_HOSTNAME}`, "heartbeat=60");
+    }
+    conn = await connPromise;
+    return conn;
+}
 
 async function initialize(queue_name) {
+    const connection = await getConnection();
+
+    if (channels[queue_name]) return;
 
-    if (conn === null) {
-        conn = await amqplib.connect(`amqp://${process.env.RABBIT_HOSTNAME}`, "heartbeat=60");
+    if (!channelPromises[queue_name]) {
+        channelPromises[queue_name] = (async () => {
+            try {
+                console.log(`[AMQP] Creating channel for queue ${queue_name}`);
+                const ch = await connection.createChannel();
+                await ch.assertQueue(queue_name, { durable: false, expires: 6000000 });
+                channels[queue_name] = ch;
+            } catch (err) {
+                delete channelPromises[queue_name]; // retry logic
+                throw err;
+            }
+        })();
     }
-    let ch = await conn.createChannel()
-    await ch.assertQueue(queue_name, {durable: false, expires: 6000000});
-    channels[queue_name] = ch
 
+    await channelPromises[queue_name];
 }
 
 function getQueueName(context) {
@@ -30,28 +55,26 @@ function getQueueName(context) {
 
 async function enqueueJobs(jobArr, taskIdArr, contextArr, customParams) {
     let context = contextArr[0];
-    let queue_name = getQueueName(context)
-    if (conn === null || !(queue_name in channels)) {
-        await initialize(queue_name)
-    }
-    let ch = channels[queue_name]
+    let queue_name = getQueueName(context);
     try {
+        await initialize(queue_name);
+        let ch = channels[queue_name];
 
-        console.log(`jobArr: ${JSON.stringify(jobArr)}, taskIdArr: ${JSON.stringify(taskIdArr)}, contextArr: ${JSON.stringify(contextArr)}, customParams: ${JSON.stringify(customParams)}`)
+        console.log(`jobArr: ${JSON.stringify(jobArr)}, taskIdArr: ${JSON.stringify(taskIdArr)}, contextArr: ${JSON.stringify(contextArr)}, customParams: ${JSON.stringify(customParams)}`);
         let tasks = [];
 
         for (let i = 0; i < jobArr.length; i++) {
             let job = jobArr[i];
             let taskId = taskIdArr[i];
             let jobMessage = createJobMessage(job.ins, job.outs, contextArr[i], taskId);
-            await context.sendMsgToJob(JSON.stringify(jobMessage), taskId) // TODO remove
-            tasks.push({"id": taskId, "message": jobMessage});
+            await context.sendMsgToJob(JSON.stringify(jobMessage), taskId); // TODO remove
+            tasks.push({ "id": taskId, "message": jobMessage });
         }
 
-        await ch.publish('', queue_name, Buffer.from(JSON.stringify({'tasks': tasks})));
+        ch.sendToQueue(queue_name, Buffer.from(JSON.stringify({ 'tasks': tasks })));
     } catch (error) {
-        console.log(error)
+        console.log(error);
     }
 }
 
-exports.enqueueJobs = enqueueJobs
+exports.enqueueJobs = enqueueJobs