stellar-stellar/src/stellar/stellar.cpp

#include <errno.h>
#include <assert.h>
#include <unistd.h>
#include <pthread.h>
#include <signal.h>
#include <sys/prctl.h>

#include "stat.h"
#include "logo.h"
#include "times.h"
#include "config.h"
#include "id_generator.h"
#include "stellar_priv.h"
#include "plugin_manager.h"

struct stellar_thread
{
    pthread_t tid;
    uint16_t idx;
    uint64_t is_runing;
    uint64_t timing_wheel_last_update_ts;
    struct ip_reassembly *ip_mgr;
    struct session_manager *sess_mgr;
    struct stellar_runtime *runtime;
};

struct stellar_runtime
{
    uint64_t need_exit;
    uint64_t stat_last_output_ts;
    struct stellar_stat *stat;
    struct packet_io *packet_io;
    struct plugin_manager_schema *plug_mgr;
    struct stellar_thread threads[MAX_THREAD_NUM];
};

struct stellar
{
    struct stellar_runtime runtime;
    struct stellar_config config;
};

static uint64_t need_exit = 0;
static thread_local uint16_t __thread_id = 0;
static const char *log_config_file = "./conf/log.toml";
static const char *main_config_file = "./conf/stellar.toml";
static const char *plugin_config_file = "./plugin/spec.toml";

/******************************************************************************
 * Stellar Thread Main Loop
 ******************************************************************************/

static void update_session_stat(struct session *sess, struct packet *pkt)
{
    if (sess)
    {
        enum flow_direction dir = session_get_current_flow_direction(sess);
        assert(dir != FLOW_DIRECTION_NONE);
        int is_ctrl = packet_is_ctrl(pkt);
        switch (packet_get_action(pkt))
        {
        case PACKET_ACTION_DROP:
            session_inc_stat(sess, dir, (is_ctrl ? STAT_CONTROL_PACKETS_DROPPED : STAT_RAW_PACKETS_DROPPED), 1);
            session_inc_stat(sess, dir, (is_ctrl ? STAT_CONTROL_BYTES_DROPPED : STAT_RAW_BYTES_DROPPED), packet_get_len(pkt));
            break;
        case PACKET_ACTION_FORWARD:
            session_inc_stat(sess, dir, (is_ctrl ? STAT_CONTROL_PACKETS_TRANSMITTED : STAT_RAW_PACKETS_TRANSMITTED), 1);
            session_inc_stat(sess, dir, (is_ctrl ? STAT_CONTROL_BYTES_TRANSMITTED : STAT_RAW_BYTES_TRANSMITTED), packet_get_len(pkt));
            break;
        default:
            assert(0);
            break;
        }

        session_set_current_packet(sess, NULL);
        session_set_current_flow_direction(sess, FLOW_DIRECTION_NONE);
    }
}

static inline void free_evicted_sessions(struct session_manager *sess_mgr, uint64_t max_free)
{
    void *plugin_ctx = NULL;
    struct session *sess = NULL;
    for (uint64_t i = 0; i < max_free; i++)
    {
        sess = session_manager_get_evicted_session(sess_mgr);
        if (sess)
        {
            plugin_ctx = session_get_user_data(sess);
            plugin_manager_session_runtime_free((struct plugin_manager_runtime *)plugin_ctx);
            session_manager_free_session(sess_mgr, sess);
        }
        else
        {
            break;
        }
    }
}

static inline void free_expired_sessions(struct session_manager *sess_mgr, uint64_t max_free, uint64_t now)
{
    void *plugin_ctx = NULL;
    struct session *sess = NULL;
    for (uint64_t i = 0; i < max_free; i++)
    {
        sess = session_manager_get_expired_session(sess_mgr, now);
        if (sess)
        {
            plugin_ctx = session_get_user_data(sess);
            plugin_manager_session_runtime_free((struct plugin_manager_runtime *)plugin_ctx);
            session_manager_free_session(sess_mgr, sess);
        }
        else
        {
            break;
        }
    }
}

static inline void merge_thread_stat(struct stellar_thread *thread, uint64_t now)
{
    struct stellar_runtime *runtime = thread->runtime;
    struct thread_stat thr_stat = {
        packet_io_stat(runtime->packet_io, thread->idx),
        ip_reassembly_stat(thread->ip_mgr),
        session_manager_stat(thread->sess_mgr),
    };
    stellar_stat_merge(runtime->stat, &thr_stat, thread->idx);
}

