blob: 99b3d96a35aad64796f6f5dea2f05e143264d4ad [file]
/*
* 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 "redis_db.h"
#include <ctime>
#include <utility>
#include "cluster/redis_slot.h"
#include "common/scope_exit.h"
#include "common/string_util.h"
#include "db_util.h"
#include "parse_util.h"
#include "rocksdb/iterator.h"
#include "rocksdb/status.h"
#include "storage/iterator.h"
#include "storage/redis_metadata.h"
#include "storage/storage.h"
#include "time_util.h"
#include "types/redis_hash.h"
#include "types/redis_list.h"
#include "types/redis_set.h"
#include "types/redis_string.h"
#include "types/redis_zset.h"
namespace redis {
Database::Database(engine::Storage *storage, std::string ns)
: storage_(storage),
metadata_cf_handle_(storage->GetCFHandle(ColumnFamilyID::Metadata)),
namespace_(std::move(ns)) {}
rocksdb::Status Database::ParseMetadata(RedisTypes types, Slice *bytes, Metadata *metadata) {
std::string old_metadata;
metadata->Encode(&old_metadata);
auto s = metadata->Decode(bytes);
// delay InvalidArgument error check after type match check
if (!s.ok() && !s.IsInvalidArgument()) return s;
if (metadata->Expired()) {
// error discarded here since it already failed
auto _ [[maybe_unused]] = metadata->Decode(old_metadata);
return rocksdb::Status::NotFound(kErrMsgKeyExpired);
}
// if type is not matched, we still need to check if the metadata is valid.
if (!types.Contains(metadata->Type()) && (metadata->size > 0 || metadata->IsEmptyableType())) {
// error discarded here since it already failed
auto _ [[maybe_unused]] = metadata->Decode(old_metadata);
return rocksdb::Status::InvalidArgument(kErrMsgWrongType);
}
if (s.IsInvalidArgument()) return s;
if (metadata->size == 0 && !metadata->IsEmptyableType()) {
// error discarded here since it already failed
auto _ [[maybe_unused]] = metadata->Decode(old_metadata);
return rocksdb::Status::NotFound("no element found");
}
return s;
}
// Some data types may support reading multiple types of metadata.
// For example, bitmap supports reading string metadata and bitmap metadata.
rocksdb::Status Database::ParseMetadataWithStats(RedisTypes types, Slice *bytes, Metadata *metadata) {
auto s = ParseMetadata(types, bytes, metadata);
storage_->RecordStat(s.ok() ? engine::StatType::KeyspaceHits : engine::StatType::KeyspaceMisses, 1);
return s;
}
rocksdb::Status Database::GetMetadata(engine::Context &ctx, RedisTypes types, const Slice &ns_key, Metadata *metadata) {
std::string raw_value;
Slice rest;
return GetMetadata(ctx, types, ns_key, &raw_value, metadata, &rest);
}
rocksdb::Status Database::GetMetadata(engine::Context &ctx, RedisTypes types, const Slice &ns_key,
std::string *raw_value, Metadata *metadata, Slice *rest) {
auto s = GetRawMetadata(ctx, ns_key, raw_value);
*rest = *raw_value;
if (!s.ok()) return s;
return ParseMetadataWithStats(types, rest, metadata);
}
rocksdb::Status Database::GetRawMetadata(engine::Context &ctx, const Slice &ns_key, std::string *bytes) {
return storage_->Get(ctx, ctx.GetReadOptions(), metadata_cf_handle_, ns_key, bytes);
}
rocksdb::Status Database::Expire(engine::Context &ctx, const Slice &user_key, uint64_t timestamp) {
std::string ns_key = AppendNamespacePrefix(user_key);
std::string value;
Metadata metadata(kRedisNone, false);
rocksdb::Status s = storage_->Get(ctx, ctx.GetReadOptions(), metadata_cf_handle_, ns_key, &value);
if (!s.ok()) return s;
s = metadata.Decode(value);
if (!s.ok()) return s;
if (metadata.Expired()) {
return rocksdb::Status::NotFound(kErrMsgKeyExpired);
}
if (!metadata.IsEmptyableType() && metadata.size == 0) {
return rocksdb::Status::NotFound("no elements");
}
if (metadata.expire == timestamp) return rocksdb::Status::OK();
// +1 to skip the flags
if (metadata.Is64BitEncoded()) {
EncodeFixed64(value.data() + 1, timestamp);
} else {
EncodeFixed32(value.data() + 1, Metadata::ExpireMsToS(timestamp));
}
auto batch = storage_->GetWriteBatchBase();
WriteBatchLogData log_data(kRedisNone, {std::to_string(kRedisCmdExpire)});
s = batch->PutLogData(log_data.Encode());
if (!s.ok()) {
return s;
}
s = batch->Put(metadata_cf_handle_, ns_key, value);
if (!s.ok()) {
return s;
}
s = storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch());
return s;
}
rocksdb::Status Database::Del(engine::Context &ctx, const Slice &user_key) {
std::string ns_key = AppendNamespacePrefix(user_key);
std::string value;
rocksdb::Status s = storage_->Get(ctx, ctx.GetReadOptions(), metadata_cf_handle_, ns_key, &value);
if (!s.ok()) return s;
Metadata metadata(kRedisNone, false);
s = metadata.Decode(value);
if (!s.ok()) return s;
if (metadata.Expired()) {
return rocksdb::Status::NotFound(kErrMsgKeyExpired);
}
return storage_->Delete(ctx, storage_->DefaultWriteOptions(), metadata_cf_handle_, ns_key);
}
rocksdb::Status Database::MDel(engine::Context &ctx, const std::vector<Slice> &keys, uint64_t *deleted_cnt) {
*deleted_cnt = 0;
std::vector<std::string> ns_keys;
ns_keys.reserve(keys.size());
for (const auto &key : keys) {
std::string ns_key = AppendNamespacePrefix(key);
ns_keys.emplace_back(std::move(ns_key));
}
auto batch = storage_->GetWriteBatchBase();
WriteBatchLogData log_data(kRedisNone);
auto s = batch->PutLogData(log_data.Encode());
if (!s.ok()) {
return s;
}
std::vector<Slice> slice_keys;
slice_keys.reserve(ns_keys.size());
for (const auto &ns_key : ns_keys) {
slice_keys.emplace_back(ns_key);
}
std::vector<rocksdb::Status> statuses(slice_keys.size());
std::vector<rocksdb::PinnableSlice> pin_values(slice_keys.size());
storage_->MultiGet(ctx, ctx.DefaultMultiGetOptions(), metadata_cf_handle_, slice_keys.size(), slice_keys.data(),
pin_values.data(), statuses.data());
for (size_t i = 0; i < slice_keys.size(); i++) {
if (!statuses[i].ok() && !statuses[i].IsNotFound()) return statuses[i];
if (statuses[i].IsNotFound()) continue;
Metadata metadata(kRedisNone, false);
// Explicit construct a rocksdb::Slice to avoid the implicit conversion from
// PinnableSlice to Slice.
auto s = metadata.Decode(rocksdb::Slice(pin_values[i].data(), pin_values[i].size()));
if (!s.ok()) continue;
if (metadata.Expired()) continue;
s = batch->Delete(metadata_cf_handle_, ns_keys[i]);
if (!s.ok()) return s;
*deleted_cnt += 1;
}
if (*deleted_cnt == 0) return rocksdb::Status::OK();
return storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch());
}
rocksdb::Status Database::Exists(engine::Context &ctx, const std::vector<Slice> &keys, uint32_t *ret) {
std::vector<std::string> ns_keys;
ns_keys.reserve(keys.size());
for (const auto &key : keys) {
ns_keys.emplace_back(AppendNamespacePrefix(key));
}
return existsInternal(ctx, ns_keys, ret);
}
rocksdb::Status Database::TTL(engine::Context &ctx, const Slice &user_key, int64_t *ttl) {
std::string ns_key = AppendNamespacePrefix(user_key);
*ttl = -2; // ttl is -2 when the key does not exist or expired
std::string value;
rocksdb::Status s = storage_->Get(ctx, ctx.GetReadOptions(), metadata_cf_handle_, ns_key, &value);
if (!s.ok()) return s.IsNotFound() ? rocksdb::Status::OK() : s;
Metadata metadata(kRedisNone, false);
s = metadata.Decode(value);
if (!s.ok()) return s;
*ttl = metadata.TTL();
return rocksdb::Status::OK();
}
rocksdb::Status Database::GetExpireTime(engine::Context &ctx, const Slice &user_key, uint64_t *timestamp) {
std::string ns_key = AppendNamespacePrefix(user_key);
Metadata metadata(kRedisNone, false);
auto s = GetMetadata(ctx, RedisTypes::All(), ns_key, &metadata);
if (!s.ok()) return s;
*timestamp = metadata.expire;
return rocksdb::Status::OK();
}
rocksdb::Status Database::GetKeyNumStats(engine::Context &ctx, const std::string &prefix, KeyNumStats *stats) {
return Keys(ctx, prefix, "*", nullptr, stats);
}
rocksdb::Status Database::Keys(engine::Context &ctx, const std::string &prefix, const std::string &suffix_glob,
std::vector<std::string> *keys, KeyNumStats *stats) {
uint16_t slot_id = 0;
std::string ns_prefix;
if (namespace_ != kDefaultNamespace || keys != nullptr) {
if (storage_->IsSlotIdEncoded()) {
ns_prefix = ComposeNamespaceKey(namespace_, "", false);
if (!prefix.empty()) {
PutFixed16(&ns_prefix, slot_id);
ns_prefix.append(prefix);
}
} else {
ns_prefix = AppendNamespacePrefix(prefix);
}
}
uint64_t ttl_sum = 0;
auto iter = util::UniqueIterator(ctx, ctx.GetReadOptions(), metadata_cf_handle_);
while (true) {
ns_prefix.empty() ? iter->SeekToFirst() : iter->Seek(ns_prefix);
for (; iter->Valid(); iter->Next()) {
if (!ns_prefix.empty() && !iter->key().starts_with(ns_prefix)) {
break;
}
auto [_, user_key] = ExtractNamespaceKey(iter->key(), storage_->IsSlotIdEncoded());
if (!util::StringMatch(suffix_glob, user_key.ToString().substr(prefix.size()))) {
continue;
}
Metadata metadata(kRedisNone, false);
auto s = metadata.Decode(iter->value());
if (!s.ok()) continue;
if (metadata.Expired()) {
if (stats) stats->n_expired++;
continue;
}
if (stats) {
int64_t ttl = metadata.TTL();
stats->n_key++;
if (ttl != -1) {
stats->n_expires++;
if (ttl > 0) ttl_sum += ttl;
}
}
if (keys) {
keys->emplace_back(user_key.ToString());
}
}
if (auto s = iter->status(); !s.ok()) {
return s;
}
if (!storage_->IsSlotIdEncoded()) break;
if (prefix.empty()) break;
if (++slot_id >= HASH_SLOTS_SIZE) break;
ns_prefix = ComposeNamespaceKey(namespace_, "", false);
PutFixed16(&ns_prefix, slot_id);
ns_prefix.append(prefix);
}
if (stats && stats->n_expires > 0) {
stats->avg_ttl = ttl_sum / stats->n_expires / 1000;
}
return rocksdb::Status::OK();
}
rocksdb::Status Database::Scan(engine::Context &ctx, const std::string &cursor, uint64_t limit,
const std::string &prefix, const std::string &suffix_glob,
std::vector<std::string> *keys, std::string *end_cursor, RedisType type,
std::optional<int> scan_slot) {
end_cursor->clear();
uint64_t cnt = 0;
uint16_t slot_start = 0;
std::string ns_prefix;
std::string user_key;
auto iter = util::UniqueIterator(ctx, ctx.GetReadOptions(), metadata_cf_handle_);
std::string ns_cursor = AppendNamespacePrefix(cursor);
if (storage_->IsSlotIdEncoded()) {
if (scan_slot.has_value()) {
slot_start = scan_slot.value();
} else {
slot_start = cursor.empty() ? 0 : GetSlotIdFromKey(cursor);
}
ns_prefix = ComposeNamespaceKey(namespace_, "", false);
if (!prefix.empty()) {
PutFixed16(&ns_prefix, slot_start);
ns_prefix.append(prefix);
}
} else {
ns_prefix = AppendNamespacePrefix(prefix);
}
if (!cursor.empty()) {
if (storage_->IsSlotIdEncoded() && !ns_prefix.empty() &&
metadata_cf_handle_->GetComparator()->Compare(rocksdb::Slice(ns_prefix), rocksdb::Slice(ns_cursor)) > 0) {
iter->Seek(ns_prefix);
} else {
iter->Seek(ns_cursor);
if (iter->Valid()) {
iter->Next();
}
}
} else if (ns_prefix.empty()) {
iter->SeekToFirst();
} else {
iter->Seek(ns_prefix);
}
uint16_t slot_id = slot_start;
while (true) {
for (; iter->Valid() && cnt < limit; iter->Next()) {
if (!ns_prefix.empty() && !iter->key().starts_with(ns_prefix)) {
break;
}
Metadata metadata(kRedisNone, false);
auto s = metadata.Decode(iter->value());
if (!s.ok()) continue;
if (type != kRedisNone && type != metadata.Type()) continue;
if (metadata.Expired()) continue;
std::tie(std::ignore, user_key) = ExtractNamespaceKey<std::string>(iter->key(), storage_->IsSlotIdEncoded());
if (!util::StringMatch(suffix_glob, user_key.substr(prefix.size()))) {
continue;
}
keys->emplace_back(user_key);
cnt++;
}
if (auto s = iter->status(); !s.ok()) {
return s;
}
if (!storage_->IsSlotIdEncoded() || prefix.empty()) {
if (!keys->empty() && cnt >= limit) {
end_cursor->append(user_key);
}
break;
}
if (cnt >= limit) {
end_cursor->append(user_key);
break;
}
if (++slot_id >= HASH_SLOTS_SIZE || scan_slot.has_value()) {
break;
}
if (slot_id > slot_start + HASH_SLOTS_MAX_ITERATIONS) {
if (keys->empty()) {
if (iter->Valid()) {
std::tie(std::ignore, user_key) = ExtractNamespaceKey<std::string>(iter->key(), storage_->IsSlotIdEncoded());
auto res = std::mismatch(prefix.begin(), prefix.end(), user_key.begin());
if (res.first == prefix.end() && util::StringMatch(suffix_glob, user_key.substr(prefix.size()))) {
keys->emplace_back(user_key);
}
end_cursor->append(user_key);
}
} else {
end_cursor->append(user_key);
}
break;
}
ns_prefix = ComposeNamespaceKey(namespace_, "", false);
PutFixed16(&ns_prefix, slot_id);
ns_prefix.append(prefix);
iter->Seek(ns_prefix);
}
return rocksdb::Status::OK();
}
rocksdb::Status Database::RandomKey(engine::Context &ctx, const std::string &cursor, std::string *key) {
key->clear();
std::string end_cursor;
std::vector<std::string> keys;
auto s = Scan(ctx, cursor, RANDOM_KEY_SCAN_LIMIT, "", "*", &keys, &end_cursor);
if (!s.ok()) {
return s;
}
if (keys.empty() && !cursor.empty()) {
// if reach the end, restart from beginning
s = Scan(ctx, "", RANDOM_KEY_SCAN_LIMIT, "", "*", &keys, &end_cursor);
if (!s.ok()) {
return s;
}
}
if (!keys.empty()) {
unsigned int seed = util::GetTimeStamp();
*key = keys.at(rand_r(&seed) % keys.size());
}
return rocksdb::Status::OK();
}
rocksdb::Status Database::FlushDB(engine::Context &ctx) {
auto begin_key = ComposeNamespaceKey(namespace_, "", false);
auto end_key = util::StringNext(begin_key);
return storage_->DeleteRange(ctx, begin_key, end_key);
}
rocksdb::Status Database::FlushAll(engine::Context &ctx) {
auto iter = util::UniqueIterator(ctx, ctx.GetReadOptions(), metadata_cf_handle_);
iter->SeekToFirst();
if (!iter->Valid()) {
return rocksdb::Status::OK();
}
auto first_key = iter->key().ToString();
iter->SeekToLast();
if (!iter->Valid()) {
return rocksdb::Status::OK();
}
auto last_key = util::StringNext(iter->key().ToString());
return storage_->DeleteRange(ctx, first_key, last_key);
}
rocksdb::Status Database::Dump(engine::Context &ctx, const Slice &user_key, std::vector<std::string> *infos) {
infos->clear();
std::string ns_key = AppendNamespacePrefix(user_key);
std::string value;
rocksdb::Status s = storage_->Get(ctx, ctx.GetReadOptions(), metadata_cf_handle_, ns_key, &value);
if (!s.ok()) return s.IsNotFound() ? rocksdb::Status::OK() : s;
Metadata metadata(kRedisNone, false);
s = metadata.Decode(value);
if (!s.ok()) return s;
infos->emplace_back("namespace");
infos->emplace_back(namespace_);
infos->emplace_back("type");
infos->emplace_back(metadata.TypeName());
infos->emplace_back("version");
infos->emplace_back(std::to_string(metadata.version));
infos->emplace_back("expire");
infos->emplace_back(std::to_string(Metadata::ExpireMsToS(metadata.expire)));
infos->emplace_back("pexpire");
infos->emplace_back(std::to_string(metadata.expire));
infos->emplace_back("size");
infos->emplace_back(std::to_string(metadata.size));
infos->emplace_back("is_64bit_common_field");
infos->emplace_back(std::to_string(metadata.Is64BitEncoded()));
infos->emplace_back("created_at");
timeval created_at = metadata.Time();
std::time_t tm = created_at.tv_sec;
char time_str[25];
if (!std::strftime(time_str, sizeof(time_str), "%Y-%m-%d %H:%M:%S", std::localtime(&tm))) {
return rocksdb::Status::TryAgain("Fail to format local time_str");
}
std::string created_at_str(time_str);
infos->emplace_back(created_at_str + "." + std::to_string(created_at.tv_usec));
if (metadata.Type() == kRedisList) {
ListMetadata list_metadata(false);
s = GetMetadata(ctx, {kRedisList}, ns_key, &list_metadata);
if (!s.ok()) return s.IsNotFound() ? rocksdb::Status::OK() : s;
infos->emplace_back("head");
infos->emplace_back(std::to_string(list_metadata.head));
infos->emplace_back("tail");
infos->emplace_back(std::to_string(list_metadata.tail));
}
return rocksdb::Status::OK();
}
rocksdb::Status Database::Type(engine::Context &ctx, const Slice &key, RedisType *type) {
std::string ns_key = AppendNamespacePrefix(key);
return typeInternal(ctx, ns_key, type);
}
std::string Database::AppendNamespacePrefix(const Slice &user_key) {
return ComposeNamespaceKey(namespace_, user_key, storage_->IsSlotIdEncoded());
}
rocksdb::Status Database::ClearKeysOfSlotRange(engine::Context &ctx, const rocksdb::Slice &ns,
const SlotRange &slot_range) {
if (!storage_->IsSlotIdEncoded()) {
return rocksdb::Status::Aborted("It is not in cluster mode");
}
std::string prefix = ComposeSlotKeyPrefix(ns, slot_range.start);
std::string prefix_end = ComposeSlotKeyPrefix(ns, slot_range.end + 1);
return storage_->DeleteRange(ctx, prefix, prefix_end);
}
rocksdb::Status Database::KeyExist(engine::Context &ctx, const std::string &key) {
uint32_t cnt = 0;
std::vector<rocksdb::Slice> keys{key};
auto s = Exists(ctx, keys, &cnt);
if (!s.ok()) {
return s;
}
if (cnt == 0) {
return rocksdb::Status::NotFound();
}
return rocksdb::Status::OK();
}
rocksdb::Status SubKeyScanner::Scan(engine::Context &ctx, RedisType type, const Slice &user_key,
const std::string &cursor, uint64_t limit, const std::string &subkey_prefix,
std::vector<std::string> *keys, std::vector<std::string> *values) {
uint64_t cnt = 0;
std::string ns_key = AppendNamespacePrefix(user_key);
Metadata metadata(type, false);
rocksdb::Status s = GetMetadata(ctx, {type}, ns_key, &metadata);
if (!s.ok()) return s;
auto iter = util::UniqueIterator(ctx, ctx.DefaultScanOptions());
std::string match_prefix_key =
InternalKey(ns_key, subkey_prefix, metadata.version, storage_->IsSlotIdEncoded()).Encode();
std::string start_key;
if (!cursor.empty()) {
start_key = InternalKey(ns_key, cursor, metadata.version, storage_->IsSlotIdEncoded()).Encode();
} else {
start_key = match_prefix_key;
}
for (iter->Seek(start_key); iter->Valid(); iter->Next()) {
if (!cursor.empty() && iter->key() == start_key) {
// if cursor is not empty, then we need to skip start_key
// because we already return that key in the last scan
continue;
}
if (!iter->key().starts_with(match_prefix_key)) {
break;
}
InternalKey ikey(iter->key(), storage_->IsSlotIdEncoded());
keys->emplace_back(ikey.GetSubKey().ToString());
if (values != nullptr) {
values->emplace_back(iter->value().ToString());
}
cnt++;
if (limit > 0 && cnt >= limit) {
break;
}
}
return iter->status();
}
RedisType WriteBatchLogData::GetRedisType() const { return type_; }
std::vector<std::string> *WriteBatchLogData::GetArguments() { return &args_; }
std::string WriteBatchLogData::Encode() const {
std::string ret = std::to_string(type_);
for (const auto &arg : args_) {
ret += " " + arg;
}
return ret;
}
Status WriteBatchLogData::Decode(const rocksdb::Slice &blob) {
const std::string &log_data = blob.ToString();
std::vector<std::string> args = util::Split(log_data, " ");
auto parse_result = ParseInt<int>(args[0], 10);
if (!parse_result) {
return parse_result.ToStatus();
}
type_ = static_cast<RedisType>(*parse_result);
args_ = std::vector<std::string>(args.begin() + 1, args.end());
return Status::OK();
}
rocksdb::Status Database::existsInternal(engine::Context &ctx, const std::vector<std::string> &keys, uint32_t *ret) {
*ret = 0;
rocksdb::Status s;
std::string value;
for (const auto &key : keys) {
s = storage_->Get(ctx, ctx.GetReadOptions(), metadata_cf_handle_, key, &value);
if (!s.ok() && !s.IsNotFound()) return s;
if (s.ok()) {
Metadata metadata(kRedisNone, false);
s = metadata.Decode(value);
if (!s.ok()) return s;
if (!metadata.Expired()) *ret += 1;
}
}
return rocksdb::Status::OK();
}
rocksdb::Status Database::typeInternal(engine::Context &ctx, const Slice &key, RedisType *type) {
*type = kRedisNone;
std::string value;
rocksdb::Status s = storage_->Get(ctx, ctx.GetReadOptions(), metadata_cf_handle_, key, &value);
if (!s.ok()) return s.IsNotFound() ? rocksdb::Status::OK() : s;
Metadata metadata(kRedisNone, false);
s = metadata.Decode(value);
if (!s.ok()) return s;
if (metadata.Expired()) {
*type = kRedisNone;
} else {
*type = metadata.Type();
}
return rocksdb::Status::OK();
}
rocksdb::Status Database::Copy(engine::Context &ctx, const std::string &key, const std::string &new_key, bool nx,
bool delete_old, CopyResult *res) {
RedisType type = kRedisNone;
auto s = typeInternal(ctx, key, &type);
if (!s.ok()) return s;
if (type == kRedisNone) {
*res = CopyResult::KEY_NOT_EXIST;
return rocksdb::Status::OK();
}
if (nx) {
uint32_t exist = 0;
if (s = existsInternal(ctx, {new_key}, &exist), !s.ok()) return s;
if (exist > 0) {
*res = CopyResult::KEY_ALREADY_EXIST;
return rocksdb::Status::OK();
}
}
*res = CopyResult::DONE;
if (key == new_key) return rocksdb::Status::OK();
auto batch = storage_->GetWriteBatchBase();
WriteBatchLogData log_data(type);
s = batch->PutLogData(log_data.Encode());
if (!s.ok()) {
return s;
}
engine::DBIterator iter(ctx, ctx.GetReadOptions());
iter.Seek(key);
if (delete_old) {
s = batch->Delete(metadata_cf_handle_, key);
if (!s.ok()) {
return s;
}
}
// copy metadata
s = batch->Put(metadata_cf_handle_, new_key, iter.Value());
if (!s.ok()) {
return s;
}
auto subkey_iter = iter.GetSubKeyIterator();
if (subkey_iter != nullptr) {
auto zset_score_cf = type == kRedisZSet ? storage_->GetCFHandle(ColumnFamilyID::SecondarySubkey) : nullptr;
for (subkey_iter->Seek(); subkey_iter->Valid(); subkey_iter->Next()) {
InternalKey from_ikey(subkey_iter->Key(), storage_->IsSlotIdEncoded());
std::string to_ikey =
InternalKey(new_key, from_ikey.GetSubKey(), from_ikey.GetVersion(), storage_->IsSlotIdEncoded()).Encode();
// copy sub key
auto s = batch->Put(subkey_iter->ColumnFamilyHandle(), to_ikey, subkey_iter->Value());
if (!s.ok()) {
return s;
}
// The ZSET type stores an extra score and member field inside `zset_score` column family
// while compared to other composed data structures. The purpose is to allow to seek by score.
if (type == kRedisZSet) {
std::string score_bytes = subkey_iter->Value().ToString();
score_bytes.append(from_ikey.GetSubKey().ToString());
// copy score key
std::string score_key =
InternalKey(new_key, score_bytes, from_ikey.GetVersion(), storage_->IsSlotIdEncoded()).Encode();
auto s = batch->Put(zset_score_cf, score_key, Slice());
if (!s.ok()) {
return s;
}
}
}
}
return storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch());
}
std::optional<std::string> Database::lookupKeyByPattern(engine::Context &ctx, const std::string &pattern,
const std::string &subst) {
if (pattern == "#") {
return subst;
}
auto match_pos = pattern.find('*');
if (match_pos == std::string::npos) {
return std::nullopt;
}
// hash field
std::string field;
auto arrow_pos = pattern.find("->", match_pos + 1);
if (arrow_pos != std::string::npos && arrow_pos + 2 < pattern.size()) {
field = pattern.substr(arrow_pos + 2);
}
std::string key = pattern.substr(0, match_pos + 1);
key.replace(match_pos, 1, subst);
std::string value;
RedisType type = RedisType::kRedisNone;
if (!field.empty()) {
auto hash_db = redis::Hash(storage_, namespace_);
if (auto s = hash_db.Type(ctx, key, &type); !s.ok() || type != RedisType::kRedisHash) {
return std::nullopt;
}
if (auto s = hash_db.Get(ctx, key, field, &value); !s.ok()) {
return std::nullopt;
}
} else {
auto string_db = redis::String(storage_, namespace_);
if (auto s = string_db.Type(ctx, key, &type); !s.ok() || type != RedisType::kRedisString) {
return std::nullopt;
}
if (auto s = string_db.Get(ctx, key, &value); !s.ok()) {
return std::nullopt;
}
}
return value;
}
rocksdb::Status Database::Sort(engine::Context &ctx, RedisType type, const std::string &key, const SortArgument &args,
std::vector<std::optional<std::string>> *elems, SortResult *res) {
// Obtain the length of the object to sort.
const std::string ns_key = AppendNamespacePrefix(key);
Metadata metadata(type, false);
auto s = GetMetadata(ctx, {type}, ns_key, &metadata);
if (!s.ok()) return s;
if (metadata.size > SORT_LENGTH_LIMIT) {
*res = SortResult::LIMIT_EXCEEDED;
return rocksdb::Status::OK();
}
auto vectorlen = static_cast<int>(metadata.size);
// Adjust the offset and count of the limit
int offset = args.offset >= vectorlen ? 0 : std::clamp(args.offset, 0, vectorlen - 1);
int count = args.offset >= vectorlen ? 0 : std::clamp(args.count, -1, vectorlen - offset);
if (count == -1) count = vectorlen - offset;
// Get the elements that need to be sorted
std::vector<std::string> str_vec;
if (count != 0) {
if (type == RedisType::kRedisList) {
auto list_db = redis::List(storage_, namespace_);
if (args.dontsort) {
if (args.desc) {
s = list_db.Range(ctx, key, -count - offset, -1 - offset, &str_vec);
if (!s.ok()) return s;
std::reverse(str_vec.begin(), str_vec.end());
} else {
s = list_db.Range(ctx, key, offset, offset + count - 1, &str_vec);
if (!s.ok()) return s;
}
} else {
s = list_db.Range(ctx, key, 0, -1, &str_vec);
if (!s.ok()) return s;
}
} else if (type == RedisType::kRedisSet) {
auto set_db = redis::Set(storage_, namespace_);
s = set_db.Members(ctx, key, &str_vec);
if (!s.ok()) return s;
if (args.dontsort) {
str_vec = std::vector(std::make_move_iterator(str_vec.begin() + offset),
std::make_move_iterator(str_vec.begin() + offset + count));
}
} else if (type == RedisType::kRedisZSet) {
auto zset_db = redis::ZSet(storage_, namespace_);
std::vector<MemberScore> member_scores;
if (args.dontsort) {
RangeRankSpec spec;
spec.start = offset;
spec.stop = offset + count - 1;
spec.reversed = args.desc;
s = zset_db.RangeByRank(ctx, key, spec, &member_scores, nullptr);
if (!s.ok()) return s;
for (auto &member_score : member_scores) {
str_vec.emplace_back(std::move(member_score.member));
}
} else {
s = zset_db.GetAllMemberScores(ctx, key, &member_scores);
if (!s.ok()) return s;
for (auto &member_score : member_scores) {
str_vec.emplace_back(std::move(member_score.member));
}
}
} else {
*res = SortResult::UNKNOWN_TYPE;
return s;
}
}
std::vector<RedisSortObject> sort_vec(str_vec.size());
for (size_t i = 0; i < str_vec.size(); ++i) {
sort_vec[i].obj = str_vec[i];
}
// Sort by BY, ALPHA, ASC/DESC
if (!args.dontsort) {
for (size_t i = 0; i < sort_vec.size(); ++i) {
std::string byval;
if (!args.sortby.empty()) {
auto lookup = lookupKeyByPattern(ctx, args.sortby, str_vec[i]);
if (!lookup.has_value()) continue;
byval = std::move(lookup.value());
} else {
byval = str_vec[i];
}
if (args.alpha && !args.sortby.empty()) {
sort_vec[i].v = byval;
} else if (!args.alpha && !byval.empty()) {
auto double_byval = ParseFloat<double>(byval);
if (!double_byval) {
*res = SortResult::DOUBLE_CONVERT_ERROR;
return rocksdb::Status::OK();
}
sort_vec[i].v = *double_byval;
}
}
std::sort(sort_vec.begin(), sort_vec.end(), [&args](const RedisSortObject &a, const RedisSortObject &b) {
return RedisSortObject::SortCompare(a, b, args);
});
// Gets the element specified by Limit
if (offset != 0 || count != vectorlen) {
sort_vec = std::vector(std::make_move_iterator(sort_vec.begin() + offset),
std::make_move_iterator(sort_vec.begin() + offset + count));
}
}
// Perform storage
for (auto &elem : sort_vec) {
if (args.getpatterns.empty()) {
elems->emplace_back(elem.obj);
}
for (const std::string &pattern : args.getpatterns) {
std::optional<std::string> val = lookupKeyByPattern(ctx, pattern, elem.obj);
if (val.has_value()) {
elems->emplace_back(val.value());
} else {
elems->emplace_back(std::nullopt);
}
}
}
if (!args.storekey.empty()) {
std::vector<std::string> store_elems;
store_elems.reserve(elems->size());
for (const auto &e : *elems) {
store_elems.emplace_back(e.value_or(""));
}
redis::List list_db(storage_, namespace_);
s = list_db.Trim(ctx, args.storekey, -1, 0);
if (!s.ok()) return s;
uint64_t new_size = 0;
s = list_db.Push(ctx, args.storekey, std::vector<Slice>(store_elems.cbegin(), store_elems.cend()), false,
&new_size);
if (!s.ok()) return s;
}
return rocksdb::Status::OK();
}
bool RedisSortObject::SortCompare(const RedisSortObject &a, const RedisSortObject &b, const SortArgument &args) {
if (!args.alpha) {
double score_a = std::get<double>(a.v);
double score_b = std::get<double>(b.v);
return !args.desc ? score_a < score_b : score_a > score_b;
} else {
if (!args.sortby.empty()) {
std::string cmp_a = std::get<std::string>(a.v);
std::string cmp_b = std::get<std::string>(b.v);
return !args.desc ? cmp_a < cmp_b : cmp_a > cmp_b;
} else {
return !args.desc ? a.obj < b.obj : a.obj > b.obj;
}
}
}
} // namespace redis