path: root/src/packet_io/under_ddos.cpp

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/sysinfo.h>
#include <pcap/pcap.h>
#include "cpu_limit.h"
#include "sapp_api.h"
#include "sapp_private_api.h"
#include "sapp_declaration.h"

#ifdef __cplusplus
extern "C" {
#endif

typedef enum
{
    TCK_USER = 0,
    TCK_NICE,
    TCK_SYSTEM,
    TCK_IDLE,
    TCK_IOWAIT,
    TCK_IRQ,
    TCK_SOFTIRQ,
    TCK_STEAL,
    TCK_GUEST,
    TCK_GUEST_NICE,
    NUM_TCK_TYPES
} cpu_tck_type;
/* https://www.linuxhowtos.org/System/procstat.htm */

typedef struct
{
    char name[16];
    uint64_t tcks[NUM_TCK_TYPES];
} cpu_tck_t;

#ifndef MAX_CORE_NUM
#define MAX_CORE_NUM 256
#endif

typedef struct{
	unsigned long long last_create_stream_new_sum[MAX_CORE_NUM];
	double last_time_cpu_total[MAX_CORE_NUM];
	double last_time_cpu_idle[MAX_CORE_NUM];
	double ewma_cpu_usage[MAX_CORE_NUM];
	cpu_tck_t all_cpu_usage[MAX_CORE_NUM];
}under_sapp_user_args_t;


/* 
	��ȡcpuռ������ͨ��core_id����, ͨ���߳�indexת����core_id.
*/
int sapp_thread_index_to_core_id(int thread_seq)
{
	int i;
    cpu_set_t current_cpu_mask;
	
    if(sapp_global_val->individual_fixed.thread_obtain_id[thread_seq] <= 0){
		return -1;
    }	

    CPU_ZERO(&current_cpu_mask);
	if(pthread_getaffinity_np(sapp_global_val->individual_fixed.thread_obtain_id[thread_seq], sizeof(cpu_set_t), &current_cpu_mask) < 0){
		return -1;
	}
	
	int tot_cpu_core = get_nprocs();
	
	for(i = 1; i <= tot_cpu_core; i++){
		if(CPU_ISSET(i, &current_cpu_mask)){
			break;
		}
	}
	
	return i;

}

static inline uint64_t calc_total_ticks(cpu_tck_t *stat)
{
    uint64_t total = 0;
    int i;
    for(i = 0; i < NUM_TCK_TYPES; i++)
        total += stat->tcks[i];
    return total;
}

static void read_cpu_usage_from_proc(cpu_tck_t *per_cpu_core_stat, int max_cpu_num)
{
    FILE *stat_fp = fopen("/proc/stat", "r");

    int i = 0;
    char no_use_string[NAME_MAX];

	if(NULL == stat_fp){
		memset(per_cpu_core_stat, 0, sizeof(cpu_tck_t)*max_cpu_num);
		return;
	}

	/* ��һ����ϵͳ����cpuռ����, Ȼ����ÿ��CPU���ĵ� */
    fgets(no_use_string, sizeof(no_use_string), stat_fp);

    for(i = 0; i < max_cpu_num; i++)
	{
        fscanf(
            stat_fp,
            "%16s %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu\n",
            per_cpu_core_stat[i].name,
            &(per_cpu_core_stat[i].tcks[TCK_USER]),
            &(per_cpu_core_stat[i].tcks[TCK_NICE]),
            &(per_cpu_core_stat[i].tcks[TCK_SYSTEM]),
            &(per_cpu_core_stat[i].tcks[TCK_IDLE]),
            &(per_cpu_core_stat[i].tcks[TCK_IOWAIT]),
            &(per_cpu_core_stat[i].tcks[TCK_IRQ]),
            &(per_cpu_core_stat[i].tcks[TCK_SOFTIRQ]),
            &(per_cpu_core_stat[i].tcks[TCK_STEAL]),
            &(per_cpu_core_stat[i].tcks[TCK_GUEST]),
            &(per_cpu_core_stat[i].tcks[TCK_GUEST_NICE])
        );
    }

    fclose(stat_fp);
}

static double sapp_get_cpu_usage_cb(cpu_limit_handle h, int _thread_index, void *_void_user_arg)
{
	double current_cpu_usage;
	uint64_t this_total_tcks, this_idle_tcks;
	int cpu_core_id;
	int sys_actual_cpu_core_num = get_nprocs();
	under_sapp_user_args_t *ud_usr_arg = (under_sapp_user_args_t *)_void_user_arg;

	/* �˺����DZ�cl_limit��̨�̵߳���, ��һ�������п���sapp IO�̻߳�û����������, ���Դ˴��ж�һ��״̬ */
	if(SAPP_STATE_PROCESSING != sapp_global_val->individual_volatile->current_state){
		return 0.0;
	}
	
	cpu_core_id = sapp_thread_index_to_core_id(_thread_index);
	if(cpu_core_id < 0){
		return 0.0;
	}

	/* �ⲿģ�����, thread_index�Ǵ�0-N˳�����ε���,
	   ���ڵ�һ�δ�/proc/stat�л�ȡ, ������ֱ�Ӵ��ڴ����ȡ, �����δ�/proc/stat */
	if(0 == _thread_index){
		read_cpu_usage_from_proc(ud_usr_arg->all_cpu_usage, sys_actual_cpu_core_num);
	}
	
	this_total_tcks = calc_total_ticks(&ud_usr_arg->all_cpu_usage[cpu_core_id]);
	this_idle_tcks = ud_usr_arg->all_cpu_usage[cpu_core_id].tcks[TCK_IDLE];

	/* ���һ��ʱ��, ��������used�IJ�ֵ�������������IJ�ֵ, 
	   ����ֱ��������, ���DZ�ʾ�Ի����ӵ���������cpuռ���ʵ���ƽ��ֵ, 
	*/
	current_cpu_usage = 100.0 * ((this_total_tcks - this_idle_tcks)-(ud_usr_arg->last_time_cpu_total[cpu_core_id] - ud_usr_arg->last_time_cpu_idle[cpu_core_id]))/(this_total_tcks - ud_usr_arg->last_time_cpu_total[cpu_core_id]);

#define EWMA_FACTOR 0.8
	if(ud_usr_arg->ewma_cpu_usage[cpu_core_id] == 0)
	{
		ud_usr_arg->ewma_cpu_usage[cpu_core_id] = current_cpu_usage;
	}
	else
	{
		ud_usr_arg->ewma_cpu_usage[cpu_core_id] = (EWMA_FACTOR * current_cpu_usage) + ((1 - EWMA_FACTOR) * ud_usr_arg->ewma_cpu_usage[cpu_core_id]);
	}
	ud_usr_arg->last_time_cpu_total[cpu_core_id] = this_total_tcks;
	ud_usr_arg->last_time_cpu_idle[cpu_core_id]  = this_idle_tcks;	

	return ud_usr_arg->ewma_cpu_usage[cpu_core_id];

}

static double sapp_get_create_stream_rate_cb(cpu_limit_handle h, int thread_index,void *_void_user_arg)
{
	under_sapp_user_args_t *ud_usr_arg = (under_sapp_user_args_t *)_void_user_arg;
	unsigned long long cur_create_stream_sum;
	unsigned long long cur_create_stream_rate;

	/* TCP��UDP�ܼ� */
	cur_create_stream_sum = sapp_global_val->mthread_volatile[thread_index]->sys_stat.count[SAPP_STAT_TCP_STREAM_NEW] + sapp_global_val->mthread_volatile[thread_index]->sys_stat.count[SAPP_STAT_UDP_STREAM_NEW];

	/* ��һ��Ϊ0, Ϊ�˱������, �˴β����� */
	if(0 == ud_usr_arg->last_create_stream_new_sum[thread_index]){
		ud_usr_arg->last_create_stream_new_sum[thread_index] = cur_create_stream_sum;
		return 0.0;
	}
	cur_create_stream_rate = cur_create_stream_sum-ud_usr_arg->last_create_stream_new_sum[thread_index];
	ud_usr_arg->last_create_stream_new_sum[thread_index] = cur_create_stream_sum;

	return (double)cur_create_stream_rate;
}

int packet_io_under_ddos_global_status(void)
{
	int global_bypass_state;
	int opt_len = sizeof(int);

	if(0 == sapp_global_val->config.packet_io.under_ddos_config.enabled){
		return 0;
	}
	
	cpu_limit_get_opt(sapp_global_val->individual_fixed.under_ddos_handle, CL_OPT_GLOBAL_BYPASS_STATE, &global_bypass_state, &opt_len);

	return global_bypass_state;
}

int packet_io_under_ddos_should_bypass(int thread_index)
{
	int ret;
	const sapp_under_ddos_config_t *p_ddos_cfg = &sapp_global_val->config.packet_io.under_ddos_config;

	if(0 == p_ddos_cfg->enabled){
		return 0;
	}

	ret = cpu_limit_can_i_do(sapp_global_val->individual_fixed.under_ddos_handle, thread_index);
	/* ����ֵҪȡ�� */
	return !ret;
}

void packet_io_under_ddos_run(void)
{
	const sapp_under_ddos_config_t *p_ddos_cfg = &sapp_global_val->config.packet_io.under_ddos_config;

	if(0 == p_ddos_cfg->enabled){
		return;
	}

	if(cpu_limit_start(sapp_global_val->individual_fixed.under_ddos_handle) < 0){
		sapp_runtime_log(RLOG_LV_FATAL, "cpu_limit_start() error !\n");
	}
}

void packet_io_under_ddos_limit_arg_free(cpu_limit_handle h, void *user_arg)
{
	if(user_arg)
	{
		free(user_arg);
	}
	return;
}

int packet_io_under_ddos_init(void)
{
	sapp_under_ddos_config_t *p_ddos_cfg = &sapp_global_val->config.packet_io.under_ddos_config;
	cpu_limit_handle ud_handle;
	int iopt_value;
	
	if(0 == p_ddos_cfg->enabled){
		return 0;
	}

	if(sapp_global_val->config.cpu.bind_mask_array_num == 0){
		sapp_log(RLOG_LV_FATAL, ~0, ~0, "cpu_limit_create() error, must set [CPU]->bind_mask when stream_bypass_enabled=1\n");
		return -1;
	}

	ud_handle = cpu_limit_create();
	if(NULL == ud_handle){
		sapp_log(RLOG_LV_FATAL, ~0, ~0, "cpu_limit_create() error!\n");
		return -1;		
	}	

	cpu_limit_set_opt(ud_handle, CL_OPT_RES_TRIGGER_THRESHOLD, &p_ddos_cfg->bypass_trigger_cpu_usage, sizeof(double));

	cpu_limit_set_opt(ud_handle, CL_OPT_FACTOR_DECREASE_RATIO, &p_ddos_cfg->factor_decrease_ratio, sizeof(double));
	cpu_limit_set_opt(ud_handle, CL_OPT_FACTOR_INCREASE_RATIO, &p_ddos_cfg->factor_increase_ratio, sizeof(double));

	iopt_value = g_packet_io_thread_num;
	cpu_limit_set_opt(ud_handle, CL_OPT_THREAD_COUNT, &iopt_value, sizeof(int));

	cpu_limit_set_opt(ud_handle, CL_OPT_STAT_INTERVAL, &p_ddos_cfg->get_cpu_usage_interval, sizeof(int));
	cpu_limit_set_opt(ud_handle, CL_OPT_OBSERVE_TIME, &p_ddos_cfg->recovery_observe_time, sizeof(int));
	cpu_limit_set_opt(ud_handle, CL_OPT_RES_SMOOTH_SCOPE, &p_ddos_cfg->smooth_avg_window, sizeof(int));
	cpu_limit_set_opt(ud_handle, CL_OPT_RES_GET_FUN, (void *)&sapp_get_cpu_usage_cb, sizeof(void *));	
	cpu_limit_set_opt(ud_handle, CL_OPT_FACTOR_GET_FUN, (void *)&sapp_get_create_stream_rate_cb, sizeof(void *));	
	under_sapp_user_args_t *ud_user_arg = (under_sapp_user_args_t *)calloc(1, sizeof(under_sapp_user_args_t));
	cpu_limit_set_opt(ud_handle, CL_OPT_USER_ARG, ud_user_arg, sizeof(void *));	
	cpu_limit_set_opt(ud_handle, CL_OPT_DESTROY_CALLBCAK_FUN, (void *)packet_io_under_ddos_limit_arg_free, sizeof(void *));	
	cpu_limit_set_opt(ud_handle, CL_OPT_DEBUG_LOG_FILE_NAME, (void *)sapp_global_val->config.data_file_path.data_under_ddos_stat_log_absolute, strlen(sapp_global_val->config.data_file_path.data_under_ddos_stat_log_absolute));		

	sapp_global_val->individual_fixed.under_ddos_handle = ud_handle;

	return 0;
}

void packet_io_under_ddos_destroy(void)
{
	if(0 == sapp_global_val->config.packet_io.under_ddos_config.enabled){
		return;
	}

	cpu_limit_destroy(sapp_global_val->individual_fixed.under_ddos_handle);	
}

#ifdef __cplusplus
}
#endif