Skip to content

Commit

Permalink
fix: Implement singleflight for segcore ChunkCache (#34250)
Browse files Browse the repository at this point in the history
See also #34249

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
  • Loading branch information
congqixia committed Jul 1, 2024
1 parent f7ecafe commit 14e827d
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 10 deletions.
40 changes: 32 additions & 8 deletions internal/core/src/storage/ChunkCache.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -15,32 +15,56 @@
// limitations under the License.

#include "ChunkCache.h"
#include <future>
#include <memory>
#include "common/Types.h"
#include "mmap/Utils.h"

namespace milvus::storage {
std::shared_ptr<ColumnBase>
ChunkCache::Read(const std::string& filepath,
const MmapChunkDescriptorPtr& descriptor) {
// use rlock to get future
{
std::shared_lock lck(mutex_);
auto it = columns_.find(filepath);
if (it != columns_.end()) {
AssertInfo(it->second, "unexpected null column, file={}", filepath);
return it->second;
lck.unlock();
auto result = it->second.second.get();
AssertInfo(result, "unexpected null column, file={}", filepath);
return result;
}
}

auto field_data = DownloadAndDecodeRemoteFile(cm_.get(), filepath);

// lock for mutation
std::unique_lock lck(mutex_);
// double check no-futurn
auto it = columns_.find(filepath);
if (it != columns_.end()) {
return it->second;
lck.unlock();
auto result = it->second.second.get();
AssertInfo(result, "unexpected null column, file={}", filepath);
return result;
}

std::promise<std::shared_ptr<ColumnBase>> p;
std::shared_future<std::shared_ptr<ColumnBase>> f = p.get_future();
columns_.emplace(filepath, std::make_pair(std::move(p), f));
lck.unlock();

// release lock and perform download and decode
// other thread request same path shall get the future.
auto field_data = DownloadAndDecodeRemoteFile(cm_.get(), filepath);
auto column = Mmap(field_data->GetFieldData(), descriptor);

// set promise value to notify the future
lck.lock();
it = columns_.find(filepath);
if (it != columns_.end()) {
// check pair exists then set value
it->second.first.set_value(column);
}
lck.unlock();
AssertInfo(column, "unexpected null column, file={}", filepath);
columns_.emplace(filepath, column);
return column;
}

Expand All @@ -58,7 +82,7 @@ ChunkCache::Prefetch(const std::string& filepath) {
return;
}

auto column = it->second;
auto column = it->second.second.get();
auto ok = madvise(
reinterpret_cast<void*>(const_cast<char*>(column->MmappedData())),
column->ByteSize(),
Expand Down
8 changes: 6 additions & 2 deletions internal/core/src/storage/ChunkCache.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
// limitations under the License.

#pragma once
#include <future>
#include <unordered_map>
#include "storage/MmapChunkManager.h"
#include "mmap/Column.h"

Expand Down Expand Up @@ -60,8 +62,10 @@ class ChunkCache {
CachePath(const std::string& filepath);

private:
using ColumnTable =
std::unordered_map<std::string, std::shared_ptr<ColumnBase>>;
using ColumnTable = std::unordered_map<
std::string,
std::pair<std::promise<std::shared_ptr<ColumnBase>>,
std::shared_future<std::shared_ptr<ColumnBase>>>>;

private:
mutable std::shared_mutex mutex_;
Expand Down

0 comments on commit 14e827d

Please sign in to comment.