/* -*- Mode: C; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*  Fluent Bit

 *  ==========

 *  Copyright (C) 2015-2022 The Fluent Bit Authors

 *

 *  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_output.h>

#include <fluent-bit/flb_input_plugin.h>

#include <fluent-bit/flb_input_chunk.h>

#include <fluent-bit/flb_storage.h>

#include <fluent-bit/flb_utils.h>

#include <chunkio/chunkio.h>

#include <chunkio/cio_error.h>

#include <sys/types.h>

#include <sys/stat.h>

#ifndef FLB_SYSTEM_WINDOWS

#include <unistd.h>

#endif

struct sb_out_chunk {

    struct cio_chunk  *chunk;

    struct cio_stream *stream;

    size_t             size;

    struct mk_list    _head;

};

struct sb_out_queue {

    struct flb_output_instance *ins;

    struct mk_list              chunks; /* head for every sb_out_chunk */

    struct mk_list              _head;

};

struct flb_sb {

    int coll_fd;                    /* collector id */

    size_t mem_limit;               /* memory limit */

    struct flb_input_instance *ins; /* input instance */

    struct cio_ctx *cio;            /* chunk i/o instance */

    struct mk_list backlogs;        /* list of all pending chunks segregated by output plugin */

};

static inline struct flb_sb *sb_get_context(struct flb_config *config);

static struct sb_out_chunk *sb_allocate_chunk(struct cio_chunk *chunk,

                                              struct cio_stream *stream,

                                              size_t size);

static void sb_destroy_chunk(struct sb_out_chunk *chunk);

static void sb_destroy_backlog(struct sb_out_queue *backlog, struct flb_sb *context);

static int  sb_allocate_backlogs(struct flb_sb *ctx);

static void sb_destroy_backlogs(struct flb_sb *ctx);

static struct sb_out_queue *sb_find_segregated_backlog_by_output_plugin_instance(

                                struct flb_output_instance *output_plugin,

                                struct flb_sb              *context);

static void sb_remove_chunk_from_segregated_backlog(struct cio_chunk    *target_chunk,

                                                    struct sb_out_queue *backlog,

                                                    int                  destroy);

static void sb_remove_chunk_from_segregated_backlogs(struct cio_chunk *chunk,

                                                     struct flb_sb    *context);

static int sb_append_chunk_to_segregated_backlog(struct cio_chunk    *target_chunk,

                                                 struct cio_stream   *stream,

                                                 size_t               target_chunk_size,

                                                 struct sb_out_queue *backlog);

static int sb_append_chunk_to_segregated_backlogs(struct cio_chunk  *target_chunk,

                                                  struct cio_stream *stream,

                                                  struct flb_sb     *context);

int sb_segregate_chunks(struct flb_config *config);

int sb_release_output_queue_space(struct flb_output_instance *output_plugin,

                                  size_t                      required_space);

ssize_t sb_get_releasable_output_queue_space(struct flb_output_instance *output_plugin,

                                             size_t                      required_space);

static inline struct flb_sb *sb_get_context(struct flb_config *config)

{

    if (config == NULL) {

        return NULL;

    }

    if (config->storage_input_plugin == NULL) {

        return NULL;

    }

    return ((struct flb_input_instance *) config->storage_input_plugin)->context;

}

static struct sb_out_chunk *sb_allocate_chunk(struct cio_chunk *chunk,

                                              struct cio_stream *stream,

                                              size_t size)

{

    struct sb_out_chunk *result;

    result = (struct sb_out_chunk *) flb_calloc(1, sizeof(struct sb_out_chunk));

    if (result != NULL) {

        result->chunk  = chunk;

        result->stream = stream;

        result->size   = size;

    }

    return result;

}

static void sb_destroy_chunk(struct sb_out_chunk *chunk)

{

    flb_free(chunk);

}

static void sb_destroy_backlog(struct sb_out_queue *backlog, struct flb_sb *context)

{

    struct mk_list      *chunk_iterator_tmp;

    struct mk_list      *chunk_iterator;

    struct sb_out_chunk *chunk;

    mk_list_foreach_safe(chunk_iterator, chunk_iterator_tmp, &backlog->chunks) {

        chunk = mk_list_entry(chunk_iterator, struct sb_out_chunk, _head);

        sb_remove_chunk_from_segregated_backlogs(chunk->chunk, context);

    }

}

static int sb_allocate_backlogs(struct flb_sb *context)