static void *work_thread(void *arg)
{
    int nr_recv;
    uint64_t now = 0;
    char thd_name[16] = {0};
    void *plugin_ctx = NULL;
    struct packet *pkt = NULL;
    struct packet *defraged_pkt = NULL;
    struct packet packets[RX_BURST_MAX];
    struct session *sess = NULL;
    struct stellar_thread *thread = (struct stellar_thread *)arg;
    struct ip_reassembly *ip_reass = thread->ip_mgr;
    struct session_manager *sess_mgr = thread->sess_mgr;
    struct stellar_runtime *runtime = thread->runtime;
    struct packet_io *packet_io = runtime->packet_io;
    struct plugin_manager_schema *plug_mgr = runtime->plug_mgr;

    uint16_t thr_idx = thread->idx;
    __thread_id = thr_idx;

    snprintf(thd_name, sizeof(thd_name), "stellar:%d", thr_idx);
    prctl(PR_SET_NAME, (unsigned long long)thd_name, NULL, NULL, NULL);

    if (packet_io_init(packet_io, thr_idx) != 0)
    {
        STELLAR_LOG_ERROR("unable to init marsio thread");
        return NULL;
    }

    ATOMIC_SET(&thread->is_runing, 1);
    STELLAR_LOG_STATE("worker thread %d runing", thr_idx);

    while (ATOMIC_READ(&need_exit) == 0)
    {
        now = stellar_get_monotonic_time_msec();
        memset(packets, 0, sizeof(packets));
        nr_recv = packet_io_ingress(packet_io, thr_idx, packets, RX_BURST_MAX);
        if (nr_recv == 0)
        {
            goto idle_tasks;
        }

        for (int i = 0; i < nr_recv; i++)
        {
            sess = NULL;
            defraged_pkt = NULL;
            pkt = &packets[i];

            plugin_manager_on_packet(plug_mgr, pkt);
            if (packet_is_fragment(pkt))
            {
                defraged_pkt = ip_reassembly_packet(ip_reass, pkt, now);
                if (defraged_pkt == NULL)
                {
                    goto fast_path;
                }
                else
                {
                    pkt = defraged_pkt;
                    plugin_manager_on_packet(plug_mgr, pkt);
                }
            }

            sess = session_manager_lookup_session(sess_mgr, pkt, now);
            if (sess == NULL)
            {
                sess = session_manager_new_session(sess_mgr, pkt, now);
                if (sess == NULL)
                {
                    goto fast_path;
                }
                plugin_ctx = plugin_manager_session_runtime_new(runtime->plug_mgr, sess);
                session_set_user_data(sess, plugin_ctx);
            }
            else
            {
                if (session_manager_update_session(sess_mgr, sess, pkt, now) == -1)
                {
                    goto fast_path;
                }
            }
            if (packet_get_session_id(pkt) == 0)
            {
                packet_set_session_id(pkt, session_get_id(sess));
            }
            plugin_manager_on_session_ingress(sess, pkt);

        fast_path:
            plugin_manager_on_session_egress(sess, pkt);
            if (sess && session_get_current_state(sess) == SESSION_STATE_DISCARD)
            {
                packet_set_action(pkt, PACKET_ACTION_DROP);
            }
            update_session_stat(sess, pkt);
            if (packet_get_action(pkt) == PACKET_ACTION_DROP)
            {
                if (pkt == defraged_pkt)
                {
                    packet_io_drop(packet_io, thr_idx, &packets[i], 1);
                    packet_free(defraged_pkt);
                }
                else
                {
                    packet_io_drop(packet_io, thr_idx, pkt, 1);
                }
            }
            else
            {
                if (pkt == defraged_pkt)
                {
                    packet_io_egress(packet_io, thr_idx, &packets[i], 1);
                    packet_free(defraged_pkt);
                }
                else
                {
                    packet_io_egress(packet_io, thr_idx, pkt, 1);
                }
            }
        }

    idle_tasks:
        // nr_recv packet atmost trigger nr_recv session evicted
        free_evicted_sessions(sess_mgr, nr_recv);

        // per 5 ms, atmost free 8 expired session
        if (now - thread->timing_wheel_last_update_ts > 5)
        {
            free_expired_sessions(sess_mgr, 8, now);
            thread->timing_wheel_last_update_ts = now;
        }

        merge_thread_stat(thread, now);
        ip_reassembly_expire(ip_reass, now);
        plugin_manager_on_polling(runtime->plug_mgr);

        if (nr_recv == 0)
        {
            packet_io_yield(packet_io, thr_idx, 10);
        }
    }

    ATOMIC_SET(&thread->is_runing, 0);
    STELLAR_LOG_STATE("worker thread %d exit !!!", thr_idx);

    return NULL;
}

