#Centipede Add `ResourcePool`: gate concurrent execution on resource availability The caller sets the total available resource pool upfront. Each contending thread should determine a conservative estimate of its lifetime's peak resource requirements, and requests to lease that amount from the pool. The call blocks until that amount becomes available, as the concurrent threads finish and release their own leased resources. The currently supported resources (by choice) are `RUsageMemory` and `RUsageTiming` (CPU etc.). PiperOrigin-RevId: 602827176
diff --git a/centipede/BUILD b/centipede/BUILD index b2ad61f..c6c0511 100644 --- a/centipede/BUILD +++ b/centipede/BUILD
@@ -352,6 +352,22 @@ ) cc_library( + name = "resource_pool", + srcs = ["resource_pool.cc"], + hdrs = ["resource_pool.h"], + deps = [ + ":rusage_stats", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/time", + ], +) + +cc_library( name = "stats", srcs = ["stats.cc"], hdrs = ["stats.h"], @@ -1747,3 +1763,18 @@ "@com_google_googletest//:gtest_main", ], ) + +cc_test( + name = "resource_pool_test", + srcs = ["resource_pool_test.cc"], + deps = [ + ":logging", + ":resource_pool", + ":rusage_stats", + ":thread_pool", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + ], +)
diff --git a/centipede/resource_pool.cc b/centipede/resource_pool.cc new file mode 100644 index 0000000..10c79a4 --- /dev/null +++ b/centipede/resource_pool.cc
@@ -0,0 +1,174 @@ +// Copyright 2024 The Centipede 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. + +#include "./centipede/resource_pool.h" + +#include <string> +#include <utility> + +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" +#include "absl/synchronization/mutex.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "./centipede/rusage_stats.h" + +namespace centipede::perf { + +template <typename ResourceT> +ResourcePool<ResourceT>::LeaseToken::LeaseToken( // + ResourcePool& leaser, LeaseRequest request) + : leaser_{leaser}, request_{std::move(request)} {} + +template <typename ResourceT> +ResourcePool<ResourceT>::LeaseToken::LeaseToken( // + ResourcePool& leaser, LeaseRequest request, absl::Status error) + : leaser_{leaser}, + request_{std::move(request)}, + status_{std::move(error)} {} + +template <typename ResourceT> +ResourcePool<ResourceT>::LeaseToken::~LeaseToken() { + CHECK(status_checked_) // + << "status() was never consulted by caller: " << *this; + if (status_.ok()) { + leaser_.ReturnLease(*this); + } +} + +template <typename ResourceT> +const typename ResourcePool<ResourceT>::LeaseRequest& +ResourcePool<ResourceT>::LeaseToken::request() const { + return request_; +} + +template <typename ResourceT> +const absl::Status& ResourcePool<ResourceT>::LeaseToken::status() const { + status_checked_ = true; + return status_; +} + +template <typename ResourceT> +std::string ResourcePool<ResourceT>::LeaseToken::id() const { + return absl::StrCat("lease_tid_", thread_id_, "_rid_", request_.id); +} + +template <typename ResourceT> +pid_t ResourcePool<ResourceT>::LeaseToken::thread_id() const { + return thread_id_; +} + +template <typename ResourceT> +absl::Time ResourcePool<ResourceT>::LeaseToken::created_at() const { + return created_at_; +} + +template <typename ResourceT> +absl::Duration ResourcePool<ResourceT>::LeaseToken::age() const { + return absl::Now() - created_at_; +} + +template <typename ResourceT> +ResourcePool<ResourceT>::ResourcePool(const ResourceT& quota) + : quota_{quota}, pool_{quota} { + LOG(INFO) << "Creating pool with quota=[" << quota.ShortStr() << "]"; +} + +template <typename ResourceT> +typename ResourcePool<ResourceT>::LeaseToken +ResourcePool<ResourceT>::AcquireLeaseBlocking(LeaseRequest&& request) { + if (VLOG_IS_ON(1)) { + absl::ReaderMutexLock lock{&pool_mu_}; + VLOG(1) << "Received lease request " << request.id // + << "\nrequested: " << request.amount.FormattedStr() // + << "\nquota: " << quota_.FormattedStr() // + << "\navailable: " << pool_.FormattedStr(); + } + + if (request.amount == ResourceT::Zero()) { + absl::Status error = // + absl::InvalidArgumentError(absl::StrCat( // + "Invalid lease request ", request.id, ": amount is zero")); + return LeaseToken{*this, std::move(request), std::move(error)}; + } + // NOTE: Using `amount > quota` would be semantically wrong, because it is + // true only when _all_ components of `amount` are strictly greater than their + // counterparts in `quota_`. + if (!(request.amount <= quota_)) { + absl::Status error = // + absl::ResourceExhaustedError(absl::StrCat( // + "Invalid lease request ", request.id, ": amount exceeds quota: [", + request.amount.ShortStr(), "] vs [", quota_.ShortStr(), "]")); + return LeaseToken{*this, std::move(request), std::move(error)}; + } + + const auto got_enough_free_pool = [this, &request]() { + pool_mu_.AssertReaderHeld(); + const bool got_pool = request.amount <= pool_; + if (!got_pool) { + VLOG(10) // + << "Pending lease '" << request.id << "':" // + << "\nreq age : " << request.age() // + << "\navailable : " << pool_.FormattedStr() // + << "\nrequested : " << (-request.amount).FormattedStr() // + << "\nmissing : " << (pool_ - request.amount).FormattedStr(); + } + return got_pool; + }; + + // Block and wait until enough of the pool becomes available to satisfy + // this request, then acquire the mutex and proceed to the true-branch. If + // the timeout is reached, proceed to the else-branch. + if (pool_mu_.LockWhenWithTimeout( // + absl::Condition{&got_enough_free_pool}, request.timeout)) { + VLOG(1) // + << "Granting lease " << request.id // + << "\nreq age : " << request.age() // + << "\nbefore : " << pool_.FormattedStr() // + << "\nleased : " << (-request.amount).FormattedStr() // + << "\nafter : " << (pool_ - request.amount).FormattedStr(); + pool_ = pool_ - request.amount; + pool_mu_.Unlock(); + return LeaseToken{*this, std::move(request)}; + } else { + absl::Status error = // + absl::DeadlineExceededError(absl::StrCat( // + "Lease request ", request.id, " timed out; timeout: ", + request.timeout, " requested: [", request.amount.ShortStr(), + "] current pool: [", pool_.ShortStr(), "]")); + pool_mu_.Unlock(); + return LeaseToken{*this, std::move(request), std::move(error)}; + } +} + +template <typename ResourceT> +void ResourcePool<ResourceT>::ReturnLease(const LeaseToken& lease) { + absl::WriterMutexLock lock{&pool_mu_}; + VLOG(1) // + << "Returning lease " << lease.request().id // + << "\nreq age : " << lease.request().age() // + << "\nlease age : " << lease.age() // + << "\nbefore : " << pool_.FormattedStr() // + << "\nreturned : " << (+lease.request().amount).FormattedStr() // + << "\nafter : " << (pool_ + lease.request().amount).FormattedStr(); + pool_ = pool_ + lease.request().amount; +} + +// Explicit instantiations for the currently supported `ResourceT`s. +template class ResourcePool<RUsageMemory>; +template class ResourcePool<RUsageTiming>; + +} // namespace centipede::perf
diff --git a/centipede/resource_pool.h b/centipede/resource_pool.h new file mode 100644 index 0000000..097a3d9 --- /dev/null +++ b/centipede/resource_pool.h
@@ -0,0 +1,203 @@ +// Copyright 2024 The Centipede 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. + +#ifndef FUZZTEST_CENTIPEDE_RESOURCE_RESOURCE_POOL_H_ +#define FUZZTEST_CENTIPEDE_RESOURCE_RESOURCE_POOL_H_ + +#include <sys/syscall.h> +#include <sys/types.h> +#include <unistd.h> + +#include <ostream> +#include <string> + +#include "absl/base/thread_annotations.h" +#include "absl/status/status.h" +#include "absl/synchronization/mutex.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" + +namespace centipede::perf { + +//------------------------------------------------------------------------------ +// ResourcePool +// +// `ResourcePool` is an accounting mechanism to effectively share a limited +// resource between concurrent consumer threads, never exceeding a quota while +// maximizing resource utilization, and thus parallelism. +// +// The quota amount is picked by the client. It can be arbitrary, or it can +// reflect an actual amount of the resource on the system (e.g. the available +// RAM). +// +// Each of the consumer threads determines a conservative estimate of its peak +// resource utilization, and requests that amount from the pool. The request +// blocks until a sufficient amount becomes available. The amount is then +// "leased" to the thread for as long as it holds the lease token, and +// auto-returned back to the pool via RAII. +// +// Notes on using in combination with `ThreadPool`: +// 1. The requested number of concurrent threads in a `ThreadPool` is often an +// attempt to indirectly control the resource usage. `ResourcePool` enables a +// more direct way of controlling it, and therefore `ThreadPool`'s thread +// count can be made as high as necessary for other purposes. +// 2. The `ResourcePool` object must is defined before the `ThreadPool` one to +// avoid dangling references to a destructed pool in the threads. +// +// The currently supported (and explicitly instantiated in the .cc) types of +// the `ResourceT` template argument are `RUsageMemory` and `RUsageTiming`. +// +// Example: +// +// { +// constexpr RUsageMemory kRssQuota{.mem_rss = RLimits::FreeRss() * 0.75}; +// ResourcePool rss_pool{kRssQuota}; +// ThreadPool threads{100}; +// for (...) { +// threads.Schedule([&rss_pool]() { +// // The thread blocks here until either the requested amount of RSS +// // becomes available as the peer threads return their leases, or the +// // 10-minute timeout expires. +// const ResourcePool::LeaseToken rss_lease = +// rss_pool.AcquireLeaseBlocking({ +// .id = absl::StrCat("rss_", shard_id), +// .amount = RUsageMemory{.mem_rss = EstimateShardPeakRss()}, +// .timeout = absl::Minutes(10), +// }); +// CHECK_OK(rss_lease.status()); +// ... +// } +// // `rss_lease` dtor returns the leased RSS to `rss_pool` and unblocks +// // other waiting threads. +// ); +// } +// } // `threads` dtor runs and joins the threads; then `rss_pool` dtor runs. +// +// TODO(ussuri): Add monitoring of claimed vs actual use by each leaser and +// a final report of over- and underutilization (possibly via RUsageProfiler). +//------------------------------------------------------------------------------ +template <typename ResourceT> +class ResourcePool { + public: + //---------------------------------------------------------------------------- + // Request + // + // Specifies a projected resource consumption between the time this request is + // submitted and the time the acquired LeaseToken goes out of scope. A + // convenient way to construct Requests is by using designated initializers + // (cf. ResourcePool's top-level doc just above). + struct LeaseRequest { + // Optional. Used in the debug logging and always included in returned + // failure statuses. + std::string id = ""; + // Mandatory. Must be > `ResourceT::Zero()`; otherwise, + // `AcquireLeaseBlocking()` immediately returns a failure. + ResourceT amount; + // Optional. `AcquireLeaseBlocking()` waits for up to this long for other + // resource consumers to free up enough of it to satisfy this request. If + // the required amount is still unavailable, `absl::DeadlineExceededError` + // is returned. The default is to acquire or fail immediately. + absl::Duration timeout = absl::ZeroDuration(); + // Should not normally be overridden by clients (but can be). Used for + // logging only. + absl::Time created_at = absl::Now(); + + // The age of this request. + absl::Duration age() const { return absl::Now() - created_at; } + }; + + //---------------------------------------------------------------------------- + // LeaseToken + // + // A RAII-based resource lock, similar to `MutexLock`. Must be held by a + // client that called `AcquireLeaseBlocking()` for as long as it continues to + // use the leased amount of the resource. Returns the resource to the leaser + // `ResourcePool` in the dtor. + class [[nodiscard]] LeaseToken { + public: + // Move-copyable only. + LeaseToken(const LeaseToken&) = delete; + LeaseToken& operator=(const LeaseToken&) = delete; + LeaseToken(LeaseToken&&) noexcept = default; + LeaseToken& operator=(LeaseToken&&) noexcept = delete; + + // Automatically returns itself to the leaser (the issuing ResourcePool). + ~LeaseToken(); + + // The outcome of resource acquisition (ie. of + // `ResourcePool::AcquireLeaseBlocking()`). Must be consulted by the client + // at least once, otherwise the dtor will CHECK. + const absl::Status& status() const; + // The originating request. + const LeaseRequest& request() const; + // A short description that can be used in logs. + std::string id() const; + // The thread ID that submitted the request. + pid_t thread_id() const; + // The creation time and the age of the lease. + absl::Time created_at() const; + absl::Duration age() const; + + private: + // Only ResourcePool can create. + friend class ResourcePool; + + // Constructs a token for a successfully acquired resource. + LeaseToken(ResourcePool& leaser, LeaseRequest request); + // Constructs a token for a resource that couldn't be acquired. + LeaseToken(ResourcePool& leaser, LeaseRequest request, absl::Status error); + + friend std::ostream& operator<<(std::ostream& os, const LeaseToken& lt) { + return os << lt.id() << ": " << lt.request().amount.ShortStr(); + } + + ResourcePool& leaser_; + LeaseRequest request_ = {}; + absl::Status status_ = absl::OkStatus(); + mutable bool status_checked_ = false; + pid_t thread_id_ = ::syscall(__NR_gettid); + absl::Time created_at_ = absl::Now(); + }; + + // `quota` is the initially available amount of the resource to be shared + // between all concurrent consumers. + // Example: `ResourcePool pool{RUsageMemory{.mem_rss = ComputeFreeRss()}};`. + explicit ResourcePool(const ResourceT& quota); + + // Blocks the current thread and waits until `request.amount` of the resources + // becomes available in the pool or until `request.timeout` expires, whichever + // comes first. When the returned object goes out of scope, the leased + // resource gets automatically returned to the pool via RAII. + // Example: `const auto lease = pool.AcquireLeaseBlocking({.mem_rss = 100});`. + LeaseToken AcquireLeaseBlocking(LeaseRequest&& request); + + private: + // `LeaseToken`'s dtor calls this to return the leased resource to the pool. + void ReturnLease(const LeaseToken& lease); + + // The total pool capacity. + const ResourceT quota_; + + // The currently available amount. + absl::Mutex pool_mu_; + ResourceT pool_ ABSL_GUARDED_BY(pool_mu_); +}; + +// An explicit deduction guide to allow `ResourcePool pool{RUsageMemory{...}}`. +template <typename R> +ResourcePool(R r) -> ResourcePool<R>; + +} // namespace centipede::perf + +#endif // FUZZTEST_CENTIPEDE_RESOURCE_RESOURCE_POOL_H_
diff --git a/centipede/resource_pool_test.cc b/centipede/resource_pool_test.cc new file mode 100644 index 0000000..05860cb --- /dev/null +++ b/centipede/resource_pool_test.cc
@@ -0,0 +1,139 @@ +// Copyright 2024 The Centipede 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. + +#include "./centipede/resource_pool.h" + +#include <array> +#include <cstddef> +#include <string> +#include <string_view> + +#include "gtest/gtest.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "./centipede/logging.h" +#include "./centipede/rusage_stats.h" +#include "./centipede/thread_pool.h" + +namespace centipede::perf { +namespace { + +TEST(ResourcePoolTest, InvalidLeaseRequests) { + const RUsageMemory kQuota = {.mem_rss = 1000}; + const RUsageMemory kZero = {.mem_rss = 0}; + const RUsageMemory kEpsilon = {.mem_rss = 1}; + ResourcePool pool{kQuota}; + { + const auto lease = pool.AcquireLeaseBlocking({.amount = kZero}); + EXPECT_EQ(lease.status().code(), absl::StatusCode::kInvalidArgument) + << VV(lease.status()); + } + { + const auto lease = pool.AcquireLeaseBlocking({.amount = kQuota - kEpsilon}); + EXPECT_EQ(lease.status().code(), absl::StatusCode::kOk) + << VV(lease.status()); + } + { + const auto lease = pool.AcquireLeaseBlocking({.amount = kQuota}); + EXPECT_EQ(lease.status().code(), absl::StatusCode::kOk) + << VV(lease.status()); + } + { + const auto lease = pool.AcquireLeaseBlocking({.amount = kQuota + kEpsilon}); + EXPECT_EQ(lease.status().code(), absl::StatusCode::kResourceExhausted) + << VV(lease.status()); + } +} + +TEST(ResourcePoolTest, Dynamic) { + struct TaskSpec { + std::string_view id; + RUsageMemory ram_chunk; + // The times are relative to time zero, when all the tasks roughly start. + int request_at_secs; + int timeout_at_secs; + int release_at_secs; + absl::StatusCode expected_lease_status; + }; + + constexpr RUsageMemory kRssQuota = {.mem_rss = 5}; + constexpr int kNumTasks = 9; + constexpr std::array<TaskSpec, kNumTasks> kTaskSpecs = {{ + // Can't request 0 amount. + {"0", {.mem_rss = 0}, 0, 1, 3, absl::StatusCode::kInvalidArgument}, + // Exceeds the initial pool capacity. + {"1", {.mem_rss = 10}, 0, 1, 3, absl::StatusCode::kResourceExhausted}, + // "2" gets the resource first. + {"2", {.mem_rss = 2}, 0, 0, 2, absl::StatusCode::kOk}, + // "1" gets the resource immediately after "2" and runs concurrently. + {"3", {.mem_rss = 2}, 0, 0, 4, absl::StatusCode::kOk}, + // "4" can't get the resource right away - 1 sec later than "2" and "3" - + // because they almost exhaust the pool; but it waits long enough for "2" + // to finish (while "3" is still running) and free up enough of the pool; + // then "4" gets the resource and runs fine. + {"4", {.mem_rss = 1}, 1, 3, 4, absl::StatusCode::kOk}, + // "5" starts while "2" and "3", and later on "3" and "4", are still + // running. They all continuously hold enough of the pool to prevent "5" + // from ever getting its resource. Eventually, "5" runs out of time. + {"5", {.mem_rss = 4}, 2, 3, 5, absl::StatusCode::kDeadlineExceeded}, + // "6" is like "5", but it waits long enough for "3" and "4" to free up + // the pool; then "6" gets the resource and runs fine. + {"6", {.mem_rss = 4}, 2, 5, 6, absl::StatusCode::kOk}, + // "7" is also like "5", but is less greedy, so although it starts 1 sec + // later, it is allowed in front of "5" and "6" and runs fine, partially + // sharing the pool with "3" and "4". + {"7", {.mem_rss = 1}, 3, 3, 5, absl::StatusCode::kOk}, + // "8" starts waiting for the maximum available amount when other + // consumers already use some of the pool. It waits long enough for all of + // them to finish, then finally grabs the entire quota and runs. + {"8", {.mem_rss = 5}, 2, 9, 10, absl::StatusCode::kOk}, + }}; + std::array<absl::Status, kNumTasks> task_lease_statuses; + + { + ResourcePool pool{kRssQuota}; + ThreadPool threads{kNumTasks}; + for (size_t i = 0; i < kNumTasks; ++i) { + const auto& t = kTaskSpecs[i]; + auto& lease_status = task_lease_statuses[i]; + threads.Schedule([&t, &pool, &lease_status]() { + // All the tasks start roughly at the same time (because there are just + // as many threads, and scheduling is fast), so they are on roughly the + // same relative timetable. + absl::SleepFor(absl::Seconds(t.request_at_secs)); + const auto lease = pool.AcquireLeaseBlocking({ + .id = std::string(t.id), + .amount = t.ram_chunk, + .timeout = absl::Seconds(t.timeout_at_secs - t.request_at_secs), + }); + lease_status = lease.status(); + if (lease_status.ok()) { + absl::SleepFor(absl::Seconds(t.release_at_secs - t.request_at_secs)); + } + }); + } + } // Threads join here. + + for (size_t i = 0; i < kNumTasks; ++i) { + const auto& task = kTaskSpecs[i]; + auto& lease_status = task_lease_statuses[i]; + EXPECT_EQ(lease_status.code(), task.expected_lease_status) + << VV(task.id) << VV(lease_status); + } +} + +} // namespace +} // namespace centipede::perf