tests: internal: task_map route data preservation across retries

Signed-off-by: Eduardo Silva <eduardo@chronosphere.io>

Author: edsiper

tests/internal/task_map.c

@@ -5,13 +5,17 @@
 #include <fluent-bit/flb_error.h>
 #include <fluent-bit/flb_socket.h>
 #include <fluent-bit/flb_task.h>
+#include <fluent-bit/flb_output.h>
+#include <fluent-bit/flb_scheduler.h>
+#include <string.h>
 
 #include "flb_tests_internal.h"
 
 #define TASK_COUNT_LIMIT (FLB_CONFIG_DEFAULT_TASK_MAP_SIZE_LIMIT + 1)
 
 struct test_ctx {
     struct flb_config     *config;
+    struct mk_event_loop  *evl;
 };
 
 struct test_ctx* test_ctx_create()
@@ -32,6 +36,22 @@ struct test_ctx* test_ctx_create()
         return NULL;
     }
 
+    ret_ctx->evl = mk_event_loop_create(8);
+    if(!TEST_CHECK(ret_ctx->evl != NULL)) {
+        flb_config_exit(ret_ctx->config);
+        flb_free(ret_ctx);
+        return NULL;
+    }
+
+    ret_ctx->config->evl = ret_ctx->evl;
+    ret_ctx->config->sched = flb_sched_create(ret_ctx->config, ret_ctx->evl);
+    if(!TEST_CHECK(ret_ctx->config->sched != NULL)) {
+        mk_event_loop_destroy(ret_ctx->evl);
+        flb_config_exit(ret_ctx->config);
+        flb_free(ret_ctx);
+        return NULL;
+    }
+
     return ret_ctx;
 }
 
@@ -88,7 +108,91 @@ void test_task_map_limit()
     test_ctx_destroy(ctx);
 }
 
+void test_task_route_data_preserved_across_retry()
+{
+    int ret;
+    int records;
+    size_t bytes;
+    struct test_ctx *ctx;
+    struct flb_task *task;
+    struct flb_output_instance out_a;
+    struct flb_output_instance out_b;
+    struct flb_task_route *route_a;
+    struct flb_task_route *route_b;
+    struct flb_task_retry *retry;
+
+    ctx = test_ctx_create();
+    if (!TEST_CHECK(ctx != NULL)) {
+        return;
+    }
+
+    task = task_alloc(ctx->config);
+    if (!TEST_CHECK(task != NULL)) {
+        test_ctx_destroy(ctx);
+        return;
+    }
+
+    /* Avoid input chunk up/down side effects in retry creation. */
+    task->users = 2;
+
+    memset(&out_a, 0, sizeof(out_a));
+    memset(&out_b, 0, sizeof(out_b));
+    out_a.retry_limit = 5;
+    out_b.retry_limit = 5;
+
+    route_a = flb_calloc(1, sizeof(struct flb_task_route));
+    route_b = flb_calloc(1, sizeof(struct flb_task_route));
+    TEST_CHECK(route_a != NULL);
+    TEST_CHECK(route_b != NULL);
+    if (!route_a || !route_b) {
+        flb_free(route_a);
+        flb_free(route_b);
+        flb_task_destroy(task, FLB_TRUE);
+        test_ctx_destroy(ctx);
+        return;
+    }
+
+    route_a->status = FLB_TASK_ROUTE_ACTIVE;
+    route_a->out = &out_a;
+    route_a->records = 3;
+    route_a->bytes = 300;
+    route_b->status = FLB_TASK_ROUTE_ACTIVE;
+    route_b->out = &out_b;
+    route_b->records = 7;
+    route_b->bytes = 700;
+    mk_list_add(&route_a->_head, &task->routes);
+    mk_list_add(&route_b->_head, &task->routes);
+
+    flb_task_set_route_data(task, &out_a, 1, 111);
+    ret = flb_task_get_route_data(task, &out_a, &records, &bytes);
+    TEST_CHECK(ret == 0);
+    TEST_CHECK(records == 1);
+    TEST_CHECK(bytes == 111);
+
+    retry = flb_task_retry_create(task, &out_a);
+    TEST_CHECK(retry != NULL);
+    if (retry != NULL) {
+        TEST_CHECK(retry->attempts == 1);
+    }
+
+    retry = flb_task_retry_create(task, &out_a);
+    TEST_CHECK(retry != NULL);
+    if (retry != NULL) {
+        TEST_CHECK(retry->attempts == 2);
+    }
+
+    ret = flb_task_get_route_data(task, &out_a, &records, &bytes);
+    TEST_CHECK(ret == 0);
+    TEST_CHECK(records == 1);
+    TEST_CHECK(bytes == 111);
+
+    flb_task_retry_clean(task, &out_a);
+    flb_task_destroy(task, FLB_TRUE);
+    test_ctx_destroy(ctx);
+}
+
 TEST_LIST = {
     { "task_map_limit" , test_task_map_limit},
+    { "task_route_data_preserved_across_retry", test_task_route_data_preserved_across_retry},
     { 0 }
 };