{

    struct mk_list             *output_plugin_iterator;

    struct flb_output_instance *output_plugin;

    struct sb_out_queue        *backlog;

    mk_list_foreach(output_plugin_iterator, &context->ins->config->outputs) {

        output_plugin = mk_list_entry(output_plugin_iterator,

                                      struct flb_output_instance,

                                      _head);

        backlog = (struct sb_out_queue *) \

                        flb_calloc(1, sizeof(struct sb_out_queue));

        if (backlog == NULL) {

            sb_destroy_backlogs(context);

            return -1;

        }

        backlog->ins = output_plugin;

        mk_list_init(&backlog->chunks);

        mk_list_add(&backlog->_head, &context->backlogs);

    }

    return 0;

}

static void sb_destroy_backlogs(struct flb_sb *context)

{

    struct mk_list      *backlog_iterator_tmp;

    struct mk_list      *backlog_iterator;

    struct sb_out_queue *backlog;

    mk_list_foreach_safe(backlog_iterator, backlog_iterator_tmp, &context->backlogs) {

        backlog = mk_list_entry(backlog_iterator, struct sb_out_queue, _head);

        mk_list_del(&backlog->_head);

        sb_destroy_backlog(backlog, context);

        flb_free(backlog);

    }

}

static struct sb_out_queue *sb_find_segregated_backlog_by_output_plugin_instance(

                                struct flb_output_instance *output_plugin,

                                struct flb_sb              *context)

{

    struct mk_list      *backlog_iterator;

    struct sb_out_queue *backlog;

    mk_list_foreach(backlog_iterator, &context->backlogs) {

        backlog = mk_list_entry(backlog_iterator, struct sb_out_queue, _head);

        if (output_plugin == backlog->ins) {

            return backlog;

        }

    }

    return NULL;

}

static void sb_remove_chunk_from_segregated_backlog(struct cio_chunk    *target_chunk,

                                                    struct sb_out_queue *backlog,

                                                    int                  destroy)

{

    struct mk_list      *chunk_iterator_tmp;

    struct mk_list      *chunk_iterator;

    struct sb_out_chunk *chunk;

    mk_list_foreach_safe(chunk_iterator, chunk_iterator_tmp, &backlog->chunks) {

        chunk = mk_list_entry(chunk_iterator, struct sb_out_chunk, _head);

        if (chunk->chunk == target_chunk) {

            mk_list_del(&chunk->_head);

            backlog->ins->fs_backlog_chunks_size -= cio_chunk_get_real_size(target_chunk);

            if (destroy) {

                sb_destroy_chunk(chunk);

            }

            break;

        }

    }

}

static void sb_remove_chunk_from_segregated_backlogs(struct cio_chunk *target_chunk,

                                                     struct flb_sb    *context)

{

    struct mk_list      *backlog_iterator;

    struct sb_out_queue *backlog;

    mk_list_foreach(backlog_iterator, &context->backlogs) {

        backlog = mk_list_entry(backlog_iterator, struct sb_out_queue, _head);

        sb_remove_chunk_from_segregated_backlog(target_chunk, backlog, FLB_TRUE);

    }

}

static int sb_append_chunk_to_segregated_backlog(struct cio_chunk    *target_chunk,

                                                 struct cio_stream   *stream,

                                                 size_t               target_chunk_size,

                                                 struct sb_out_queue *backlog)

{

    struct sb_out_chunk *chunk;

    chunk = sb_allocate_chunk(target_chunk, stream, target_chunk_size);

    if (chunk == NULL) {

        flb_errno();

        return -1;

    }

    mk_list_add(&chunk->_head, &backlog->chunks);

    backlog->ins->fs_backlog_chunks_size += target_chunk_size;

    return 0;

}

static int sb_append_chunk_to_segregated_backlogs(struct cio_chunk  *target_chunk,

                                                  struct cio_stream *stream,

                                                  struct flb_sb     *context)

{

    struct flb_input_chunk  dummy_input_chunk;

    struct mk_list         *tmp;

    struct mk_list         *head;

    size_t                  chunk_size;

    struct sb_out_queue    *backlog;

    int                     tag_len;

    const char *            tag_buf;

    int                     result;

    memset(&dummy_input_chunk, 0, sizeof(struct flb_input_chunk));

    dummy_input_chunk.in    = context->ins;

    dummy_input_chunk.chunk = target_chunk;

    chunk_size = cio_chunk_get_real_size(target_chunk);

