Merge pull request #225 from santiment/fixPrimaryKey

Option to disable Kafka producer sticky partition

Author: YordanPavlov

src/e2e/producer-transaction.spec.ts

@@ -117,7 +117,7 @@ describe('Producer transactions', function () {
     // consumer but there is something I am missing.
     setTimeout(async function () {
       for (let i = 0; i < num_messages_test; i++) {
-        exporter.sendDataWithKey({
+        exporter.sendData({
           timestamp: 10000000,
           iso_date: new Date().toISOString(),
           key: 1

src/lib/constants.ts

@@ -13,4 +13,5 @@ export const START_PRIMARY_KEY = getIntEnvVariable('START_PRIMARY_KEY', -1);
 export const WRITE_SIGNAL_RECORDS_KAFKA = getBoolEnvVariable('WRITE_SIGNAL_RECORDS_KAFKA', false);
 export const KAFKA_TOPIC = process.env.KAFKA_TOPIC;
 export const TEST_ENV = getBoolEnvVariable('TEST_ENV', false);
+export const KAFKA_PRODUCER_DISABLE_STICKY_PARTITION = getBoolEnvVariable('KAFKA_PRODUCER_DISABLE_STICKY_PARTITION', false);
 

src/lib/kafka_storage.ts

@@ -43,13 +43,13 @@ class Partitioner {
 
   async init(hashFunction: ((value: object) => number), topicName: string, producer: Kafka.Producer) {
     this.hashFunction = hashFunction;
-    this.partitionCount = await this.findPartitionCount(topicName, producer);
+    this.partitionCount = await this.#findPartitionCount(topicName, producer);
     if (this.partitionCount <= 0) {
       throw new Error('Partition count should be > 0');
     }
   }
 
-  async findPartitionCount(topicName: string, producer: Kafka.Producer) {
+  async #findPartitionCount(topicName: string, producer: Kafka.Producer) {
     // We depend on the producer already being connected
     logger.info(`Finding partition count for topic ${topicName}`);
     const metadata: Kafka.Metadata = await new Promise((resolve, reject) => {
@@ -106,7 +106,7 @@ export class Exporter {
   private readonly zookeeperClient: ZookeeperClientAsync;
   private partitioner: Partitioner | null;
 
-  constructor(exporter_name: string, transactional: boolean, topicName: string) {
+  constructor(exporter_name: string, transactional: boolean, topicName: string, disableStickyPartition: boolean = false) {
     this.exporter_name = exporter_name;
 
     const producer_settings: ProducerGlobalConfig = {
@@ -133,6 +133,10 @@ export class Exporter {
       producer_settings['enable.idempotence'] = true;
     }
 
+    if (disableStickyPartition) {
+      producer_settings['sticky.partitioning.linger.ms'] = 0;
+    }
+
     this.producer = new Kafka.Producer(producer_settings);
 
     this.producer.on('event.log', function (log) {
@@ -296,42 +300,24 @@ export class Exporter {
     }
   }
 
-  async sendData(events: Array<any>) {
-    if (events.constructor !== Array) {
-      events = [events];
-    }
-
-    events = events.map(
-      event => (typeof event === 'object' ? JSON.stringify(event) : event)
-    );
-    events.forEach(event => {
-      this.producer.produce(this.topicName, null, Buffer.from(event));
-    });
-
-    return new Promise<void>((resolve, reject) =>
-      this.producer.flush(KAFKA_FLUSH_TIMEOUT, (err: LibrdKafkaError) => {
-        if (err) return reject(err);
-        resolve();
-      })
-    );
-  }
-
-  async sendDataWithKey(events: object | Array<object>, keyField: string, signalRecordData: object | null) {
+  async sendData(events: object | Array<object>, keyField: string | null, signalRecordData: object | null) {
     const arrayEvents: Array<object> = (events.constructor !== Array) ? [events] : events
 
     if (signalRecordData !== null && this.partitioner === null) {
       throw new Error('Signal record logic needs partitioner');
     }
 
     arrayEvents.forEach((event: any) => {
-      const partitionNumberPayload = this.partitioner ? this.partitioner.getPartitionNumber(event) : null;
+      const key = keyField !== null ? event[keyField] : null
+
+      const partitionNumberOfPayload = this.partitioner ? this.partitioner.getPartitionNumber(event) : null;
       const eventString = typeof event === 'object' ? JSON.stringify(event) : event;
-      this.producer.produce(this.topicName, partitionNumberPayload, Buffer.from(eventString), event[keyField]);
+      this.producer.produce(this.topicName, partitionNumberOfPayload, Buffer.from(eventString), key);
       if (signalRecordData !== null && this.partitioner !== null) {
         const signalRecordString = typeof signalRecordData === 'object' ? JSON.stringify(signalRecordData) : signalRecordData;
         for (let partitionNumber = 0; partitionNumber < this.partitioner.getPartitionCount(); ++partitionNumber) {
-          if (partitionNumber !== partitionNumberPayload) {
-            this.producer.produce(this.topicName, partitionNumber, Buffer.from(signalRecordString), event[keyField]);
+          if (partitionNumber !== partitionNumberOfPayload) {
+            this.producer.produce(this.topicName, partitionNumber, Buffer.from(signalRecordString), key);
           }
         }
       }
@@ -382,11 +368,12 @@ export class Exporter {
     const signalRecord: object | null = writeSignalRecordsKafka ? { 'santiment_signal_record': true } : null
     try {
       if (BLOCKCHAIN === 'utxo') {
-        await this.sendDataWithKey(events, 'height', signalRecord);
-      } else if (BLOCKCHAIN === 'receipts') {
-        await this.sendDataWithKey(events, 'transactionHash', signalRecord);
-      } else {
-        await this.sendDataWithKey(events, 'primaryKey', signalRecord);
+        await this.sendData(events, 'height', signalRecord);
+      } else if (BLOCKCHAIN === 'receipts' || BLOCKCHAIN === 'eth') {
+        await this.sendData(events, 'transactionHash', signalRecord);
+      }
+      else {
+        await this.sendData(events, 'primaryKey', signalRecord);
       }
       await this.commitTransaction();
     } catch (exception) {

src/main.ts

@@ -27,9 +27,9 @@ export class Main {
     ))
   }
 
-  async initExporter(exporterName: string, isTransactions: boolean, kafkaTopic: string) {
+  async initExporter(exporterName: string, isTransactions: boolean, kafkaTopic: string, disableStickyPartition: boolean) {
     const INIT_EXPORTER_ERR_MSG = 'Error when initializing exporter: ';
-    this.exporter = new Exporter(exporterName, isTransactions, kafkaTopic);
+    this.exporter = new Exporter(exporterName, isTransactions, kafkaTopic, disableStickyPartition);
     await this.exporter
       .connect()
       .then(() => this.exporter.initTransactions())
@@ -70,7 +70,7 @@ export class Main {
   async init(blockchain: string) {
     const blockchainSpecificConstants = require(`./blockchains/${blockchain}/lib/constants`);
     const mergedConstants = { ...constantsBase, ...blockchainSpecificConstants };
-    await this.initExporter(EXPORTER_NAME, true, mergedConstants.KAFKA_TOPIC);
+    await this.initExporter(EXPORTER_NAME, true, mergedConstants.KAFKA_TOPIC, mergedConstants.KAFKA_PRODUCER_DISABLE_STICKY_PARTITION);
     await this.initWorker(blockchain, mergedConstants);
     metrics.startCollection();