C语言基于令牌桶算法实现slowcat命令

一、功能概述

封装一个令牌桶模块，实现cat命令的慢输出版本，即根据指定每秒输出的字节数，依次文件内容，直到文件末尾

二、令牌桶模块

1、定义结构体指针数组，用于存放多个令牌桶指针

struct mytbf_st *job[MYTBF_MAX];

2、定义结构体，用来存储令牌桶

struct mytbf_st
{
    int cps; // 每秒产生的令牌数
    int burst; // 令牌上限
    int token; // 令牌总数
    int pos; // 保存在数组中的下标，便于检索
};

3、定义头文件 mytbf.h ，声明对外暴露的接口

#ifndef _MYTBF_H_
#define _MYTBF_H_

// 数组容量
#define MYTBF_MAX   1024

// 隐藏实现
typedef void mytbf_t;

mytbf_t *mytbf_init(int cps, int burst);

// 取令牌，返回取到的膈俞
int mytbf_fetchtoken(mytbf_t *, int);
// 归还，返回还了多少
int mytbf_returntoken(mytbf_t *, int);

int mytbf_destroy(mytbf_t *);

#endif

4、具体实现 mytbf.c

#include <stdlib.h>
#include <stdio.h>
#include <errno.h>
#include <signal.h>
#include <unistd.h>

#include "mytbf.h"

typedef void (*__sighandler_t) (int);

static struct mytbf_st *job[MYTBF_MAX];
static int inited = 0;
// 保存调用该模块前SIGALRM的默认策略，用于卸载模块后进行复原
static __sighandler_t alarm_handler_save;

struct mytbf_st
{
    int cps;
    int burst; // 上限
    int token;
    int pos;
};

// 获取数组中的空闲位置
static int get_free_pos(void)
{
    int i;
    for (i = 0; i < MYTBF_MAX; i++)
    {
        if (job[i] == NULL)
            return i;
    }
    return -1;
}

// 定义SIGALRM信号处理策略
static void alarm_handler(int s)
{
    int i;
    alarm(1);
    for (i = 0; i < MYTBF_MAX; i++)
    {
        if (job[i] == NULL)
            continue;
        job[i]->token += job[i]->cps;
        if (job[i]->token > job[i]->burst)
            job[i]->token = job[i]->burst;
    }
}

// 模块卸载
static void module_unload()
{
    int i;
    // 复原信号策略
    signal(SIGALRM, alarm_handler_save);
    // cancel
    alarm(0);
    for (i = 0; i < MYTBF_MAX; i++)
    {
        free(job[i]);
        job[i] = NULL;
    }
}

// 注册闹钟信号的策略，并构建alarm链
static void module_load()
{
    alarm_handler_save = signal(SIGALRM, alarm_handler);
    alarm(1);
    // 挂载钩子函数，便于在退出进程后进行资源释放和信号策略复原
    atexit(module_unload);
}

// 初始化令牌桶
mytbf_t *mytbf_init(int cps, int burst)
{
    struct mytbf_st *me;
    if (!inited)
    {
        module_load();
        inited = 1;
    }

    me = (struct mytbf_st *)malloc(sizeof(*me));
    if (me == NULL)
        return NULL;

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

    int pos = get_free_pos();
    if (pos < 0)
    {
        free(me);
        return NULL;
    }
    me->pos = pos;
    job[pos] = me;

    return me;
}

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

// 取令牌，返回取到的数量
int mytbf_fetchtoken(mytbf_t *ptr, int size)
{
    int n;
    struct mytbf_st *me = (struct mytbf_st *)ptr;

    if (size <= 0)
        return -EINVAL;

    while (me->token <= 0)
        pause();

    // 取其小
    n = min(me->token, size);
    me->token -= n;

    return n;
}

// 如果令牌没有消费完，则进行归还
int mytbf_returntoken(mytbf_t *ptr, int size)
{
    struct mytbf_st *me = (struct mytbf_st *)ptr;

    if (size <= 0)
        return -EINVAL;

    me->token += size;
    if (me->token > me->burst)
        me->token = me->burst;

    return size;
}

int mytbf_destroy(mytbf_t *ptr)
{
    struct mytbf_st *me = (struct mytbf_st *)ptr;

    job[me->pos] = NULL;
    free(me);

    return 0;
}

三、通过令牌桶模块，实现slowcat.c

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

#include "mytbf.h"

#define CPS     10
#define BUFSIZE 1024
#define BURST 100


int main(int argc, char **argv)
{
    int sfd, dfd = 1;
    int len = 0;
    int ret = 0;
    int pos = 0;
    int size = 0;
    char buf[BUFSIZE];
    mytbf_t *tbf;

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

    tbf = mytbf_init(CPS, BURST);
    if(tbf == NULL)
    {
        fprintf(stderr, "mytbf_init() failed!\n");
        exit(1);
    }

    do
    {
        sfd = open(argv[1], O_RDONLY);
        if(sfd < 0)
        {
            if(errno != EINTR)
            {
                perror("open()");
                exit(1);
            }
        }
    }while (sfd < 0);
    

    while(1)
    {
        size = mytbf_fetchtoken(tbf, BUFSIZE);
        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()");
            break;
        }

        if(len == 0)
            break;

        if(size - len > 0)
            mytbf_returntoken(tbf, size - len);

        pos = 0;
        while(len > 0)
        {
            // printf("pos=%d len=%d\n", pos, len);
            ret = write(dfd, buf + pos, len);
            // printf("ret=%d\n", ret);
            if(ret < 0)
            {
                if(errno == EINTR)
                    continue; 
                perror("write()");
                exit(1);
            }
            pos += ret;
            len -= ret;
        }
        
    }

    close(sfd);

    mytbf_destroy(tbf);

    exit(0);
}

附录：makefile

all:slowcat

slowcat:slowcat.o mytbf.o
	gcc $^ -o $@

clean:
	rm -rf *.o slowcat