令牌桶算法实现

令牌桶算法最初来源于计算机网络。

在网络传输数据时，为了防止网络拥塞，需限制流出网络的流量，使流量以比较均匀的速度向外发送。

令牌桶算法是网络流量整形（Traffic Shaping）和速率限制（Rate Limiting）中最常使用的一种算法。

流程
产生令牌：周期性的以一定速率往令牌桶中增加令牌。

如果桶中的令牌数达到桶容量，丢弃多余令牌。

消耗令牌：接受请求或输入数据时会消耗桶中的令牌。

以请求或消息为单位时，可以一次消耗一个令牌。

在网络传输中，消耗令牌的数量可以根据数据包的大小决定。

是否通过：桶中的令牌 >= 所需令牌时，请求或者数据包通过，否者被限流。

对于被限流的请求或者数据包，可以有不同的处理方式。

（1、直接丢弃；2、进队列等待； 3、可以通过，但需要做特殊标记）

//mytdf.h
#ifndef MYTBF_H__
#define MYTBF_H__

#define MYTBF_MAX  1024 //能够支持1024个不同速率的令牌桶

typedef void mytbf_t;

mytbf_t *mytbf_init(int, int);

int mytbf_fetchtoken(mytbf_t *, int);

int mytbf_returntoken(mytbf_t *, int);

int mytbf_destroy(mytbf_t*);

#endif

//mytbf.c
#include "mytbf.h"
#include 
#include 
#include 
#include 

struct mytbf_t
{
	int token;
	int cps;
	int burst;
	int pos;
};

//令牌桶数组，全局变量，元素全部自动初始化为NULL
static struct mytbf_t *job[MYTBF_MAX];
static int inited = 0;

typedef void(*sighandler_t)(int);
static sighandler_t alarm_handler_save;

static int get_free_pos(void)
{
	int i;
	for(i = 0; i < MYTBF_MAX; i++)
	{
		if(job[i] == NULL)
			return i;
	}
	
	return -1;
}

static int min(int a, int b)
{
	if(a < b)
		return a;
	return b;
}

static void alrm_handler(int s)
{
	int i;
	
	alarm(1);
	
	for(i = 0; i < MYTBF_MAX; i++)
	{
		if(job[i] != NULL)
		{
			job[i]->token += job[i]->cps;
			if(job[i]->token > job[i]->burst)
				job[i]->token = job[i]->burst;
		}
	}
}

//恢复现场
static void module_unload(void)
{
	int i;
	
	//恢复signal先前行为，关闭闹钟，移除模块
	signal(SIGALRM, alarm_handler_save);
	alarm(0);
	for(i = 0; i < MYTBF_MAX; i++)
		free(job[i]);
}

static void module_load(void)
{
	alarm_handler_save = signal(SIGALRM, alrm_handler);
	alarm(1);
	
	atexit(module_unload);
}

mytbf_t *mytbf_init(int cps, int burst)
{
	struct mytbf_t *me;
	
	int pos;
	
	if(!inited)
	{
		module_load();
		inited = 1;
	}
	
	pos = get_free_pos();
	if(pos < 0)
		return NULL;
	
	me = malloc(sizeof(*me));
	if(NULL == me)
		return NULL;
	
	me->token = 0;
	me->cps = cps;
	me->burst = burst;
	me->pos = pos;
	
	job[pos] = me;
	
	return me;
}

/*
从令牌桶中获取令牌数
参数：
mytbf *：要获取令牌的对象
size：要获取的令牌数
返回值：
获取到的真正令牌数
*/
int mytbf_fetchtoken(mytbf_t *ptr, int size)
{
	int n;
	
	if(size <= 0)
		return -EINVAL;
	
	struct mytbf_t *me = ptr;
	
	while(me->token <= 0)
		pause();
	
	n = min(me->token, size);
	
	me->token -= n;
	
	return n;
}

/*
参数：
mytbf *：要归还令牌的对象
size：要归还的令牌数
返回值：
归还的令牌数
*/
int mytbf_returntoken(mytbf_t *ptr, int size)
{
	struct mytbf_t *me = ptr;
	
	if(size <= 0)
		return -EINVAL;
		
	me->token += size;
	if(me->token > me->burst)
		return me->burst;
	
	return size;
}

int mytbf_destroy(mytbf_t *ptr)
{
	 struct mytbf_t *me = ptr;
	 
	 job[me->pos] = NULL;
	 free(me);
	 
	 return 0;
}

//main.c
#include  
#include  
#include  
#include  
#include  
#include 
#include 
#include 
#include "mytbf.h"
#include 

#define CPS 10
#define BUFFSIZE 1024
#define BURST 100 

int main(int argc, char* argv[]) 
{ 
	if(argc < 2) 
	{ 
		fprintf(stderr, "usage: %s ", argv[0]);
		exit(1);
    }
    
    int sfd, dfd = 1;
    int len, wet, pos, size;
    char buf[BUFFSIZE];
	mytbf_t *tbf;
	
	tbf = mytbf_init(CPS, BUFFSIZE);
	if(NULL == tbf)
	{
		fprintf(stderr, "mytbf_init() failed!\n");
		exit(1);
	}

	do
	{
		sfd = open(argv[1], O_RDONLY);
		if(sfd < 0)
		{
			if(errno != EINTR)
			{
				perror("open src_file");
				exit(1); 
			}
		}
	}while(sfd < 0);
   
	
   
    while(1)
    {
		size = mytbf_fetchtoken(tbf, BUFFSIZE);
		if(size < 0)
		{
			fprintf(stderr, "mytbf_fetchtoken():%s\n", strerror(-size));
			exit(1);
		}
		
        while((len = read(sfd, buf, size)) < 0)
		{
			if(errno == EINTR)
				continue;
            perror("read src_file");
            break;
		}
      
        if(len == 0)
        {
            break;    
        }
        
		//最后一次取的令牌数可能并没有全部用完，就把文件内容读取完毕
		if(size - len > 0)
			mytbf_returntoken(tbf, size-len);
        
		pos = 0;
		while(len > 0)
		{
			wet = write(dfd, buf+pos, len);
			if(wet < 0)
			{
				if(errno == EINTR)
					continue;
				perror("write des_file");
				exit(1);
			}
			pos += wet;
			len -= wet;
		} 	
    }
     
    close(sfd);
	
	mytbf_destroy(tbf);
	
    exit(0);
}