@514labs/moose-lib 0.6.276-ci-6-g278c5539 → 0.6.276-ci-1-gfe86cd2c

package/dist/moose-runner.js

@@ -55,7 +55,7 @@ function createProducerConfig(maxMessageBytes) {
     ...maxMessageBytes && { "message.max.bytes": maxMessageBytes }
   };
 }
-var import_http, import_client, import_kafka_javascript, Kafka, compilerLog, getClickhouseClient, cliLog, MAX_RETRIES, MAX_RETRY_TIME_MS, RETRY_INITIAL_TIME_MS, MAX_RETRIES_PRODUCER, ACKs, parseBrokerString, logError, buildSaslConfig, getKafkaClient;
+var import_http, import_client, import_kafka_javascript, Kafka, compilerLog, getClickhouseClient, cliLog, MAX_RETRIES, MAX_RETRY_TIME_MS, RETRY_INITIAL_TIME_MS, MAX_RETRIES_PRODUCER, ACKs, parseBrokerString, logError, buildSaslConfig, getKafkaClient, buildNativeSaslConfig, createNativeKafkaConsumer;
 var init_commons = __esm({
   "src/commons.ts"() {
     "use strict";
@@ -152,6 +152,59 @@ var init_commons = __esm({
         }
       });
     };
+    buildNativeSaslConfig = (logger2, cfg) => {
+      if (!cfg.saslMechanism || !cfg.saslUsername || !cfg.saslPassword) {
+        return {};
+      }
+      const mechanism = cfg.saslMechanism.toUpperCase();
+      const validMechanisms = ["PLAIN", "SCRAM-SHA-256", "SCRAM-SHA-512"];
+      if (!validMechanisms.includes(mechanism)) {
+        logger2.warn(`Unsupported SASL mechanism: ${cfg.saslMechanism}`);
+        return {};
+      }
+      return {
+        "sasl.mechanisms": mechanism,
+        "sasl.username": cfg.saslUsername,
+        "sasl.password": cfg.saslPassword
+      };
+    };
+    createNativeKafkaConsumer = (cfg, logger2, rebalanceCb) => {
+      const brokers = parseBrokerString(cfg.broker || "");
+      if (brokers.length === 0) {
+        throw new Error(`No valid broker addresses found in: "${cfg.broker}"`);
+      }
+      logger2.log(
+        `Creating native KafkaConsumer with brokers: ${brokers.join(", ")}`
+      );
+      logger2.log(`Security protocol: ${cfg.securityProtocol || "plaintext"}`);
+      logger2.log(`Client ID: ${cfg.clientId}`);
+      logger2.log(`Group ID: ${cfg.groupId}`);
+      const saslConfig = buildNativeSaslConfig(logger2, cfg);
+      const consumerConfig = {
+        // Connection
+        "bootstrap.servers": brokers.join(","),
+        "client.id": cfg.clientId,
+        // Group management
+        "group.id": cfg.groupId,
+        "session.timeout.ms": cfg.sessionTimeoutMs ?? 3e4,
+        "heartbeat.interval.ms": cfg.heartbeatIntervalMs ?? 3e3,
+        "max.poll.interval.ms": cfg.maxPollIntervalMs ?? 3e5,
+        // Offset management
+        "enable.auto.commit": cfg.autoCommit ?? true,
+        "auto.commit.interval.ms": cfg.autoCommitIntervalMs ?? 5e3,
+        // Security
+        ...cfg.securityProtocol === "SASL_SSL" && {
+          "security.protocol": "sasl_ssl"
+        },
+        ...saslConfig,
+        // Rebalance callback
+        ...rebalanceCb && { rebalance_cb: rebalanceCb }
+      };
+      const topicConfig = {
+        "auto.offset.reset": cfg.autoOffsetReset ?? "earliest"
+      };
+      return new import_kafka_javascript.KafkaConsumer(consumerConfig, topicConfig);
+    };
   }
 });
 
@@ -1879,8 +1932,6 @@ init_commons();
 var { Kafka: Kafka2 } = import_kafka_javascript2.KafkaJS;
 var HOSTNAME = process3.env.HOSTNAME;
 var AUTO_COMMIT_INTERVAL_MS = 5e3;
-var PARTITIONS_CONSUMED_CONCURRENTLY = 3;
-var MAX_RETRIES_CONSUMER = 150;
 var SESSION_TIMEOUT_CONSUMER = 3e4;
 var HEARTBEAT_INTERVAL_CONSUMER = 3e3;
 var DEFAULT_MAX_STREAMING_CONCURRENCY = 100;
