allocator_async.h

// Copyright 2024 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 <optional>
 
#include "pw_async2/dispatcher.h"  // IWYU pragma: keep
#include "pw_async2/poll.h"
#include "pw_multibuf/allocator.h"
 
namespace pw::multibuf {
 
class MultiBufAllocationFuture;
 
class MultiBufAllocatorAsync {
  // : public MultiBufAllocator::MoreMemoryDelegate {
 public:
  MultiBufAllocatorAsync(MultiBufAllocator& mbuf_allocator)
      : mbuf_allocator_(mbuf_allocator) {}
 
  MultiBufAllocatorAsync(MultiBufAllocatorAsync&) = delete;
  MultiBufAllocatorAsync& operator=(MultiBufAllocatorAsync&) = delete;
  MultiBufAllocatorAsync(MultiBufAllocatorAsync&&) = delete;
  MultiBufAllocatorAsync& operator=(MultiBufAllocatorAsync&&) = delete;
 
  MultiBufAllocationFuture AllocateAsync(size_t size);
 
  MultiBufAllocationFuture AllocateAsync(size_t min_size, size_t desired_size);
 
  MultiBufAllocationFuture AllocateContiguousAsync(size_t size);
 
  MultiBufAllocationFuture AllocateContiguousAsync(size_t min_size,
                                                   size_t desired_size);
 
 private:
  MultiBufAllocator& mbuf_allocator_;
};
 
class MultiBufAllocationFuture
    : public MultiBufAllocator::MemoryAvailableDelegate {
 public:
  constexpr explicit MultiBufAllocationFuture(MultiBufAllocator& allocator)
      : allocator_(&allocator),
        min_size_(0),
        desired_size_(0),
        contiguity_requirement_(kAllowDiscontiguous) {}
  MultiBufAllocationFuture(MultiBufAllocator& allocator,
                           size_t min_size,
                           size_t desired_size,
                           ContiguityRequirement contiguity_requirement)
      : allocator_(&allocator),
        min_size_(min_size),
        desired_size_(desired_size),
        contiguity_requirement_(contiguity_requirement) {}
 
  MultiBufAllocationFuture(MultiBufAllocationFuture&&);
  MultiBufAllocationFuture& operator=(MultiBufAllocationFuture&&);
  ~MultiBufAllocationFuture() override;
 
  void SetDesiredSize(size_t min_size) {
    SetDesiredSizes(min_size, min_size, kAllowDiscontiguous);
  }
  void SetDesiredSizes(size_t min_size,
                       size_t desired_size,
                       ContiguityRequirement contiguity_requirement);
  async2::Poll<std::optional<MultiBuf>> Pend(async2::Context& cx);
 
  MultiBufAllocator& allocator() { return *allocator_; }
  size_t min_size() const { return min_size_; }
  size_t desired_size() const { return desired_size_; }
  bool needs_contiguous() const {
    return contiguity_requirement_ == kNeedsContiguous;
  }
 
 private:
  friend class MultiBufAllocator;
 
  // Handle new memory being available. Return true if our need has been met (so
  // we can be destroyed by owner).
  bool HandleMemoryAvailable(MultiBufAllocator& alloc,
                             size_t size_available,
                             size_t contiguous_size_available) const override;
 
  async2::Poll<std::optional<MultiBuf>> TryAllocate();
 
  // The allocator this future is tied to.
  MultiBufAllocator* allocator_;
 
  // The waker to wake when a suitably-sized allocation becomes available.
  async2::Waker waker_;
 
  // The properties of the kind of allocation being waited for.
  //
  // These properties can only be mutated while holding the allocator's lock,
  // however the MultiBufAllocationFuture owner can freely read these values
  // without needing to acquire the lock.
  //
  // The allocator may read these values so long as this value is listed and
  // the allocator holds the lock.
  size_t min_size_;
  size_t desired_size_;
  ContiguityRequirement contiguity_requirement_;
};
 
}  // namespace pw::multibuf