Redis 分布式锁

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
// --------------------
#include <unistd.h>
#include <pthread.h>
// --------------------
#include <hiredis.h>
// --------------------
#include "thread_helper.h"

#define ATOMIC_UNLOCK 0
#define ATOMIC_LOCK 1

int lock = 0;
int num = 0;
int total_num = 50;
ThreadSchema* ts;

bool acquire_lock(int sn, redisContext* c) {
    redisReply* r = NULL;
    int res = 0;

    if(NULL == c)
        return false;

    while(true) {
        r = (redisReply*)redisCommand(c,"SETNX tom_lock 1");
        res = r->integer;
        freeReplyObject(r);

        if (res == 1) {
            r = (redisReply*)redisCommand(c,"EXPIRE tom_lock 60");
            freeReplyObject(r);
            return true;
        } else {

            r = (redisReply*)redisCommand(c,"TTL tom_lock");
            res = r->integer;
            //printf("acquire_lock %d %d\n", r->type, r->integer);
            freeReplyObject(r);

            if(res == -1){
                r = (redisReply*)redisCommand(c,"EXPIRE tom_lock 1");
                freeReplyObject(r);
            }
        }

        usleep(1000);
    }

    return false;
}

bool release_lock(int sn, redisContext* c) {
    redisReply* r = NULL;

    while(true) {
        r = (redisReply*)redisCommand(c,"WATCH tom_lock");
        freeReplyObject(r);

        r = (redisReply*)redisCommand(c,"GET tom_lock");

        //printf("release_lock %lld\n", r->integer);

        if (((REDIS_REPLY_INTEGER == r->type) && (1 == r->integer)) || 
            ((REDIS_REPLY_STRING == r->type) && (strcmp("1", r->str) == 0))
            ) {

            //printf("release_lock state 1\n");

            freeReplyObject(r);

            r = (redisReply*)redisCommand(c,"MULTI");
            freeReplyObject(r);

            r = (redisReply*)redisCommand(c,"DEL tom_lock");
            freeReplyObject(r);

            r = (redisReply*)redisCommand(c,"EXEC");

            return true;
        } else if(REDIS_REPLY_NIL == r->type) {
            printf("release_lock state 2\n");

            freeReplyObject(r);

            r = (redisReply*)redisCommand(c,"UNWATCH");
            freeReplyObject(r);

            return true;
        } else {
            printf("release_lock sn:%d, tate 3\n", sn);

            freeReplyObject(r);

            r = (redisReply*)redisCommand(c,"UNWATCH");
            freeReplyObject(r);
        }

        usleep(1000);
    }

    return false;
}

void* test_thread(void* arg) {
    int sn = *((int*)arg);
    redisContext* c = ts->rctx[sn];
    redisReply* r = NULL;
    int tom = 0;

    while(true) {

        //printf("En Lock sn:%d\n", sn);
        fflush(stdout);

        if(!acquire_lock(sn, c)) {
            printf("Co Lock\n");
            fflush(stdout);
            sched_yield();
            continue;
        }

        // printf("Lv Lock sn:%d\n", sn);
        // fflush(stdout);

        r = (redisReply*)redisCommand(c,"GET tom");
        tom = atoi(r->str);
        freeReplyObject(r);

        if(tom >= 50) {
            release_lock(sn, c);
            break;
        }

        tom += 1;
        printf("sn: %d, tom: %d\n", sn, tom);
        fflush(stdout);

        r = (redisReply*)redisCommand(c,"SET tom %d", tom);
        freeReplyObject(r);

        // printf("En UnLock sn:%d\n", sn);
        // fflush(stdout);

        release_lock(sn, c);

        //usleep(1);
        // printf("Lv UnLock sn:%d\n", sn);
        // fflush(stdout);
    }

    printf("Lv Test:%d\n", sn);
    fflush(stdout);

    return NULL;
}

int main(int argc, char* argv[])
{
    ts = new ThreadSchema(20);

    const char *hostname = "127.0.0.1";
    int port = 6379;
    struct timeval timeout = { 1, 500000 }; // 1.5 seconds

    for(int i=0; i<ts->num; ++i) {
        ts->rctx[i] = redisConnectWithTimeout(hostname, port, timeout);
        if (ts->rctx[i] == NULL || ts->rctx[i]->err) {
            if (ts->rctx[i]) {
                printf("Connection error: %s\n", ts->rctx[i]->errstr);
            } else {
                printf("Connection error: can't allocate redis context\n");
            }
            exit(1);
        }
    }

    redisReply* r = (redisReply*)redisCommand(ts->rctx[0],"DEL tom_lock");
    freeReplyObject(r);

    r = (redisReply*)redisCommand(ts->rctx[0],"SET tom 0");
    freeReplyObject(r);

    for(int i=0; i<ts->num; ++i)
    {
        pthread_create(&ts->pid[i], NULL, test_thread, &ts->sn[i]);
    }

    /* Disconnects and frees the context */

    while(num < total_num)
    {
        sched_yield();
    }

    return 0;
}

 

#ifndef THREAD_HELPER_H
#define THREAD_HELPER_H

#include <hiredis.h>
#include <pthread.h>

class ThreadSchema
{
public:
    ThreadSchema(int n);
    ~ThreadSchema();

    int             num;
    int*            sn;
    redisContext**  rctx;
    pthread_t*      pid;
};

#endif

 

#include "thread_helper.h"
#include <memory.h>
#include <stdlib.h>

ThreadSchema::ThreadSchema(int n)
{
    num = n;

    sn = (int*)malloc(sizeof(int) * num);
    if (NULL != sn) {
        for (int i=0; i < num; ++i) {
            sn[i] = i;
        }
    }

    rctx = (redisContext**)malloc(sizeof(redisContext*) * num);
    if (NULL != rctx) {
        memset(rctx, 0, sizeof(redisContext*) * num);
    }

    pid = (pthread_t*)malloc(sizeof(pthread_t) * num);
    if(NULL != pid) {
        memset(pid, 0, sizeof(pthread_t*) * num);
    }
}

ThreadSchema::~ThreadSchema() {

    for (int i=0; i < num; ++i) {
        if(NULL != rctx[i]) {
            redisFree(rctx[i]);
        }
    }

    if(NULL != rctx) {
        free(rctx);
        rctx = NULL;
    }

    if(NULL != pid) {
        free(pid);
        pid = NULL;
    }
}