    if (chunk_size < 0) {

        flb_warn("[storage backlog] could not get real size of chunk %s/%s",

                  stream->name, target_chunk->name);

        return -1;

    }

    result = flb_input_chunk_get_tag(&dummy_input_chunk, &tag_buf, &tag_len);

    if (result == -1) {

        flb_error("[storage backlog] could not retrieve chunk tag from %s/%s, "

                  "removing it from the queue",

                  stream->name, target_chunk->name);

        return -2;

    }

    flb_routes_mask_set_by_tag(dummy_input_chunk.routes_mask, tag_buf, tag_len,

                               context->ins);

    mk_list_foreach_safe(head, tmp, &context->backlogs) {

        backlog = mk_list_entry(head, struct sb_out_queue, _head);

        if (flb_routes_mask_get_bit(dummy_input_chunk.routes_mask,

                                    backlog->ins->id)) {

            result = sb_append_chunk_to_segregated_backlog(target_chunk, stream,

                                                           chunk_size, backlog);

            if (result) {

                return -3;

            }

        }

    }

    return 0;

}

int sb_segregate_chunks(struct flb_config *config)

{

    int                ret;

    size_t             size;

    struct mk_list    *tmp;

    struct mk_list    *stream_iterator;

    struct mk_list    *chunk_iterator;

    int                chunk_error;

    struct flb_sb     *context;

    struct cio_stream *stream;

    struct cio_chunk  *chunk;

    context = sb_get_context(config);

    if (context == NULL) {

        return 0;

    }

    ret = sb_allocate_backlogs(context);

    if (ret) {

        return -2;

    }

    mk_list_foreach(stream_iterator, &context->cio->streams) {

        stream = mk_list_entry(stream_iterator, struct cio_stream, _head);

        mk_list_foreach_safe(chunk_iterator, tmp, &stream->chunks) {

            chunk = mk_list_entry(chunk_iterator, struct cio_chunk, _head);

            if (!cio_chunk_is_up(chunk)) {

                ret = cio_chunk_up_force(chunk);

                if (ret == CIO_CORRUPTED) {

                    if (config->storage_del_bad_chunks) {

                        chunk_error = cio_error_get(chunk);

                        if (chunk_error == CIO_ERR_BAD_FILE_SIZE ||

                            chunk_error == CIO_ERR_BAD_LAYOUT)

                        {

                            flb_plg_error(context->ins, "discarding irrecoverable chunk %s/%s", stream->name, chunk->name);

                            cio_chunk_close(chunk, CIO_TRUE);

                        }

                    }

                    continue;

                }

            }

            if (!cio_chunk_is_up(chunk)) {

                return -3;

            }

            /* try to segregate a chunk */

            ret = sb_append_chunk_to_segregated_backlogs(chunk, stream, context);

            if (ret) {

                /*

                 * if the chunk could not be segregated, just remove it from the

                 * queue and continue.

                 *

                 * if content size is zero, it's safe to 'delete it'.

                 */

                size = cio_chunk_get_content_size(chunk);

                if (size <= 0) {

                    cio_chunk_close(chunk, CIO_TRUE);

                }

                else {

                    cio_chunk_close(chunk, CIO_FALSE);

                }

                continue;

            }

            /* lock the chunk */

            flb_plg_info(context->ins, "register %s/%s", stream->name, chunk->name);

            cio_chunk_lock(chunk);

            cio_chunk_down(chunk);

        }

    }

    return 0;

}

ssize_t sb_get_releasable_output_queue_space(struct flb_output_instance *output_plugin,

                                             size_t                      required_space)

{

    ssize_t              releasable_space;

    struct mk_list      *chunk_iterator;

    struct flb_sb       *context;

    struct sb_out_queue *backlog;

    struct sb_out_chunk *chunk;

    context = sb_get_context(output_plugin->config);

    if (context == NULL) {

        return 0;

    }

    backlog = sb_find_segregated_backlog_by_output_plugin_instance(

                                output_plugin, context);

    if (backlog == NULL) {

        return 0;

    }

    releasable_space = 0;

    mk_list_foreach(chunk_iterator, &backlog->chunks) {

        chunk = mk_list_entry(chunk_iterator, struct sb_out_chunk, _head);

        releasable_space += chunk->size;

        if (releasable_space >= required_space) {

            break;

        }

    }

    return releasable_space;

}

