1 /*
2 * Copyright (c) Facebook, Inc. and its affiliates.
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 #pragma once
18
19 #include <stdexcept>
20
21 #include <folly/File.h>
22 #include <folly/portability/GFlags.h>
23 #include <folly/system/MemoryMapping.h>
24 #include <thrift/lib/cpp2/frozen/Frozen.h>
25 #include <thrift/lib/cpp2/protocol/Serializer.h>
26 #include <thrift/lib/thrift/gen-cpp2/frozen_constants.h>
27
28 DECLARE_bool(thrift_frozen_util_disable_mlock);
29 DECLARE_bool(thrift_frozen_util_mlock_on_fault);
30
31 namespace apache {
32 namespace thrift {
33 namespace frozen {
34
35 class FrozenFileForwardIncompatible : public std::runtime_error {
36 public:
37 explicit FrozenFileForwardIncompatible(int fileVersion);
38
fileVersion()39 int fileVersion() const { return fileVersion_; }
supportedVersion()40 int supportedVersion() const {
41 return schema::frozen_constants::kCurrentFrozenFileVersion();
42 }
43
44 private:
45 int fileVersion_;
46 };
47
48 /**
49 * A FreezeRoot that mallocs buffers as needed.
50 */
51 class MallocFreezer final : public FreezeRoot {
52 public:
MallocFreezer()53 explicit MallocFreezer() {}
54
55 template <class T>
freeze(const Layout<T> & layout,const T & root)56 void freeze(const Layout<T>& layout, const T& root) {
57 doFreeze(layout, root);
58 }
59
appendTo(std::string & out)60 void appendTo(std::string& out) const {
61 out.reserve(size_ + out.size());
62 for (auto& segment : segments_) {
63 out.append(segment.buffer, segment.buffer + segment.size);
64 }
65 }
66
67 private:
68 size_t distanceToEnd(const byte* origin) const;
69 size_t offsetOf(const byte* origin) const;
70
71 folly::MutableByteRange appendBuffer(size_t size);
72
73 void doAppendBytes(
74 byte* origin,
75 size_t n,
76 folly::MutableByteRange& range,
77 size_t& distance,
78 size_t alignment) override;
79
80 struct Segment {
81 explicit Segment(size_t size);
SegmentSegment82 Segment(Segment&& other) : size(other.size), buffer(other.buffer) {
83 other.buffer = nullptr;
84 }
85
86 ~Segment();
87
88 size_t size{0};
89 byte* buffer{nullptr}; // owned
90 };
91 std::map<const byte*, size_t> offsets_;
92 std::vector<Segment> segments_;
93 size_t size_{0};
94 };
95
96 /**
97 * Returns an upper bound estimate of the number of bytes required to freeze
98 * this object with a minimal layout. Actual bytes required will depend on the
99 * alignment of the freeze buffer.
100 */
101 template <class T>
frozenSize(const T & v)102 size_t frozenSize(const T& v) {
103 Layout<T> layout;
104 return LayoutRoot::layout(v, layout) - LayoutRoot::kPaddingBytes;
105 }
106
107 /**
108 * Returns an upper bound estimate of the number of bytes required to freeze
109 * this object with a given layout. Actual bytes required will depend on on
110 * the alignment of the freeze buffer.
111 */
112 template <class T>
frozenSize(const T & v,const Layout<T> & fixedLayout)113 size_t frozenSize(const T& v, const Layout<T>& fixedLayout) {
114 Layout<T> layout = fixedLayout;
115 size_t size;
116 bool changed;
117 LayoutRoot::layout(v, layout, changed, size);
118 if (changed) {
119 throw LayoutException();
120 }
121 return size;
122 }
123
124 template <class T>
serializeRootLayout(const Layout<T> & layout,std::string & out)125 void serializeRootLayout(const Layout<T>& layout, std::string& out) {
126 schema::MemorySchema memSchema;
127 schema::Schema schema;
128 saveRoot(layout, memSchema);
129 schema::convert(memSchema, schema);
130
131 *schema.fileVersion_ref() =
132 schema::frozen_constants::kCurrentFrozenFileVersion();
133 out.clear();
134 CompactSerializer::serialize(schema, &out);
135 }
136
137 template <class T>
deserializeRootLayout(folly::ByteRange & range,Layout<T> & layoutOut)138 void deserializeRootLayout(folly::ByteRange& range, Layout<T>& layoutOut) {
139 schema::Schema schema;
140 size_t schemaSize = CompactSerializer::deserialize(range, schema);
141
142 if (*schema.fileVersion_ref() >
143 schema::frozen_constants::kCurrentFrozenFileVersion()) {
144 throw FrozenFileForwardIncompatible(*schema.fileVersion_ref());
145 }
146
147 schema::MemorySchema memSchema;
148 schema::convert(std::move(schema), memSchema);
149 loadRoot(layoutOut, memSchema);
150 range.advance(schemaSize);
151 }
152
153 template <class T>
freezeToFile(const T & x,folly::File file)154 void freezeToFile(const T& x, folly::File file) {
155 std::string schemaStr;
156 auto layout = std::make_unique<Layout<T>>();
157 auto contentSize = LayoutRoot::layout(x, *layout);
158
159 serializeRootLayout(*layout, schemaStr);
160
161 size_t initialBufferSize = contentSize + schemaStr.size();
162 folly::MemoryMapping mapping(
163 file.dup(), 0, initialBufferSize, folly::MemoryMapping::writable());
164 auto mappingRange = mapping.writableRange();
165 auto writeRange = mapping.writableRange();
166 std::copy(schemaStr.begin(), schemaStr.end(), writeRange.begin());
167 writeRange.advance(schemaStr.size());
168 ByteRangeFreezer::freeze(*layout, x, writeRange);
169 size_t finalBufferSize = writeRange.begin() - mappingRange.begin();
170 ftruncate(file.fd(), finalBufferSize);
171 }
172
173 template <class T>
freezeToString(const T & x,std::string & out)174 void freezeToString(const T& x, std::string& out) {
175 out.clear();
176 Layout<T> layout;
177 size_t contentSize = LayoutRoot::layout(x, layout);
178 serializeRootLayout(layout, out);
179
180 size_t schemaSize = out.size();
181 size_t bufferSize = schemaSize + contentSize;
182 out.resize(bufferSize, 0);
183 folly::MutableByteRange writeRange(
184 reinterpret_cast<byte*>(&out[schemaSize]), contentSize);
185 ByteRangeFreezer::freeze(layout, x, writeRange);
186 out.resize(out.size() - writeRange.size());
187 }
188
189 template <class T>
freezeToString(const T & x)190 std::string freezeToString(const T& x) {
191 std::string result;
192 freezeToString(x, result);
193 return result;
194 }
195
196 template <class T>
freezeDataToString(const T & x,const Layout<T> & layout)197 std::string freezeDataToString(const T& x, const Layout<T>& layout) {
198 std::string out;
199 MallocFreezer freezer;
200 freezer.freeze(layout, x);
201 freezer.appendTo(out);
202 return out;
203 }
204
205 template <class T>
freezeToStringMalloc(const T & x,std::string & out)206 void freezeToStringMalloc(const T& x, std::string& out) {
207 Layout<T> layout;
208 LayoutRoot::layout(x, layout);
209 out.clear();
210 serializeRootLayout(layout, out);
211 MallocFreezer freezer;
212 freezer.freeze(layout, x);
213 freezer.appendTo(out);
214 }
215
216 /**
217 * mapFrozen<T>() returns an owned reference to a frozen object which can be
218 * used as a Frozen view of the object. All overloads of this function return
219 * MappedFrozen<T>, which is an alias for a type which bundles the view with its
220 * associated resources. This type may be used directly to hold references to
221 * mapped frozen objects.
222 *
223 * Depending on which overload is used, this bundle will hold references to
224 * different associated data:
225 *
226 * - mapFrozen<T>(ByteRange): Only the layout tree associated with the object.
227 * - mapFrozen<T>(StringPiece): Same as mapFrozen<T>(ByteRange).
228 * - mapFrozen<T>(MemoryMapping): Takes ownership of the memory mapping
229 * in addition to the layout tree.
230 * - mapFrozen<T>(File): Owns the memory mapping created from the File (which,
231 * in turn, takes ownership of the File) in addition to the layout tree.
232 */
233 template <class T>
234 using MappedFrozen = Bundled<typename Layout<T>::View>;
235
236 template <class T>
mapFrozen(folly::ByteRange range)237 MappedFrozen<T> mapFrozen(folly::ByteRange range) {
238 auto layout = std::make_unique<Layout<T>>();
239 deserializeRootLayout(range, *layout);
240 MappedFrozen<T> ret(layout->view({range.begin(), 0}));
241 ret.hold(std::move(layout));
242 return ret;
243 }
244
245 template <class T>
mapFrozen(folly::StringPiece range)246 MappedFrozen<T> mapFrozen(folly::StringPiece range) {
247 return mapFrozen<T>(folly::ByteRange(range));
248 }
249
250 template <class T>
mapFrozen(folly::MemoryMapping mapping)251 MappedFrozen<T> mapFrozen(folly::MemoryMapping mapping) {
252 auto ret = mapFrozen<T>(mapping.range());
253 ret.hold(std::move(mapping));
254 return ret;
255 }
256
257 /**
258 * Maps from the given string, taking ownership of it and bundling it with the
259 * return object to ensure its lifetime.
260 * @param trim Trims the serialized layout from the input string.
261 */
262 template <class T>
263 MappedFrozen<T> mapFrozen(std::string&& str, bool trim = true) {
264 auto layout = std::make_unique<Layout<T>>();
265 auto holder = std::make_unique<HolderImpl<std::string>>(std::move(str));
266 auto& ownedStr = holder->t_;
267 folly::ByteRange rangeBefore = folly::StringPiece(ownedStr);
268 folly::ByteRange range = rangeBefore;
269 deserializeRootLayout(range, *layout);
270 if (trim) {
271 size_t trimSize = range.begin() - rangeBefore.begin();
272 ownedStr.erase(ownedStr.begin(), ownedStr.begin() + trimSize);
273 ownedStr.shrink_to_fit();
274 range = folly::StringPiece(ownedStr);
275 }
276 MappedFrozen<T> ret(layout->view({range.begin(), 0}));
277 ret.holdImpl(std::move(holder));
278 ret.hold(std::move(layout));
279 return ret;
280 }
281
282 template <class T>
283 [[deprecated(
284 "std::string values must be passed by move with std::move(str) or "
285 "passed through non-owning StringPiece")]] MappedFrozen<T>
286 mapFrozen(const std::string& str) = delete;
287
288 template <class T>
mapFrozen(folly::File file,folly::MemoryMapping::LockMode lockMode)289 MappedFrozen<T> mapFrozen(
290 folly::File file, folly::MemoryMapping::LockMode lockMode) {
291 folly::MemoryMapping mapping(std::move(file), 0);
292 folly::MemoryMapping::LockFlags flags{};
293 flags.lockOnFault = FLAGS_thrift_frozen_util_mlock_on_fault;
294 mapping.mlock(lockMode, flags);
295 return mapFrozen<T>(std::move(mapping));
296 }
297
298 template <class T>
mapFrozen(folly::File file)299 MappedFrozen<T> mapFrozen(folly::File file) {
300 return mapFrozen<T>(
301 std::move(file), folly::MemoryMapping::LockMode::TRY_LOCK);
302 }
303
304 } // namespace frozen
305 } // namespace thrift
306 } // namespace apache
307