RWLock::RWLock(GetZooKeeper get_zookeeper_, const std::string & path_)
: get_zookeeper{get_zookeeper_}, path{path_}
{
if (!get_zookeeper)
throw DB::Exception{"No ZooKeeper accessor specified", DB::ErrorCodes::LOGICAL_ERROR};
int32_t code = get_zookeeper()->tryCreate(path, "", CreateMode::Persistent);
if ((code != ZOK) && (code != ZNODEEXISTS))
{
if (code == ZNONODE)
throw DB::Exception{"No such lock", DB::ErrorCodes::RWLOCK_NO_SUCH_LOCK};
else
throw KeeperException{code};
}
}
void RWLock::release()
{
__sync_synchronize();
if (!*this)
{
owns_lock = false;
return;
}
if (key.empty())
throw DB::Exception{"RWLock: no lock is held", DB::ErrorCodes::LOGICAL_ERROR};
get_zookeeper()->tryRemoveEphemeralNodeWithRetries(path + "/" + key);
key.clear();
owns_lock = false;
}
std::optional<std::string> ZooKeeperNodeCache::get(const std::string & path)
{
zkutil::ZooKeeperPtr zookeeper;
std::unordered_set<std::string> invalidated_paths;
{
std::lock_guard<std::mutex> lock(context->mutex);
if (!context->zookeeper)
{
/// Possibly, there was a previous session and it has expired. Clear the cache.
nonexistent_nodes.clear();
node_cache.clear();
context->zookeeper = get_zookeeper();
}
zookeeper = context->zookeeper;
invalidated_paths.swap(context->invalidated_paths);
}
if (!zookeeper)
throw DB::Exception("Could not get znode: `" + path + "'. ZooKeeper not configured.", DB::ErrorCodes::NO_ZOOKEEPER);
for (const auto & path : invalidated_paths)
{
nonexistent_nodes.erase(path);
node_cache.erase(path);
}
if (nonexistent_nodes.count(path))
return std::nullopt;
auto watch_callback = [context=context](const ZooKeeperImpl::ZooKeeper::WatchResponse & response)
{
if (!(response.type != ZooKeeperImpl::ZooKeeper::SESSION || response.state == ZooKeeperImpl::ZooKeeper::EXPIRED_SESSION))
return;
bool changed = false;
{
std::lock_guard<std::mutex> lock(context->mutex);
if (response.type != ZooKeeperImpl::ZooKeeper::SESSION)
changed = context->invalidated_paths.emplace(response.path).second;
else if (response.state == ZooKeeperImpl::ZooKeeper::EXPIRED_SESSION)
{
context->zookeeper = nullptr;
context->invalidated_paths.clear();
changed = true;
}
}
if (changed)
context->changed_event.set();
};
std::string contents;
auto cache_it = node_cache.find(path);
if (cache_it != node_cache.end())
{
return cache_it->second;
}
if (zookeeper->tryGetWatch(path, contents, /* stat = */nullptr, watch_callback))
{
node_cache.emplace(path, contents);
return contents;
}
/// Node doesn't exist. Create a watch on node creation.
nonexistent_nodes.insert(path);
if (!zookeeper->existsWatch(path, /* stat = */nullptr, watch_callback))
return std::nullopt;
/// Node was created between the two previous calls, try again. Watch is already set.
if (zookeeper->tryGet(path, contents))
{
nonexistent_nodes.erase(path);
node_cache.emplace(path, contents);
return contents;
}
return std::nullopt;
}
void RWLock::acquireImpl(Mode mode)
{
static_assert((lock_type == RWLock::Read) || (lock_type == RWLock::Write), "Invalid RWLock type");
__sync_synchronize();
if (!*this)
{
owns_lock = true;
return;
}
if (!key.empty())
throw DB::Exception{"RWLock: lock already held", DB::ErrorCodes::RWLOCK_ALREADY_HELD};
try
{
/// Enqueue a new request for a lock.
int32_t code = get_zookeeper()->tryCreate(path + "/" + Prefix<lock_type>::name,
"", CreateMode::EphemeralSequential, key);
if (code == ZNONODE)
throw DB::Exception{"No such lock", DB::ErrorCodes::RWLOCK_NO_SUCH_LOCK};
else if (code != ZOK)
throw KeeperException(code);
key = key.substr(path.length() + 1);
while (true)
{
auto zookeeper = get_zookeeper();
std::vector<std::string> children;
int32_t code = zookeeper->tryGetChildren(path, children);
if (code == ZNONODE)
throw DB::Exception{"No such lock", DB::ErrorCodes::RWLOCK_NO_SUCH_LOCK};
else if (code != ZOK)
throw KeeperException{code};
std::sort(children.begin(), children.end(), nodeQueueCmp);
auto it = std::lower_bound(children.cbegin(), children.cend(), key, nodeQueueCmp);
/// This should never happen.
if ((it == children.cend()) || !nodeQueueEquals(*it, key))
throw DB::Exception{"RWLock: corrupted lock request queue. Own request not found.",
DB::ErrorCodes::LOGICAL_ERROR};
const std::string * observed_key = nullptr;
if (lock_type == RWLock::Read)
{
/// Look for the first write lock request that is older than us.
auto it2 = std::find_if(
std::make_reverse_iterator(it),
children.crend(),
[](const std::string & child)
{
return startsWith(child, Prefix<RWLock::Write>::name);
});
if (it2 != children.crend())
observed_key = &*it2;
}
else if (lock_type == RWLock::Write)
{
if (it != children.cbegin())
{
/// Take the next lock request that is older than us.
auto it2 = std::prev(it);
observed_key = &*it2;
}
}
if (observed_key == nullptr)
{
/// We hold the lock.
owns_lock = true;
break;
}
if (mode == NonBlocking)
{
int32_t code = zookeeper->tryRemoveEphemeralNodeWithRetries(path + "/" + key);
if (code == ZNONODE)
throw DB::Exception{"No such lock", DB::ErrorCodes::RWLOCK_NO_SUCH_LOCK};
else if (code != ZOK)
throw KeeperException{code};
key.clear();
break;
}
abortIfRequested();
/// Wait for our turn to come.
if (zookeeper->exists(path + "/" + *observed_key, nullptr, event))
{
do
{
abortIfRequested();
}
while (!event->tryWait(wait_duration));
}
}
}
catch (...)
{
try
{
if (!key.empty())
get_zookeeper()->tryRemoveEphemeralNodeWithRetries(path + "/" + key);
}
catch (...)
{
DB::tryLogCurrentException(__PRETTY_FUNCTION__);
}
throw;
}
}
ZooKeeperNodeCache::ZNode ZooKeeperNodeCache::get(const std::string & path, Coordination::WatchCallback caller_watch_callback)
{
std::unordered_set<std::string> invalidated_paths;
{
std::lock_guard lock(context->mutex);
if (context->all_paths_invalidated)
{
/// Possibly, there was a previous session and it has expired. Clear the cache.
path_to_cached_znode.clear();
context->all_paths_invalidated = false;
}
invalidated_paths.swap(context->invalidated_paths);
}
zkutil::ZooKeeperPtr zookeeper = get_zookeeper();
if (!zookeeper)
throw DB::Exception("Could not get znode: `" + path + "'. ZooKeeper not configured.", DB::ErrorCodes::NO_ZOOKEEPER);
for (const auto & invalidated_path : invalidated_paths)
path_to_cached_znode.erase(invalidated_path);
auto cache_it = path_to_cached_znode.find(path);
if (cache_it != path_to_cached_znode.end())
return cache_it->second;
std::weak_ptr<Context> weak_context(context);
auto watch_callback = [weak_context, caller_watch_callback](const Coordination::WatchResponse & response)
{
if (!(response.type != Coordination::SESSION || response.state == Coordination::EXPIRED_SESSION))
return;
auto owned_context = weak_context.lock();
if (!owned_context)
return;
bool changed = false;
{
std::lock_guard lock(owned_context->mutex);
if (response.type != Coordination::SESSION)
changed = owned_context->invalidated_paths.emplace(response.path).second;
else if (response.state == Coordination::EXPIRED_SESSION)
{
owned_context->invalidated_paths.clear();
owned_context->all_paths_invalidated = true;
changed = true;
}
}
if (changed && caller_watch_callback)
caller_watch_callback(response);
};
ZNode result;
result.exists = zookeeper->tryGetWatch(path, result.contents, &result.stat, watch_callback);
if (result.exists)
{
path_to_cached_znode.emplace(path, result);
return result;
}
/// Node doesn't exist. We must set a watch on node creation (because it wasn't set by tryGetWatch).
result.exists = zookeeper->existsWatch(path, &result.stat, watch_callback);
if (!result.exists)
{
path_to_cached_znode.emplace(path, result);
return result;
}
/// Node was created between the two previous calls, try again. Watch is already set.
result.exists = zookeeper->tryGet(path, result.contents, &result.stat);
path_to_cached_znode.emplace(path, result);
return result;
}
