Batch Kafka messages and drop oversized payloads

src/lib/kafka.ts

@@ -8,11 +8,35 @@ const log = debug('umami:kafka');
 const CONNECT_TIMEOUT = 5000;
 const SEND_TIMEOUT = 3000;
 const ACKS = 1;
+const DEFAULT_MAX_MESSAGE_BYTES = 900_000;
 
 let kafka: Kafka;
 let producer: Producer;
 const enabled = Boolean(process.env.KAFKA_URL && process.env.KAFKA_BROKER);
 
+type KafkaMessage = Record<string, unknown>;
+type KafkaProducerMessage = { value: string };
+
+function getMaxMessageBytes() {
+  const size = Number(process.env.KAFKA_MAX_MESSAGE_BYTES);
+
+  return Number.isFinite(size) && size > 0 ? size : DEFAULT_MAX_MESSAGE_BYTES;
+}
+
+function getMessageSize(value: string) {
+  return Buffer.byteLength(value, 'utf8');
+}
+
+function getMessages(message: KafkaMessage | KafkaMessage[]) {
+  const items = Array.isArray(message) ? message : [message];
+
+  return items.map(item => {
+    const value = JSON.stringify(item);
+
+    return { value, size: getMessageSize(value) };
+  });
+}
+
 function getClient() {
   const { username, password } = new URL(process.env.KAFKA_URL);
   const brokers = process.env.KAFKA_BROKER.split(',');
@@ -65,28 +89,57 @@ async function getProducer(): Promise<Producer> {
 
 async function sendMessage(
   topic: string,
-  message: Record<string, string | number> | Record<string, string | number>[],
+  message: KafkaMessage | KafkaMessage[],
 ): Promise<RecordMetadata[]> {
   try {
     await connect();
 
-    return producer.send({
-      topic,
-      messages: Array.isArray(message)
-        ? message.map(a => {
-            return { value: JSON.stringify(a) };
-          })
-        : [
-            {
-              value: JSON.stringify(message),
-            },
-          ],
-      timeout: SEND_TIMEOUT,
-      acks: ACKS,
-    });
+    const maxMessageBytes = getMaxMessageBytes();
+    const messages = getMessages(message);
+    const result: RecordMetadata[] = [];
+    let batch: KafkaProducerMessage[] = [];
+    let batchSize = 0;
+
+    for (const { value, size } of messages) {
+      if (size > maxMessageBytes) {
+        log('Kafka message dropped: topic=%s size=%d max=%d', topic, size, maxMessageBytes);
+        continue;
+      }
+
+      if (batch.length && batchSize + size > maxMessageBytes) {
+        result.push(
+          ...(await producer.send({
+            topic,
+            messages: batch,
+            timeout: SEND_TIMEOUT,
+            acks: ACKS,
+          })),
+        );
+        batch = [];
+        batchSize = 0;
+      }
+
+      batch.push({ value });
+      batchSize += size;
+    }
+
+    if (batch.length) {
+      result.push(
+        ...(await producer.send({
+          topic,
+          messages: batch,
+          timeout: SEND_TIMEOUT,
+          acks: ACKS,
+        })),
+      );
+    }
+
+    return result;
   } catch (e) {
     // eslint-disable-next-line no-console
     console.log('KAFKA ERROR:', serializeError(e));
+
+    return [];
   }
 }