1 | #pragma once |
2 | |
3 | #if defined(__linux__) || defined(__FreeBSD__) |
4 | |
5 | #include <IO/WriteBufferFromFileBase.h> |
6 | #include <IO/WriteBuffer.h> |
7 | #include <IO/BufferWithOwnMemory.h> |
8 | #include <IO/AIO.h> |
9 | #include <Core/Defines.h> |
10 | #include <Common/CurrentMetrics.h> |
11 | |
12 | #include <string> |
13 | #include <unistd.h> |
14 | #include <fcntl.h> |
15 | |
16 | |
17 | namespace CurrentMetrics |
18 | { |
19 | extern const Metric OpenFileForWrite; |
20 | } |
21 | |
22 | namespace DB |
23 | { |
24 | |
25 | /** Class for asynchronous data writing. |
26 | */ |
27 | class WriteBufferAIO : public WriteBufferFromFileBase |
28 | { |
29 | public: |
30 | WriteBufferAIO(const std::string & filename_, size_t buffer_size_ = DBMS_DEFAULT_BUFFER_SIZE, int flags_ = -1, mode_t mode_ = 0666, |
31 | char * existing_memory_ = nullptr); |
32 | ~WriteBufferAIO() override; |
33 | |
34 | WriteBufferAIO(const WriteBufferAIO &) = delete; |
35 | WriteBufferAIO & operator=(const WriteBufferAIO &) = delete; |
36 | |
37 | off_t getPositionInFile() override; |
38 | void sync() override; |
39 | std::string getFileName() const override { return filename; } |
40 | int getFD() const override { return fd; } |
41 | |
42 | private: |
43 | void nextImpl() override; |
44 | off_t doSeek(off_t off, int whence) override; |
45 | void doTruncate(off_t length) override; |
46 | |
47 | /// If there's still data in the buffer, we'll write them. |
48 | void flush(); |
49 | /// Wait for the end of the current asynchronous task. |
50 | bool waitForAIOCompletion(); |
51 | /// Prepare an asynchronous request. |
52 | void prepare(); |
53 | /// |
54 | void finalize() override; |
55 | |
56 | private: |
57 | /// Buffer for asynchronous data writes. |
58 | BufferWithOwnMemory<WriteBuffer> flush_buffer; |
59 | |
60 | /// Description of the asynchronous write request. |
61 | iocb request{}; |
62 | iocb * request_ptr{&request}; |
63 | |
64 | AIOContext aio_context{1}; |
65 | |
66 | const std::string filename; |
67 | |
68 | /// The number of bytes to be written to the disk. |
69 | off_t bytes_to_write = 0; |
70 | /// Number of bytes written with the last request. |
71 | off_t bytes_written = 0; |
72 | /// The number of zero bytes to be cut from the end of the file |
73 | /// after the data write operation completes. |
74 | off_t truncation_count = 0; |
75 | |
76 | /// The current position in the file. |
77 | off_t pos_in_file = 0; |
78 | /// The maximum position reached in the file. |
79 | off_t max_pos_in_file = 0; |
80 | |
81 | /// The starting position of the aligned region of the disk to which the data is written. |
82 | off_t region_aligned_begin = 0; |
83 | /// The size of the aligned region of the disk. |
84 | size_t region_aligned_size = 0; |
85 | |
86 | /// The file descriptor for writing. |
87 | int fd = -1; |
88 | |
89 | /// The data buffer that we want to write to the disk. |
90 | Position buffer_begin = nullptr; |
91 | |
92 | /// Is the asynchronous write operation still in progress? |
93 | bool is_pending_write = false; |
94 | /// Did the asynchronous operation fail? |
95 | bool aio_failed = false; |
96 | |
97 | CurrentMetrics::Increment metric_increment{CurrentMetrics::OpenFileForWrite}; |
98 | }; |
99 | |
100 | } |
101 | |
102 | #endif |
103 | |