Merge pull request #956 from nabilwadih:nwadih/resolveCmakeWarning PiperOrigin-RevId: 602869243
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