/******************************************************************************
 * Stellar Main Function
 ******************************************************************************/

static void signal_handler(int signo)
{
    if (signo == SIGINT)
    {
        STELLAR_LOG_STATE("SIGINT received, notify threads to exit !!!");
        ATOMIC_SET(&need_exit, 1);
    }

    if (signo == SIGQUIT)
    {
        STELLAR_LOG_STATE("SIGQUIT received, notify threads to exit !!!");
        ATOMIC_SET(&need_exit, 1);
    }

    if (signo == SIGTERM)
    {
        STELLAR_LOG_STATE("SIGTERM received, notify threads to exit !!!");
        ATOMIC_SET(&need_exit, 1);
    }

    if (signo == SIGHUP)
    {
        STELLAR_LOG_STATE("SIGHUP received, reload log level !!!");
        log_reload_level(log_config_file);
    }
}

static int all_session_have_freed(struct stellar_runtime *runtime, struct stellar_config *config)
{
    for (int i = 0; i < config->io_opts.nr_threads; i++)
    {
        struct session_manager *sess_mgr = runtime->threads[i].sess_mgr;
        struct session_manager_stat *sess_stat = session_manager_stat(sess_mgr);

        if (ATOMIC_READ(&sess_stat->curr_nr_tcp_sess_used) != 0)
        {
            return 0;
        }

        if (ATOMIC_READ(&sess_stat->curr_nr_udp_sess_used) != 0)
        {
            return 0;
        }
    }

    return 1;
}

static int all_stat_have_output(struct stellar_runtime *runtime, struct stellar_config *config)
{
    static int count = 0;
    if (runtime->stat_last_output_ts == stellar_get_monotonic_time_msec())
    {
        count++;
    }

    return count == 2;
}

static int stellar_thread_init(struct stellar_runtime *runtime, struct stellar_config *config)
{
    uint64_t now = stellar_get_monotonic_time_msec();
    for (uint16_t i = 0; i < config->io_opts.nr_threads; i++)
    {
        struct stellar_thread *thread = &runtime->threads[i];
        thread->idx = i;
        thread->is_runing = 0;
        thread->timing_wheel_last_update_ts = now;
        thread->sess_mgr = session_manager_new(&config->sess_mgr_opts, now);
        if (thread->sess_mgr == NULL)
        {
            STELLAR_LOG_ERROR("unable to create session manager");
            return -1;
        }
        thread->ip_mgr = ip_reassembly_new(&config->ip_opts);
        if (thread->ip_mgr == NULL)
        {
            STELLAR_LOG_ERROR("unable to create ip reassemble manager");
            return -1;
        }
        thread->runtime = runtime;
    }

    return 0;
}

static void stellar_thread_clean(struct stellar_runtime *runtime, struct stellar_config *config)
{
    for (uint16_t i = 0; i < config->io_opts.nr_threads; i++)
    {
        struct stellar_thread *thread = &runtime->threads[i];
        if (ATOMIC_READ(&thread->is_runing) == 0)
        {
            STELLAR_LOG_STATE("wait worker thread %d free context", i);
            session_manager_free(thread->sess_mgr);
            ip_reassembly_free(thread->ip_mgr);
        }
    }
}

static int stellar_thread_run(struct stellar_runtime *runtime, struct stellar_config *config)
{
    for (uint16_t i = 0; i < config->io_opts.nr_threads; i++)
    {
        struct stellar_thread *thread = &runtime->threads[i];
        if (pthread_create(&thread->tid, NULL, work_thread, (void *)thread) < 0)
        {
            STELLAR_LOG_ERROR("unable to create worker thread, error %d: %s", errno, strerror(errno));
            return -1;
        }
    }

    return 0;
}

static void stellar_thread_join(struct stellar_runtime *runtime, struct stellar_config *config)
{
    for (uint16_t i = 0; i < config->io_opts.nr_threads; i++)
    {
        struct stellar_thread *thread = &runtime->threads[i];
        while (ATOMIC_READ(&thread->is_runing) == 1)
        {
            STELLAR_LOG_STATE("wait worker thread %d stop", i);
            sleep(1);
        }
    }
}