int sb_release_output_queue_space(struct flb_output_instance *output_plugin,

                                  size_t                      required_space)

{

    struct mk_list      *chunk_iterator_tmp;

    struct mk_list      *chunk_iterator;

    size_t               released_space;

    struct flb_sb       *context;

    struct sb_out_queue *backlog;

    struct sb_out_chunk *chunk;

    context = sb_get_context(output_plugin->config);

    if (context == NULL) {

        return -1;

    }

    backlog = sb_find_segregated_backlog_by_output_plugin_instance(

                                                        output_plugin, context);

    if (backlog == NULL) {

        return -2;

    }

    released_space = 0;

    mk_list_foreach_safe(chunk_iterator, chunk_iterator_tmp, &backlog->chunks) {

        chunk = mk_list_entry(chunk_iterator, struct sb_out_chunk, _head);

        released_space += chunk->size;

        cio_chunk_close(chunk->chunk, FLB_TRUE);

        sb_remove_chunk_from_segregated_backlogs(chunk->chunk, context);

        if (released_space >= required_space) {

            break;

        }

    }

    if (released_space < required_space) {

        return -3;

    }

    return 0;

}

/* Collection callback */

static int cb_queue_chunks(struct flb_input_instance *in,

                           struct flb_config *config, void *data)

{

    size_t                  empty_output_queue_count;

    struct mk_list         *output_queue_iterator;

    struct sb_out_queue    *output_queue_instance;

    struct sb_out_chunk    *chunk_instance;

    struct flb_sb          *ctx;

    struct flb_input_chunk *ic;

    struct flb_input_chunk  tmp_ic;

    void                   *ch;

    size_t                  total = 0;

    ssize_t                 size;

    int                     ret;

    int                     event_type;

    /* Get context */

    ctx = (struct flb_sb *) data;

    /* Get the total number of bytes already enqueued */

    total = flb_input_chunk_total_size(in);

    /* If we already hitted our limit, just wait and re-check later */

    if (total >= ctx->mem_limit) {

        return 0;

    }

    empty_output_queue_count = 0;

    while (total < ctx->mem_limit &&

           empty_output_queue_count < mk_list_size(&ctx->backlogs)) {

        empty_output_queue_count = 0;

        mk_list_foreach(output_queue_iterator, &ctx->backlogs) {

            output_queue_instance = mk_list_entry(output_queue_iterator,

                                                  struct sb_out_queue,

                                                  _head);

            if (mk_list_is_empty(&output_queue_instance->chunks) != 0) {

                chunk_instance = mk_list_entry_first(&output_queue_instance->chunks,

                                                     struct sb_out_chunk,

                                                     _head);

                /* Try to enqueue one chunk */

                /*

                 * All chunks on this backlog are 'file' based, always try to set

                 * them up. We validate the status.

                 */

                ret = cio_chunk_is_up(chunk_instance->chunk);

                if (ret == CIO_FALSE) {

                    ret = cio_chunk_up_force(chunk_instance->chunk);

                    if (ret == CIO_CORRUPTED) {

                        flb_plg_error(ctx->ins, "removing corrupted chunk from the "

                                      "queue %s:%s",

                                      chunk_instance->stream->name, chunk_instance->chunk->name);

                        cio_chunk_close(chunk_instance->chunk, FLB_FALSE);

                        sb_remove_chunk_from_segregated_backlogs(chunk_instance->chunk, ctx);

                        /* This function will indirecly release chunk_instance so it has to be

                         * called last.

                         */

                        continue;

                    }

                    else if (ret == CIO_ERROR || ret == CIO_RETRY) {

                        continue;

                    }

                }

                /*

                 * Map the chunk file context into a temporary buffer since the

                 * flb_input_chunk_get_event_type() interface needs an

                 * struct fb_input_chunk argument.

                 */

                tmp_ic.chunk = chunk_instance->chunk;

                /* Retrieve the event type: FLB_INPUT_LOGS, FLB_INPUT_METRICS of FLB_INPUT_TRACES */

                ret = flb_input_chunk_get_event_type(&tmp_ic);

                if (ret == -1) {

                    flb_plg_error(ctx->ins, "removing chunk with wrong metadata "

                                  "from the queue %s:%s",

                                  chunk_instance->stream->name,

                                  chunk_instance->chunk->name);

                    cio_chunk_close(chunk_instance->chunk, FLB_TRUE);

                    sb_remove_chunk_from_segregated_backlogs(chunk_instance->chunk,

                                                             ctx);

                    continue;

                }

                event_type = ret;

                /* get the number of bytes being used by the chunk */

                size = cio_chunk_get_content_size(chunk_instance->chunk);

                if (size <= 0) {

                    flb_plg_error(ctx->ins, "removing empty chunk from the "

                                  "queue %s:%s",

                                  chunk_instance->stream->name, chunk_instance->chunk->name);

                    cio_chunk_close(chunk_instance->chunk, FLB_TRUE);

                    sb_remove_chunk_from_segregated_backlogs(chunk_instance->chunk, ctx);

                    /* This function will indirecly release chunk_instance so it has to be

                     * called last.

                     */

                    continue;

                }

                ch = chunk_instance->chunk;

                /* Associate this backlog chunk to this instance into the engine */

                ic = flb_input_chunk_map(in, event_type, ch);

                if (!ic) {

                    flb_plg_error(ctx->ins, "removing chunk %s:%s from the queue",

                                  chunk_instance->stream->name, chunk_instance->chunk->name);

                    cio_chunk_down(chunk_instance->chunk);

                    /*

                     * If the file cannot be mapped, just drop it. Failures are all

                     * associated with data corruption.

                     */

                    cio_chunk_close(chunk_instance->chunk, FLB_TRUE);

                    sb_remove_chunk_from_segregated_backlogs(chunk_instance->chunk, ctx);

                    /* This function will indirecly release chunk_instance so it has to be

                     * called last.

                     */

                    continue;

                }

                flb_plg_info(ctx->ins, "queueing %s:%s",

                             chunk_instance->stream->name, chunk_instance->chunk->name);

                /* We are removing this chunk reference from this specific backlog

                 * queue but we need to leave it in the remainder queues.

                 */

                sb_remove_chunk_from_segregated_backlogs(chunk_instance->chunk, ctx);

                /* check our limits */

                total += size;

            }

            else {

                empty_output_queue_count++;

            }

        }

    }

