missive/resources/resource_manager.cc

// Copyright 2021 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "missive/resources/resource_manager.h"

#include <atomic>
#include <cstdint>
#include <optional>
#include <utility>

#include <base/logging.h>
#include <base/memory/ref_counted.h>
#include <base/memory/scoped_refptr.h>
#include <base/sequence_checker.h>
#include <base/task/bind_post_task.h>
#include <base/task/sequenced_task_runner.h>
#include <base/task/task_traits.h>
#include <base/task/thread_pool.h>

namespace reporting {

ResourceManager::ResourceManager(uint64_t total_size)
    : total_(total_size),
      sequenced_task_runner_(base::ThreadPool::CreateSequencedTaskRunner(
          {base::TaskPriority::BEST_EFFORT})) {
  DETACH_FROM_SEQUENCE(sequence_checker_);
}

ResourceManager::~ResourceManager() = default;

bool ResourceManager::Reserve(uint64_t size) {
  uint64_t old_used = used_.fetch_add(size);
  if (old_used + size > total_) {
    used_.fetch_sub(size);
    return false;
  }
  return true;
}

void ResourceManager::Discard(uint64_t size) {
  CHECK_LE(size, used_.load());
  used_.fetch_sub(size);

  sequenced_task_runner_->PostTask(
      FROM_HERE, base::BindOnce(&ResourceManager::FlushCallbacks,
                                base::WrapRefCounted(this)));
}

uint64_t ResourceManager::GetTotal() const {
  return total_;
}

uint64_t ResourceManager::GetUsed() const {
  return used_.load();
}

void ResourceManager::Test_SetTotal(uint64_t test_total) {
  total_ = test_total;
}

void ResourceManager::RegisterCallback(uint64_t size, base::OnceClosure cb) {
  sequenced_task_runner_->PostTask(
      FROM_HERE,
      base::BindOnce(
          [](scoped_refptr<ResourceManager> self, uint64_t size,
             base::OnceClosure cb) {
            DCHECK_CALLED_ON_VALID_SEQUENCE(self->sequence_checker_);
            self->resource_callbacks_.emplace(size, std::move(cb));

            // Attempt to apply remaining callbacks
            // (this is especially important if the new callback is registered
            // with no allocations to be released - we don't want the callback
            // to wait indefinitely in this case).
            self->FlushCallbacks();
          },
          base::WrapRefCounted(this), size,
          base::BindPostTask(base::SequencedTaskRunner::GetCurrentDefault(),
                             std::move(cb))));
}

void ResourceManager::FlushCallbacks() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  int64_t remained = GetTotal() - GetUsed();  // No synchronization whatsoever.
  while (remained > 0 && !resource_callbacks_.empty()) {
    auto& cb_pair = resource_callbacks_.front();
    if (cb_pair.first > static_cast<uint64_t>(remained)) {
      break;
    }
    remained -= cb_pair.first;
    std::move(cb_pair.second).Run();
    resource_callbacks_.pop();
  }
}

ScopedReservation::ScopedReservation() noexcept = default;

ScopedReservation::ScopedReservation(
    uint64_t size, scoped_refptr<ResourceManager> resource_manager) noexcept
    : resource_manager_(resource_manager) {
  if (size == 0uL || !resource_manager->Reserve(size)) {
    return;
  }
  size_ = size;
}

ScopedReservation::ScopedReservation(
    uint64_t size, const ScopedReservation& other_reservation) noexcept
    : resource_manager_(other_reservation.resource_manager_) {
  if (size == 0uL || !resource_manager_.get() ||
      !resource_manager_->Reserve(size)) {
    return;
  }
  size_ = size;
}

ScopedReservation::ScopedReservation(ScopedReservation&& other) noexcept
    : resource_manager_(other.resource_manager_),
      size_(std::exchange(other.size_, std::nullopt)) {}

bool ScopedReservation::reserved() const {
  return size_.has_value();
}

bool ScopedReservation::Reduce(uint64_t new_size) {
  if (!reserved()) {
    return false;
  }
  if (new_size < 0 || size_.value() < new_size) {
    return false;
  }
  resource_manager_->Discard(size_.value() - new_size);
  if (new_size > 0) {
    size_ = new_size;
  } else {
    size_ = std::nullopt;
  }
  return true;
}

void ScopedReservation::HandOver(ScopedReservation& other) {
  if (resource_manager_.get()) {
    CHECK_EQ(resource_manager_.get(), other.resource_manager_.get())
        << "Reservations are not related";
  } else {
    CHECK(!reserved()) << "Unattached reservation may not have size";
    resource_manager_ = other.resource_manager_;
  }
  if (!other.reserved()) {
    return;  // Nothing changes.
  }
  const uint64_t old_size = (reserved() ? size_.value() : 0uL);
  size_ = old_size + std::exchange(other.size_, std::nullopt).value();
}

ScopedReservation::~ScopedReservation() {
  if (reserved()) {
    resource_manager_->Discard(size_.value());
  }
}
}  // namespace reporting