tracking_allocator.h

// Copyright 2023 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <cstring>
 
#include "pw_allocator/allocator.h"
#include "pw_allocator/capability.h"
#include "pw_allocator/metrics.h"
#include "pw_assert/assert.h"
#include "pw_metric/metric.h"
#include "pw_preprocessor/compiler.h"
#include "pw_result/result.h"
#include "pw_status/status.h"
#include "pw_status/status_with_size.h"
 
namespace pw::allocator {
 
 
static constexpr struct AddTrackingAllocatorAsChild {
} kAddTrackingAllocatorAsChild = {};
 
template <typename MetricsType>
class TrackingAllocator : public pw::Allocator {
 public:
  TrackingAllocator(metric::Token token, Allocator& allocator)
      : Allocator(allocator.capabilities() | kImplementsGetRequestedLayout),
        allocator_(allocator),
        metrics_(token) {}
 
  template <typename OtherMetrics>
  TrackingAllocator(metric::Token token,
                    TrackingAllocator<OtherMetrics>& parent,
                    const AddTrackingAllocatorAsChild&)
      : TrackingAllocator(token, parent) {
    parent.metric_group().Add(metric_group());
  }
 
  const metric::Group& metric_group() const { return metrics_.group(); }
  metric::Group& metric_group() { return metrics_.group(); }
 
  const MetricsType& metrics() const { return metrics_.metrics(); }
 
  void UpdateDeferred() const { metrics_.UpdateDeferred(allocator_); }
 
 private:
  void* DoAllocate(Layout layout) override;
 
  void DoDeallocate(void* ptr) override;
 
  bool DoResize(void* ptr, size_t new_size) override;
 
  void* DoReallocate(void* ptr, Layout new_layout) override;
 
  size_t DoGetAllocated() const override { return allocator_.GetAllocated(); }
 
  std::optional<allocator::Fragmentation> DoMeasureFragmentation()
      const override {
    return allocator_.MeasureFragmentation();
  }
 
  Result<Layout> DoGetInfo(InfoType info_type, const void* ptr) const override {
    return GetInfo(allocator_, info_type, ptr);
  }
 
  Allocator& allocator_;
  mutable internal::Metrics<MetricsType> metrics_;
};
 
// Template method implementation.
 
template <typename MetricsType>
void* TrackingAllocator<MetricsType>::DoAllocate(Layout layout) {
  if constexpr (internal::AnyEnabled<MetricsType>()) {
    Layout requested = layout;
    size_t allocated = allocator_.GetAllocated();
    void* new_ptr = allocator_.Allocate(requested);
    if (new_ptr == nullptr) {
      metrics_.RecordFailure(requested.size());
      return nullptr;
    }
    metrics_.IncrementAllocations();
    metrics_.ModifyRequested(requested.size(), 0);
    metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
    return new_ptr;
  } else {
    return allocator_.Allocate(layout);
  }
}
 
template <typename MetricsType>
void TrackingAllocator<MetricsType>::DoDeallocate(void* ptr) {
  if constexpr (internal::AnyEnabled<MetricsType>()) {
    Layout requested = Layout::Unwrap(GetRequestedLayout(ptr));
    size_t allocated = allocator_.GetAllocated();
    allocator_.Deallocate(ptr);
    metrics_.IncrementDeallocations();
    metrics_.ModifyRequested(0, requested.size());
    metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
  } else {
    allocator_.Deallocate(ptr);
  }
}
 
template <typename MetricsType>
bool TrackingAllocator<MetricsType>::DoResize(void* ptr, size_t new_size) {
  if constexpr (internal::AnyEnabled<MetricsType>()) {
    Layout requested = Layout::Unwrap(GetRequestedLayout(ptr));
    size_t allocated = allocator_.GetAllocated();
    if (!allocator_.Resize(ptr, new_size)) {
      metrics_.RecordFailure(new_size);
      return false;
    }
    metrics_.IncrementResizes();
    metrics_.ModifyRequested(new_size, requested.size());
    metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
    return true;
  } else {
    return allocator_.Resize(ptr, new_size);
  }
}
 
template <typename MetricsType>
void* TrackingAllocator<MetricsType>::DoReallocate(void* ptr,
                                                   Layout new_layout) {
  if constexpr (internal::AnyEnabled<MetricsType>()) {
    // Check if possible to resize in place with no additional overhead.
    Layout requested = Layout::Unwrap(GetRequestedLayout(ptr));
    size_t allocated = allocator_.GetAllocated();
    size_t new_size = new_layout.size();
    if (allocator_.Resize(ptr, new_size)) {
      metrics_.IncrementReallocations();
      metrics_.ModifyRequested(new_size, requested.size());
      metrics_.ModifyAllocated(allocator_.GetAllocated(), allocated);
      return ptr;
    }
 
    // Retrieve the allocated layout of the old pointer to calculate the peak
    // memory usage during reallocation if a move occurs. Note that if the
    // underlying allocator does not have the `kImplementsGetAllocatedLayout`
    // capability, this will return an error and the peak metric may be lower
    // than the actual peak allocation value.
    Result<Layout> old_allocated_layout = GetAllocatedLayout(ptr);
 
    void* new_ptr = allocator_.Reallocate(ptr, new_layout);
    if (new_ptr == nullptr) {
      metrics_.RecordFailure(new_size);
      return nullptr;
    }
    metrics_.IncrementReallocations();
    metrics_.ModifyRequested(new_size, requested.size());
 
    size_t current_allocated = allocator_.GetAllocated();
    if (new_ptr != ptr && old_allocated_layout.ok()) {
      size_t peak_allocated = current_allocated + old_allocated_layout->size();
      metrics_.ModifyAllocated(peak_allocated, allocated);
      metrics_.ModifyAllocated(current_allocated, peak_allocated);
    } else {
      metrics_.ModifyAllocated(current_allocated, allocated);
    }
    return new_ptr;
  } else {
    return allocator_.Reallocate(ptr, new_layout);
  }
}
 
 
}  // namespace pw::allocator