Use Temporary Sequence Nodes
https://www.51cto.com/article/687086.html
#include <vector>
#include <iostream>
#include "gtest/gtest.h"
#include <zookeeper.h>
#include "glog/logging.h"
#include "fmt/format.h"
#include <thread>
#include <chrono>
#include <mutex>
#include <condition_variable>
class Status{
public:
enum ErrorCode {
SUCCESS = 0,
FAIL = 1,
};
Status(){}
Status(int c = 0):code(c) {}
template<typename ...Args>
Status(Args ...args) : code(FAIL), msg(fmt::format(std::forward<Args>(args)...)) {
}
template<typename ...Args>
Status(int c, Args ...args) : code(c), msg(fmt::format(std::forward<Args>(args)...)) {
}
operator bool() const {
return code == SUCCESS;
}
int code = 0;
std::string msg;
~Status(){
if(code != SUCCESS) {
LOG(INFO) << "code = " << code << " msg = " << msg;
}
}
};
#define CHECK_COND_RETURN(cond, ...) \
do{\
if(!(cond)){\
return Status(__VA_ARGS__);\
}\
}while(0)
class ZooLock {
public:
ZooLock(const std::string &addr, const std::string &path)
: zk_addr(addr), lock_path(path)
{
}
Status Init() {
handle = zookeeper_init(zk_addr.c_str(), watcher_fn, recv_timeout, nullptr, this, 0);
CHECK_COND_RETURN(handle != nullptr, "zk init fail, return null handle");
LOG(INFO) << "zk init ok, handle is: " << handle;
return 0;
}
Status Lock() {
// int zoo_create(zhandle_t *zh, const char *path, const char *value, int valuelen,
// const struct ACL_vector *acl, int flags, char *path_buffer,
// int path_buffer_len)
//create father path
auto ret = zoo_create(
handle, "/lock",
nullptr, 0,
&ZOO_OPEN_ACL_UNSAFE,
0,
nullptr, 0);
CHECK_COND_RETURN(
ret == ZOK || ret == ZNODEEXISTS,
"zk create /lock fail, return null handle");
//create tmp and seq node
char path_buff[256] = {0};
CHECK_COND_RETURN(
zoo_create(
handle, lock_path.c_str(),
nullptr, 0,
&ZOO_OPEN_ACL_UNSAFE,
ZOO_EPHEMERAL | ZOO_SEQUENCE,
path_buff, 256) == ZOK,
"zk create fail, return null handle");
/*int zoo_get_children(zhandle_t *zh, const char *path, int watch,
struct String_vector *strings)
*/
self_node_path = std::string(path_buff);
self_node_path = self_node_path.substr(6, self_node_path.size());
//check self is min?
do{
struct String_vector strings;
CHECK_COND_RETURN(
zoo_get_children(handle, "/lock", 0, &strings) == ZOK,
"zk get children fail, return null handle");
std::vector<std::string> all_path;
for(int i = 0; i < strings.count; i++) {
all_path.emplace_back(strings.data[i]);
delete strings.data[i];
}
std::sort(all_path.begin(), all_path.end());
// for(auto &path : all_path) {
// LOG(INFO) << "all path: " << path;
// }
CHECK_COND_RETURN(all_path.size() > 0, "no path, create failed ?");
if(all_path[0] == self_node_path) { //self is min, so get the lock
break;
}
bool found = false;
for(int i = 0; i < all_path.size(); i++) {
if(all_path[i] == self_node_path) {
found = true;
min_node_path = "/lock/" + all_path[i-1];
break;
}
}
// zoo_get(zhandle_t *zh, const char *path, int watch, char *buffer,
// int *buffer_len, struct Stat *stat)
struct Stat stat;
int len = 256;
//get min node, add watch. int watc = 1, will call watcher_fn
auto ret = zoo_get(handle, min_node_path.c_str(), 1, path_buff, &len, &stat);
if(ret == ZNONODE){
continue;
}
CHECK_COND_RETURN(ret == ZOK, "add watch to {}, failed code: {}", min_node_path, ret);
CHECK_COND_RETURN(found, "no path found, create failed ?");
std::unique_lock<std::mutex> guard(_mtx);
cv.wait(guard, [this](){return min_node_deleted;});
}while(true);
LOG(INFO) << "create success: " << path_buff;
return 0;
}
Status UnLock() {
auto path = "/lock/" + self_node_path;
CHECK_COND_RETURN(zoo_delete(handle, path.c_str() , -1) == ZOK, "close zk fail");
// handle = nullptr;
return 0;
}
static void watcher_fn(zhandle_t *zh, int type,
int state, const char *path,void *watcherCtx)
{
ZooLock *self = (ZooLock*) watcherCtx;
LOG(INFO) << "zk state, self " << self
<< " type: " << type
<< " state: " << state
<< " path: " << path
<< " min_node_path: " << self->min_node_path;
if(path == self->min_node_path) {
if(type == ZOO_DELETED_EVENT) { //node deleted, so awake
LOG(INFO) << "min node deleted";
std::unique_lock<std::mutex> guard(self->_mtx);
self->min_node_deleted = true;
self->cv.notify_all();
}
}
}
~ZooLock() {
zookeeper_close(handle);
}
private:
std::string zk_addr;
std::string lock_path;
zhandle_t *handle = nullptr;
int recv_timeout = 5000;
std::string self_node_path;
std::condition_variable cv;
std::mutex _mtx;
bool min_node_deleted = false;
std::string min_node_path;
};
TEST(zk, zk) {
{
std::vector<std::thread> ths;
int counter = 0;
int target = 20;
int num = 30;
for(int i = 0; i < target; i++) {
ths.emplace_back([i, &counter, num](){
ZooLock lock2("localhost:2181", "/lock/ckpt-");
ASSERT_TRUE(lock2.Init());
for(int j = 0; j < num; j++) {
ASSERT_TRUE(lock2.Lock());
auto tmp = counter;
tmp++;
// std::this_thread::sleep_for(std::chrono::milliseconds(std::rand() % 1000));
counter = tmp;
LOG(INFO) << "after lock got " << i << " counter: " << counter;
ASSERT_TRUE(lock2.UnLock());
}
});
}
for(auto &t : ths) {
t.join();
}
ASSERT_EQ(counter, target * num);
}
}