| 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 |  | 
|---|