@@ -1921,23 +1972,31 @@ var stopProducer = async (logger2, producer) => {
 var stopConsumer = async (logger2, consumer, sourceTopic) => {
   try {
     logger2.log("Pausing consumer...");
-    const partitionNumbers = Array.from(
+    const topicPartitions = Array.from(
       { length: sourceTopic.partitions },
-      (_, i) => i
-    );
-    await consumer.pause([
-      {
+      (_, i) => ({
         topic: sourceTopic.name,
-        partitions: partitionNumbers
-      }
-    ]);
+        partition: i
+      })
+    );
+    consumer.pause(topicPartitions);
     logger2.log("Disconnecting consumer...");
-    await consumer.disconnect();
+    await new Promise((resolve2, reject) => {
+      consumer.disconnect((err) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve2();
+        }
+      });
+    });
     logger2.log("Consumer is shutting down...");
   } catch (error) {
     logger2.error(`Error during consumer shutdown: ${error}`);
     try {
-      await consumer.disconnect();
+      await new Promise((resolve2) => {
+        consumer.disconnect(() => resolve2());
+      });
       logger2.log("Consumer disconnected after error");
     } catch (disconnectError) {
       logger2.error(`Failed to disconnect consumer: ${disconnectError}`);
@@ -2177,7 +2236,15 @@ var startConsumer = async (args, logger2, metrics, _parallelism, consumer, produ
   }
   try {
     logger2.log("Connecting consumer...");
-    await consumer.connect();
+    await new Promise((resolve2, reject) => {
+      consumer.connect({}, (err) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve2();
+        }
+      });
+    });
     logger2.log("Consumer connected successfully");
   } catch (error) {
     logger2.error("Failed to connect consumer:");
@@ -2202,61 +2269,94 @@ var startConsumer = async (args, logger2, metrics, _parallelism, consumer, produ
     streamingFunctions = [[loadStreamingFunction(args.functionFilePath), {}]];
     fieldMutations = void 0;
   }
-  await consumer.subscribe({
-    topics: [args.sourceTopic.name]
-    // Use full topic name for Kafka operations
-  });
-  await consumer.run({
-    eachBatchAutoResolve: true,
-    // Enable parallel processing of partitions
-    partitionsConsumedConcurrently: PARTITIONS_CONSUMED_CONCURRENTLY,
-    // To be adjusted
-    eachBatch: async ({ batch, heartbeat, isRunning, isStale }) => {
-      if (!isRunning() || isStale()) {
-        return;
-      }
-      metrics.count_in += batch.messages.length;
-      cliLog({
-        action: "Received",
-        message: `${logger2.logPrefix} ${batch.messages.length} message(s)`
-      });
-      logger2.log(`Received ${batch.messages.length} message(s)`);
-      let index = 0;
-      const readableStream = import_node_stream2.Readable.from(batch.messages);
-      const processedMessages = await readableStream.map(
-        async (message) => {
-          index++;
-          if (batch.messages.length > DEFAULT_MAX_STREAMING_CONCURRENCY && index % DEFAULT_MAX_STREAMING_CONCURRENCY || index - 1 === batch.messages.length) {
-            await heartbeat();
+  consumer.subscribe([args.sourceTopic.name]);
+  consumer.setDefaultConsumeTimeout(1e3);
+  let isRunning = true;
+  const consumeLoop = async () => {
+    while (isRunning && consumer.isConnected()) {
+      try {
+        const messages = await new Promise(
+          (resolve2, reject) => {
+            consumer.consume(CONSUMER_MAX_BATCH_SIZE, (err, messages2) => {
+              if (err) {
+                reject(err);
+              } else {
+                resolve2(messages2 || []);
+              }
+            });
+          }
+        );
+        if (messages.length === 0) {
+          continue;
+        }
+        metrics.count_in += messages.length;
+        cliLog({
+          action: "Received",
+          message: `${logger2.logPrefix} ${messages.length} message(s)`
+        });
+        logger2.log(`Received ${messages.length} message(s)`);
+        const readableStream = import_node_stream2.Readable.from(messages);
+        const processedMessages = await readableStream.map(
+          async (message) => {
+            const kafkaMessage = {
+              value: message.value,
+              key: message.key,
+              partition: message.partition,
+              offset: message.offset,
+              timestamp: message.timestamp,
+              headers: message.headers
+            };
+            return handleMessage(
+              logger2,
+              streamingFunctions,
+              kafkaMessage,
+              producer,
+              fieldMutations
+            );
+          },
+          {
+            concurrency: MAX_STREAMING_CONCURRENCY
           }
-          return handleMessage(
+        ).toArray();
+        const filteredMessages = processedMessages.flat().filter((msg) => msg !== void 0 && msg.value !== void 0);
+        if (args.targetTopic === void 0 || processedMessages.length === 0) {
+          continue;
+        }
+        if (filteredMessages.length > 0) {
+          await sendMessages(
             logger2,
-            streamingFunctions,
-            message,
+            metrics,
+            args.targetTopic,
             producer,
-            fieldMutations
+            filteredMessages
           );
-        },
-        {
-          concurrency: MAX_STREAMING_CONCURRENCY
         }
-      ).toArray();
-      const filteredMessages = processedMessages.flat().filter((msg) => msg !== void 0 && msg.value !== void 0);
-      if (args.targetTopic === void 0 || processedMessages.length === 0) {
-        return;
-      }
-      await heartbeat();
-      if (filteredMessages.length > 0) {
-        await sendMessages(
-          logger2,
-          metrics,
-          args.targetTopic,
-          producer,
-          filteredMessages
-        );
+      } catch (error) {
+        if (error && typeof error === "object" && "code" in error) {
+          const kafkaError = error;
+          if (kafkaError.code === import_kafka_javascript2.CODES.ERRORS.ERR__TIMED_OUT) {
+            continue;
+          }
+          if (kafkaError.code === import_kafka_javascript2.CODES.ERRORS.ERR__PARTITION_EOF) {
+            continue;
+          }
+        }
+        logger2.error(`Error consuming messages: ${error}`);
+        if (error instanceof Error) {
+          logError(logger2, error);
+        }
+        await new Promise((resolve2) => setTimeout(resolve2, 100));
       }
     }
+  };
+  consumeLoop().catch((err) => {
+    logger2.error(`Consumer loop crashed: ${err}`);
+    isRunning = false;
   });
+  consumer._isRunning = isRunning;
+  consumer._stopConsuming = () => {
+    isRunning = false;
+  };
   logger2.log("Consumer is running...");
 };
 var buildLogger = (args, workerId) => {
@@ -2336,6 +2436,33 @@ var runStreamingFunctions = async (args) => {
       setTimeout(() => sendMessageMetrics(logger2, metrics), 1e3);
       const clientIdPrefix = HOSTNAME ? `${HOSTNAME}-` : "";
       const processId = `${clientIdPrefix}${streamingFuncId}-ts-${worker.id}`;
+      const consumer = createNativeKafkaConsumer(
+        {
+          clientId: processId,
+          broker: args.broker,
+          groupId: streamingFuncId,
+          securityProtocol: args.securityProtocol,
+          saslUsername: args.saslUsername,
+          saslPassword: args.saslPassword,
+          saslMechanism: args.saslMechanism,
+          sessionTimeoutMs: SESSION_TIMEOUT_CONSUMER,
+          heartbeatIntervalMs: HEARTBEAT_INTERVAL_CONSUMER,
+          autoCommit: true,
+          autoCommitIntervalMs: AUTO_COMMIT_INTERVAL_MS,
+          autoOffsetReset: "earliest",
+          maxBatchSize: CONSUMER_MAX_BATCH_SIZE
+        },
+        logger2,
+        (err, assignments) => {
+          if (err.code === import_kafka_javascript2.CODES.ERRORS.ERR__ASSIGN_PARTITIONS) {
+            logger2.log(`Assigned partitions: ${JSON.stringify(assignments)}`);
+          } else if (err.code === import_kafka_javascript2.CODES.ERRORS.ERR__REVOKE_PARTITIONS) {
+            logger2.log(`Revoked partitions: ${JSON.stringify(assignments)}`);
+          } else {
+            logger2.error(`Rebalance error: ${err.message}`);
+          }
+        }
+      );
       const kafka = await getKafkaClient(
         {
           clientId: processId,
@@ -2347,20 +2474,6 @@ var runStreamingFunctions = async (args) => {
         },
         logger2
       );
-      const consumer = kafka.consumer({
-        kafkaJS: {
-          groupId: streamingFuncId,
-          sessionTimeout: SESSION_TIMEOUT_CONSUMER,
-          heartbeatInterval: HEARTBEAT_INTERVAL_CONSUMER,
-          retry: {
-            retries: MAX_RETRIES_CONSUMER
-          },
-          autoCommit: true,
-          autoCommitInterval: AUTO_COMMIT_INTERVAL_MS,
-          fromBeginning: true
-        },
-        "js.consumer.max.batch.size": CONSUMER_MAX_BATCH_SIZE
-      });
       const maxMessageBytes = args.targetTopic?.max_message_bytes || 1024 * 1024;
       const producer = kafka.producer(
         createProducerConfig(maxMessageBytes)
@@ -2397,6 +2510,9 @@ var runStreamingFunctions = async (args) => {
     },
     workerStop: async ([logger2, producer, consumer]) => {
       logger2.log(`Received SIGTERM, shutting down gracefully...`);
+      if (consumer._stopConsuming) {
+        consumer._stopConsuming();
+      }
       logger2.log("Stopping consumer first...");
       await stopConsumer(logger2, consumer, args.sourceTopic);
       logger2.log("Waiting for in-flight messages to complete...");