令牌桶的实现 使用多线程互斥量实现

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

mytbf.c

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <unistd.h>
#include <string.h>
#include <pthread.h>
#include "mytbf.h"


struct mytbf_st
{
    
    int cps;
    int token;
    int burst;
    int pos;
        pthread_mutex_t mut;
};
static struct mytbf_st* job[MYTBF_MAX];
static pthread_mutex_t mut_job = PTHREAD_MUTEX_INITIALIZER;


static sighandler_t alrm_handler_save;
static pthread_t tid_alrm;
static pthread_once_t init_once = PTHREAD_ONCE_INIT;
static void *thr_alrm(void *p)
{
        while(1)
        {
                pthread_mutex_lock(&mut_job);
                for( int i = 0; i < MYTBF_MAX; i++)
                {
                        if( job[i] != NULL)
                        {
                                pthread_mutex_lock(&job[i]->mut);
                                job[i]->token += job[i]->cps;
                                if( job[i]->token > job[i]->burst)
                                        job[i]->token = job[i]->burst;
                                pthread_mutex_unlock(&job[i]->mut);
                        }
                }
                pthread_mutex_unlock(&mut_job);
                sleep(1);
        }
   
}
static void module_unload()
{
  
        pthread_cancel(tid_alrm);
        pthread_join(tid_alrm,NULL);
    for( int i = 0; i < MYTBF_MAX;i++)
    {
                if( job[i] != NULL)
                {
                        mytbf_destroy(job[i]);
                }
    }
    pthread_mutex_destroy(&mut_job);
}
static void module_load()
{
   
   int err = pthread_create(&tid_alrm,NULL,thr_alrm,NULL);
   if( err)
   {
           fprintf(stderr,"pthread_create() err  %s \n",strerror(err));
           exit(1);
   }

    atexit(module_unload);
}
static int get_free_pos_unlocked(void)
{
    int i = 0;
    for( ;i < MYTBF_MAX; i++)
    {
        if( job[i] == NULL)
        {
            
            return i;
        }
    }
    return -1;
}
mytbf_t * mytbf_init(int cps,int burst)
{
    struct mytbf_st* me;
    pthread_once(&init_once,module_load);
    me = malloc(sizeof(*me));
    if( me == NULL)
    {
        return NULL;
    }


    me->token= 0;
    me->cps = cps;
    me->burst = burst;
        pthread_mutex_init(&me->mut,NULL);
        pthread_mutex_lock(&mut_job);
    int pos = get_free_pos_unlocked();
    if(pos <0)
        {
                pthread_mutex_unlock(&mut_job);
                free(me);
                return NULL;
        }

    me->pos = pos;
    job[pos] = me;
        pthread_mutex_unlock(&mut_job);
    return me;
}
static int min(int a,int b)
{
    return a > b? b: a;
}
int mytbf_fetchtoken(mytbf_t*ptr ,int size)
{
    struct mytbf_st * me = ptr;
       if( size <=0)
       {
           return -EINVAL;
       }
           pthread_mutex_lock(&me->mut);
        while(me->token <=0)
       {
                   pthread_mutex_unlock(&me->mut);
                   sched_yield();
                   pthread_mutex_lock(&me->mut);

           }

          
        int n = min(me->token , size);
        me->token -= n;
                pthread_mutex_unlock(&me->mut);
        return n;
}

int mytbf_returntoken(mytbf_t*ptr,int size)
{
    struct mytbf_st * me = ptr;
    if( size <= 0)
    {
        return -EINVAL;
    }
        pthread_mutex_lock(&me->mut);
    me->token += size;
    if( me->token > me->burst)
        me->token = me->burst;
        pthread_mutex_unlock(&me->mut);
    return size;
}
int mytbf_destroy(mytbf_t *ptr)
{
     struct mytbf_st * me = ptr;
         pthread_mutex_lock(&mut_job);
         job[me->pos] = NULL;
         pthread_mutex_unlock(&mut_job);
         pthread_mutex_destroy(&me->mut);
     free(ptr);
     return 0;
}

main.c

#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <unistd.h>
#include <stdbool.h>
#include <errno.h>
#include <fcntl.h>
#define CPS     100
#define BUFFSIZE  1024
#define  BURST   1000
volatile static int token =0;

#include "mytbf.h"
int main( int argc, char** argv)
{


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

        exit(0);
    }
    

    int sfd,dfd = 1;
    char buff[BUFFSIZE];
    mytbf_t *tbf;
    int size;
    tbf = mytbf_init(CPS,BURST);
    if( tbf == NULL)
    {   
        fprintf(stderr,"mytbf_init() error");
        exit(0);
    }
    if( (sfd = open(argv[1],O_RDONLY))<0)
    {
        perror("open()");
        exit(1);
    }
/*      if( (dfd = open("/tmp/out",O_RDWR)) <0)
        {
                perror("open()");
                exit(1);
        }
        */
    while(1)
    {
      size = mytbf_fetchtoken(tbf,BUFFSIZE);
      ssize_t read_size;
      errno =0;
      while((read_size =  read(sfd,buff,size)) < 0)
      {
       
            if( errno == EINTR)
          {
                continue;
         }
      }
      if( size - read_size > 0)
           mytbf_returntoken(tbf,size - read_size);
  
    
//      read_size =  read(sfd,buff,BUFFSIZE);
      
      if( read_size == 0)
          break;
      if( read_size < 0)
          break;
      ssize_t write_size = 0;
      errno = 0;
      while( read_size)
      {
          ssize_t ret = write(dfd,buff+write_size,read_size);
          if( ret < 0 && errno != EINTR)
          {
              break;
          }
          errno =0;
          read_size -= ret;
      }
    }


    mytbf_destroy(tbf);
    close(sfd);
    exit(0);

}