Skip to content

perf(umbp): take the peer pool off the resolve critical path #629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
TianDi101 wants to merge 1 commit into from
Closed

perf(umbp): take the peer pool off the resolve critical path #629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 67 additions & 34 deletions src/umbp/distributed/pool/peer_pool.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,10 +70,12 @@ void PeerPool::StopTransitionWorker() {

void PeerPool::TouchLocked(const std::string& key) {
auto access = last_access_.find(key);
if (access != last_access_.end()) access_order_.erase(access->second);
const uint64_t stamp = ++access_clock_;
last_access_[key] = stamp;
access_order_[stamp] = key;
if (access != last_access_.end()) {
access_order_.splice(access_order_.begin(), access_order_, access->second);
return;
}
access_order_.push_front(key);
last_access_.emplace(key, access_order_.begin());
}

void PeerPool::ForgetAccessLocked(const std::string& key) {
Expand Down Expand Up @@ -130,10 +132,10 @@ size_t PeerPool::EnqueueWatermarkCandidatesLocked(LogicalTierGraph::TierIndex ti
candidates.push_back({TierTransitionKind::kOffload, key, backend_id, tier_index});
};

// Coldest first: access_order_ is keyed by access stamp, so the keys least
// likely to be read again leave the tier before the rest.
for (const auto& [stamp, key] : access_order_) {
(void)stamp;
// Coldest first: access_order_ keeps the hottest key at the front, so walking
// it backwards reaches the keys least likely to be read again first.
for (auto it = access_order_.rbegin(); it != access_order_.rend(); ++it) {
const std::string& key = *it;
auto placement = placements_.find(key);
if (placement == placements_.end()) continue;
consider(key, placement->second);
Expand Down Expand Up @@ -539,29 +541,52 @@ std::vector<bool> PeerPool::BatchAbort(const std::vector<PoolSlotRef>& slots) {

std::vector<PoolResolvedEntry> PeerPool::BatchResolve(const std::vector<std::string>& keys,
bool include_descs) {
std::lock_guard<std::mutex> operation_lock(operation_mutex_);
// Three phases, and only the first and last hold operation_mutex_. The backend
// lookups in between are the expensive part of a resolve, and every backend
// guards itself, so holding the pool lock across them only serializes callers
// that could have run side by side. Under TP that is every rank at once: the
// ranged-call timers showed resolve queueing into a 0 -> 30 ms staircase
// across a wave of concurrent batches, 84% of ranged-get time.
//
// placements_ is a hint, so acting on a stale one costs at most one extra
// backend probe, and phase 3 repairs it exactly as the miss path already did.
std::vector<PoolResolvedEntry> out(keys.size());
if (backends_ == nullptr || policy_ == nullptr || keys.empty()) return out;
if (keys.empty()) return out;

std::vector<std::optional<uint32_t>> preferred(keys.size());
for (size_t i = 0; i < keys.size(); ++i) {
auto it = placements_.find(keys[i]);
if (it != placements_.end()) preferred[i] = it->second;
std::vector<uint32_t> read_order;
{
std::lock_guard<std::mutex> operation_lock(operation_mutex_);
if (backends_ == nullptr || policy_ == nullptr) return out;
for (size_t i = 0; i < keys.size(); ++i) {
auto it = placements_.find(keys[i]);
if (it != placements_.end()) preferred[i] = it->second;
}
// Snapshot the order under the lock so a concurrent backend add or remove
// cannot be observed halfway through the walk below.
read_order = policy_->ReadOrder(*backends_);
}

std::vector<std::vector<bool>> attempted(
keys.size(), std::vector<bool>(BackendRegistry::kMaxBackends, false));
// One bit per (key, backend) in a flat word per key. A vector<bool> per key
// heap-allocated once per key, and a 1023-key batch is the common shape here.
static_assert(BackendRegistry::kMaxBackends <= 16, "attempted mask is 16 bits wide");
std::vector<uint16_t> attempted(keys.size(), 0);

const auto resolve = [&](uint32_t backend_id, const std::vector<size_t>& indices) {
auto* backend = backends_->Get(backend_id);
if (backend == nullptr || indices.empty()) return;
// The single-backend tier -- the common deployment -- sends every key to one
// backend in order, so hand it `keys` directly. Copying it first cost an
// allocation and a ~128-byte string copy per key for nothing.
bool identity = indices.size() == keys.size();
for (size_t i = 0; identity && i < indices.size(); ++i) identity = indices[i] == i;
std::vector<std::string> backend_keys;
backend_keys.reserve(indices.size());
if (!identity) backend_keys.reserve(indices.size());
for (size_t index : indices) {
attempted[index][backend_id] = true;
backend_keys.push_back(keys[index]);
attempted[index] |= static_cast<uint16_t>(1u << backend_id);
if (!identity) backend_keys.push_back(keys[index]);
}
auto results = backend->BatchResolve(backend_keys, include_descs);
auto results = backend->BatchResolve(identity ? keys : backend_keys, include_descs);
for (size_t i = 0; i < indices.size() && i < results.size(); ++i) {
const size_t index = indices[i];
const ResolveOutcome outcome = EffectiveResolveOutcome(results[i]);
Expand Down Expand Up @@ -594,12 +619,13 @@ std::vector<PoolResolvedEntry> PeerPool::BatchResolve(const std::vector<std::str
}
for (const auto& [backend_id, indices] : preferred_groups) resolve(backend_id, indices);

const auto read_order = policy_->ReadOrder(*backends_);
for (uint32_t backend_id : read_order) {
if (backend_id >= BackendRegistry::kMaxBackends) continue;
std::vector<size_t> unresolved;
for (size_t i = 0; i < keys.size(); ++i) {
if (!out[i].resolved.found && !attempted[i][backend_id]) unresolved.push_back(i);
if (!out[i].resolved.found && (attempted[i] & (1u << backend_id)) == 0) {
unresolved.push_back(i);
}
}
resolve(backend_id, unresolved);
}
Expand All @@ -612,6 +638,22 @@ std::vector<PoolResolvedEntry> PeerPool::BatchResolve(const std::vector<std::str
});
if (batch_busy) return out;

// Probe the backends for the keys that resolved to nothing before retaking the
// lock: Contains() is another backend call, and phase 3 exists to mutate maps,
// not to wait on media.
std::vector<std::optional<uint32_t>> repaired(keys.size());
for (size_t i = 0; i < keys.size(); ++i) {
if (out[i].resolved.found) continue;
for (uint32_t backend_id : read_order) {
auto* backend = backends_->Get(backend_id);
if (backend != nullptr && backend->Contains(keys[i])) {
repaired[i] = backend_id;
break;
}
}
}

std::lock_guard<std::mutex> operation_lock(operation_mutex_);
for (size_t i = 0; i < keys.size(); ++i) {
if (out[i].resolved.found) {
placements_[keys[i]] = out[i].backend_id;
Expand All @@ -628,22 +670,13 @@ std::vector<PoolResolvedEntry> PeerPool::BatchResolve(const std::vector<std::str
{TierTransitionKind::kPromotion, keys[i], out[i].backend_id, *source_tier});
}
}
} else if (repaired[i].has_value()) {
placements_[keys[i]] = *repaired[i];
} else {
bool exists = false;
for (uint32_t backend_id : read_order) {
auto* backend = backends_->Get(backend_id);
if (backend != nullptr && backend->Contains(keys[i])) {
placements_[keys[i]] = backend_id;
exists = true;
break;
}
}
// A key the backends no longer hold must leave the access index too, or
// it accumulates there for the life of the process and skews LRU order.
if (!exists) {
placements_.erase(keys[i]);
ForgetAccessLocked(keys[i]);
}
placements_.erase(keys[i]);
ForgetAccessLocked(keys[i]);
}
}
return out;
Expand Down
23 changes: 14 additions & 9 deletions src/umbp/include/umbp/distributed/pool/peer_pool.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <list>
#include <map>
#include <memory>
#include <mutex>
Expand Down Expand Up @@ -156,9 +157,11 @@ class PeerPool {
// Pool favors correctness over concurrent dispatch; later policies can shard
// this by key without changing the public interface.
mutable std::mutex operation_mutex_;
// Ordered so watermark offload can walk candidates in key order without
// materializing and sorting every placement first.
std::map<std::string, uint32_t> placements_;
// Unordered on purpose: every use is a point lookup, insert or erase, and
// BatchResolve publishes one entry per key with the pool lock held, where an
// ordered tree spends that critical section on string comparisons for an
// ordering nothing reads.
std::unordered_map<std::string, uint32_t> placements_;
// A migration may install its target while an independent read lease delays
// source deletion. Eviction retries must drain only that source, never fan
// out and delete the durable target.
Expand All @@ -170,16 +173,18 @@ class PeerPool {
std::condition_variable transition_idle_cv_;
std::unordered_map<std::string, PendingPlacement> pending_keys_;
std::map<std::pair<uint32_t, uint64_t>, std::string> pending_slots_;
std::unordered_map<std::string, uint64_t> last_access_;
std::unordered_map<std::string, std::list<std::string>::iterator> last_access_;
// Reads served for a key from a tier whose trigger is kOnHits. Separate from
// last_access_, which is an LRU stamp and carries no count. Cleared when the
// last_access_, which is an LRU position and carries no count. Cleared when the
// promotion is queued and whenever the key leaves the access index, so it
// cannot outlive the key it counts for.
std::unordered_map<std::string, uint32_t> promote_hit_counts_;
// Reverse index of last_access_, so the least recently used key is the first
// element rather than the result of sorting every placement.
std::map<uint64_t, std::string> access_order_;
uint64_t access_clock_ = 0;
// Reverse index of last_access_, hottest at the front, so the least recently
// used key is reached without sorting every placement. An intrusive list
// rather than a stamp-keyed tree: BatchResolve touches one entry per resolved
// key with the pool lock held, and a tree insert paid a node allocation plus
// a copy of the ~128-byte key each time. A splice pays neither.
std::list<std::string> access_order_;
// Earliest next scan per tier, moved forward only when a scan found nothing
// to queue. A pressured tier with no candidates must not make every commit
// batch pay for another walk.
Expand Down
5 changes: 5 additions & 0 deletions src/umbp/include/umbp/standalone/standalone_process_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,11 @@ class StandaloneProcessClient : public IUMBPClient {
// Resolves `ptr` against the registered host regions. On success writes the
// region-relative `offset` and the matched region's worker VA `region_base`.
bool OffsetFor(uintptr_t ptr, size_t size, uint64_t* offset, uint64_t* region_base) const;
// The batch paths translate one pointer per RANGE -- a 1023-key ranged get
// carries roughly 5k of them -- so they take registration_mu_ once for the
// whole request and call this instead of a lock round trip per range.
bool OffsetForLocked(uintptr_t ptr, size_t size, uint64_t* offset,
uint64_t* region_base) const;
// Not const: a successful Ping is where the server's backend mode and ranged
// capability become known, and they are cached on the client.
bool WaitReady(int timeout_ms);
Expand Down
61 changes: 36 additions & 25 deletions src/umbp/standalone/standalone_process_client.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,11 @@ void StandaloneProcessClient::MaybeAutoStart() {
bool StandaloneProcessClient::OffsetFor(uintptr_t ptr, size_t size, uint64_t* offset,
uint64_t* region_base) const {
std::lock_guard<std::mutex> lock(registration_mu_);
return OffsetForLocked(ptr, size, offset, region_base);
}

bool StandaloneProcessClient::OffsetForLocked(uintptr_t ptr, size_t size, uint64_t* offset,
uint64_t* region_base) const {
for (const auto& region : regions_) {
if (ptr < region.base) continue;
uintptr_t rel = ptr - region.base;
Expand Down Expand Up @@ -436,18 +441,21 @@ std::vector<bool> StandaloneProcessClient::BatchGetRanges(

::umbp::BatchRangeDataRequest req;
req.set_client_id(ClientId());
for (size_t i = 0; i < keys.size(); ++i) {
if (dsts[i].size() > std::numeric_limits<uint32_t>::max()) return failed;
req.add_keys(keys[i]);
req.add_range_counts(static_cast<uint32_t>(dsts[i].size()));
for (size_t j = 0; j < dsts[i].size(); ++j) {
uint64_t shm_offset = 0;
uint64_t region_base = 0;
if (!OffsetFor(dsts[i][j], sizes[i][j], &shm_offset, &region_base)) return failed;
req.add_shm_offsets(shm_offset);
req.add_region_bases(region_base);
req.add_sizes(sizes[i][j]);
req.add_object_offsets(src_offsets[i][j]);
{
std::lock_guard<std::mutex> registration_lock(registration_mu_);
for (size_t i = 0; i < keys.size(); ++i) {
if (dsts[i].size() > std::numeric_limits<uint32_t>::max()) return failed;
req.add_keys(keys[i]);
req.add_range_counts(static_cast<uint32_t>(dsts[i].size()));
for (size_t j = 0; j < dsts[i].size(); ++j) {
uint64_t shm_offset = 0;
uint64_t region_base = 0;
if (!OffsetForLocked(dsts[i][j], sizes[i][j], &shm_offset, &region_base)) return failed;
req.add_shm_offsets(shm_offset);
req.add_region_bases(region_base);
req.add_sizes(sizes[i][j]);
req.add_object_offsets(src_offsets[i][j]);
}
}
}

Expand All @@ -472,19 +480,22 @@ std::vector<bool> StandaloneProcessClient::BatchPutRanges(

::umbp::BatchRangeDataRequest req;
req.set_client_id(ClientId());
for (size_t i = 0; i < keys.size(); ++i) {
if (srcs[i].size() > std::numeric_limits<uint32_t>::max()) return failed;
req.add_keys(keys[i]);
req.add_object_sizes(object_sizes[i]);
req.add_range_counts(static_cast<uint32_t>(srcs[i].size()));
for (size_t j = 0; j < srcs[i].size(); ++j) {
uint64_t shm_offset = 0;
uint64_t region_base = 0;
if (!OffsetFor(srcs[i][j], sizes[i][j], &shm_offset, &region_base)) return failed;
req.add_shm_offsets(shm_offset);
req.add_region_bases(region_base);
req.add_sizes(sizes[i][j]);
req.add_object_offsets(dst_offsets[i][j]);
{
std::lock_guard<std::mutex> registration_lock(registration_mu_);
for (size_t i = 0; i < keys.size(); ++i) {
if (srcs[i].size() > std::numeric_limits<uint32_t>::max()) return failed;
req.add_keys(keys[i]);
req.add_object_sizes(object_sizes[i]);
req.add_range_counts(static_cast<uint32_t>(srcs[i].size()));
for (size_t j = 0; j < srcs[i].size(); ++j) {
uint64_t shm_offset = 0;
uint64_t region_base = 0;
if (!OffsetForLocked(srcs[i][j], sizes[i][j], &shm_offset, &region_base)) return failed;
req.add_shm_offsets(shm_offset);
req.add_region_bases(region_base);
req.add_sizes(sizes[i][j]);
req.add_object_offsets(dst_offsets[i][j]);
}
}
}

Expand Down