From 377ae9300ff7cb4917fdbbc09b0b3c153df67a71 Mon Sep 17 00:00:00 2001
From: Eduardo Silva <eduardo@treasure-data.com>
Date: Wed, 19 Oct 2016 08:49:57 -0600
Subject: [PATCH] core: use new memory handler wrappers

Signed-off-by: Eduardo Silva <eduardo@treasure-data.com>
---
 include/fluent-bit/flb_mem.h | 18 ++++++++++++------
 src/flb_buffer.c             |  9 +++++----
 src/flb_buffer_chunk.c       |  7 ++++---
 src/flb_buffer_qchunk.c      |  5 +++--
 src/flb_config.c             |  3 ++-
 src/flb_http_client.c        | 10 ++++++----
 src/flb_input.c              | 16 +++++++++-------
 src/flb_io_tls.c             |  6 ++++--
 src/flb_kernel.c             |  6 ++++--
 src/flb_lib.c                |  5 +++--
 src/flb_log.c                |  5 +++--
 src/flb_network.c            |  3 ++-
 src/flb_output.c             |  5 +++--
 src/flb_pack.c               |  8 +++++---
 src/flb_router.c             |  4 +++-
 src/flb_scheduler.c          |  3 ++-
 src/flb_stats.c              | 12 ++++++------
 src/flb_task.c               | 11 ++++++-----
 src/flb_upstream.c           |  6 ++++--
 src/flb_uri.c                |  5 ++++-
 src/flb_worker.c             |  3 ++-
 21 files changed, 92 insertions(+), 58 deletions(-)

diff --git a/include/fluent-bit/flb_mem.h b/include/fluent-bit/flb_mem.h
index 1c902b313b9..5a1c1267511 100644
--- a/include/fluent-bit/flb_mem.h
+++ b/include/fluent-bit/flb_mem.h
@@ -22,7 +22,16 @@
 
 #include <fluent-bit/flb_info.h>
 #include <fluent-bit/flb_macros.h>
-#include <fluent-bit/flb_log.h>
+
+#include <stdlib.h>
+
+/*
+ * The following memory handling wrappers, aims to simplify the way to use
+ * the default memory allocator from the libc or an alternative one as Jemalloc.
+ *
+ * Here there is no error logging in case of failures, we defer that task to the
+ * caller.
+ */
 
 #if ((__GNUC__ * 100 + __GNUC__MINOR__) > 430)  /* gcc version > 4.3 */
 # define ALLOCSZ_ATTR(x,...) __attribute__ ((alloc_size(x, ##__VA_ARGS__)))