int stellar_main(int argc, char **argv)
{
    static struct stellar st = {0};
    static struct stellar_runtime *runtime = &st.runtime;
    static struct stellar_config *config = &st.config;

    stellar_update_time_cache();

    signal(SIGINT, signal_handler);
    signal(SIGQUIT, signal_handler);
    signal(SIGTERM, signal_handler);
    signal(SIGHUP, signal_handler);

    if (log_init(log_config_file) != 0)
    {
        STELLAR_LOG_ERROR("unable to init log");
        goto error_out;
    }
    STELLAR_LOG_STATE("start stellar (version: %s)\n %s", __stellar_version, logo_str);
    STELLAR_LOG_STATE("log config file    : %s", log_config_file);
    STELLAR_LOG_STATE("main config file   : %s", main_config_file);
    STELLAR_LOG_STATE("plugin config file : %s", plugin_config_file);

    if (stellar_load_config(main_config_file, config) != 0)
    {
        STELLAR_LOG_ERROR("unable to load config file");
        goto error_out;
    }
    stellar_print_config(config);
    STELLAR_LOG_DEBUG("sizeof(struct session) = %lu bytes", sizeof(struct session));

    if (id_generator_init(config->dev_opts.base, config->dev_opts.offset) != 0)
    {
        STELLAR_LOG_ERROR("unable to init id generator");
        goto error_out;
    }

    runtime->stat = stellar_stat_new(config->io_opts.nr_threads);
    if (runtime->stat == NULL)
    {
        STELLAR_LOG_ERROR("unable to create stellar stat");
        goto error_out;
    }
    runtime->plug_mgr = plugin_manager_init(&st, plugin_config_file);
    if (runtime->plug_mgr == NULL)
    {
        STELLAR_LOG_ERROR("unable to create plugin manager");
        goto error_out;
    }

    runtime->packet_io = packet_io_new(&config->io_opts);
    if (runtime->packet_io == NULL)
    {
        STELLAR_LOG_ERROR("unable to create packet io");
        goto error_out;
    }

    if (stellar_thread_init(runtime, config) != 0)
    {
        STELLAR_LOG_ERROR("unable to init thread context");
        goto error_out;
    }

    if (stellar_thread_run(runtime, config) != 0)
    {
        STELLAR_LOG_ERROR("unable to create worker thread");
        goto error_out;
    }

    runtime->stat_last_output_ts = stellar_get_monotonic_time_msec();
    while (!ATOMIC_READ(&need_exit))
    {
        stellar_update_time_cache();
        if (stellar_get_monotonic_time_msec() - runtime->stat_last_output_ts > 2000)
        {
            runtime->stat_last_output_ts = stellar_get_monotonic_time_msec();
            stellar_stat_output(runtime->stat);
        }
        usleep(1000); // 1ms

        // Only available in dump file mode, automatically exits when all sessions have been released
        if (packet_io_wait_exit(runtime->packet_io) && all_session_have_freed(runtime, config) && all_stat_have_output(runtime, config))
        {
            STELLAR_LOG_STATE("all sessions have been released and all stat have been output, notify threads to exit !!!");
            ATOMIC_SET(&need_exit, 1);
        }
    }

error_out:
    stellar_thread_join(runtime, config);
    stellar_thread_clean(runtime, config);
    packet_io_free(runtime->packet_io);
    plugin_manager_exit(runtime->plug_mgr);
    stellar_stat_free(runtime->stat);
    STELLAR_LOG_STATE("stellar exit !!!\n");
    log_free();

    return 0;
}

/******************************************************************************
 * Stellar Utility Function
 ******************************************************************************/

struct packet_io *stellar_get_packet_io(const struct stellar *st)
{
    return st->runtime.packet_io;
}

struct plugin_manager_schema *stellar_get_plugin_manager(const struct stellar *st)
{
    return st->runtime.plug_mgr;
}

void stellar_set_plugin_manger(struct stellar *st, struct plugin_manager_schema *plug_mgr)
{
    st->runtime.plug_mgr = plug_mgr;
}

struct session_manager *stellar_get_session_manager(const struct stellar *st)
{
    uint16_t idx = stellar_get_current_thread_index();
    return st->runtime.threads[idx].sess_mgr;
}

uint16_t stellar_get_current_thread_index()
{
    return __thread_id;
}