@cryptexlabs/codex-nodejs-common 0.2.3 → 0.3.0

package/src/service/kafka/kafka.service.ts

@@ -2,6 +2,7 @@ import { Inject, Injectable, LoggerService } from "@nestjs/common";
 import {
   Admin,
   Consumer,
+  EachMessagePayload,
   ITopicConfig,
   Kafka,
   Message,
@@ -115,6 +116,39 @@ export class KafkaService implements ConsumerServiceDelegateInterface {
     });
   }
 
+  public async startManualConsumer(
+    topics: (string | RegExp)[],
+    callback: (topic: string, message: Message) => Promise<void>
+  ): Promise<void> {
+    await this._setManualConsume();
+
+    const promises = [];
+    for (const topic of topics) {
+      promises.push(
+        this._kafkaConsumer.subscribe({ topic, fromBeginning: true })
+      );
+    }
+    try {
+      await Promise.all(promises);
+
+      await this._ensureTopicsExist(topics);
+
+      const results = [
+        new Promise((resolve) => {
+          this._kafkaConsumer.run({
+            eachMessage: async (messagePayload: EachMessagePayload) => {
+              const { topic, partition, message } = messagePayload;
+              await callback(topic, message);
+            },
+          });
+        }),
+      ];
+      await Promise.all(results);
+    } catch (e) {
+      this.logger.error(e.message, e.trace);
+    }
+  }
+
   private async _ensureTopicsExist(topics: (string | RegExp)[]): Promise<void> {
     const stringTopics = [];
     for (const topic of topics) {
@@ -165,4 +199,115 @@ export class KafkaService implements ConsumerServiceDelegateInterface {
       messages,
     });
   }
+
+  private async _setManualConsume(): Promise<void> {
+    /* Based on solution provided here:
+     * https://github.com/tulios/kafkajs/issues/825#issuecomment-674106799
+     *
+     * 1. We need to know which partitions we are assigned.
+     * 2. Which partitions have we consumed the last offset for
+     * 3. If all partitions have 0 lag, we exit.
+     */
+
+    /*
+     * `consumedTopicPartitions` will be an object of all topic-partitions
+     * and a boolean indicating whether or not we have consumed all
+     * messages in that topic-partition. For example:
+     *
+     * {
+     *   "topic-test-0": false,
+     *   "topic-test-1": false,
+     *   "topic-test-2": false
+     * }
+     */
+    let consumedTopicPartitions = {};
+    this._kafkaConsumer.on(
+      this._kafkaConsumer.events.GROUP_JOIN,
+      async ({ payload }) => {
+        const { memberAssignment } = payload;
+
+        consumedTopicPartitions = Object.entries(memberAssignment).reduce(
+          (topics, [topic, partitions]) => {
+            for (const partition in partitions) {
+              this.logger.debug("topic-partition: ", `${topic}-${partition}`);
+              topics[`${topic}-${partition}`] = false;
+            }
+            return topics;
+          },
+          {}
+        );
+      }
+    );
+
+    /*
+     * This is extremely unergonomic, but if we are currently caught up to the head
+     * of all topic-partitions, we won't actually get any batches, which means we'll
+     * never find out that we are actually caught up. So as a workaround, what we can do
+     * is to check in `FETCH_START` if we have previously made a fetch without
+     * processing any batches in between. If so, it means that we received empty
+     * fetch responses, meaning there was no more data to fetch.
+     *
+     * We need to initially set this to true, or we would immediately exit.
+     */
+    let processedBatch = true;
+    this._kafkaConsumer.on(this._kafkaConsumer.events.FETCH_START, async () => {
+      if (processedBatch === false) {
+        await this._kafkaConsumer.disconnect();
+        process.exit(0);
+      }
+
+      processedBatch = false;
+    });
+
+    /*
+     * Now whenever we have finished processing a batch, we'll update `consumedTopicPartitions`
+     * and exit if all topic-partitions have been consumed,
+     */
+    this._kafkaConsumer.on(
+      this._kafkaConsumer.events.END_BATCH_PROCESS,
+      async ({ payload }) => {
+        const { topic, partition, offsetLag } = payload;
+
+        consumedTopicPartitions[`${topic}-${partition}`] = offsetLag === "0";
+
+        if (
+          Object.values(consumedTopicPartitions).every((consumed) =>
+            Boolean(consumed)
+          )
+        ) {
+          await this._kafkaConsumer.disconnect();
+          process.exit(0);
+        }
+
+        processedBatch = true;
+      }
+    );
+
+    // Graceful shutdown
+    const errorTypes = ["unhandledRejection", "uncaughtException"];
+    const signalTraps: NodeJS.Signals[] = ["SIGTERM", "SIGINT", "SIGUSR2"];
+
+    errorTypes.map((type) => {
+      process.on(type, async (e) => {
+        try {
+          this.logger.debug(`process.on ${type}`);
+          this.logger.error(e);
+          await this.disconnect();
+          process.exit(0);
+        } catch (_) {
+          process.exit(1);
+        }
+      });
+    });
+
+    signalTraps.map((type) => {
+      process.once(type, async () => {
+        try {
+          await this.disconnect();
+        } finally {
+          process.kill(process.pid, type);
+        }
+      });
+    });
+  }
 }