@@ -36,7 +45,6 @@ void *flb_malloc(const size_t size) {
 
     aux = malloc(size);
     if (flb_unlikely(!aux && size)) {
-        flb_errno();
         return NULL;
     }
 
@@ -48,8 +56,7 @@ void *flb_calloc(size_t n, const size_t size) {
     void *buf;
 
     buf = calloc(n, size);
-    if (mk_unlikely(!buf)) {
-        flb_errno();
+    if (flb_unlikely(!buf)) {
         return NULL;
     }
 
@@ -62,8 +69,7 @@ void *flb_realloc(void *ptr, const size_t size)
     void *aux;
 
     aux = realloc(ptr, size);
-    if (mk_unlikely(!aux && size)) {
-        flb_errno();
+    if (flb_unlikely(!aux && size)) {
         return NULL;
     }
 
diff --git a/src/flb_buffer.c b/src/flb_buffer.c
index 890388f1bbb..d981c7935e6 100644
--- a/src/flb_buffer.c
+++ b/src/flb_buffer.c
@@ -29,6 +29,7 @@
 #include <inttypes.h>
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 
 #ifdef FLB_HAVE_BUFFERING
 
@@ -362,7 +363,7 @@ struct flb_buffer *flb_buffer_create(char *path, int workers,
     }
 
     /* Main buffer context */
-    ctx = malloc(sizeof(struct flb_buffer));
+    ctx = flb_malloc(sizeof(struct flb_buffer));
     if (!ctx) {
         return NULL;
     }
@@ -371,7 +372,7 @@ struct flb_buffer *flb_buffer_create(char *path, int workers,
 
     path_len = strlen(path);
     if (path[path_len - 1] != '/') {
-        ctx->path = malloc(path_len + 2);
+        ctx->path = flb_malloc(path_len + 2);
         memcpy(ctx->path, path, path_len);
         ctx->path[path_len++] = '/';
         ctx->path[path_len++] = '\0';
@@ -391,7 +392,7 @@ struct flb_buffer *flb_buffer_create(char *path, int workers,
 
     for (i = 0; i < ctx->workers_n; i++) {
         /* Allocate worker context */
-        worker = calloc(1, sizeof(struct flb_buffer_worker));
+        worker = flb_calloc(1, sizeof(struct flb_buffer_worker));
         if (!worker) {
             flb_buffer_destroy(ctx);
             return NULL;
@@ -450,7 +451,7 @@ struct flb_buffer *flb_buffer_create(char *path, int workers,
     ctx->workers_n = i;
 
     /* Generate pseudo input plugin and instance */
-    ctx->i_ins = calloc(1, sizeof(struct flb_input_instance));
+    ctx->i_ins = flb_calloc(1, sizeof(struct flb_input_instance));
     if (!ctx->i_ins) {
         flb_errno();
         flb_buffer_destroy(ctx);
diff --git a/src/flb_buffer_chunk.c b/src/flb_buffer_chunk.c
index b92c152a19e..7b857f66b85 100644
--- a/src/flb_buffer_chunk.c
+++ b/src/flb_buffer_chunk.c
@@ -18,6 +18,7 @@
  */
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 
 #ifdef FLB_HAVE_BUFFERING
 
@@ -218,7 +219,7 @@ static int chunk_find(char *root_path, char *hash,
         return -1;
     }
 
-    target = malloc(PATH_MAX);
+    target = flb_malloc(PATH_MAX);
     if (!target) {
         free(file);
         return -1;
@@ -267,7 +268,7 @@ static int chunk_remove_route(char *root_path, char *abs_path,
         return -1;
     }
 
-    to = malloc(PATH_MAX);
+    to = flb_malloc(PATH_MAX);
     if (!to) {
         flb_errno();
         return -1;
@@ -385,7 +386,7 @@ int flb_buffer_chunk_add(struct flb_buffer_worker *worker,
      *
      *     SHA1(chunk.data).routes_id.wID.tag
      */
-    fchunk = malloc(PATH_MAX);
+    fchunk = flb_malloc(PATH_MAX);
     if (!fchunk) {
         flb_errno();
         return -1;
diff --git a/src/flb_buffer_qchunk.c b/src/flb_buffer_qchunk.c
index c39e7405c07..746c7b9d9c2 100644
--- a/src/flb_buffer_qchunk.c
+++ b/src/flb_buffer_qchunk.c
@@ -18,6 +18,7 @@
  */
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_buffer.h>
 #include <fluent-bit/flb_buffer_qchunk.h>
 #include <fluent-bit/flb_engine_dispatch.h>
@@ -57,7 +58,7 @@ struct flb_buffer_qchunk *flb_buffer_qchunk_add(struct flb_buffer_qworker *qw,
     int len;
     struct flb_buffer_qchunk *qchunk;
 
-    qchunk = malloc(sizeof(struct flb_buffer_qchunk));
+    qchunk = flb_malloc(sizeof(struct flb_buffer_qchunk));
     if (!qchunk) {
         perror("malloc");
         return NULL;
@@ -343,7 +344,7 @@ int flb_buffer_qchunk_create(struct flb_buffer *ctx)
     struct flb_buffer_qworker *qw;
 
     /* Allocate context */
-    qw = malloc(sizeof(struct flb_buffer_qworker));
+    qw = flb_malloc(sizeof(struct flb_buffer_qworker));
     if (!qw) {
         perror("malloc");
         return -1;
diff --git a/src/flb_config.c b/src/flb_config.c
index 446cbb1d8d3..8b6cb79a595 100644
--- a/src/flb_config.c
+++ b/src/flb_config.c
@@ -24,6 +24,7 @@
 #include <stddef.h>
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_macros.h>
 #include <fluent-bit/flb_config.h>
 #include <fluent-bit/flb_plugins.h>
@@ -77,7 +78,7 @@ struct flb_config *flb_config_init()
 {
     struct flb_config *config;
 
-    config = calloc(1, sizeof(struct flb_config));
+    config = flb_calloc(1, sizeof(struct flb_config));
     if (!config) {
         perror("malloc");
         return NULL;
diff --git a/src/flb_http_client.c b/src/flb_http_client.c
index 5227530d129..5261bf035d6 100644
--- a/src/flb_http_client.c
+++ b/src/flb_http_client.c
@@ -29,6 +29,8 @@
  * - Get return Status, Headers and Body content if found.
  */
 
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_http_client.h>
 
 /* check if there is enough space in the client header buffer */
@@ -156,7 +158,7 @@ struct flb_http_client *flb_http_client(struct flb_upstream_conn *u_conn,
         break;
     };
 
-    buf = calloc(1, FLB_HTTP_BUF_SIZE);
+    buf = flb_calloc(1, FLB_HTTP_BUF_SIZE);
     if (!buf) {
         perror("malloc");
         return NULL;
@@ -190,7 +192,7 @@ struct flb_http_client *flb_http_client(struct flb_upstream_conn *u_conn,
         return NULL;
     }
 
-    c = calloc(1, sizeof(struct flb_http_client));
+    c = flb_calloc(1, sizeof(struct flb_http_client));
     if (!c) {
         free(buf);
         return NULL;
@@ -246,7 +248,7 @@ int flb_http_add_header(struct flb_http_client *c,
         else {
             new_size = c->header_size + required;
         }
-        tmp = realloc(c->header_buf, new_size);
+        tmp = flb_realloc(c->header_buf, new_size);
         if (!tmp) {
             perror("realloc");
             return -1;
@@ -290,7 +292,7 @@ int flb_http_do(struct flb_http_client *c, size_t *bytes)
     /* check enough space for the ending CRLF */
     if (header_available(c, crlf) != 0) {
         new_size = c->header_size + 2;
-        tmp = realloc(c->header_buf, new_size);
+        tmp = flb_realloc(c->header_buf, new_size);
         if (!tmp) {
             return -1;
         }
diff --git a/src/flb_input.c b/src/flb_input.c
index 9441718e9ce..d0584beedee 100644
--- a/src/flb_input.c
+++ b/src/flb_input.c
@@ -21,6 +21,8 @@
 #include <string.h>
 #include <unistd.h>
 
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_macros.h>
 #include <fluent-bit/flb_input.h>
 #include <fluent-bit/flb_error.h>
@@ -99,7 +101,7 @@ struct flb_input_instance *flb_input_new(struct flb_config *config,
         }
 
         /* Create plugin instance */
-        instance = malloc(sizeof(struct flb_input_instance));
+        instance = flb_malloc(sizeof(struct flb_input_instance));
         if (!instance) {
             perror("malloc");
             return NULL;
@@ -179,7 +181,7 @@ int flb_input_set_property(struct flb_input_instance *in, char *k, char *v)
     }
     else {
         /* Append any remaining configuration key to prop list */
-        prop = malloc(sizeof(struct flb_config_prop));
+        prop = flb_malloc(sizeof(struct flb_config_prop));
         if (!prop) {
             return -1;
         }
@@ -377,7 +379,7 @@ int flb_input_set_collector_time(struct flb_input_instance *in,
 {
     struct flb_input_collector *collector;
 
-    collector = malloc(sizeof(struct flb_input_collector));
+    collector = flb_malloc(sizeof(struct flb_input_collector));
     collector->type        = FLB_COLLECT_TIME;
     collector->cb_collect  = cb_collect;
     collector->fd_event    = -1;
@@ -397,7 +399,7 @@ int flb_input_set_collector_event(struct flb_input_instance *in,
 {
     struct flb_input_collector *collector;
 
-    collector = malloc(sizeof(struct flb_input_collector));
+    collector = flb_malloc(sizeof(struct flb_input_collector));
     collector->type        = FLB_COLLECT_FD_EVENT;
     collector->cb_collect  = cb_collect;
     collector->fd_event    = fd;
@@ -417,7 +419,7 @@ int flb_input_set_collector_socket(struct flb_input_instance *in,
 {
     struct flb_input_collector *collector;
 
-    collector = malloc(sizeof(struct flb_input_collector));
+    collector = flb_malloc(sizeof(struct flb_input_collector));
     collector->type        = FLB_COLLECT_FD_SERVER;
     collector->cb_collect  = cb_new_connection;
     collector->fd_event    = fd;
@@ -441,14 +443,14 @@ struct flb_input_dyntag *flb_input_dyntag_create(struct flb_input_instance *in,
     }
 
     /* Allocate node and reset fields */
-    dt = malloc(sizeof(struct flb_input_dyntag));
+    dt = flb_malloc(sizeof(struct flb_input_dyntag));
     if (!dt) {
         return NULL;
     }
     dt->busy = FLB_FALSE;
     dt->lock = FLB_FALSE;
     dt->in   = in;
-    dt->tag  = malloc(tag_len + 1);
+    dt->tag  = flb_malloc(tag_len + 1);
     memcpy(dt->tag, tag, tag_len);
     dt->tag[tag_len] = '\0';
     dt->tag_len = tag_len;
diff --git a/src/flb_io_tls.c b/src/flb_io_tls.c
index db012623659..886d46f41db 100644
--- a/src/flb_io_tls.c
+++ b/src/flb_io_tls.c
@@ -25,6 +25,8 @@
 #include <mbedtls/debug.h>
 #include <mbedtls/error.h>
 
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_io.h>
 #include <fluent-bit/flb_tls.h>
 #include <fluent-bit/flb_io_tls.h>
@@ -75,7 +77,7 @@ struct flb_tls_context *flb_tls_context_new(int verify,
     int ret;
     struct flb_tls_context *ctx;
 
-    ctx = malloc(sizeof(struct flb_tls_context));
+    ctx = flb_malloc(sizeof(struct flb_tls_context));
     if (!ctx) {
         perror("malloc");
         return NULL;
@@ -157,7 +159,7 @@ struct flb_tls_session *flb_tls_session_new(struct flb_tls_context *ctx)
     int ret;
     struct flb_tls_session *session;
 
-    session = malloc(sizeof(struct flb_tls_session));
+    session = flb_malloc(sizeof(struct flb_tls_session));
     if (!session) {
         return NULL;
     }
diff --git a/src/flb_kernel.c b/src/flb_kernel.c
index a924bc8031b..9c70972ddc1 100644
--- a/src/flb_kernel.c
+++ b/src/flb_kernel.c
@@ -20,6 +20,8 @@
 #include <ctype.h>
 #include <sys/utsname.h>
 
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_kernel.h>
 #include <fluent-bit/flb_utils.h>
 
@@ -77,14 +79,14 @@ struct flb_kernel *flb_kernel_info()
     c = atoi(tmp);
     free(tmp);
 
-    kernel = malloc(sizeof(struct flb_kernel));
+    kernel = flb_malloc(sizeof(struct flb_kernel));
     if (!kernel) {
         return NULL;
     }
     kernel->minor = a;
     kernel->major = b;
     kernel->patch = c;
-    kernel->s_version.data = malloc(16);
+    kernel->s_version.data = flb_malloc(16);
 
     len = snprintf(kernel->s_version.data, 16, "%i.%i.%i", a, b, c);
     if (len == -1) {
diff --git a/src/flb_lib.c b/src/flb_lib.c
index 1e18a53fb3f..a14ceecd228 100644
--- a/src/flb_lib.c
+++ b/src/flb_lib.c
@@ -22,6 +22,7 @@
 #include <stdarg.h>
 
 #include <fluent-bit/flb_lib.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_engine.h>
 #include <fluent-bit/flb_input.h>
 #include <fluent-bit/flb_output.h>
@@ -82,7 +83,7 @@ flb_ctx_t *flb_create()
     mtrace();
 #endif
 
-    ctx = calloc(1, sizeof(flb_ctx_t));
+    ctx = flb_calloc(1, sizeof(flb_ctx_t));
     if (!ctx) {
         perror("malloc");
         return NULL;
@@ -115,7 +116,7 @@ flb_ctx_t *flb_create()
     config->ch_evl = ctx->event_loop;
 
     /* Prepare the notification channels */
-    ctx->event_channel = calloc(1, sizeof(struct mk_event));
+    ctx->event_channel = flb_calloc(1, sizeof(struct mk_event));
     ret = mk_event_channel_create(config->ch_evl,
                                   &config->ch_notif[0],
                                   &config->ch_notif[1],
diff --git a/src/flb_log.c b/src/flb_log.c
index 8b6a2ca89ad..2fad3837a6e 100644
--- a/src/flb_log.c
+++ b/src/flb_log.c
@@ -33,6 +33,7 @@
 #include <fluent-bit/flb_log.h>
 #include <fluent-bit/flb_config.h>
 #include <fluent-bit/flb_worker.h>
+#include <fluent-bit/flb_mem.h>
 
 FLB_TLS_DEFINE(struct flb_log, flb_log_ctx)
 
@@ -170,7 +171,7 @@ struct flb_log *flb_log_init(struct flb_config *config, int type,
     struct flb_worker *worker;
     struct mk_event_loop *evl;
 
-    log = malloc(sizeof(struct flb_log));
+    log = flb_malloc(sizeof(struct flb_log));
     if (!log) {
         perror("malloc");
         return NULL;
@@ -216,7 +217,7 @@ struct flb_log *flb_log_init(struct flb_config *config, int type,
      * it will need a 'worker-like' context, here we create a fake worker
      * context just for messaging purposes.
      */
-    worker = malloc(sizeof(struct flb_worker));
+    worker = flb_malloc(sizeof(struct flb_worker));
     if (!worker) {
         flb_errno();
         free(log);
diff --git a/src/flb_network.c b/src/flb_network.c
index a7a32bd191d..104dbeaea85 100644
--- a/src/flb_network.c
+++ b/src/flb_network.c
@@ -34,6 +34,7 @@
 #include <arpa/inet.h>
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_network.h>
 #include <fluent-bit/flb_utils.h>
 #include <fluent-bit/flb_macros.h>
@@ -47,7 +48,7 @@ static char *copy_substr(char *str, int s)
 {
     char *buf;
 
-    buf = malloc(s + 1);
+    buf = flb_malloc(s + 1);
     strncpy(buf, str, s);
     buf[s] = '\0';
 
diff --git a/src/flb_output.c b/src/flb_output.c
index edbadb9db9f..01a9863eab7 100644
--- a/src/flb_output.c
+++ b/src/flb_output.c
@@ -22,6 +22,7 @@
 #include <string.h>
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_thread.h>
 #include <fluent-bit/flb_output.h>
 
@@ -174,7 +175,7 @@ struct flb_output_instance *flb_output_new(struct flb_config *config,
         }
 
         /* Output instance */
-        instance = calloc(1, sizeof(struct flb_output_instance));
+        instance = flb_calloc(1, sizeof(struct flb_output_instance));
         if (!instance) {
             perror("malloc");
             return NULL;
@@ -295,7 +296,7 @@ int flb_output_set_property(struct flb_output_instance *out, char *k, char *v)
 #endif
     else {
         /* Append any remaining configuration key to prop list */
-        prop = malloc(sizeof(struct flb_config_prop));
+        prop = flb_malloc(sizeof(struct flb_config_prop));
         if (!prop) {
             return -1;
         }
diff --git a/src/flb_pack.c b/src/flb_pack.c
index d4fbc368ace..b6cc22fa081 100644
--- a/src/flb_pack.c
+++ b/src/flb_pack.c
@@ -20,6 +20,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_error.h>
 #include <fluent-bit/flb_utils.h>
 #include <fluent-bit/flb_pack.h>
@@ -40,7 +42,7 @@ static int json_tokenise(char *js, size_t len,
 
     while (ret == JSMN_ERROR_NOMEM) {
         n = state->tokens_size += 256;
-        tmp = realloc(state->tokens, sizeof(jsmntok_t) * n);
+        tmp = flb_realloc(state->tokens, sizeof(jsmntok_t) * n);
         if (!tmp) {
             perror("realloc");
             return -1;
@@ -142,7 +144,7 @@ static char *tokens_to_msgpack(char *js,
 
     /* dump data back to a new buffer */
     *out_size = sbuf.size;
-    buf = malloc(sbuf.size);
+    buf = flb_malloc(sbuf.size);
     memcpy(buf, sbuf.data, sbuf.size);
     msgpack_sbuffer_destroy(&sbuf);
 
@@ -194,7 +196,7 @@ int flb_pack_state_init(struct flb_pack_state *s)
     int size = 256;
 
     jsmn_init(&s->parser);
-    s->tokens = calloc(1, sizeof(jsmntok_t) * size);
+    s->tokens = flb_calloc(1, sizeof(jsmntok_t) * size);
     if (!s->tokens) {
         perror("calloc");
         return -1;
diff --git a/src/flb_router.c b/src/flb_router.c
index 1e9bde414ab..c04c6994821 100644
--- a/src/flb_router.c
+++ b/src/flb_router.c
@@ -17,6 +17,8 @@
  *  limitations under the License.
  */
 
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_input.h>
 #include <fluent-bit/flb_output.h>
 #include <fluent-bit/flb_config.h>
@@ -74,7 +76,7 @@ static int flb_router_connect(struct flb_input_instance *in,
 {
     struct flb_router_path *p;
 
-    p = malloc(sizeof(struct flb_router_path));
+    p = flb_malloc(sizeof(struct flb_router_path));
     if (!p) {
         perror("malloc");
         return -1;
diff --git a/src/flb_scheduler.c b/src/flb_scheduler.c
index d94a5668bea..e7887a0e1b6 100644
--- a/src/flb_scheduler.c
+++ b/src/flb_scheduler.c
@@ -18,6 +18,7 @@
  */
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_thread.h>
 #include <fluent-bit/flb_config.h>
 #include <fluent-bit/flb_scheduler.h>
@@ -114,7 +115,7 @@ int flb_sched_request_create(struct flb_config *config,
     struct flb_sched_request *request;
 
     /* Allocate request node */
-    request = malloc(sizeof(struct flb_sched_request));
+    request = flb_malloc(sizeof(struct flb_sched_request));
     if (!request) {
         perror("malloc");
         return -1;
diff --git a/src/flb_stats.c b/src/flb_stats.c
index d44683ac4b4..7474b6d8b15 100644
--- a/src/flb_stats.c
+++ b/src/flb_stats.c
@@ -66,7 +66,7 @@ static int stats_userver_timer(struct flb_stats *stats)
     struct flb_stats_userver *userver;
 
     userver = stats->userver;
-    timer = malloc(sizeof(struct flb_stats_userver_t));
+    timer = flb_malloc(sizeof(struct flb_stats_userver_t));
     if (!timer) {
         return -1;
     }
@@ -204,7 +204,7 @@ static int stats_userver_add(int fd, struct flb_stats *stats)
     struct flb_stats_userver *userver = stats->userver;
 
     /* Allocate connection node */
-    client = calloc(1, sizeof(struct flb_stats_userver_c));
+    client = flb_calloc(1, sizeof(struct flb_stats_userver_c));
     if (!client) {
         return -1;
     }
@@ -342,7 +342,7 @@ static int flb_stats_userver(struct flb_stats *stats)
     struct flb_stats_userver *userver;
 
     /* Create a TCP server based on unix sockets */
-    userver = malloc(sizeof(struct flb_stats_userver));
+    userver = flb_malloc(sizeof(struct flb_stats_userver));
     if (!userver) {
         flb_error("[stats_usrv] no mem!");
         return -1;
@@ -458,7 +458,7 @@ static int register_input_plugin(struct flb_input_plugin *plugin,
     struct flb_stats_in_plugin *sp;
 
     /* Allocate stats object for this stats entry */
-    sp = malloc(sizeof(struct flb_stats_in_plugin));
+    sp = flb_malloc(sizeof(struct flb_stats_in_plugin));
     if (!sp) {
         return -1;
     }
@@ -504,7 +504,7 @@ static int register_output_plugin(struct flb_output_plugin *plugin,
     struct flb_stats_out_plugin *sp;
 
     /* Allocate stats object for this stats entry */
-    sp = malloc(sizeof(struct flb_stats_out_plugin));
+    sp = flb_malloc(sizeof(struct flb_stats_out_plugin));
     if (!sp) {
         return -1;
     }
@@ -582,7 +582,7 @@ int flb_stats_init(struct flb_config *config)
     int ret;
     struct flb_stats *stats;
 
-    stats = malloc(sizeof(struct flb_stats));
+    stats = flb_malloc(sizeof(struct flb_stats));
     if (!stats) {
         flb_error("[stats] could not initialize");
         return -1;
diff --git a/src/flb_task.c b/src/flb_task.c
index fb20eee02a3..2ba27f75f1e 100644
--- a/src/flb_task.c
+++ b/src/flb_task.c
@@ -26,6 +26,7 @@
 #include <fluent-bit/flb_output.h>
 #include <fluent-bit/flb_router.h>
 #include <fluent-bit/flb_task.h>
+#include <fluent-bit/flb_mem.h>
 
 #ifdef FLB_HAVE_BUFFERING
 #include <fluent-bit/flb_sha1.h>
@@ -93,7 +94,7 @@ struct flb_task_retry *flb_task_retry_create(struct flb_task *task,
 
     if (!retry) {
         /* Create a new re-try instance */
-        retry = malloc(sizeof(struct flb_task_retry));
+        retry = flb_malloc(sizeof(struct flb_task_retry));
         if (!retry) {
             perror("malloc");
             return NULL;
@@ -140,7 +141,7 @@ static struct flb_task *task_alloc(struct flb_config *config)
     struct flb_task *task;
 
     /* Allocate the new task */
-    task = (struct flb_task *) calloc(1, sizeof(struct flb_task));
+    task = (struct flb_task *) flb_calloc(1, sizeof(struct flb_task));
     if (!task) {
         perror("malloc");
         return NULL;
@@ -208,7 +209,7 @@ struct flb_task *flb_task_create(uint64_t ref_id,
             router_path = mk_list_entry(head, struct flb_router_path, _head);
             o_ins = router_path->ins;
 
-            route = malloc(sizeof(struct flb_task_route));
+            route = flb_malloc(sizeof(struct flb_task_route));
             if (!route) {
                 perror("malloc");
                 continue;
@@ -228,7 +229,7 @@ struct flb_task *flb_task_create(uint64_t ref_id,
                                   struct flb_output_instance, _head);
 
             if (flb_router_match(tag, o_ins->match)) {
-                route = malloc(sizeof(struct flb_task_route));
+                route = flb_malloc(sizeof(struct flb_task_route));
                 if (!route) {
                     perror("malloc");
                     continue;
@@ -326,7 +327,7 @@ struct flb_task *flb_task_create_direct(uint64_t ref_id,
     mk_list_foreach(head, &config->outputs) {
         o_ins = mk_list_entry(head, struct flb_output_instance, _head);
         if (o_ins->mask_id & routes) {
-            route = malloc(sizeof(struct flb_task_route));
+            route = flb_malloc(sizeof(struct flb_task_route));
             if (!route) {
                 perror("malloc");
                 continue;
diff --git a/src/flb_upstream.c b/src/flb_upstream.c
index 7902ab1969a..5213cc7d8fa 100644
--- a/src/flb_upstream.c
+++ b/src/flb_upstream.c
@@ -20,6 +20,8 @@
 #include <unistd.h>
 
 #include <mk_core.h>
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_upstream.h>
 #include <fluent-bit/flb_io.h>
 #include <fluent-bit/flb_io_tls.h>
@@ -32,7 +34,7 @@ struct flb_upstream *flb_upstream_create(struct flb_config *config,
 {
     struct flb_upstream *u;
 
-    u = calloc(1, sizeof(struct flb_upstream));
+    u = flb_calloc(1, sizeof(struct flb_upstream));
     if (!u) {
         perror("malloc");
         return NULL;
@@ -103,7 +105,7 @@ static struct flb_upstream_conn *create_conn(struct flb_upstream *u)
     void *th = NULL;
 #endif
 
-    conn = malloc(sizeof(struct flb_upstream_conn));
+    conn = flb_malloc(sizeof(struct flb_upstream_conn));
     if (!conn) {
         perror("malloc");
         return NULL;
diff --git a/src/flb_uri.c b/src/flb_uri.c
index 8e3957af93c..1053b9bb866 100644
--- a/src/flb_uri.c
+++ b/src/flb_uri.c
@@ -19,6 +19,9 @@
 
 #include <stdlib.h>
 #include <mk_core.h>
+
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_uri.h>
 #include <fluent-bit/flb_utils.h>
 
@@ -53,7 +56,7 @@ struct flb_uri *flb_uri_create(char *full_uri)
     uri_size  = sizeof(struct flb_uri);
     uri_size += (sizeof(struct flb_uri_field) * FLB_URI_MAX);
 
-    p  = calloc(1, uri_size);
+    p  = flb_calloc(1, uri_size);
     if (!p) {
         perror("malloc");
         return NULL;
diff --git a/src/flb_worker.c b/src/flb_worker.c
index 9b4f09ac6f3..5bb4056351a 100644
--- a/src/flb_worker.c
+++ b/src/flb_worker.c
@@ -18,6 +18,7 @@
  */
 
 #include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_utils.h>
 #include <fluent-bit/flb_config.h>
 #include <fluent-bit/flb_worker.h>
@@ -68,7 +69,7 @@ int flb_worker_create(void (*func) (void *), void *arg, pthread_t *tid,
     int ret;
     struct flb_worker *worker;
 
-    worker = malloc(sizeof(struct flb_worker));
+    worker = flb_malloc(sizeof(struct flb_worker));
     if (!worker) {
         perror("malloc");
         return -1;