令牌桶封装成库示例

token_bucket.h文件(主要用来函数声明)

#ifndef __TOKEN_BUCKET_H
#define __TOKEN_BUCKET_H
/*
 * 实现令牌桶
 *      token  令牌
 *      cps    速率
 *      burst  上限
 */

#define TBF_MAX    1024

//创建令牌桶
int tbf_init(int cps, int burst);

//取令牌
int tbf_fetch(int id, int size);

//还令牌
int tbf_return(int id, int size);

//销毁令牌桶
int tbf_destory(int id);

#endif

token_bucket.c文件

#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <unistd.h>
#include <sys/time.h>
#include "token_bucket.h"

//定义桶的类型
struct tbf_st
{
    int cps;
    int burst;
    int token;
};

//定义1024的数组存桶的结构体地址
static struct tbf_st *line[TBF_MAX];
static int inited = 0;
struct sigaction oldact;
struct itimerval olditv;

//模块卸载，信号行为和时钟恢复
static void  mouder_unload(void)
{
    sigaction(SIGALRM, &oldact, NULL);
    setitimer(ITIMER_REAL, &olditv, NULL);
}

/******************************************/
static void alrm_handler(int s)
{
    for(int i=0; i<TBF_MAX; i++)
    {
        if(line[i] != NULL)
        {
            line[i] -> token += line[i] -> cps;
            if(line[i]->token >= line[i] -> burst)
                line[i] -> token = line[i] -> burst;
        }
    }
}

static void mouder_load(void)
{
    struct sigaction act;
    struct itimerval itv;

    act.sa_handler = alrm_handler;
    act.sa_flags = 0;
    sigemptyset(&act.sa_mask);
    
    sigaction(SIGALRM, &act, &oldact);
    
    itv.it_interval.tv_sec = 1;
    itv.it_interval.tv_usec = 0;

    itv.it_value.tv_sec = 1;
    itv.it_value.tv_usec = 0;

    setitimer(ITIMER_REAL, &itv, &olditv);

    atexit(mouder_unload);
}


//创建令牌桶
int tbf_init(int cps, int burst)
{
    struct tbf_st *new;
    int i;

    if(inited == 0)
    {
        mouder_load();
        inited = 1;
    }

    //找空位
    for(i=0; i<TBF_MAX; i++)
    {
        if(line[i] == NULL)
            break;
        return -1;
    }

    new = malloc(sizeof(*new));
    if( new == NULL)
        return -1;

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

    line[i] = new;

    return i;
}

/**********************************************/
//找最小值
static int min(int a,int b)
{
    return a>b?a:b;
}

//取令牌
int tbf_fetch(int id, int size)
{
    int n;
    if(size <= 0)
        return -1;
    while(line[id]->token <= 0)
    {
        pause();
    }

    //取两者最小
    n = min(line[id]->token,size);

    line[id]->token -= n;

    return 0;
}
/**********************************************/

//还令牌
int tbf_return(int id, int size)
{
    if(size <= 0)
        return -1;
    line[id]->token += size;
    if(line[id]->token > line[id]->burst)
        line[id]->token = line[id]->burst;
    return size;
}

//销毁令牌桶
int tbf_destory(int id)
{
    if(id < 0 || id > TBF_MAX-1)
    {
        return -1;
    }
    free(line[id]);
    line[id] = NULL;
    return 0;
}

调试文件main.c

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

#include "token_bucket.h"

#define BUFSIZE 10
#define CPS BUFSIZE 
#define BURST CPS*10

int main(int argc, char **argv)
{
    int rfd, wfd;
    char buf[BUFSIZE] = {};
    int cnt, ret, pos;
    int tb,n;

    if (argc < 2) 
    {
        fprintf(stderr, "Usage....\n");
        return 1;
    }

    while (1) 
    {
        rfd = open(argv[1], O_RDONLY);
        if (rfd == -1) 
        {
            if (errno == EINTR)
            {
                // 信号打断---》假错
                continue;
            }
            perror("open()"); // 输出到标准错误输出
            return 1;
        }
        break;
    }

    tb = tbf_init(CPS, BURST);

    if(tb < 0)
    {
        fprintf(stderr, "tbf_init():");
        exit(1);
    }

    wfd = 1;

    while (1)
    {
        n = tbf_fetch(tb, CPS);
        if(n < 0)
        {
            fprintf(stderr, "tbf_fetch():");
            goto ERROR1;
        }

        cnt = read(rfd, buf, CPS); // 阻塞
        if (-1 == cnt) {
            if (errno == EINTR)
                continue;
            perror("read()");
            goto ERROR1;
        }
        if (0 == cnt) {
            // rfd结束标志
            break;
        }
        pos = 0;
        while (1) {
            ret = write(wfd, buf+pos, cnt);
            if (ret < 0) {
                if (errno == EINTR)
                    continue;
                perror("write()");
                goto ERROR1;
            }
            cnt -= ret;
            if (cnt > 0) {
                // 没写完
                pos += ret;
            } else {
                break;
            }
        }
    }

    tbf_destory(tb);

    close(rfd);

    return 0;
ERROR1:
    close(rfd);
    return 1;
}

makefile文件

SRC=main.o token_bucket.o
OBJ=tbf

$(OBJ):$(SRC)
    gcc -o $@ $^

clean:
    rm *.o $(OBJ)