令牌桶库实现

实现及管理一个最大1024个不同令牌桶(通过数组管理也可通过链表管理)

头文件：mytbf.h

#ifndef __MYTBF_H__
#define __MYTBF_H__

#define MAXCOUNT 1024　　//设置最大令牌桶数

typedef void tbf_st;　　//定义数据类型

tbf_st *mytbf_init(int cps,int burst);　　//初始化令牌桶
int mytbf_reversSzie(tbf_st *,int);　　//申请令牌
int mytbf_returnSzie(tbf_st *,int);　　//归还令牌
int mytbf_destory(tbf_st *);　　　　//销毁令牌桶

#endif

库文件：mytbf.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <errno.h>
#include <unistd.h>
#include “mytbf.h”

typedef void (*sighandler_t)(int);　　//定义数据类型
sighandler_t sig_save;　　　　　　//保存信号原始状态

static int boot = 0;　　　　//启动标志

struct mytbf_st　　　　//令牌桶结构体
{
　　int cps;　　   //每秒增加值
　　int burst;　　//资源上线
　　int token;　 //当前资源
　　int pos;　　//数组中存放位置
};

static struct mytbf_st *tbf[MAXCOUNT];  //令牌桶指针数组

static void alrm_handler(int s)　　//闹钟处理函数
{
　　int i;
　　alarm(1);　　//循环闹钟设置
　　for(i=0;i<MAXCOUNT;i++)　　//闹钟时间到查询令牌桶数组为每个令牌桶增加令牌资源并检测是否到资源上线
　　{
　　　　if(tbf[i] !=NULL)
　　　　{
　　　　　　tbf[i]->token +=tbf[i]->cps;
　　　　　　if(tbf[i]->token >= tbf[i]->burst)
　　　　　　　　tbf[i]->token = tbf[i]->burst;
　　　　}
　　}
}

static void modle_load(void)　　//加载模块时执行
{
　　sig_save = signal(SIGALRM,alrm_handler);　　//注册闹钟处理函数并保存原始处理函数
　　alarm(1);
}

static void modle_unload(void)　　//卸载模块时执行
{
　　int i;

　　signal(SIGALRM,sig_save);　　//恢复闹钟原始处理函数
　　alarm(0);　　　　　　//关闭闹钟

　　for(i=0;i<MAXCOUNT;i++)　　//循环释放申请的令牌桶空间
　　{
　　　　if(tbf[i] != NULL)
　　　　　　free(tbf[i]);
　　}
}

static int get_tbf_pos(void)　　//查询令牌桶指针数组为空地址
{
　　int i=0;
　　for(i=0;i<MAXCOUNT;i++)
　　{
　　　　if(tbf[i] == NULL)
　　　　return i;
　　}

　　return -1;
}

tbf_st *mytbf_init(int cps,int burst)　　//初始化令牌桶
{
　　int pos;
　　struct mytbf_st *new;

　　if(boot ==0)　　//如果首次执行那么加载闹钟模块
　　{
　　　　modle_load();
　　　　boot =1;
　　}
　　atexit(modle_unload);　　//挂上钩子函数出错时释放资源

　　pos = get_tbf_pos();　　//获取令牌桶指针数组空地址
　　if(pos <0)
　　　　return NULL;

　　new = malloc(sizeof(mytbf));　　//申请空间创建令牌桶并初始化
　　if(new ==NULL)
　　　　return NULL;

　　new->cps = cps;
　　new->burst = burst;
　　new->token =0;
　　new->pos = pos;

　　tbf[pos] = new;　　//将令牌桶添加到令牌桶指针数组中
　　return new;
}

static int min(int a,int b)　　//计算两个数的最小值
{
　　if(a>b)
　　　　return b;
　　return a;
}

int mytbf_reversSzie(tbf_st *ptr,int size)　　//向指定令牌桶申请令牌
{
　　struct mytbf_st *me = ptr;
　　int ret;

　　if(size <=0)
　　　　return -EINVAL;

　　ret = min(size,me->token);　　//申请令牌数与实际持有令牌数取最小值
　　me->token -=ret;

　　return ret;
}

int mytbf_returnSzie(tbf_st *ptr,int size)　　//向指定令牌桶归返令牌
{
　　struct mytbf_st *me = ptr;

　　if(size<=0)
　　　　return -EINVAL;

　　if((me->token+size) >= me->burst )　　//当持有令牌+归返令牌大于资源上线时为资源上线
　　　　me->token = me->burst;
　　else
　　　　me->token +=size;

　　return size;
}

int mytbf_destory(tbf_st *ptr)　　//销毁执行令牌桶
{
　　struct mytbf_st *me = ptr;

　　tbf[me->pos] = NULL;　　//将令牌指针数组指定位置空
　　free(me);　　//释放令牌桶申请的空间

　　return 0;
}

主函数：main.c

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <errno.h>
#include “mytbf.h”

#define CPS 10　　　　　　//每秒增加的令牌数
#define MAXSZIE 1024　　//资源上线

int main(int argc, char**argv)
{
　　int fd;
　　int readcnt;
　　int writecnt;
　　int pos;
　　char readbuff[MAXSZIE];
　　tbf_st *my_tbf;　　　　　　//创建令牌指针
　　int getcnt;

　　if(argc<2)
　　{
　　　　fprintf(stderr,”Usage:%s <filename>/n”,argv[0]);
　　　　exit(1);
　　}

　　my_tbf = mytbf_init(CPS,MAXSZIE);　　//初始化令牌桶

　　while(1)
　　{
　　　　fd = open(argv[1],O_RDONLY);
　　　　if(fd <0 && errno !=EINTR)
　　　　{
　　　　　　perror(“open()”);
　　　　　　exit(1);
　　　　}
　　　　else if(fd >=0)
　　　　　　break;
　　}

　　while(1)
　　{
　　　　while((getcnt = mytbf_reversSzie(my_tbf,MAXSZIE)) <= 0)　　//当申请的令牌资源小于等于0时
　　　　　　pause();　　//等待

　　　　while((readcnt = read(fd, readbuff, getcnt))<0)
　　　　{
　　　　　　if(errno!=EINTR)
　　　　　　{
　　　　　　　　perror(“read()”);
　　　　　　　　exit(1);
　　　　　　}
　　　　　　else
　　　　　　　　continue;
　　　　}
　　　　if(readcnt ==0)
　　　　　　break;

　　　　if(readcnt< getcnt)　　//当申请的令牌未使用完时
　　　　　　mytbf_returnSzie(my_tbf,getcnt-readcnt);　　//归返未使用完的令牌

　　　　pos = 0;
　　　　while(readcnt>0)
　　　　{
　　　　　　writecnt = write(1,readbuff+pos,readcnt);
　　　　　　if(writecnt<0)
　　　　　　{
　　　　　　　　if(errno == EINTR)
　　　　　　　　　　continue;
　　　　　　　　perror(“write()”);
　　　　　　　　exit(1);
　　　　　　}
　　　　　　pos +=writecnt;
　　　　　　readcnt -=writecnt;
　　　　}
　　}

　　mytbf_destory(my_tbf);　　//销毁当前令牌桶
　　close(fd);　　
　　exit(0);
}

编译文件：Makefile

mytbf:mytbf.o main.o
　　$(CC) $^ -o $@

clean:
　　$(RM) *.o mytbf