QoS令牌桶2 ------- 完整实现代码

根据前面的第一篇文章的TB的第一种实现的方式，我们实现的TB的完成代码，如下：


#ifndef __INCLUDE_QOS_METER_H__
#define __INCLUDE_QOS_METER_H__

#ifdef __cplusplus
extern "C" {
#endif

/**
 * @file
 * qos Traffic Metering
 *
 * Traffic metering algorithms:
 *    1. Single Rate Three Color Marker (srTCM): defined by IETF RFC 2697
 *    2. Two Rate Three Color Marker (trTCM): defined by IETF RFC 2698
 *
 ***/

#include <stdint.h>

/*
 * Application Programmer's Interface (API)
 *
 */

#define QOS_METER_TIME_SHIFT		8	/* 扩大因子，防止HZ/PIR出现小数 */
#define QOS_METER_CYCLE_SCALE       1

/* Packet Color Set */
enum qos_meter_color {
	QOS_METER_GREEN = 0, 	/* Green */
	QOS_METER_YELLOW,    	/* Yellow */
	QOS_METER_RED,       	/* Red */
	QOS_METER_COLORS     	/* Number of available colors */
};

/*
 * srTCM parameters per metered traffic flow. The CIR, CBS and EBS parameters only
 * count bytes of IP packets and do not include link specific headers. At least one of
 * the CBS or EBS parameters has to be greater than zero. 
 */
struct qos_meter_srtcm_params {
	uint32_t cir; 			/* CIR. Measured in Kbps per second. */
	uint32_t cbs; 			/* CBS. Measured in bytes. */
	uint32_t ebs; 			/* EBS. Measured in bytes. */
};

struct qos_meter_srtcm {
	uint32_t cir;
	uint32_t cbs_bytes;
	
	int64_t cbs;			
	int64_t ebs;
	int64_t tc;   				
	int64_t te;
	uint64_t time; 						/* Time of latest update of C and E token buckets */
			
	/* 时间周期，简单理解为XX字节对应的cycle数.目前使用的是1 << QOS_METER_TIME_SHIFT*/	
	uint64_t cir_period;					
};

/*
 * trTCM parameters per metered traffic flow. The CIR, PIR, CBS and PBS parameters
 * only count bytes of IP packets and do not include link specific headers. PIR has to
 * be greater than or equal to CIR. Both CBS or EBS have to be greater than zero. 
 */
struct qos_meter_trtcm_params {
	uint32_t cir; 			/* CIR. Measured in Kbps per second. */
	uint32_t pir; 			/* PIR. Measured in Kbps per second. */
	uint32_t cbs; 			/* CBS. Measured in byes. */
	uint32_t pbs; 			/* PBS. Measured in bytes. */
};

/**
 * Internal data structure storing the trTCM run-time context per metered
 * traffic flow.
 */
struct qos_meter_trtcm {
	uint32_t cir;
	uint32_t pir;
	uint32_t cbs_bytes;
	uint32_t pbs_bytes;

	int64_t cbs;
	int64_t pbs;
	int64_t tc;					
	int64_t tp;	
	uint64_t time;

	uint64_t cir_period;					
	uint64_t pir_period;
};

int qos_meter_srtcm_config(struct qos_meter_srtcm *srtcm, struct qos_meter_srtcm_params *params);

int qos_meter_trtcm_config(struct qos_meter_trtcm *trtcm, struct qos_meter_trtcm_params *params);

char const *qos_meter_color_str(int color);

static inline int64_t qos_meter_get_time_diff(uint64_t now, uint64_t last_time)
{
	#ifdef QOS_METER_CYCLE_SCALE
	return (now - last_time) << QOS_METER_TIME_SHIFT;
	#else
	return (now - last_time);
	#endif
}

static inline int64_t qos_meter_l2t(uint64_t period, uint32_t bytes)
{
	if (bytes == 0) {
		return 0;
	}

	#ifdef QOS_METER_CYCLE_SCALE
	return (period * bytes);
	#else
	return (period * bytes) >> qos_METER_TIME_SHIFT;
	#endif
}

static inline int64_t qos_meter_t2l(uint64_t period, uint64_t time)
{
	if (period == 0) {
		return 0;
	}

	#ifdef QOS_METER_CYCLE_SCALE
	return (time) / period;
	#else
	return (time << QOS_METER_TIME_SHIFT) / period;
	#endif
}

static inline enum qos_meter_color 
qos_meter_srtcm_color_blind_check(struct qos_meter_srtcm *m, uint32_t pkt_len)
{
	uint64_t now;
	int64_t time_diff, b2t;

	now = qos_get_tsc_cycles();
	
	time_diff = qos_meter_get_time_diff(now, m->time);
	b2t = qos_meter_l2t(m->cir_period, pkt_len);

	/* Put the tokens overflowing from tc into te bucket */
	m->tc += time_diff;
	if (m->tc > m->cbs) {
		m->te += (m->tc - m->cbs);
		if (m->te > m->ebs) {
			m->te = m->ebs;
		}
		
		m->tc = m->cbs;
	}

	m->time = now;

	/* Color logic */
	if (m->tc > 0) {
		m->tc -= b2t;
		return QOS_METER_GREEN;
	}

	if (m->te > 0) {
		m->te -= b2t;
		return QOS_METER_YELLOW;
	}

	return QOS_METER_RED;
}

static inline enum qos_meter_color 
qos_meter_srtcm_color_aware_check(struct qos_meter_srtcm *m, uint32_t pkt_len, enum qos_meter_color pkt_color)
{
	uint64_t now, b2t;
	int64_t time_diff;

	now = qos_get_tsc_cycles();
	if (now < m->time)
    {
		;
	}
	
	/* Bucket update */
	time_diff = qos_meter_get_time_diff(now, m->time);
	m->time = now;

	b2t = qos_meter_l2t(m->cir_period, pkt_len);
		
	/* Put the tokens overflowing from tc into te bucket */
	m->tc += time_diff;
	if (m->tc > m->cbs) {
		m->te += (m->tc - m->cbs);
		if (m->te > m->ebs) {
			m->te = m->ebs;
		}
		
		m->tc = m->cbs;
	}

	/* Color logic */
	
	/* 如果报文已被标记为绿色且tc > 0，则报文被标记为绿色，且tc=tc-B */
	if ((pkt_color == QOS_METER_GREEN) && (m->tc > 0)) {
		m->tc -= b2t;
		return QOS_METER_GREEN;
	}

	/*
 	 * 1、如果报文已被标记为绿色且tc < 0但te > 0，则报文被标记为黄色，te=te-B
 	 * 2、如果报文已被标记为黄色但te > 0，则报文被标记为黄色，且te=te-B
 	 */
	if ((pkt_color != QOS_METER_RED) && (m->te > 0)) {
		m->te -= b2t;
		return QOS_METER_YELLOW;
	}

	/*
 	 * 1、如果报文已被标记为黄色且te < 0，则报文被标记为红色，且te保持不变；
 	 * 2、如果报文已被标记为红色，直接将报文标记为红色，tc和te不变。
 	 */
	return QOS_METER_RED;
}

static inline enum qos_meter_color
qos_meter_trtcm_color_blind_check(struct qos_meter_trtcm *m, uint32_t pkt_len)
{
	uint64_t now;
	int64_t cb2t, pb2t;
	uint64_t time_diff_tc, time_diff_tp;

	now = qos_get_tsc_cycles();
	if (now < m->time)
    {
		;
	}
	
	/* Bucket update */
	time_diff_tc = qos_meter_get_time_diff(now, m->time);
	time_diff_tp = time_diff_tc;
	m->time = now;

	cb2t = qos_meter_l2t(m->cir_period, pkt_len);
	pb2t = qos_meter_l2t(m->pir_period, pkt_len);
	
	m->tc += time_diff_tc;
	if (m->tc > m->cbs) {
		m->tc = m->cbs;
	}

	m->tp += time_diff_tp;
	if (m->tp > m->pbs) {
		m->tp = m->pbs;
	}

	/* Color logic */
	if (m->tp > 0) {
		if (m->tc > 0) {
			m->tc -= cb2t;
			m->tp -= pb2t;
			return QOS_METER_GREEN;
		} else {
			m->tp -= pb2t;
			return QOS_METER_YELLOW;
		}
	} else {
		return QOS_METER_RED;
	}

}

static inline enum qos_meter_color
qos_meter_trtcm_color_aware_check(struct qos_meter_trtcm *m, uint32_t pkt_len, enum qos_meter_color pkt_color)
{
	uint64_t now, cb2t, pb2t;
	uint64_t time_diff_tc, time_diff_tp;

	now = qos_get_tsc_cycles();
	if (now < m->time)
    {
		;
	}
	
	/* Bucket update */
	time_diff_tc = qos_meter_get_time_diff(now, m->time);
	time_diff_tp = time_diff_tc;
	m->time = now;

	cb2t = qos_meter_l2t(m->cir_period, pkt_len);
	pb2t = qos_meter_l2t(m->pir_period, pkt_len);
	
	m->tc += time_diff_tc;
	if (m->tc > m->cbs) {
		m->tc = m->cbs;
	}

	m->tp += time_diff_tp;
	if (m->tp > m->pbs) {
		m->tp = m->pbs;
	}

	/* Color logic */
	
	/* 非色盲模式下
	 * 1、如果报文标记为红色,标识为红色
	 * 2、如果报文是绿色或者黄色,但tp < 0, 标识为红色
	 */
	if ((pkt_color == QOS_METER_RED) || (m->tp < 0)) {
		return QOS_METER_RED;
	}

	/* 非色盲模式下
	 * 1、如果报文已被标记为黄色,且tp > 0，标识为黄色
	 * 2、如果报文已被标记为绿色,且tp > 0 && tc < 0，则报文被标记为黄色
	 */
	if ((pkt_color == QOS_METER_YELLOW) || (m->tc < 0)) {
		m->tp -= pb2t;
		return QOS_METER_YELLOW;
	}

	/* 非色盲模式下,如果报文已被标记为绿色且Tc > 0，则报文被标记为绿色 */
	m->tc -= cb2t;
	m->tp -= pb2t;

	return QOS_METER_GREEN;
}

#ifdef __cplusplus
}
#endif

