| 1 | #pragma once | 
|---|---|
| 2 | |
| 3 | #include <Common/ZooKeeper/ZooKeeper.h> | 
| 4 | #include <Common/Exception.h> | 
| 5 | #include <IO/ReadHelpers.h> | 
| 6 | |
| 7 | |
| 8 | namespace DB | 
| 9 | { | 
| 10 | |
| 11 | namespace 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). | 
| 20 | class EphemeralLockInZooKeeper : public boost::noncopyable | 
| 21 | { | 
| 22 | public: | 
| 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 | |
| 81 | private: | 
| 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. | 
| 90 | class EphemeralLocksInAllPartitions | 
| 91 | { | 
| 92 | public: | 
| 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 | |
| 112 | private: | 
| 113 | zkutil::ZooKeeper & zookeeper; | 
| 114 | std::vector<LockInfo> locks; | 
| 115 | }; | 
| 116 | |
| 117 | } | 
| 118 | 
