xref: /aosp_15_r20/external/pigweed/pw_allocator/public/pw_allocator/tracking_allocator.h (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1 // Copyright 2023 The Pigweed Authors
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
4 // use this file except in compliance with the License. You may obtain a copy of
5 // the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12 // License for the specific language governing permissions and limitations under
13 // the License.
14 #pragma once
15 
16 #include <cstddef>
17 #include <cstdint>
18 #include <cstring>
19 
20 #include "pw_allocator/allocator.h"
21 #include "pw_allocator/capability.h"
22 #include "pw_allocator/metrics.h"
23 #include "pw_assert/assert.h"
24 #include "pw_metric/metric.h"
25 #include "pw_preprocessor/compiler.h"
26 #include "pw_result/result.h"
27 #include "pw_status/status.h"
28 #include "pw_status/status_with_size.h"
29 
30 namespace pw::allocator {
31 
32 /// This tag type is used to explicitly select the constructor which adds
33 /// the tracking allocator's metrics group as a child of the info
34 /// allocator it is wrapping.
35 static constexpr struct AddTrackingAllocatorAsChild {
36 } kAddTrackingAllocatorAsChild = {};
37 
38 /// Wraps an `Allocator` and records details of its usage.
39 ///
40 /// Metric collection is performed using the provided template parameter type.
41 /// Callers can not instantiate this class directly, as it lacks a public
42 /// constructor. Instead, callers should use derived classes which provide the
43 /// template parameter type, such as `TrackingAllocator` which uses the
44 /// default metrics implementation, or `TrackingAllocatorForTest` which
45 /// always uses the real metrics implementation.
46 template <typename MetricsType>
47 class TrackingAllocator : public Allocator {
48  public:
TrackingAllocator(metric::Token token,Allocator & allocator)49   TrackingAllocator(metric::Token token, Allocator& allocator)
50       : Allocator(allocator.capabilities() | kImplementsGetRequestedLayout),
51         allocator_(allocator),
52         metrics_(token) {}
53 
54   template <typename OtherMetrics>
TrackingAllocator(metric::Token token,TrackingAllocator<OtherMetrics> & parent,const AddTrackingAllocatorAsChild &)55   TrackingAllocator(metric::Token token,
56                     TrackingAllocator<OtherMetrics>& parent,
57                     const AddTrackingAllocatorAsChild&)
58       : TrackingAllocator(token, parent) {
59     parent.metric_group().Add(metric_group());
60   }
61 
metric_group()62   const metric::Group& metric_group() const { return metrics_.group(); }
metric_group()63   metric::Group& metric_group() { return metrics_.group(); }
64 
metrics()65   const MetricsType& metrics() const { return metrics_.metrics(); }
66 
67  private:
68   /// @copydoc Allocator::Allocate
69   void* DoAllocate(Layout layout) override;
70 
71   /// @copydoc Allocator::Deallocate
72   void DoDeallocate(void* ptr) override;
73 
74   /// @copydoc Allocator::Deallocate
DoDeallocate(void * ptr,Layout)75   void DoDeallocate(void* ptr, Layout) override { DoDeallocate(ptr); }
76 
77   /// @copydoc Allocator::Resize
78   bool DoResize(void* ptr, size_t new_size) override;
79 
80   /// @copydoc Allocator::Reallocate
81   void* DoReallocate(void* ptr, Layout new_layout) override;
82 
83   /// @copydoc Allocator::GetAllocated
DoGetAllocated()84   size_t DoGetAllocated() const override { return allocator_.GetAllocated(); }
85 
86   /// @copydoc Deallocator::GetInfo
DoGetInfo(InfoType info_type,const void * ptr)87   Result<Layout> DoGetInfo(InfoType info_type, const void* ptr) const override {
88     return GetInfo(allocator_, info_type, ptr);
89   }
90 
91   Allocator& allocator_;
92   internal::Metrics<MetricsType> metrics_;
93 };
94 
95 // Template method implementation.
96 
97 template <typename MetricsType>
DoAllocate(Layout layout)98 void* TrackingAllocator<MetricsType>::DoAllocate(Layout layout) {
99   Layout requested = layout;
100   size_t allocated = allocator_.GetAllocated();
101   void* new_ptr = allocator_.Allocate(requested);
102   if (new_ptr == nullptr) {
103     metrics_.RecordFailure(requested.size());
104     return nullptr;
105   }
106   metrics_.IncrementAllocations();
107   metrics_.ModifyRequested(requested.size(), 0);
108   metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
109   return new_ptr;
110 }
111 
112 template <typename MetricsType>
DoDeallocate(void * ptr)113 void TrackingAllocator<MetricsType>::DoDeallocate(void* ptr) {
114   Layout requested = Layout::Unwrap(GetRequestedLayout(ptr));
115   size_t allocated = allocator_.GetAllocated();
116   allocator_.Deallocate(ptr);
117   metrics_.IncrementDeallocations();
118   metrics_.ModifyRequested(0, requested.size());
119   metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
120 }
121 
122 template <typename MetricsType>
DoResize(void * ptr,size_t new_size)123 bool TrackingAllocator<MetricsType>::DoResize(void* ptr, size_t new_size) {
124   Layout requested = Layout::Unwrap(GetRequestedLayout(ptr));
125   size_t allocated = allocator_.GetAllocated();
126   Layout new_requested(new_size, requested.alignment());
127   if (!allocator_.Resize(ptr, new_requested.size())) {
128     metrics_.RecordFailure(new_size);
129     return false;
130   }
131   metrics_.IncrementResizes();
132   metrics_.ModifyRequested(new_requested.size(), requested.size());
133   metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
134   return true;
135 }
136 
137 template <typename MetricsType>
DoReallocate(void * ptr,Layout new_layout)138 void* TrackingAllocator<MetricsType>::DoReallocate(void* ptr,
139                                                    Layout new_layout) {
140   // Check if possible to resize in place with no additional overhead.
141   Layout requested = Layout::Unwrap(GetRequestedLayout(ptr));
142   size_t allocated = allocator_.GetAllocated();
143   Layout new_requested(new_layout.size(), requested.alignment());
144   if (allocator_.Resize(ptr, new_layout.size())) {
145     metrics_.IncrementReallocations();
146     metrics_.ModifyRequested(new_requested.size(), requested.size());
147     metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
148     return ptr;
149   }
150 
151   // Need to move data to a brand new allocation.
152   // In order to properly record the peak allocation, this method needs to
153   // perform the steps of allocating, copying, and deallocating memory, and
154   // recording metrics in the interim steps.
155   Result<Layout> old_layout = GetUsableLayout(ptr);
156   if (!old_layout.ok()) {
157     metrics_.RecordFailure(new_layout.size());
158     return nullptr;
159   }
160   void* new_ptr = allocator_.Allocate(new_layout);
161   if (new_ptr == nullptr) {
162     metrics_.RecordFailure(new_layout.size());
163     return nullptr;
164   }
165   // Update with transient allocation to ensure peak metrics are correct.
166   size_t transient_allocated = allocator_.GetAllocated();
167   metrics_.ModifyAllocated(transient_allocated, allocated);
168   if (ptr != nullptr) {
169     std::memcpy(new_ptr, ptr, std::min(new_layout.size(), old_layout->size()));
170     allocator_.Deallocate(ptr);
171   }
172   metrics_.IncrementReallocations();
173   metrics_.ModifyRequested(new_requested.size(), requested.size());
174   metrics_.ModifyAllocated(allocator_.GetAllocated(), transient_allocated);
175   return new_ptr;
176 }
177 
178 // TODO(b/326509341): This is an interim alias to facilitate refactoring
179 // downstream consumers of `TrackingAllocator` to add a template parameter.
180 //
181 // The following migration steps are complete:
182 // 1. Downstream consumers will be updated to use `TrackingAllocatorImpl<...>`.
183 // 2. The iterim `TrackingAllocator` class will be removed.
184 // 3. `TrackingAllocatorImpl<...>` will be renamed to `TrackingAllocator<...>`,
185 //    with a `TrackingAllocatorImpl<...>` alias pointing to it.
186 //
187 // The following migration steps remain:
188 // 4. Downstream consumers will be updated to use `TrackingAllocator<...>`.
189 // 5. The `TrackingAllocatorImpl<...>` alias will be removed.
190 template <typename MetricsType>
191 using TrackingAllocatorImpl = TrackingAllocator<MetricsType>;
192 
193 }  // namespace pw::allocator
194