#endif /* __INCLUDE_QOS_METER_H__ */

char const *qos_meter_color_str(int color)
{
	if (color == QOS_METER_GREEN) {
		return "GREEN";
	} else if (color == QOS_METER_YELLOW) {
		return "YELLOW";
	} else if (color == QOS_METER_RED) {
		return "RED";
	} else {
		return "NONE";
	}
}

static void qos_meter_get_tb_params(uint64_t rate_bytes, uint64_t *tb_period)
{
	uint64_t hz;

	if (rate_bytes == 0) {
		return;
	}

	hz = qos_get_tsc_hz();

	*tb_period = (hz << QOS_METER_TIME_SHIFT) / rate_bytes;

	return;
}

int qos_meter_srtcm_config(struct qos_meter_srtcm *srtcm, struct qos_meter_srtcm_params *params)
{
	uint64_t rate_bytes;
	
	/* Check input parameters */
	if ((srtcm == NULL) ||
		(params == NULL) ||
		(params->cir == 0)) {
		return -EINVAL;
   	}

	/* Initialize srTCM run-time structure */
	srtcm->cir = params->cir;	
	rate_bytes = ((uint64_t)params->cir * 1000) >> 3;
	if (params->cbs == 0) {
		params->cbs = rate_bytes / 10; /* 100ms */
	}

	srtcm->cbs_bytes = params->cbs;
	
	qos_meter_get_tb_params(rate_bytes, &srtcm->cir_period);
	
	srtcm->cbs = srtcm->tc = qos_meter_l2t(srtcm->cir_period, params->cbs);
	srtcm->ebs = srtcm->te = qos_meter_l2t(srtcm->cir_period, params->ebs);

	srtcm->time = qos_get_tsc_cycles();
	
	return 0;
}

