src/write_kafka.c

   1 /**
   2  * collectd - src/write_kafka.c
   3  * Copyright (C) 2014       Pierre-Yves Ritschard
   4  *
   5  * Permission is hereby granted, free of charge, to any person obtaining a
   6  * copy of this software and associated documentation files (the "Software"),
   7  * to deal in the Software without restriction, including without limitation
   8  * the rights to use, copy, modify, merge, publish, distribute, sublicense,
   9  * and/or sell copies of the Software, and to permit persons to whom the
  10  * Software is furnished to do so, subject to the following conditions:
  11  *
  12  * The above copyright notice and this permission notice shall be included in
  13  * all copies or substantial portions of the Software.
  14  *
  15  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  18  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  20  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
  21  * DEALINGS IN THE SOFTWARE.
  22  *
  23  * Authors:
  24  *   Pierre-Yves Ritschard <pyr at spootnik.org>
  25  */
  26
  27 #include "collectd.h"
  28
  29 #include "plugin.h"
  30 #include "common.h"
  31 #include "configfile.h"
  32 #include "utils_cmd_putval.h"
  33 #include "utils_format_graphite.h"
  34 #include "utils_format_json.h"
  35
  36 #include <stdint.h>
  37 #include <librdkafka/rdkafka.h>
  38 #include <errno.h>
  39
  40 struct kafka_topic_context {
  41 #define KAFKA_FORMAT_JSON        0
  42 #define KAFKA_FORMAT_COMMAND     1
  43 #define KAFKA_FORMAT_GRAPHITE    2
  44     uint8_t                      format;
  45     unsigned int                 graphite_flags;
  46     _Bool                        store_rates;
  47     rd_kafka_topic_conf_t       *conf;
  48     rd_kafka_topic_t            *topic;
  49     rd_kafka_conf_t             *kafka_conf;
  50     rd_kafka_t                  *kafka;
  51     char                        *key;
  52     char                        *prefix;
  53     char                        *postfix;
  54     char                         escape_char;
  55     char                        *topic_name;
  56     pthread_mutex_t              lock;
  57 };
  58
  59 static int kafka_handle(struct kafka_topic_context *);
  60 static int kafka_write(const data_set_t *, const value_list_t *, user_data_t *);
  61 static int32_t kafka_partition(const rd_kafka_topic_t *, const void *, size_t,
  62                                int32_t, void *, void *);
  63
  64 #if defined HAVE_LIBRDKAFKA_LOGGER || defined HAVE_LIBRDKAFKA_LOG_CB
  65 static void kafka_log(const rd_kafka_t *, int, const char *, const char *);
  66
  67 static void kafka_log(const rd_kafka_t *rkt, int level,
  68                       const char *fac, const char *msg)
  69 {
  70     plugin_log(level, "%s", msg);
  71 }
  72 #endif
  73
  74 static uint32_t kafka_hash(const char *keydata, size_t keylen)
  75 {
  76     uint32_t hash = 5381;
  77     for (; keylen > 0; keylen--)
  78         hash = ((hash << 5) + hash) + keydata[keylen - 1];
  79     return hash;
  80 }
  81
  82 static int32_t kafka_partition(const rd_kafka_topic_t *rkt,
  83                                const void *keydata, size_t keylen,
  84                                int32_t partition_cnt, void *p, void *m)
  85 {
  86     uint32_t key = kafka_hash(keydata, keylen);
  87     uint32_t target = key % partition_cnt;
  88     int32_t  i = partition_cnt;
  89
  90     while (--i > 0 && !rd_kafka_topic_partition_available(rkt, target)) {
  91         target = (target + 1) % partition_cnt;
  92     }
  93     return target;
  94 }
  95
  96 static int kafka_handle(struct kafka_topic_context *ctx) /* {{{ */
  97 {
  98     char                         errbuf[1024];
  99     rd_kafka_conf_t             *conf;
 100     rd_kafka_topic_conf_t       *topic_conf;
 101
 102     if (ctx->kafka != NULL && ctx->topic != NULL)
 103         return(0);
 104
 105     if (ctx->kafka == NULL) {
 106         if ((conf = rd_kafka_conf_dup(ctx->kafka_conf)) == NULL) {
 107             ERROR("write_kafka plugin: cannot duplicate kafka config");
 108             return(1);
 109         }
 110
 111         if ((ctx->kafka = rd_kafka_new(RD_KAFKA_PRODUCER, conf,
 112                                     errbuf, sizeof(errbuf))) == NULL) {
 113             ERROR("write_kafka plugin: cannot create kafka handle.");
 114             return 1;
 115         }
 116
 117         rd_kafka_conf_destroy(ctx->kafka_conf);
 118         ctx->kafka_conf = NULL;
 119
 120         INFO ("write_kafka plugin: created KAFKA handle : %s", rd_kafka_name(ctx->kafka));
 121
 122 #if defined(HAVE_LIBRDKAFKA_LOGGER) && !defined(HAVE_LIBRDKAFKA_LOG_CB)
 123         rd_kafka_set_logger(ctx->kafka, kafka_log);
 124 #endif
 125     }
 126
 127     if (ctx->topic == NULL ) {
 128         if ((topic_conf = rd_kafka_topic_conf_dup(ctx->conf)) == NULL) {
 129             ERROR("write_kafka plugin: cannot duplicate kafka topic config");
 130             return 1;
 131         }
 132
 133         if ((ctx->topic = rd_kafka_topic_new(ctx->kafka, ctx->topic_name,
 134                                             topic_conf)) == NULL) {
 135             ERROR("write_kafka plugin: cannot create topic : %s\n",
 136             rd_kafka_err2str(rd_kafka_errno2err(errno)));
 137             return errno;
 138         }
 139
 140         rd_kafka_topic_conf_destroy(ctx->conf);
 141         ctx->conf = NULL;
 142
 143         INFO ("write_kafka plugin: handle created for topic : %s", rd_kafka_topic_name(ctx->topic));
 144     }
 145
 146     return(0);
 147
 148 } /* }}} int kafka_handle */
 149
 150 static int kafka_write(const data_set_t *ds, /* {{{ */
 151           const value_list_t *vl,
 152           user_data_t *ud)
 153 {
 154     int      status = 0;
 155     void    *key;
 156     size_t   keylen = 0;
 157     char     buffer[8192];
 158     size_t   bfree = sizeof(buffer);
 159     size_t   bfill = 0;
 160     size_t   blen = 0;
 161     struct   kafka_topic_context  *ctx = ud->data;
 162
 163     if ((ds == NULL) || (vl == NULL) || (ctx == NULL))
 164         return EINVAL;
 165
 166     pthread_mutex_lock (&ctx->lock);
 167     status = kafka_handle(ctx);
 168     pthread_mutex_unlock (&ctx->lock);
 169     if( status != 0 )
 170         return status;
 171
 172     bzero(buffer, sizeof(buffer));
 173
 174     switch (ctx->format) {
 175     case KAFKA_FORMAT_COMMAND:
 176         status = create_putval(buffer, sizeof(buffer), ds, vl);
 177         if (status != 0) {
 178             ERROR("write_kafka plugin: create_putval failed with status %i.",
 179                   status);
 180             return status;
 181         }
 182         blen = strlen(buffer);
 183         break;
 184     case KAFKA_FORMAT_JSON:
 185         format_json_initialize(buffer, &bfill, &bfree);
 186         format_json_value_list(buffer, &bfill, &bfree, ds, vl,
 187                                ctx->store_rates);
 188         format_json_finalize(buffer, &bfill, &bfree);
 189         blen = strlen(buffer);
 190         break;
 191     case KAFKA_FORMAT_GRAPHITE:
 192         status = format_graphite(buffer, sizeof(buffer), ds, vl,
 193                                  ctx->prefix, ctx->postfix, ctx->escape_char,
 194                                  ctx->graphite_flags);
 195         if (status != 0) {
 196             ERROR("write_kafka plugin: format_graphite failed with status %i.",
 197                   status);
 198             return status;
 199         }
 200         blen = strlen(buffer);
 201         break;
 202     default:
 203         ERROR("write_kafka plugin: invalid format %i.", ctx->format);
 204         return -1;
 205     }
 206
 207     key = ctx->key;
 208     if (key != NULL)
 209         keylen = strlen (key);
 210     else
 211         keylen = 0;
 212
 213     rd_kafka_produce(ctx->topic, RD_KAFKA_PARTITION_UA,
 214                      RD_KAFKA_MSG_F_COPY, buffer, blen,
 215                      key, keylen, NULL);
 216
 217     return status;
 218 } /* }}} int kafka_write */
 219
 220 static void kafka_topic_context_free(void *p) /* {{{ */
 221 {
 222     struct kafka_topic_context *ctx = p;
 223
 224     if (ctx == NULL)
 225         return;
 226
 227     if (ctx->topic_name != NULL)
 228         sfree(ctx->topic_name);
 229     if (ctx->topic != NULL)
 230         rd_kafka_topic_destroy(ctx->topic);
 231     if (ctx->conf != NULL)
 232         rd_kafka_topic_conf_destroy(ctx->conf);
 233     if (ctx->kafka_conf != NULL)
 234         rd_kafka_conf_destroy(ctx->kafka_conf);
 235     if (ctx->kafka != NULL)
 236         rd_kafka_destroy(ctx->kafka);
 237
 238     sfree(ctx);
 239 } /* }}} void kafka_topic_context_free */
 240
 241 static void kafka_config_topic(rd_kafka_conf_t *conf, oconfig_item_t *ci) /* {{{ */
 242 {
 243     int                          status;
 244     int                          i;
 245     struct kafka_topic_context  *tctx;
 246     char                        *key = NULL;
 247     char                        *val;
 248     char                         callback_name[DATA_MAX_NAME_LEN];
 249     char                         errbuf[1024];
 250     user_data_t                  ud;
 251     oconfig_item_t              *child;
 252     rd_kafka_conf_res_t          ret;
 253
 254     if ((tctx = calloc(1, sizeof (*tctx))) == NULL) {
 255         ERROR ("write_kafka plugin: calloc failed.");
 256         return;
 257     }
 258
 259     tctx->escape_char = '.';
 260     tctx->store_rates = 1;
 261     tctx->format = KAFKA_FORMAT_JSON;
 262     tctx->key = NULL;
 263
 264     if ((tctx->kafka_conf = rd_kafka_conf_dup(conf)) == NULL) {
 265         sfree(tctx);
 266         ERROR("write_kafka plugin: cannot allocate memory for kafka config");
 267         return;
 268     }
 269
 270 #ifdef HAVE_LIBRDKAFKA_LOG_CB
 271     rd_kafka_conf_set_log_cb(tctx->kafka_conf, kafka_log);
 272 #endif
 273
 274     if ((tctx->conf = rd_kafka_topic_conf_new()) == NULL) {
 275         rd_kafka_conf_destroy(tctx->kafka_conf);
 276         sfree(tctx);
 277         ERROR ("write_kafka plugin: cannot create topic configuration.");
 278         return;
 279     }
 280
 281     if (ci->values_num != 1) {
 282         WARNING("kafka topic name needed.");
 283         goto errout;
 284     }
 285
 286     if (ci->values[0].type != OCONFIG_TYPE_STRING) {
 287         WARNING("kafka topic needs a string argument.");
 288         goto errout;
 289     }
 290
 291     if ((tctx->topic_name = strdup(ci->values[0].value.string)) == NULL) {
 292         ERROR("write_kafka plugin: cannot copy topic name.");
 293         goto errout;
 294     }
 295
 296     for (i = 0; i < ci->children_num; i++) {
 297         /*
 298          * The code here could be simplified but makes room
 299          * for easy adding of new options later on.
 300          */
 301         child = &ci->children[i];
 302         status = 0;
 303
 304         if (strcasecmp ("Property", child->key) == 0) {
 305             if (child->values_num != 2) {
 306                 WARNING("kafka properties need both a key and a value.");
 307                 goto errout;
 308             }
 309             if (child->values[0].type != OCONFIG_TYPE_STRING ||
 310                 child->values[1].type != OCONFIG_TYPE_STRING) {
 311                 WARNING("kafka properties needs string arguments.");
 312                 goto errout;
 313             }
 314             key = child->values[0].value.string;
 315             val = child->values[1].value.string;
 316             ret = rd_kafka_topic_conf_set(tctx->conf,key, val,
 317                                           errbuf, sizeof(errbuf));
 318             if (ret != RD_KAFKA_CONF_OK) {
 319                 WARNING("cannot set kafka topic property %s to %s: %s.",
 320                         key, val, errbuf);
 321                 goto errout;
 322             }
 323
 324         } else if (strcasecmp ("Key", child->key) == 0)  {
 325             cf_util_get_string (child, &tctx->key);
 326             assert (tctx->key != NULL);
 327         } else if (strcasecmp ("Format", child->key) == 0) {
 328             status = cf_util_get_string(child, &key);
 329             if (status != 0)
 330                 goto errout;
 331
 332             assert(key != NULL);
 333
 334             if (strcasecmp(key, "Command") == 0) {
 335                 tctx->format = KAFKA_FORMAT_COMMAND;
 336
 337             } else if (strcasecmp(key, "Graphite") == 0) {
 338                 tctx->format = KAFKA_FORMAT_GRAPHITE;
 339
 340             } else if (strcasecmp(key, "Json") == 0) {
 341                 tctx->format = KAFKA_FORMAT_JSON;
 342
 343             } else {
 344                 WARNING ("write_kafka plugin: Invalid format string: %s",
 345                          key);
 346             }
 347
 348             sfree(key);
 349
 350         } else if (strcasecmp ("StoreRates", child->key) == 0) {
 351             status = cf_util_get_boolean (child, &tctx->store_rates);
 352             (void) cf_util_get_flag (child, &tctx->graphite_flags,
 353                                      GRAPHITE_STORE_RATES);
 354
 355         } else if (strcasecmp ("GraphiteSeparateInstances", child->key) == 0) {
 356             status = cf_util_get_flag (child, &tctx->graphite_flags,
 357                                        GRAPHITE_SEPARATE_INSTANCES);
 358
 359         } else if (strcasecmp ("GraphiteAlwaysAppendDS", child->key) == 0) {
 360             status = cf_util_get_flag (child, &tctx->graphite_flags,
 361                                        GRAPHITE_ALWAYS_APPEND_DS);
 362
 363         } else if (strcasecmp ("GraphitePrefix", child->key) == 0) {
 364             status = cf_util_get_string (child, &tctx->prefix);
 365         } else if (strcasecmp ("GraphitePostfix", child->key) == 0) {
 366             status = cf_util_get_string (child, &tctx->postfix);
 367         } else if (strcasecmp ("GraphiteEscapeChar", child->key) == 0) {
 368             char *tmp_buff = NULL;
 369             status = cf_util_get_string (child, &tmp_buff);
 370             if (strlen (tmp_buff) > 1)
 371                 WARNING ("write_kafka plugin: The option \"GraphiteEscapeChar\" handles "
 372                         "only one character. Others will be ignored.");
 373             tctx->escape_char = tmp_buff[0];
 374             sfree (tmp_buff);
 375         } else {
 376             WARNING ("write_kafka plugin: Invalid directive: %s.", child->key);
 377         }
 378
 379         if (status != 0)
 380             break;
 381     }
 382
 383     rd_kafka_topic_conf_set_partitioner_cb(tctx->conf, kafka_partition);
 384     rd_kafka_topic_conf_set_opaque(tctx->conf, tctx);
 385
 386     ssnprintf(callback_name, sizeof(callback_name),
 387               "write_kafka/%s", tctx->topic_name);
 388
 389     ud.data = tctx;
 390     ud.free_func = kafka_topic_context_free;
 391
 392     status = plugin_register_write (callback_name, kafka_write, &ud);
 393     if (status != 0) {
 394         WARNING ("write_kafka plugin: plugin_register_write (\"%s\") "
 395                 "failed with status %i.",
 396                 callback_name, status);
 397         goto errout;
 398     }
 399
 400     pthread_mutex_init (&tctx->lock, /* attr = */ NULL);
 401
 402     return;
 403  errout:
 404     if (tctx->topic_name != NULL)
 405         free(tctx->topic_name);
 406     if (tctx->conf != NULL)
 407         rd_kafka_topic_conf_destroy(tctx->conf);
 408     if (tctx->kafka_conf != NULL)
 409         rd_kafka_conf_destroy(tctx->kafka_conf);
 410     sfree(tctx);
 411 } /* }}} int kafka_config_topic */
 412
 413 static int kafka_config(oconfig_item_t *ci) /* {{{ */
 414 {
 415     int                          i;
 416     oconfig_item_t              *child;
 417     rd_kafka_conf_t             *conf;
 418     rd_kafka_conf_res_t          ret;
 419     char                         errbuf[1024];
 420
 421     if ((conf = rd_kafka_conf_new()) == NULL) {
 422         WARNING("cannot allocate kafka configuration.");
 423         return -1;
 424     }
 425     for (i = 0; i < ci->children_num; i++)  {
 426         child = &ci->children[i];
 427
 428         if (strcasecmp("Topic", child->key) == 0) {
 429             kafka_config_topic (conf, child);
 430         } else if (strcasecmp(child->key, "Property") == 0) {
 431             char *key = NULL;
 432             char *val = NULL;
 433
 434             if (child->values_num != 2) {
 435                 WARNING("kafka properties need both a key and a value.");
 436                 goto errout;
 437             }
 438             if (child->values[0].type != OCONFIG_TYPE_STRING ||
 439                 child->values[1].type != OCONFIG_TYPE_STRING) {
 440                 WARNING("kafka properties needs string arguments.");
 441                 goto errout;
 442             }
 443             if ((key = strdup(child->values[0].value.string)) == NULL) {
 444                 WARNING("cannot allocate memory for attribute key.");
 445                 goto errout;
 446             }
 447             if ((val = strdup(child->values[1].value.string)) == NULL) {
 448                 WARNING("cannot allocate memory for attribute value.");
 449                 sfree(key);
 450                 goto errout;
 451             }
 452             ret = rd_kafka_conf_set(conf, key, val, errbuf, sizeof(errbuf));
 453             if (ret != RD_KAFKA_CONF_OK) {
 454                 WARNING("cannot set kafka property %s to %s: %s",
 455                         key, val, errbuf);
 456                 sfree(key);
 457                 sfree(val);
 458                 goto errout;
 459             }
 460             sfree(key);
 461             sfree(val);
 462         } else {
 463             WARNING ("write_kafka plugin: Ignoring unknown "
 464                  "configuration option \"%s\" at top level.",
 465                  child->key);
 466         }
 467     }
 468     if (conf != NULL)
 469         rd_kafka_conf_destroy(conf);
 470     return (0);
 471  errout:
 472     if (conf != NULL)
 473         rd_kafka_conf_destroy(conf);
 474     return -1;
 475 } /* }}} int kafka_config */
 476
 477 void module_register(void)
 478 {
 479     plugin_register_complex_config ("write_kafka", kafka_config);
 480 }