1#pragma once
2
3#include <Common/ZooKeeper/ZooKeeper.h>
4#include <Common/Exception.h>
5#include <IO/ReadHelpers.h>
6
7
8namespace DB
9{
10
11namespace ErrorCodes
12{
13 extern const int LOGICAL_ERROR;
14}
15
16/// A class that is used for locking a block number in a partition.
17/// It creates a secondary ephemeral node in `temp_path` and a main ephemeral node with `path_prefix`
18/// that references the secondary node. The reasons for this two-level scheme are historical (of course
19/// it would be simpler to allocate block numbers for all partitions in one ZK directory).
20class EphemeralLockInZooKeeper : public boost::noncopyable
21{
22public:
23 EphemeralLockInZooKeeper(
24 const String & path_prefix_, const String & temp_path, zkutil::ZooKeeper & zookeeper_, Coordination::Requests * precheck_ops = nullptr);
25
26 EphemeralLockInZooKeeper() = default;
27
28 EphemeralLockInZooKeeper(EphemeralLockInZooKeeper && rhs) noexcept
29 {
30 *this = std::move(rhs);
31 }
32
33 EphemeralLockInZooKeeper & operator=(EphemeralLockInZooKeeper && rhs) noexcept
34 {
35 zookeeper = rhs.zookeeper;
36 rhs.zookeeper = nullptr;
37 path_prefix = std::move(rhs.path_prefix);
38 path = std::move(rhs.path);
39 holder_path = std::move(rhs.holder_path);
40 return *this;
41 }
42
43 bool isCreated() const
44 {
45 return zookeeper && !holder_path.empty() && !path.empty();
46 }
47
48 String getPath() const
49 {
50 checkCreated();
51 return path;
52 }
53
54 /// Parse the number at the end of the path.
55 UInt64 getNumber() const
56 {
57 checkCreated();
58 return parse<UInt64>(path.c_str() + path_prefix.size(), path.size() - path_prefix.size());
59 }
60
61 void unlock();
62
63 /// Adds actions equivalent to `unlock()` to the list.
64 void getUnlockOps(Coordination::Requests & ops);
65
66 /// Do not delete nodes in destructor. You may call this method after 'getUnlockOps' and successful execution of these ops,
67 /// because the nodes will be already deleted.
68 void assumeUnlocked()
69 {
70 holder_path.clear();
71 }
72
73 void checkCreated() const
74 {
75 if (!isCreated())
76 throw Exception("EphemeralLock is not created", ErrorCodes::LOGICAL_ERROR);
77 }
78
79 ~EphemeralLockInZooKeeper();
80
81private:
82 zkutil::ZooKeeper * zookeeper = nullptr;
83 String path_prefix;
84 String path;
85 String holder_path;
86};
87
88
89/// Acquires block number locks in all partitions.
90class EphemeralLocksInAllPartitions
91{
92public:
93 EphemeralLocksInAllPartitions(
94 const String & block_numbers_path, const String & path_prefix, const String & temp_path,
95 zkutil::ZooKeeper & zookeeper_);
96
97 struct LockInfo
98 {
99 String path;
100 String holder_path;
101
102 String partition_id;
103 UInt64 number = 0;
104 };
105
106 const std::vector<LockInfo> & getLocks() const { return locks; }
107
108 void unlock();
109
110 ~EphemeralLocksInAllPartitions();
111
112private:
113 zkutil::ZooKeeper & zookeeper;
114 std::vector<LockInfo> locks;
115};
116
117}
118