be/src/io/cache/shard_mem_cache.cpp

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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
//
//   http://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 "io/cache/shard_mem_cache.h"

#include <bvar/bvar.h>

#include <cstring>
#include <memory>
#include <mutex>

#include "common/logging.h"

namespace doris::io {

bvar::Adder<uint64_t> g_mem_file_cache_duplicate_key_replace_num(
        "file_cache", "memory_duplicate_key_replace_num");

Status ShardMemHashTable::remove(const FileCacheKey& key) {
    std::unique_lock<std::shared_mutex> lock(_shard_mt);
    auto map_key = std::make_pair(key.hash, key.offset);
    auto iter = _cache_map.find(map_key);
    if (iter == _cache_map.end()) {
        LOG_WARNING("key not found in cache map")
                .tag("hash", key.hash.to_string())
                .tag("offset", key.offset);
        return Status::IOError("key not found in in-memory cache map when remove");
    }
    _cache_map.erase(iter);
    return Status::OK();
}

Status ShardMemHashTable::append(const FileCacheKey& key, const Slice& value) {
    return appendv(key, &value, 1);
}

Status ShardMemHashTable::appendv(const FileCacheKey& key, const Slice* values, size_t value_cnt) {
    std::unique_lock<std::shared_mutex> lock(_shard_mt);

    auto map_key = std::make_pair(key.hash, key.offset);
    auto iter = _cache_map.find(map_key);
    size_t total_size = 0;
    for (size_t idx = 0; idx < value_cnt; ++idx) {
        total_size += values[idx].size;
    }
    if (total_size == 0) {
        return Status::InvalidArgument("appendv requires non-empty slices");
    }

    MemBlock mem_block;
    mem_block.addr = std::shared_ptr<char[]>(new char[total_size], std::default_delete<char[]>());
    mem_block.size = total_size;
    DCHECK(mem_block.addr != nullptr);
    char* dst = mem_block.addr.get();
    for (size_t idx = 0; idx < value_cnt; ++idx) {
        memcpy(dst, values[idx].data, values[idx].size);
        dst += values[idx].size;
    }
    if (iter != _cache_map.end()) {
        g_mem_file_cache_duplicate_key_replace_num << 1;
        LOG_WARNING("replace duplicate key in in-memory cache map")
                .tag("hash", key.hash.to_string())
                .tag("offset", key.offset)
                .tag("old_size", iter->second.size)
                .tag("new_size", total_size);
        // Replacing the payload keeps existing readers alive while reusing the same key.
        iter->second = std::move(mem_block);
        return Status::OK();
    }

    _cache_map.emplace(map_key, std::move(mem_block));
    return Status::OK();
}

Status ShardMemHashTable::append_iobuf(const FileCacheKey& key, const butil::IOBuf& value) {
    std::unique_lock<std::shared_mutex> lock(_shard_mt);

    auto map_key = std::make_pair(key.hash, key.offset);
    auto iter = _cache_map.find(map_key);

    const size_t total_size = value.length();
    if (total_size == 0) {
        return Status::InvalidArgument("append_iobuf requires non-empty payload");
    }

    MemBlock mem_block;
    mem_block.payload.append(value);
    mem_block.size = total_size;
    mem_block.use_iobuf = true;
    if (iter != _cache_map.end()) {
        g_mem_file_cache_duplicate_key_replace_num << 1;
        LOG_WARNING("replace duplicate key in in-memory cache map")
                .tag("hash", key.hash.to_string())
                .tag("offset", key.offset)
                .tag("old_size", iter->second.size)
                .tag("new_size", total_size);
        iter->second = std::move(mem_block);
        return Status::OK();
    }

    _cache_map.emplace(map_key, std::move(mem_block));
    return Status::OK();
}

Status ShardMemHashTable::read(const FileCacheKey& key, size_t value_offset, Slice buffer) {
    std::shared_lock<std::shared_mutex> lock(_shard_mt);
    auto map_key = std::make_pair(key.hash, key.offset);
    auto iter = _cache_map.find(map_key);
    if (iter == _cache_map.end()) {
        LOG_WARNING("key not found in cache map")
                .tag("hash", key.hash.to_string())
                .tag("offset", key.offset);
        return Status::IOError("key not found in in-memory cache map when read");
    }

    const auto& mem_block = iter->second;
    if (value_offset > mem_block.size || buffer.size > mem_block.size - value_offset) {
        return Status::InternalError(
                "read buffer exceeds cached block, key={}, offset={}, request_offset={}, "
                "request_size={}, cached_size={}",
                key.hash.to_string(), key.offset, value_offset, buffer.size, mem_block.size);
    }
    if (mem_block.use_iobuf) {
        const size_t copied = mem_block.payload.copy_to(buffer.data, buffer.size, value_offset);
        if (copied != buffer.size) {
            return Status::InternalError(
                    "short read from iobuf cache block, key={}, offset={}, request_offset={}, "
                    "request_size={}, actual_size={}",
                    key.hash.to_string(), key.offset, value_offset, buffer.size, copied);
        }
        return Status::OK();
    }
    DCHECK(mem_block.addr != nullptr);
    memcpy(buffer.data, mem_block.addr.get() + value_offset, buffer.size);
    return Status::OK();
}

Status ShardMemHashTable::read_to_iobuf(const FileCacheKey& key, size_t value_offset,
                                        size_t bytes_req, butil::IOBuf* out, size_t* bytes_read) {
    if (out == nullptr || bytes_read == nullptr) {
        return Status::InvalidArgument("read_to_iobuf requires non-null out and bytes_read");
    }
    std::shared_lock<std::shared_mutex> lock(_shard_mt);
    auto map_key = std::make_pair(key.hash, key.offset);
    auto iter = _cache_map.find(map_key);
    if (iter == _cache_map.end()) {
        LOG_WARNING("key not found in cache map")
                .tag("hash", key.hash.to_string())
                .tag("offset", key.offset);
        return Status::IOError("key not found in in-memory cache map when read_to_iobuf");
    }
    const auto& mem_block = iter->second;
    if (value_offset > mem_block.size || bytes_req > mem_block.size - value_offset) {
        return Status::InternalError(
                "read iobuf exceeds cached block, key={}, offset={}, request_offset={}, "
                "request_size={}, cached_size={}",
                key.hash.to_string(), key.offset, value_offset, bytes_req, mem_block.size);
    }
    if (mem_block.use_iobuf) {
        const size_t appended = mem_block.payload.append_to(out, bytes_req, value_offset);
        if (appended != bytes_req) {
            return Status::InternalError(
                    "short append_to from iobuf cache block, key={}, offset={}, request_offset={}, "
                    "request_size={}, actual_size={}",
                    key.hash.to_string(), key.offset, value_offset, bytes_req, appended);
        }
    } else {
        DCHECK(mem_block.addr != nullptr);
        out->append(mem_block.addr.get() + value_offset, bytes_req);
    }
    *bytes_read = bytes_req;
    return Status::OK();
}

Status ShardMemHashTable::clear() {
    std::unique_lock<std::shared_mutex> lock(_shard_mt);
    _cache_map.clear();
    return Status::OK();
}

} // namespace doris::io