    return 0;

}

/* Initialize plugin */

static int cb_sb_init(struct flb_input_instance *in,

                      struct flb_config *config, void *data)

{

    int ret;

    char mem[32];

    struct flb_sb *ctx;

    ctx = flb_malloc(sizeof(struct flb_sb));

    if (!ctx) {

        flb_errno();

        return -1;

    }

    ctx->cio = data;

    ctx->ins = in;

    ctx->mem_limit = flb_utils_size_to_bytes(config->storage_bl_mem_limit);

    mk_list_init(&ctx->backlogs);

    flb_utils_bytes_to_human_readable_size(ctx->mem_limit, mem, sizeof(mem) - 1);

    flb_plg_info(ctx->ins, "queue memory limit: %s", mem);

    /* export plugin context */

    flb_input_set_context(in, ctx);

    /* Set a collector to trigger the callback to queue data every second */

    ret = flb_input_set_collector_time(in, cb_queue_chunks, 1, 0, config);

    if (ret < 0) {

        flb_plg_error(ctx->ins, "could not create collector");

        flb_free(ctx);

        return -1;

    }

    ctx->coll_fd = ret;

    return 0;

}

static void cb_sb_pause(void *data, struct flb_config *config)

{

    struct flb_sb *ctx = data;

    flb_input_collector_pause(ctx->coll_fd, ctx->ins);

}

static void cb_sb_resume(void *data, struct flb_config *config)

{

    struct flb_sb *ctx = data;

    flb_input_collector_resume(ctx->coll_fd, ctx->ins);

}

static int cb_sb_exit(void *data, struct flb_config *config)

{

    struct flb_sb *ctx = data;

    flb_input_collector_pause(ctx->coll_fd, ctx->ins);

    sb_destroy_backlogs(ctx);

    flb_free(ctx);

    return 0;

}

/* Plugin reference */

struct flb_input_plugin in_storage_backlog_plugin = {

    .name         = "storage_backlog",

    .description  = "Storage Backlog",

    .cb_init      = cb_sb_init,

    .cb_pre_run   = NULL,

    .cb_collect   = NULL,

    .cb_ingest    = NULL,

    .cb_flush_buf = NULL,

    .cb_pause     = cb_sb_pause,

    .cb_resume    = cb_sb_resume,

    .cb_exit      = cb_sb_exit,

    /* This plugin can only be configured and invoked by the Engine */

    .flags        = FLB_INPUT_PRIVATE

};