Merging upstream version 1.44.3.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 382 insertions, 0 deletions

diff --git a/fluent-bit/plugins/in_kafka/in_kafka.c b/fluent-bit/plugins/in_kafka/in_kafka.c
new file mode 100644
index 000000000..972ae4170
--- /dev/null
+++ b/fluent-bit/plugins/in_kafka/in_kafka.c
@@ -0,0 +1,382 @@
+/* -*- Mode: C; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
+
+/*  Fluent Bit
+ *  ==========
+ *  Copyright (C) 2019-2021 The Fluent Bit Authors
+ *  Copyright (C) 2015-2018 Treasure Data Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include <fluent-bit/flb_input_plugin.h>
+#include <fluent-bit/flb_pack.h>
+#include <fluent-bit/flb_engine.h>
+#include <fluent-bit/flb_time.h>
+#include <fluent-bit/flb_parser.h>
+#include <fluent-bit/flb_error.h>
+#include <fluent-bit/flb_utils.h>
+#include <mpack/mpack.h>
+#include <stddef.h>
+#include <stdio.h>
+
+#include "fluent-bit/flb_input.h"
+#include "fluent-bit/flb_kafka.h"
+#include "fluent-bit/flb_mem.h"
+#include "in_kafka.h"
+#include "rdkafka.h"
+
+static int try_json(struct flb_log_event_encoder *log_encoder,
+                    rd_kafka_message_t *rkm)
+{
+    int root_type;
+    char *buf = NULL;
+    size_t bufsize;
+    int ret;
+
+    ret = flb_pack_json(rkm->payload, rkm->len, &buf, &bufsize, &root_type, NULL);
+    if (ret) {
+        if (buf) {
+            flb_free(buf);
+        }
+        return ret;
+    }
+    flb_log_event_encoder_append_body_raw_msgpack(log_encoder, buf, bufsize);
+    flb_free(buf);
+    return 0;
+}
+
+static int process_message(struct flb_in_kafka_config *ctx,
+                           rd_kafka_message_t *rkm)
+{
+    struct flb_log_event_encoder *log_encoder = ctx->log_encoder;
+    int ret;
+
+    ret = flb_log_event_encoder_begin_record(log_encoder);
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_set_current_timestamp(log_encoder);
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "topic");
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        if (rkm->rkt) {
+            ret = flb_log_event_encoder_append_body_cstring(log_encoder,
+                                                            rd_kafka_topic_name(rkm->rkt));
+        }
+        else {
+            ret = flb_log_event_encoder_append_body_null(log_encoder);
+        }
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_append_body_values(log_encoder,
+                                                       FLB_LOG_EVENT_CSTRING_VALUE("partition"),
+                                                       FLB_LOG_EVENT_INT32_VALUE(rkm->partition));
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_append_body_values(log_encoder,
+                                                       FLB_LOG_EVENT_CSTRING_VALUE("offset"),
+                                                       FLB_LOG_EVENT_INT64_VALUE(rkm->offset));
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "error");
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        if (rkm->err) {
+            ret = flb_log_event_encoder_append_body_cstring(log_encoder,
+                                                            rd_kafka_message_errstr(rkm));
+        }
+        else {
+            ret = flb_log_event_encoder_append_body_null(log_encoder);
+        }
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "key");
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        if (rkm->key) {
+            ret = flb_log_event_encoder_append_body_string(log_encoder,
+                                                           rkm->key,
+                                                           rkm->key_len);
+        }
+        else {
+            ret = flb_log_event_encoder_append_body_null(log_encoder);
+        }
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "payload");
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        if (rkm->payload) {
+            if (ctx->format != FLB_IN_KAFKA_FORMAT_JSON ||
+                    try_json(log_encoder, rkm)) {
+                ret = flb_log_event_encoder_append_body_string(log_encoder,
+                                                               rkm->payload,
+                                                               rkm->len);
+            }
+        }
+        else {
+            ret = flb_log_event_encoder_append_body_null(log_encoder);
+        }
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        ret = flb_log_event_encoder_commit_record(log_encoder);
+    }
+
+    if (ret != FLB_EVENT_ENCODER_SUCCESS) {
+        flb_log_event_encoder_rollback_record(log_encoder);
+    }
+
+    return ret;
+}
+
+static int in_kafka_collect(struct flb_input_instance *ins,
+                            struct flb_config *config, void *in_context)
+{
+    int ret;
+    struct flb_in_kafka_config *ctx = in_context;
+    rd_kafka_message_t *rkm;
+
+    ret = FLB_EVENT_ENCODER_SUCCESS;
+
+    while (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        rkm = rd_kafka_consumer_poll(ctx->kafka.rk, 1);
+
+        if (!rkm) {
+            break;
+        }
+
+        flb_plg_debug(ins, "kafka message received");
+
+        ret = process_message(ctx, rkm);
+
+        rd_kafka_message_destroy(rkm);
+
+        /* TO-DO: commit the record based on `ret` */
+        rd_kafka_commit(ctx->kafka.rk, NULL, 0);
+    }
+
+    if (ret == FLB_EVENT_ENCODER_SUCCESS) {
+        flb_input_log_append(ins, NULL, 0,
+                             ctx->log_encoder->output_buffer,
+                             ctx->log_encoder->output_length);
+        ret = 0;
+    }
+    else {
+        flb_plg_error(ins, "Error encoding record : %d", ret);
+        ret = -1;
+    }
+
+    flb_log_event_encoder_reset(ctx->log_encoder);
+
+    return ret;
+}
+
+/* Initialize plugin */
+static int in_kafka_init(struct flb_input_instance *ins,
+                         struct flb_config *config, void *data)
+{
+    int ret;
+    const char *conf;
+    struct flb_in_kafka_config *ctx;
+    rd_kafka_conf_t *kafka_conf = NULL;
+    rd_kafka_topic_partition_list_t *kafka_topics = NULL;
+    rd_kafka_resp_err_t err;
+    char errstr[512];
+    (void) data;
+
+    /* Allocate space for the configuration context */
+    ctx = flb_malloc(sizeof(struct flb_in_kafka_config));
+    if (!ctx) {
+        return -1;
+    }
+    ctx->ins = ins;
+
+    ret = flb_input_config_map_set(ins, (void*) ctx);
+    if (ret == -1) {
+        flb_plg_error(ins, "unable to load configuration.");
+        flb_free(ctx);
+        return -1;
+    }
+
+    kafka_conf = flb_kafka_conf_create(&ctx->kafka, &ins->properties, 1);
+    if (!kafka_conf) {
+        flb_plg_error(ins, "Could not initialize kafka config object");
+        goto init_error;
+    }
+
+    ctx->kafka.rk = rd_kafka_new(RD_KAFKA_CONSUMER, kafka_conf, errstr,
+            sizeof(errstr));
+
+    /* Create Kafka consumer handle */
+    if (!ctx->kafka.rk) {
+        flb_plg_error(ins, "Failed to create new consumer: %s", errstr);
+        goto init_error;
+    }
+
+    conf = flb_input_get_property("topics", ins);
+    if (!conf) {
+        flb_plg_error(ins, "config: no topics specified");
+        goto init_error;
+    }
+
+    kafka_topics = flb_kafka_parse_topics(conf);
+    if (!kafka_topics) {
+        flb_plg_error(ins, "Failed to parse topic list");
+        goto init_error;
+    }
+
+    if (strcasecmp(ctx->format_str, "none") == 0) {
+        ctx->format = FLB_IN_KAFKA_FORMAT_NONE;
+    }
+    else if (strcasecmp(ctx->format_str, "json") == 0) {
+        ctx->format = FLB_IN_KAFKA_FORMAT_JSON;
+    }
+    else {
+        flb_plg_error(ins, "config: invalid format \"%s\"", ctx->format_str);
+        goto init_error;
+    }
+
+    if ((err = rd_kafka_subscribe(ctx->kafka.rk, kafka_topics))) {
+        flb_plg_error(ins, "Failed to start consuming topics: %s", rd_kafka_err2str(err));
+        goto init_error;
+    }
+    rd_kafka_topic_partition_list_destroy(kafka_topics);
+    kafka_topics = NULL;
+
+    /* Set the context */
+    flb_input_set_context(ins, ctx);
+    /* Collect upon data available on the pipe read fd */
+
+    int poll_seconds = ctx->poll_ms / 1000;
+    int poll_milliseconds = ctx->poll_ms % 1000;
+
+    ret = flb_input_set_collector_time(ins,
+                                       in_kafka_collect,
+                                       poll_seconds, poll_milliseconds * 1e6,
+                                       config);
+    if (ret) {
+        flb_plg_error(ctx->ins, "could not set collector for kafka input plugin");
+        goto init_error;
+    }
+
+    ctx->log_encoder = flb_log_event_encoder_create(FLB_LOG_EVENT_FORMAT_DEFAULT);
+
+    if (ctx->log_encoder == NULL) {
+        flb_plg_error(ins, "could not initialize log encoder");
+        goto init_error;
+    }
+
+    return 0;
+
+init_error:
+    if (kafka_topics) {
+        rd_kafka_topic_partition_list_destroy(kafka_topics);
+    }
+    if (ctx->kafka.rk) {
+        rd_kafka_destroy(ctx->kafka.rk);
+    }
+    else if (kafka_conf) {
+        /* conf is already destroyed when rd_kafka is initialized */
+        rd_kafka_conf_destroy(kafka_conf);
+    }
+    flb_free(ctx);
+
+    return -1;
+}
+
+/* Cleanup serial input */
+static int in_kafka_exit(void *in_context, struct flb_config *config)
+{
+    struct flb_in_kafka_config *ctx;
+
+    if (!in_context) {
+        return 0;
+    }
+
+    ctx = in_context;
+    rd_kafka_destroy(ctx->kafka.rk);
+    flb_free(ctx->kafka.brokers);
+
+    if (ctx->log_encoder){
+        flb_log_event_encoder_destroy(ctx->log_encoder);
+    }
+
+    flb_free(ctx);
+
+    return 0;
+}
+
+static struct flb_config_map config_map[] = {
+   {
+    FLB_CONFIG_MAP_INT, "poll_ms", FLB_IN_KAFKA_DEFAULT_POLL_MS,
+    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, poll_ms),
+    "Interval in milliseconds to check for new messages."
+   },
+   {
+    FLB_CONFIG_MAP_STR, "topics", (char *)NULL,
+    0, FLB_FALSE, 0,
+    "Set the kafka topics, delimited by commas."
+   },
+   {
+    FLB_CONFIG_MAP_STR, "format", FLB_IN_KAFKA_DEFAULT_FORMAT,
+    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, format_str),
+    "Set the data format which will be used for parsing records."
+   },
+   {
+    FLB_CONFIG_MAP_STR, "brokers", (char *)NULL,
+    0, FLB_FALSE, 0,
+    "Set the kafka brokers, delimited by commas."
+   },
+   {
+    FLB_CONFIG_MAP_STR, "client_id", (char *)NULL,
+    0, FLB_FALSE, 0,
+    "Set the kafka client_id."
+   },
+   {
+    FLB_CONFIG_MAP_STR, "group_id", (char *)NULL,
+    0, FLB_FALSE, 0,
+    "Set the kafka group_id."
+   },
+   {
+    FLB_CONFIG_MAP_STR_PREFIX, "rdkafka.", NULL,
+    /* FLB_CONFIG_MAP_MULT, FLB_TRUE, offsetof(struct flb_in_kafka_config, rdkafka_opts), */
+    0,  FLB_FALSE, 0,
+    "Set the librdkafka options"
+   },
+   /* EOF */
+   {0}
+};
+
+/* Plugin reference */
+struct flb_input_plugin in_kafka_plugin = {
+    .name         = "kafka",
+    .description  = "Kafka consumer input plugin",
+    .cb_init      = in_kafka_init,
+    .cb_pre_run   = NULL,
+    .cb_collect   = in_kafka_collect,
+    .cb_flush_buf = NULL,
+    .cb_exit      = in_kafka_exit,
+    .config_map   = config_map
+};