Conversation
Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>
Summary of ChangesHello @XucSh, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a crucial local caching mechanism for Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a local caching mechanism to improve performance for get operations, covering get_buffer, get_into, and their batch versions, and includes coordination with the master service to manage local replicas, along with a thorough test suite. However, critical security vulnerabilities have been identified, including a potential heap buffer overflow in the CacheLocal function due to inconsistent object size assumptions, a risk of arbitrary memory access from the Python wrapper's use of raw memory pointers, and a race condition in the caching logic that could lead to resource leaks on the master server. Furthermore, there is an inefficiency in cache population for get_into operations, some code duplication, and a misleading comment that need addressing.
| uintptr_t buffer_ptr, size_t size, bool local_cache = false) { | ||
| // Get data directly into user-provided buffer | ||
| void *buffer = reinterpret_cast<void *>(buffer_ptr); |
There was a problem hiding this comment.
The get_into, batch_get_into, put_from, and put_from_with_metadata functions accept raw memory addresses as uintptr_t from Python and cast them directly to void* (e.g., line 1336). This allows a Python script to read from or write to arbitrary memory locations within the process. In shared or multi-tenant environments, this can lead to sensitive information disclosure, process crashes (if NULL or invalid pointers are passed), or Remote Code Execution. The library should ideally validate that these pointers belong to registered memory regions or use safer abstractions like BufferHandle objects instead of raw pointers.
| auto& mem_desc = target_desc.get_memory_descriptor(); | ||
| void* target_addr = | ||
| reinterpret_cast<void*>(mem_desc.buffer_descriptor.buffer_address_); | ||
| memcpy(target_addr, data_ptr, data_size); |
There was a problem hiding this comment.
Potential heap buffer overflow in Client::CacheLocal. The memcpy operation at line 1518 uses data_size provided by the caller without verifying it against the actual size of the allocated buffer (mem_desc.buffer_descriptor.size_). In a distributed environment, the object size on the master could change between the initial query and the CacheLocalStart call, leading to an allocation smaller than the data being copied. Additionally, the destination address target_addr should be validated to ensure it resides within a memory segment previously mounted by the client to prevent arbitrary memory writes if the master is compromised.
| auto end_result = master_client_.CacheLocalEnd(key, replica_id); | ||
| if (!end_result) { | ||
| LOG(ERROR) << "CacheLocal: CacheLocalEnd failed for key=" << key | ||
| << ", error=" << toString(end_result.error()); | ||
| return tl::unexpected(end_result.error()); | ||
| } |
There was a problem hiding this comment.
There is a potential resource leak here. If master_client_.CacheLocalEnd(key, replica_id) fails, the replica allocated by CacheLocalStart is not revoked. This will leave the replica in an INITIALIZED state on the master, consuming resources indefinitely. You should call master_client_.CacheLocalRevoke(key, replica_id) in the failure case to clean up the allocated replica.
| auto end_result = master_client_.CacheLocalEnd(key, replica_id); | |
| if (!end_result) { | |
| LOG(ERROR) << "CacheLocal: CacheLocalEnd failed for key=" << key | |
| << ", error=" << toString(end_result.error()); | |
| return tl::unexpected(end_result.error()); | |
| } | |
| auto end_result = master_client_.CacheLocalEnd(key, replica_id); | |
| if (!end_result) { | |
| LOG(ERROR) << "CacheLocal: CacheLocalEnd failed for key=" << key | |
| << ", error=" << toString(end_result.error()); | |
| master_client_.CacheLocalRevoke(key, replica_id); | |
| return tl::unexpected(end_result.error()); | |
| } |
| if (client_ && !local_cache_.Contains(key)) { | ||
| auto cache_result = | ||
| client_->CacheLocal(key, buffer, result.value()); | ||
| if (cache_result) { | ||
| auto alloc_result = | ||
| client_buffer_allocator_->allocate(result.value()); | ||
| if (alloc_result) { | ||
| memcpy(alloc_result->ptr(), buffer, result.value()); | ||
| auto handle = std::make_shared<BufferHandle>( | ||
| std::move(*alloc_result)); | ||
| local_cache_.Insert(key, handle, cache_result.value(), | ||
| result.value()); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
The logic for populating the cache on a get_into miss seems inefficient. It performs two separate memcpy operations from the user-provided buffer: one inside client_->CacheLocal to create a segment replica, and another to populate a buffer from client_buffer_allocator_. This results in reading from the user buffer twice.
To improve this and make the data flow clearer, consider a more linear approach:
- Allocate a buffer from
client_buffer_allocator_. - Copy data from the user buffer into this new buffer once.
- Use this new buffer as the single source to create the segment replica via
client_->CacheLocal.
if (client_ && !local_cache_.Contains(key)) {
auto alloc_result =
client_buffer_allocator_->allocate(result.value());
if (alloc_result) {
// First, copy data to our managed buffer
memcpy(alloc_result->ptr(), buffer, result.value());
// Then, create the segment replica from our managed buffer
auto cache_result =
client_->CacheLocal(key, alloc_result->ptr(), result.value());
if (cache_result) {
auto handle = std::make_shared<BufferHandle>(
std::move(*alloc_result));
local_cache_.Insert(key, handle, cache_result.value(),
result.value());
}
}
}| // Cache the result locally via master-coordinated allocation | ||
| if (client_) { | ||
| auto cache_result = | ||
| client_->CacheLocal(key, result->ptr(), result->size()); |
There was a problem hiding this comment.
Race condition in get_buffer, get_into, and batch_get_into when local_cache=true. Multiple threads requesting the same key simultaneously may all miss the cache and call client_->CacheLocal. Since MasterService::CacheLocalStart allocates a new replica if the existing one is not yet COMPLETE, this can lead to multiple redundant replicas being allocated on the master for the same key and segment. If eviction is disabled (as noted in the PR description), this results in a permanent resource leak on the master, potentially leading to Denial of Service due to memory exhaustion.
| * @return pointer to CacheEntry if found, nullptr otherwise. | ||
| * The pointer is valid only while the caller holds no | ||
| * write lock (i.e., use under shared lock externally, | ||
| * or copy the entry). | ||
| */ |
There was a problem hiding this comment.
The documentation for LookupEntry is misleading. It states that the function returns a pointer to CacheEntry, but the implementation returns a std::optional<CacheEntry>, which is a value type (a copy of the entry). This could confuse future developers about the lifetime of the returned object. The comment should be updated to reflect that it returns a copy of the cache entry.
Suggested change:
* @return An optional containing a copy of the CacheEntry if found,
* std::nullopt otherwise.| // Cache locally on success only when local_cache=true | ||
| if (local_cache && internal_results[i] && client_ && | ||
| !local_cache_.Contains(miss_keys[i])) { | ||
| int64_t data_size = internal_results[i].value(); | ||
| auto cache_result = | ||
| client_->CacheLocal(miss_keys[i], miss_buffers[i], data_size); | ||
| if (cache_result) { | ||
| auto alloc_result = | ||
| client_buffer_allocator_->allocate(data_size); | ||
| if (alloc_result) { | ||
| memcpy(alloc_result->ptr(), miss_buffers[i], data_size); | ||
| auto handle = std::make_shared<BufferHandle>( | ||
| std::move(*alloc_result)); | ||
| local_cache_.Insert(miss_keys[i], handle, | ||
| cache_result.value(), data_size); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
The cache population logic for a cache miss in batch_get_into is very similar to the logic in get_into. This code duplication could be avoided by extracting the caching logic into a private helper function. This would improve maintainability and ensure that any future bug fixes or optimizations to the caching logic are applied in both places.
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
2 participants
Description
Now this feature has to disable evict.
Test result:
--- Benchmark: Data size sweep (get_buffer) ---
Size | Avg Miss | Avg Hit | Speedup | Throughput
---------+------------+------------+----------+-------------
4 KB | 0.276ms | 0.001ms | 191.4x | 2.65 GB/s
64 KB | 0.174ms | 0.001ms | 150.1x | 52.50 GB/s
256 KB | 0.245ms | 0.002ms | 109.2x | 108.75 GB/s
1 MB | 0.590ms | 0.002ms | 277.5x | 459.56 GB/s
4 MB | 1.959ms | 0.002ms | 824.8x | 1644.83 GB/s
--- Benchmark: Data size sweep (get_into) ---
Size | Avg Miss | Avg Hit | Speedup | Throughput
---------+------------+------------+----------+-------------
4 KB | 0.042ms | 0.019ms | 2.2x | 0.20 GB/s
64 KB | 0.031ms | 0.027ms | 1.1x | 2.23 GB/s
256 KB | 0.053ms | 0.049ms | 1.1x | 4.94 GB/s
1 MB | 0.165ms | 0.163ms | 1.0x | 6.00 GB/s
4 MB | 0.604ms | 0.602ms | 1.0x | 6.48 GB/s
Type of Change
How Has This Been Tested?
Checklist
./scripts/code_format.shbefore submitting.