@514labs/moose-lib 0.6.276-ci-1-gfe86cd2c → 0.6.276-ci-3-g7a7bb18a

package/dist/moose-runner.js

@@ -55,13 +55,13 @@ 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, buildNativeSaslConfig, createNativeKafkaConsumer;
+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 init_commons = __esm({
   "src/commons.ts"() {
     "use strict";
     import_http = __toESM(require("http"));
     import_client = require("@clickhouse/client");
-    import_kafka_javascript = require("@confluentinc/kafka-javascript");
+    import_kafka_javascript = require("@514labs/kafka-javascript");
     ({ Kafka } = import_kafka_javascript.KafkaJS);
     compilerLog = (message) => {
       if (!isTruthy(process.env.MOOSE_DISABLE_COMPILER_LOGS)) {
@@ -152,59 +152,6 @@ 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);
-    };
   }
 });
 
@@ -1924,7 +1871,7 @@ var runBlocks = async (config) => {
 
 // src/streaming-functions/runner.ts
 var import_node_stream2 = require("stream");
-var import_kafka_javascript2 = require("@confluentinc/kafka-javascript");
+var import_kafka_javascript2 = require("@514labs/kafka-javascript");
 var import_node_buffer = require("buffer");
 var process3 = __toESM(require("process"));
 var http3 = __toESM(require("http"));
@@ -1932,6 +1879,8 @@ 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;
@@ -1972,31 +1921,23 @@ var stopProducer = async (logger2, producer) => {
 var stopConsumer = async (logger2, consumer, sourceTopic) => {
   try {
     logger2.log("Pausing consumer...");
-    const topicPartitions = Array.from(
+    const partitionNumbers = Array.from(
       { length: sourceTopic.partitions },
-      (_, i) => ({
-        topic: sourceTopic.name,
-        partition: i
-      })
+      (_, i) => i
     );
-    consumer.pause(topicPartitions);
+    await consumer.pause([
+      {
+        topic: sourceTopic.name,
+        partitions: partitionNumbers
+      }
+    ]);
     logger2.log("Disconnecting consumer...");
-    await new Promise((resolve2, reject) => {
-      consumer.disconnect((err) => {
-        if (err) {
-          reject(err);
-        } else {
-          resolve2();
-        }
-      });
-    });
+    await consumer.disconnect();
     logger2.log("Consumer is shutting down...");
   } catch (error) {
     logger2.error(`Error during consumer shutdown: ${error}`);
     try {
-      await new Promise((resolve2) => {
-        consumer.disconnect(() => resolve2());
-      });
+      await consumer.disconnect();
       logger2.log("Consumer disconnected after error");
     } catch (disconnectError) {
       logger2.error(`Failed to disconnect consumer: ${disconnectError}`);
@@ -2236,15 +2177,7 @@ var startConsumer = async (args, logger2, metrics, _parallelism, consumer, produ
   }
   try {
     logger2.log("Connecting consumer...");
-    await new Promise((resolve2, reject) => {
-      consumer.connect({}, (err) => {
-        if (err) {
-          reject(err);
-        } else {
-          resolve2();
-        }
-      });
-    });
+    await consumer.connect();
     logger2.log("Consumer connected successfully");
   } catch (error) {
     logger2.error("Failed to connect consumer:");
@@ -2269,94 +2202,61 @@ var startConsumer = async (args, logger2, metrics, _parallelism, consumer, produ
     streamingFunctions = [[loadStreamingFunction(args.functionFilePath), {}]];
     fieldMutations = void 0;
   }
-  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
+  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();
           }
-        ).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(
+          return handleMessage(
             logger2,
-            metrics,
-            args.targetTopic,
+            streamingFunctions,
+            message,
             producer,
-            filteredMessages
+            fieldMutations
           );
+        },
+        {
+          concurrency: MAX_STREAMING_CONCURRENCY
         }
-      } 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));
+      ).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
+        );
       }
     }
-  };
-  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) => {
@@ -2436,33 +2336,6 @@ 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,
@@ -2474,6 +2347,20 @@ 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)
@@ -2510,9 +2397,6 @@ 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...");