int qos_meter_trtcm_config(struct qos_meter_trtcm *trtcm, struct qos_meter_trtcm_params *params)
{
	uint64_t rate_bytes;

	/* Check input parameters */
	if ((trtcm == NULL) ||
		(params == NULL) ||
		(params->cir == 0) ||
		(params->pir == 0) ||
		(params->pir < params->cir)) {
		return -EINVAL;
	}

	/* Initialize trTCM run-time structure */
	trtcm->cir = params->cir;
	trtcm->pir = params->pir;
	
	rate_bytes = ((uint64_t)params->cir * 1000) >> 3;
	qos_meter_get_tb_params(rate_bytes, &trtcm->cir_period);
	if (params->cbs == 0) {
		params->cbs = rate_bytes / 10; 		
	}
	rate_bytes = ((uint64_t)params->pir * 1000) >> 3;
	qos_meter_get_tb_params(rate_bytes, &trtcm->pir_period);
	if (params->pbs == 0) {
		params->pbs = (rate_bytes >> 3) / 10;
	}

	trtcm->pbs_bytes = params->pbs;
	trtcm->cbs_bytes = params->cbs;
	
	trtcm->cbs = trtcm->tc = qos_meter_l2t(trtcm->cir_period, params->cbs);
	trtcm->pbs = trtcm->tp = qos_meter_l2t(trtcm->pir_period, params->pbs);
	
	trtcm->time = qos_get_tsc_cycles();
		
	return 0;
}