1 /*
2  * Copyright (C) 2017 The Android Open Source Project
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 INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
18 #define INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
19 
20 #include <memory>
21 #include <string>
22 #include <vector>
23 
24 #include "perfetto/base/export.h"
25 #include "perfetto/base/logging.h"
26 #include "perfetto/protozero/root_message.h"
27 #include "perfetto/protozero/scattered_stream_writer.h"
28 
29 namespace protozero {
30 
31 class Message;
32 
33 class PERFETTO_EXPORT ScatteredHeapBuffer
34     : public protozero::ScatteredStreamWriter::Delegate {
35  public:
36   class PERFETTO_EXPORT Slice {
37    public:
38     Slice();
39     explicit Slice(size_t size);
40     Slice(Slice&& slice) noexcept;
41     ~Slice();
42     Slice& operator=(Slice&&);
43 
GetTotalRange()44     inline protozero::ContiguousMemoryRange GetTotalRange() const {
45       return {buffer_.get(), buffer_.get() + size_};
46     }
47 
GetUsedRange()48     inline protozero::ContiguousMemoryRange GetUsedRange() const {
49       return {buffer_.get(), buffer_.get() + size_ - unused_bytes_};
50     }
51 
start()52     uint8_t* start() const { return buffer_.get(); }
size()53     size_t size() const { return size_; }
unused_bytes()54     size_t unused_bytes() const { return unused_bytes_; }
set_unused_bytes(size_t unused_bytes)55     void set_unused_bytes(size_t unused_bytes) {
56       PERFETTO_DCHECK(unused_bytes_ <= size_);
57       unused_bytes_ = unused_bytes;
58     }
59 
60     void Clear();
61 
62    private:
63     std::unique_ptr<uint8_t[]> buffer_;
64     size_t size_;
65     size_t unused_bytes_;
66   };
67 
68   ScatteredHeapBuffer(size_t initial_slice_size_bytes = 128,
69                       size_t maximum_slice_size_bytes = 128 * 1024);
70   ~ScatteredHeapBuffer() override;
71 
72   // protozero::ScatteredStreamWriter::Delegate implementation.
73   protozero::ContiguousMemoryRange GetNewBuffer() override;
74 
75   // Stitch all the slices into a single contiguous buffer.
76   std::vector<uint8_t> StitchSlices();
77 
78   // Note that the returned ranges point back to this buffer and thus cannot
79   // outlive it.
80   std::vector<protozero::ContiguousMemoryRange> GetRanges();
81 
slices()82   const std::vector<Slice>& slices() const { return slices_; }
83 
set_writer(protozero::ScatteredStreamWriter * writer)84   void set_writer(protozero::ScatteredStreamWriter* writer) {
85     writer_ = writer;
86   }
87 
88   // Update unused_bytes() of the current |Slice| based on the writer's state.
89   void AdjustUsedSizeOfCurrentSlice();
90 
91   // Returns the total size the slices occupy in heap memory (including unused).
92   size_t GetTotalSize();
93 
94   // Reset the contents of this buffer but retain one slice allocation (if it
95   // exists) to be reused for future writes.
96   void Reset();
97 
98  private:
99   size_t next_slice_size_;
100   const size_t maximum_slice_size_;
101   protozero::ScatteredStreamWriter* writer_ = nullptr;
102   std::vector<Slice> slices_;
103 
104   // Used to keep an allocated slice around after this buffer is reset.
105   Slice cached_slice_;
106 };
107 
108 // Helper function to create heap-based protozero messages in one line.
109 // Useful when manually serializing a protozero message (primarily in
110 // tests/utilities). So instead of the following:
111 //   protozero::MyMessage msg;
112 //   protozero::ScatteredHeapBuffer shb;
113 //   protozero::ScatteredStreamWriter writer(&shb);
114 //   shb.set_writer(&writer);
115 //   msg.Reset(&writer);
116 //   ...
117 // You can write:
118 //   protozero::HeapBuffered<protozero::MyMessage> msg;
119 //   msg->set_stuff(...);
120 //   msg.SerializeAsString();
121 template <typename T = ::protozero::Message>
122 class HeapBuffered {
123  public:
HeapBuffered()124   HeapBuffered() : HeapBuffered(4096, 4096) {}
HeapBuffered(size_t initial_slice_size_bytes,size_t maximum_slice_size_bytes)125   HeapBuffered(size_t initial_slice_size_bytes, size_t maximum_slice_size_bytes)
126       : shb_(initial_slice_size_bytes, maximum_slice_size_bytes),
127         writer_(&shb_) {
128     shb_.set_writer(&writer_);
129     msg_.Reset(&writer_);
130   }
131 
132   // This can't be neither copied nor moved because Message hands out pointers
133   // to itself when creating submessages.
134   HeapBuffered(const HeapBuffered&) = delete;
135   HeapBuffered& operator=(const HeapBuffered&) = delete;
136   HeapBuffered(HeapBuffered&&) = delete;
137   HeapBuffered& operator=(HeapBuffered&&) = delete;
138 
get()139   T* get() { return &msg_; }
140   T* operator->() { return &msg_; }
141 
empty()142   bool empty() const { return shb_.slices().empty(); }
143 
SerializeAsArray()144   std::vector<uint8_t> SerializeAsArray() {
145     msg_.Finalize();
146     return shb_.StitchSlices();
147   }
148 
SerializeAsString()149   std::string SerializeAsString() {
150     auto vec = SerializeAsArray();
151     return std::string(reinterpret_cast<const char*>(vec.data()), vec.size());
152   }
153 
GetRanges()154   std::vector<protozero::ContiguousMemoryRange> GetRanges() {
155     msg_.Finalize();
156     return shb_.GetRanges();
157   }
158 
Reset()159   void Reset() {
160     shb_.Reset();
161     writer_.Reset(protozero::ContiguousMemoryRange{});
162     msg_.Reset(&writer_);
163     PERFETTO_DCHECK(empty());
164   }
165 
166  private:
167   ScatteredHeapBuffer shb_;
168   ScatteredStreamWriter writer_;
169   RootMessage<T> msg_;
170 };
171 
172 }  // namespace protozero
173 
174 #endif  // INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
175