C语言读写伯克利DB 4

因为缓存数据的buffer总是不够大（会引起段错误）索性从堆上拿了两块大内存

/*
功能说明：逐日存储来访用户（使用伯克利DB）
根据存储的用户信息确定某用户是否是首次来访用户（未被存储的伯克利DB）

调用方式1: 查询游客在2013年8月8日是否访问了指定渠道
./channeldb -s "bch2000 guest:123456789" 20130808
调用方式2：将指定文件里的用户信息写入DB，同时将该文件里的首次访问用户写入日志
./channeldb -f 20130809 ./clog/20130809.log 20130809
*/

#include <assert.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <db.h>
#include <sys/types.h>
#include <getopt.h>

#define DATABASE "/mnt/disk1/ucshuqi/touch/userlist/historydb/channel.db"
#define YES 1
#define NO  0
#define BUFFER_SIZE 1024 * 8

char *readBuffer = NULL;
char *writeBuffer = NULL;

/* ViewData 组件：记录
    用户第一次来访时需要记录的数据，例如日期
 */

struct ViewData {
        int date;
};

void setDate(struct ViewData *data,char *s)
{
        assert(data!=NULL && s!= NULL && strlen(s) == 8);
        data->date = atoi(s);
}

/*当 query.date >= stored.date 返回YES，即包含此附属数据的用户信息是历史来访用户*/
int isHistoryViewInfo(struct ViewData *query , struct ViewData *stored)
{
        assert(query != NULL && stored != NULL);
        printf("query date is %d , stored date is %d \n",query->date, stored->date);

        if(query->date >= stored->date)
        {
                return YES;
        }
        else
        {
                return NO;
        }
}

void printViewData(struct ViewData *data)
{
        assert(data != NULL);
        printf("print view date : %d\n",data->date);
}


/* string helper module */

char *trim(char *s)
{
        int i;

        assert(s!=NULL);
        i = strlen(s);
        for(;i>0;i--)
        {
                if(s[i]==' ' || s[i]=='\n' || s[i]=='\0' || s[i]=='\t')
                {
                        s[i] = '\0';
                }
                else
                {
                        break;
                }
        }
        return s;
}

/* 数据库访问 组件*/

DB *openDb()
{
        int ret;
        DB *dbp = NULL;

        ret = db_create(&dbp, NULL, 0);
        if(ret != 0)
        {
                fprintf(stderr,"create Db error!\n");
                exit(1);
        }

        ret = dbp->open(dbp, NULL, DATABASE, NULL, DB_BTREE, DB_CREATE,0664);

        if(ret != 0)
        {
                fprintf(stderr,"open Db error!\n");
                exit(1);
        }

        return dbp;
}

/* 业务逻辑 */

/* 存储用户访问信息和附属信息viewdata */
int saveViewInfo(DB *dbp, char *user, struct ViewData *data)
{
        DBT key,value;

        assert(dbp != NULL && user != NULL && data != NULL);

        memset(&key, 0, sizeof(key));
        key.data = user;
        key.size = strlen(user) + 1;

        memset(&value, 0, sizeof(value));
        value.data = data;
        value.size = sizeof(*data);
        if(dbp->put(dbp, NULL, &key, &value, 0) == 0)
        {
                //printf("save ---> %s, %d\n", key.data, key.size);
                printViewData(value.data);
                return YES;
        }
        else
        {
                return NO;
        }
}

/*查找用户访问信息，并将附属信息写入data*/
int findViewInfo(DB *dbp, char *user, struct ViewData *data)
{
        DBT key,value;

        assert(dbp != NULL && user != NULL && data != NULL);

        memset(&key, 0, sizeof(key));
        key.data = user;
        key.size = strlen(user) + 1;

        memset(&value,0,sizeof(value)); // must !!!

        if(dbp->get(dbp, NULL, &key, &value, 0) == 0)
        {
                memcpy((char*)data,(char*)value.data,value.size);
                //printViewData(data);
                return YES; //view in database 
        }
        else
        {
                return NO; //view not in database
        }
}

/*记录用户信息，如果用户是首次来访*/
void recordNewView(FILE *fp, char *user, char *buffer)
{
        //char buffer[BUFFER_SIZE] = {0};                                                                
        assert(fp != NULL && user != NULL && buffer != NULL);
        printf("%s is new view\n", user);
        sprintf(buffer, "%s found\n", user);                                               
        fwrite(buffer, sizeof(char), strlen(buffer), fp);
}

/*存储访问信息到数据库，同时写入附属信息*/
void saveViewFile(char *from, char *to, struct ViewData *writeData)
{
        DB *dbp = openDb();
        struct ViewData stored;
        FILE *fp = fopen(from,"r");
        FILE *fpResult = fopen(to,"w");

        assert(fp != NULL && fpResult != NULL && writeData != NULL);
        assert(readBuffer != NULL && writeBuffer != NULL);

        /* clear global buffer content */
        memset(readBuffer, 0, BUFFER_SIZE); 
        memset(writeBuffer, 0, BUFFER_SIZE); 

        while(fgets((char*)readBuffer, BUFFER_SIZE, fp)!=NULL)
        {
                char *user = trim((char*)readBuffer);
                memset(&stored, 0, sizeof(stored));

                if(findViewInfo(dbp, user, &stored) == YES)
                {
                        if(isHistoryViewInfo(writeData,&stored) == NO)
                        {
                                recordNewView(fpResult, user, writeBuffer);
                        }
                        continue;
                }

                recordNewView(fpResult, user, writeBuffer);
                if(saveViewInfo(dbp, user, writeData) == NO)
                {
                        printf("save %s faild\n", user);
                }

                /* clear global buffer content */
                memset(readBuffer, 0, BUFFER_SIZE); 
                memset(writeBuffer, 0, BUFFER_SIZE); 
        }

        free(readBuffer);
        free(writeBuffer);
        dbp->close(dbp, 0);
        fclose(fp);
        fclose(fpResult);
}

/*给命令行调用的接口，找茬指定的用户是否是历史用户*/
void hasViewInfo(char *user, struct ViewData *query)
{
        DB *dbp = NULL;
        struct ViewData stored;

        dbp = openDb();
        assert(user != NULL && query != NULL);
        memset(&stored, 0 ,sizeof(stored));
        user = trim(user);

        if(findViewInfo(dbp, user, &stored) == YES)
        {
                if(isHistoryViewInfo(query, &stored) == YES)
                {
                        printf("found %s\n",user);
                }
                else
                {
                        //printf("%s in db\n",user);
                        printf("not found %s\n",user);
                }
        }
        else
        {
                printf("%s not in db\n",user);
                printf("not found %s\n",user);
        }

        dbp->close(dbp, 0);
}

int main (int argc, char *argv[])
{
        int oc;
        extern char *optarg;
        extern int optind, opterr, optopt;

        char *from = NULL;
        char *to = NULL;

        struct ViewData viewData;
        memset(&viewData, 0, sizeof(viewData));

        readBuffer = malloc(BUFFER_SIZE);
        writeBuffer = malloc(BUFFER_SIZE);

        while((oc=getopt(argc,argv,"f:s:t:")) != -1)
        {
                switch(oc)
                {
                        case 's':
                                setDate(&viewData,argv[optind]);
                                hasViewInfo(optarg, &viewData);
                                break;

                        case 'f':
                                from = optarg;
                                to = argv[optind++];
                                setDate(&viewData,argv[optind]);
                                saveViewFile(from, to, &viewData);
                                break;
                }
        }

        return 0;
}