1 | #pragma once |
2 | |
3 | #include <Disks/IDisk.h> |
4 | #include <IO/ReadBufferFromFile.h> |
5 | #include <IO/WriteBufferFromFile.h> |
6 | |
7 | #include <mutex> |
8 | #include <Poco/DirectoryIterator.h> |
9 | #include <Poco/File.h> |
10 | |
11 | namespace DB |
12 | { |
13 | namespace ErrorCodes |
14 | { |
15 | extern const int LOGICAL_ERROR; |
16 | } |
17 | |
18 | class DiskLocalReservation; |
19 | |
20 | class DiskLocal : public IDisk |
21 | { |
22 | public: |
23 | friend class DiskLocalReservation; |
24 | |
25 | DiskLocal(const String & name_, const String & path_, UInt64 keep_free_space_bytes_) |
26 | : name(name_), disk_path(path_), keep_free_space_bytes(keep_free_space_bytes_) |
27 | { |
28 | if (disk_path.back() != '/') |
29 | throw Exception("Disk path must ends with '/', but '" + disk_path + "' doesn't." , ErrorCodes::LOGICAL_ERROR); |
30 | } |
31 | |
32 | const String & getName() const override { return name; } |
33 | |
34 | const String & getPath() const override { return disk_path; } |
35 | |
36 | ReservationPtr reserve(UInt64 bytes) override; |
37 | |
38 | UInt64 getTotalSpace() const override; |
39 | |
40 | UInt64 getAvailableSpace() const override; |
41 | |
42 | UInt64 getUnreservedSpace() const override; |
43 | |
44 | UInt64 getKeepingFreeSpace() const override { return keep_free_space_bytes; } |
45 | |
46 | bool exists(const String & path) const override; |
47 | |
48 | bool isFile(const String & path) const override; |
49 | |
50 | bool isDirectory(const String & path) const override; |
51 | |
52 | void createDirectory(const String & path) override; |
53 | |
54 | void createDirectories(const String & path) override; |
55 | |
56 | DiskDirectoryIteratorPtr iterateDirectory(const String & path) override; |
57 | |
58 | void moveFile(const String & from_path, const String & to_path) override; |
59 | |
60 | void copyFile(const String & from_path, const String & to_path) override; |
61 | |
62 | std::unique_ptr<ReadBuffer> readFile(const String & path) const override; |
63 | |
64 | std::unique_ptr<WriteBuffer> writeFile(const String & path) override; |
65 | |
66 | private: |
67 | bool tryReserve(UInt64 bytes); |
68 | |
69 | private: |
70 | const String name; |
71 | const String disk_path; |
72 | const UInt64 keep_free_space_bytes; |
73 | |
74 | /// Used for reservation counters modification |
75 | static std::mutex mutex; |
76 | UInt64 reserved_bytes = 0; |
77 | UInt64 reservation_count = 0; |
78 | }; |
79 | |
80 | using DiskLocalPtr = std::shared_ptr<DiskLocal>; |
81 | |
82 | |
83 | class DiskLocalDirectoryIterator : public IDiskDirectoryIterator |
84 | { |
85 | public: |
86 | explicit DiskLocalDirectoryIterator(const String & path) : iter(path) {} |
87 | |
88 | void next() override { ++iter; } |
89 | |
90 | bool isValid() const override { return iter != Poco::DirectoryIterator(); } |
91 | |
92 | String name() const override { return iter.name(); } |
93 | |
94 | private: |
95 | Poco::DirectoryIterator iter; |
96 | }; |
97 | |
98 | class DiskLocalReservation : public IReservation |
99 | { |
100 | public: |
101 | DiskLocalReservation(const DiskLocalPtr & disk_, UInt64 size_) |
102 | : disk(disk_), size(size_), metric_increment(CurrentMetrics::DiskSpaceReservedForMerge, size_) |
103 | { |
104 | } |
105 | |
106 | UInt64 getSize() const override { return size; } |
107 | |
108 | DiskPtr getDisk() const override { return disk; } |
109 | |
110 | void update(UInt64 new_size) override; |
111 | |
112 | ~DiskLocalReservation() override; |
113 | |
114 | private: |
115 | DiskLocalPtr disk; |
116 | UInt64 size; |
117 | CurrentMetrics::Increment metric_increment; |
118 | }; |
119 | |
120 | class DiskFactory; |
121 | void registerDiskLocal(DiskFactory & factory); |
122 | |
123 | } |
124 | |