1 | /* |
2 | * Copyright 2014-present Facebook, Inc. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | */ |
16 | |
17 | #ifndef FOLLY_GEN_FILE_H_ |
18 | #error This file may only be included from folly/gen/File.h |
19 | #endif |
20 | |
21 | #include <system_error> |
22 | |
23 | #include <folly/gen/String.h> |
24 | |
25 | namespace folly { |
26 | namespace gen { |
27 | namespace detail { |
28 | |
29 | class FileReader : public GenImpl<ByteRange, FileReader> { |
30 | public: |
31 | FileReader(File file, std::unique_ptr<IOBuf> buffer) |
32 | : file_(std::move(file)), buffer_(std::move(buffer)) { |
33 | buffer_->clear(); |
34 | } |
35 | |
36 | template <class Body> |
37 | bool apply(Body&& body) const { |
38 | for (;;) { |
39 | ssize_t n; |
40 | do { |
41 | n = ::read(file_.fd(), buffer_->writableTail(), buffer_->capacity()); |
42 | } while (n == -1 && errno == EINTR); |
43 | if (n == -1) { |
44 | throw std::system_error(errno, std::system_category(), "read failed" ); |
45 | } |
46 | if (n == 0) { |
47 | return true; |
48 | } |
49 | if (!body(ByteRange(buffer_->tail(), size_t(n)))) { |
50 | return false; |
51 | } |
52 | } |
53 | } |
54 | |
55 | // Technically, there could be infinite files (e.g. /dev/random), but people |
56 | // who open those can do so at their own risk. |
57 | static constexpr bool infinite = false; |
58 | |
59 | private: |
60 | File file_; |
61 | std::unique_ptr<IOBuf> buffer_; |
62 | }; |
63 | |
64 | class FileWriter : public Operator<FileWriter> { |
65 | public: |
66 | FileWriter(File file, std::unique_ptr<IOBuf> buffer) |
67 | : file_(std::move(file)), buffer_(std::move(buffer)) { |
68 | if (buffer_) { |
69 | buffer_->clear(); |
70 | } |
71 | } |
72 | |
73 | template <class Source, class Value> |
74 | void compose(const GenImpl<Value, Source>& source) const { |
75 | auto fn = [&](ByteRange v) { |
76 | if (!this->buffer_ || v.size() >= this->buffer_->capacity()) { |
77 | this->flushBuffer(); |
78 | this->write(v); |
79 | } else { |
80 | if (v.size() > this->buffer_->tailroom()) { |
81 | this->flushBuffer(); |
82 | } |
83 | memcpy(this->buffer_->writableTail(), v.data(), v.size()); |
84 | this->buffer_->append(v.size()); |
85 | } |
86 | }; |
87 | |
88 | // Iterate |
89 | source.foreach(std::move(fn)); |
90 | |
91 | flushBuffer(); |
92 | file_.close(); |
93 | } |
94 | |
95 | private: |
96 | void write(ByteRange v) const { |
97 | ssize_t n; |
98 | while (!v.empty()) { |
99 | do { |
100 | n = ::write(file_.fd(), v.data(), v.size()); |
101 | } while (n == -1 && errno == EINTR); |
102 | if (n == -1) { |
103 | throw std::system_error( |
104 | errno, std::system_category(), "write() failed" ); |
105 | } |
106 | v.advance(size_t(n)); |
107 | } |
108 | } |
109 | |
110 | void flushBuffer() const { |
111 | if (buffer_ && buffer_->length() != 0) { |
112 | write(ByteRange(buffer_->data(), buffer_->length())); |
113 | buffer_->clear(); |
114 | } |
115 | } |
116 | |
117 | mutable File file_; |
118 | std::unique_ptr<IOBuf> buffer_; |
119 | }; |
120 | |
121 | inline auto byLineImpl(File file, char delim, bool keepDelimiter) { |
122 | // clang-format off |
123 | return fromFile(std::move(file)) |
124 | | eachAs<StringPiece>() |
125 | | resplit(delim, keepDelimiter); |
126 | // clang-format on |
127 | } |
128 | |
129 | } // namespace detail |
130 | |
131 | /** |
132 | * Generator which reads lines from a file. |
133 | * Note: This produces StringPieces which reference temporary strings which are |
134 | * only valid during iteration. |
135 | */ |
136 | inline auto byLineFull(File file, char delim = '\n') { |
137 | return detail::byLineImpl(std::move(file), delim, true); |
138 | } |
139 | |
140 | inline auto byLineFull(int fd, char delim = '\n') { |
141 | return byLineFull(File(fd), delim); |
142 | } |
143 | |
144 | inline auto byLineFull(const char* f, char delim = '\n') { |
145 | return byLineFull(File(f), delim); |
146 | } |
147 | |
148 | inline auto byLine(File file, char delim = '\n') { |
149 | return detail::byLineImpl(std::move(file), delim, false); |
150 | } |
151 | |
152 | inline auto byLine(int fd, char delim = '\n') { |
153 | return byLine(File(fd), delim); |
154 | } |
155 | |
156 | inline auto byLine(const char* f, char delim = '\n') { |
157 | return byLine(File(f), delim); |
158 | } |
159 | |
160 | } // namespace gen |
161 | } // namespace folly |
162 | |