int proc_zrscan(NetworkServer *net, Link *link, const Request &req, Response *resp){
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(6);
uint64_t limit = req[5].Uint64();
uint64_t offset = 0;
if(req.size() > 6){
offset = limit;
limit = offset + req[6].Uint64();
}
ZIterator *it = serv->ssdb->zrscan(req[1], req[2], req[3], req[4], limit);
if(offset > 0){
it->skip(offset);
}
resp->push_back("ok");
uint64_t size = 0;
while(it->next()){
size += it->key.size() + it->score.size();
CHECK_SCAN_OUTPUT_LIMIT(size);
resp->push_back(it->key);
resp->push_back(it->score);
}
delete it;
return 0;
}
static int proc_zremrangebyrank(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() < 4){
resp->push_back("client_error");
return 0;
}
uint64_t offset = req[2].Uint64();
uint64_t limit = req[3].Uint64();
ZIterator *it = serv->ssdb->zrange(req[1], offset, limit);
uint64_t count = 0;
while(it->next()){
count ++;
int ret = serv->ssdb->zdel(req[1], it->key);
if(ret == -1){
resp->push_back("error");
delete it;
return 0;
}
}
delete it;
char buf[20];
snprintf(buf, sizeof(buf), "%" PRIu64 "", count);
resp->push_back("ok");
resp->push_back(buf);
return 0;
}
int proc_zclear(NetworkServer *net, Link *link, const Request &req, Response *resp){
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(2);
const Bytes &name = req[1];
int64_t count = 0;
while(1){
ZIterator *it = serv->ssdb->zrange(name, 0, 1000);
int num = 0;
while(it->next()){
int ret = serv->ssdb->zdel(name, it->key);
if(ret == -1){
resp->push_back("error");
delete it;
return 0;
}
num ++;
};
delete it;
if(num == 0){
break;
}
count += num;
}
serv->save_kv_stats();
resp->reply_int(0, count);
return 0;
}
static int proc_zclear(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() < 2){
resp->push_back("client_error");
return 0;
}
const Bytes &name = req[1];
uint64_t total = 0;
while(1){
ZIterator *it = serv->ssdb->zrange(name, 0, 1000);
int num = 0;
while(it->next()){
int ret = serv->ssdb->zdel(name, it->key);
if(ret == -1){
resp->push_back("error");
delete it;
return 0;
}
num ++;
};
delete it;
if(num == 0){
break;
}
total += num;
}
char buf[20];
snprintf(buf, sizeof(buf), "%" PRIu64 "", total);
resp->push_back("ok");
resp->push_back(buf);
return 0;
}
ZIterator* SSDB::zrrange(const Bytes &name, uint64_t offset, uint64_t limit){
if(offset + limit > limit){
limit = offset + limit;
}
ZIterator *it = ziterator(this, name, "", "", "", limit, Iterator::BACKWARD);
it->skip(offset);
return it;
}
static int proc_zkeys(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() < 6){
resp->push_back("client_error");
}else{
uint64_t limit = req[5].Uint64();
ZIterator *it = serv->ssdb->zscan(req[1], req[2], req[3], req[4], limit);
resp->push_back("ok");
while(it->next()){
resp->push_back(it->key);
}
delete it;
}
return 0;
}
int proc_zsum(NetworkServer *net, Link *link, const Request &req, Response *resp){
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(4);
int64_t sum = 0;
ZIterator *it = serv->ssdb->zscan(req[1], "", req[2], req[3], -1);
while(it->next()){
sum += str_to_int64(it->score);
}
delete it;
resp->reply_int(0, sum);
return 0;
}
int CommandProc::proc_zkeys(const Link &link, const Request &req, Response *resp){
if(req.size() < 6){
resp->push_back("client_error");
}else{
int limit = req[5].Int();
ZIterator *it = ssdb->zscan(req[1], req[2], req[3], req[4], limit);
resp->push_back("ok");
while(it->next()){
resp->push_back(it->key);
}
delete it;
}
return 0;
}
int64_t SSDB::zrrank(const Bytes &name, const Bytes &key) const{
ZIterator *it = ziterator(this, name, "", "", "", INT_MAX, Iterator::BACKWARD);
uint64_t ret = 0;
while(true){
if(it->next() == false){
ret = -1;
break;
}
if(key == it->key){
break;
}
ret ++;
}
delete it;
return ret;
}
static int proc_zrrange(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() != 4){
resp->push_back("client_error");
}else{
uint64_t offset = req[2].Uint64();
uint64_t limit = req[3].Uint64();
ZIterator *it = serv->ssdb->zrange(req[1], offset, limit);
resp->push_back("ok");
while(it->next()){
resp->push_back(it->key);
resp->push_back(it->score);
}
delete it;
}
return 0;
}
int proc_zkeys(NetworkServer *net, Link *link, const Request &req, Response *resp){
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(6);
uint64_t limit = req[5].Uint64();
ZIterator *it = serv->ssdb->zscan(req[1], req[2], req[3], req[4], limit);
resp->push_back("ok");
uint64_t size = 0;
while(it->next()){
size += it->key.size();
CHECK_SCAN_OUTPUT_LIMIT(size);
resp->push_back(it->key);
}
delete it;
return 0;
}
void ExpirationHandler::load_expiration_keys_from_db(int num){
ZIterator *it;
it = ssdb->zscan(this->list_name, "", "", "", num);
int n = 0;
while(it->next()){
n ++;
std::string &key = it->key;
int64_t score = str_to_int64(it->score);
if(score < 2000000000){
// older version compatible
score *= 1000;
}
fast_keys.add(key, score);
}
delete it;
log_debug("load %d keys into fast_keys", n);
}
static int proc_zsum(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() < 4){
resp->push_back("client_error");
return 0;
}
int64_t sum = 0;
ZIterator *it = serv->ssdb->zscan(req[1], "", req[2], req[3], -1);
while(it->next()){
sum += str_to_int64(it->score);
}
delete it;
char buf[20];
snprintf(buf, sizeof(buf), "%" PRId64 "", sum);
resp->push_back("ok");
resp->push_back(buf);
return 0;
}
static int proc_zcount(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() < 4){
resp->push_back("client_error");
return 0;
}
uint64_t count = 0;
ZIterator *it = serv->ssdb->zscan(req[1], "", req[2], req[3], -1);
while(it->next()){
count ++;
}
delete it;
char buf[20];
snprintf(buf, sizeof(buf), "%" PRIu64 "", count);
resp->push_back("ok");
resp->push_back(buf);
return 0;
}
int proc_zavg(NetworkServer *net, Link *link, const Request &req, Response *resp){
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(4);
int64_t sum = 0;
int64_t count = 0;
ZIterator *it = serv->ssdb->zscan(req[1], "", req[2], req[3], -1);
while(it->next()){
sum += str_to_int64(it->score);
count ++;
}
delete it;
double avg = (double)sum/count;
resp->push_back("ok");
resp->add(avg);
return 0;
}
int proc_zremrangebyscore(NetworkServer *net, Link *link, const Request &req, Response *resp){
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(4);
ZIterator *it = serv->ssdb->zscan(req[1], "", req[2], req[3], -1);
int64_t count = 0;
while(it->next()){
count ++;
int ret = serv->ssdb->zdel(req[1], it->key);
if(ret == -1){
delete it;
resp->push_back("error");
return 0;
}
}
delete it;
serv->save_kv_stats();
resp->reply_int(0, count);
return 0;
}
static int proc_zclear(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() < 2){
resp->push_back("client_error");
return 0;
}
const Bytes &name = req[1];
uint64_t total = 0;
while(1){
ZIterator *it = serv->ssdb->zrange(name, 0, 1000);
// we need std::string to hold the memory, because Bytes never alloc memory
std::vector<std::string> s_keys;
while(it->next()){
s_keys.push_back(it->key);
}
delete it;
if(s_keys.empty()){
break;
}
std::vector<Bytes> keys;
for(std::vector<std::string>::iterator it=s_keys.begin(); it!=s_keys.end(); it++){
keys.push_back(*it);
}
int ret = serv->ssdb->multi_zdel(name, keys, 0);
if(ret == -1){
resp->push_back("error");
break;
}else{
total += ret;
}
}
char buf[20];
snprintf(buf, sizeof(buf), "%"PRIu64"", total);
resp->push_back("ok");
resp->push_back(buf);
return 0;
}
static int proc_zrscan(Server *serv, Link *link, const Request &req, Response *resp){
if(req.size() < 6){
resp->push_back("client_error");
}else{
uint64_t limit = req[5].Uint64();
uint64_t offset = 0;
if(req.size() > 6){
offset = limit;
limit = offset + req[6].Uint64();
}
ZIterator *it = serv->ssdb->zrscan(req[1], req[2], req[3], req[4], limit);
if(offset > 0){
it->skip(offset);
}
resp->push_back("ok");
while(it->next()){
resp->push_back(it->key);
resp->push_back(it->score);
}
delete it;
}
return 0;
}
int proc_zremrangebyrank(NetworkServer *net, Link *link, const Request &req, Response *resp){
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(4);
uint64_t start = req[2].Uint64();
uint64_t end = req[3].Uint64();
ZIterator *it = serv->ssdb->zrange(req[1], start, end - start + 1);
int64_t count = 0;
while(it->next()){
count ++;
int ret = serv->ssdb->zdel(req[1], it->key);
if(ret == -1){
resp->push_back("error");
delete it;
return 0;
}
}
delete it;
serv->save_kv_stats();
resp->reply_int(0, count);
return 0;
}
void ExpirationHandler::load_expiration_keys_from_db(int64_t expired, int num){
char data[30];
int size = snprintf(data, sizeof(data), "%" PRId64, expired);
if(size <= 0){
size = 0;
}
ZIterator *it;
it = ssdb->zscan(this->list_name, "", "", Bytes(data,size), num);
int n = 0;
while(it->next()){
n ++;
std::string &key = it->key;
int64_t score = str_to_int64(it->score);
if(score < 2000000000){
// older version compatible
// TODO 新老版本共存的时候，这个会有一些bug
score *= 1000;
}
fast_keys.add(key, score);
}
delete it;
log_debug("load %d keys into fast_keys", n);
}
/*
* ZUNIONSTORE destination numkeys key [key ...] [WEIGHTS weight [weight ...]] [AGGREGATE SUM|MIN|MAX]
*/
static int proc_join_zsets(NetworkServer *net, Link *link, const Request &req, Response *resp, int type){
int offset = 3;
SSDBServer *serv = (SSDBServer *)net->data;
CHECK_NUM_PARAMS(offset);
int numkeys = req[2].Int();
int size = numkeys + offset;
if (numkeys <= 0 || req.size() < size) {
resp->reply_client_error("wrong number of arguments");
return 0;
}
std::vector<int> weights;
int aggregate_type = AGGREGATE_SUM_TYPE;
for (int i = size; i < req.size();) {
if (req[i] == "weights") {
CHECK_NUM_PARAMS(i + 1 + numkeys);
for (int j = 1; j <= numkeys; j ++) {
weights.push_back(req[i + j].Int());
}
i += numkeys + 1;
} else if (req[i] == "aggregate") {
CHECK_NUM_PARAMS(i + 2);
if (req[i+1] == "sum") {
aggregate_type = AGGREGATE_SUM_TYPE;
} else if (req[i+1] == "min") {
aggregate_type = AGGREGATE_MIN_TYPE;
} else if (req[i+1] == "max") {
aggregate_type = AGGREGATE_MAX_TYPE;
} else {
resp->reply_client_error("wrong format of arguments");
return 0;
}
i += 2;
} else {
resp->reply_client_error("wrong format of arguments");
return 0;
}
}
if (weights.empty()) {
for (int i = 0; i < numkeys; i ++) {
weights.push_back(1);
}
}
int min_len = -1, min_index = 0;
std::vector< std::map<std::string,int64_t> > vectors;
for (int i = offset; i < offset + numkeys; i ++) {
std::map<std::string,int64_t> key_scores;
// FIXME 控制大小,防止内存爆掉
ZIterator *it = serv->ssdb->zscan(req[i], "", "", "", SSDB_MAX_SCAN_LEN);
while(it->next()){
int64_t score = str_to_int64(it->score) * weights[i-offset];
key_scores[it->key] = score;
}
delete it;
if (key_scores.size() == 0 && type == INTER_TYPE) {
resp->reply_int(0, 0);
return 0;
}
vectors.push_back(key_scores);
if (min_len == -1 || min_len > key_scores.size()) {
min_len = key_scores.size();
min_index = i - offset;
}
}
std::map<std::string,int64_t> result;
if (type == INTER_TYPE) {
result = vectors[min_index]; // copy sets
for (std::map<std::string,int64_t>::iterator it = vectors[min_index].begin(); it != vectors[min_index].end(); it ++) {
std::string key = it->first;
int64_t score = it->second;
for (int i = 0; i < vectors.size(); i ++) {
if (i == min_index) {
continue;
}
std::map<std::string,int64_t>::iterator fit = vectors[i].find(key);
if (fit == vectors[i].end()){
// remove
result.erase(key);
break;
} else {
if (aggregate_type == AGGREGATE_MIN_TYPE) {
result[key] = score < fit->second ? score : fit->second;
} else if (aggregate_type == AGGREGATE_MAX_TYPE) {
result[key] = score > fit->second ? score : fit->second;
} else {
result[key] = score + fit->second;
}
}
}
}
} else if (type == UNION_TYPE) {
for (int i = 0; i < vectors.size(); i ++) {
for (std::map<std::string,int64_t>::iterator it = vectors[i].begin(); it != vectors[i].end(); it ++) {
std::string key = it->first;
int64_t score = it->second;
std::map<std::string,int64_t>::iterator rit = result.find(key);
if (rit == result.end()) { // not found
result[key] = score;
} else {
if (aggregate_type == AGGREGATE_MIN_TYPE) {
rit->second = score < rit->second ? score : rit->second;
} else if (aggregate_type == AGGREGATE_MAX_TYPE) {
rit->second = score > rit->second ? score : rit->second;
} else {
rit->second = score + rit->second;
}
}
}
}
}
for(std::map<std::string,int64_t>::iterator it = result.begin(); it != result.end(); it ++) {
int ret = serv->ssdb->zset(req[1], it->first, str(it->second));
if(ret == -1){
resp->push_back("error");
return 0;
}
}
serv->save_kv_stats();
resp->reply_int(0, result.size());
return 0;
}
