159 lines
4.1 KiB
C
159 lines
4.1 KiB
C
|
||
/**
|
||
* Request
|
||
* Cmd: | OP(1) | argc(1) | repeat { arglen(4) | arg } |
|
||
*
|
||
* Response
|
||
* Rsp: | OP(1) | status(1) | datalen(4) | data |
|
||
*/
|
||
#include <stdio.h>
|
||
#include <stdlib.h>
|
||
#include <stdint.h>
|
||
#include <string.h>
|
||
#include <sys/time.h>
|
||
#include <sys/socket.h>
|
||
#include <arpa/inet.h>
|
||
|
||
#define CMD_SIZE (1024)
|
||
#define BATCH_SIZE (65536)
|
||
#define KVS_BATCH_MAX 128
|
||
#define TIME_SUB_MS(tv1, tv2) ((tv1.tv_sec - tv2.tv_sec) * 1000 + (tv1.tv_usec - tv2.tv_usec) / 1000)
|
||
|
||
// #define PRESP print_response
|
||
#define PRESP
|
||
|
||
|
||
typedef enum {
|
||
KVS_STATUS_OK = 0,
|
||
KVS_STATUS_ERROR = 1,
|
||
KVS_STATUS_NO_EXIST = 2,
|
||
KVS_STATUS_EXIST = 3,
|
||
KVS_STATUS_BADREQ = 4
|
||
}rsp_ret_status_e;
|
||
|
||
enum {
|
||
KVS_CMD_START = 0,
|
||
// array
|
||
KVS_CMD_SET = KVS_CMD_START,
|
||
KVS_CMD_GET,
|
||
KVS_CMD_DEL,
|
||
KVS_CMD_MOD,
|
||
KVS_CMD_EXIST,
|
||
// rbtree
|
||
KVS_CMD_RSET,
|
||
KVS_CMD_RGET,
|
||
KVS_CMD_RDEL,
|
||
KVS_CMD_RMOD,
|
||
KVS_CMD_REXIST,
|
||
// hash
|
||
KVS_CMD_HSET,
|
||
KVS_CMD_HGET,
|
||
KVS_CMD_HDEL,
|
||
KVS_CMD_HMOD,
|
||
KVS_CMD_HEXIST,
|
||
KVS_CMD_SAVE,
|
||
|
||
KVS_CMD_COUNT,
|
||
};
|
||
|
||
typedef struct {
|
||
uint8_t op;
|
||
uint8_t status;
|
||
uint32_t datalen;
|
||
uint8_t *data;
|
||
} kvs_response_t;
|
||
|
||
int kvs_need(const uint8_t *p, const uint8_t *end, size_t n);
|
||
int kvs_read_u8(const uint8_t **pp, const uint8_t *end, uint8_t *out);
|
||
int kvs_read_u16(const uint8_t **pp, const uint8_t *end, uint16_t *out);
|
||
int kvs_read_u32(const uint8_t **pp, const uint8_t *end, uint32_t *out);
|
||
|
||
int kvs_write_u8(uint8_t **pp, const uint8_t *end, uint8_t v);
|
||
int kvs_write_u16(uint8_t **pp, const uint8_t *end, uint16_t v);
|
||
int kvs_write_u32(uint8_t **pp, const uint8_t *end, uint32_t v);
|
||
|
||
int getcmd(uint8_t op, const char *key, uint32_t key_len, const char *value, uint32_t value_len, uint8_t *buf);
|
||
|
||
int parse_response(const uint8_t *buf, int buflen, kvs_response_t *rsp);
|
||
void print_response(const char *cmd_name, const kvs_response_t *rsp);
|
||
int verify_response(const kvs_response_t *rsp, uint8_t expected_op,
|
||
uint8_t expected_status, const char *expected_data, uint32_t expected_len);
|
||
|
||
|
||
|
||
|
||
|
||
typedef struct {
|
||
uint8_t buf[BATCH_SIZE];
|
||
int len; // 当前已写入的 batch 字节数
|
||
int cnt; // 当前 batch 里命令条数
|
||
int cmd_len[KVS_BATCH_MAX];
|
||
} kvs_batch_t;
|
||
|
||
static void kvs_batch_init(kvs_batch_t *b)
|
||
{
|
||
b->len = 0;
|
||
b->cnt = 0;
|
||
memset(b->cmd_len, 0, sizeof(b->cmd_len));
|
||
}
|
||
|
||
/**
|
||
* 用 getcmd() 生成单条命令，然后 append 到 batch buffer
|
||
* 返回：0 成功，-1 失败（太多条 or buffer 不够）
|
||
*/
|
||
static int kvs_batch_add(kvs_batch_t *b, uint8_t op, const char *key, uint32_t key_len, const char *value, uint32_t value_len){
|
||
if (b->cnt >= KVS_BATCH_MAX) return -1;
|
||
uint8_t tmp[CMD_SIZE];
|
||
int n = getcmd(op, key, key_len, value, value_len, tmp); // 你提供的函数
|
||
if (n <= 0) return -1;
|
||
|
||
if (b->len + n > (int)sizeof(b->buf)) return -1;
|
||
|
||
memcpy(b->buf + b->len, tmp, n);
|
||
b->cmd_len[b->cnt] = n;
|
||
b->cnt++;
|
||
b->len += n;
|
||
return 0;
|
||
}
|
||
|
||
/**
|
||
* 一次性发送 batch
|
||
* 返回：发送字节数，<0 表示失败
|
||
*/
|
||
static int kvs_batch_send(int fd, const kvs_batch_t *b)
|
||
{
|
||
// printf("send : %d\n", b->len);
|
||
return (int)send(fd, b->buf, b->len, 0);
|
||
}
|
||
|
||
/**
|
||
* 一次 recv 收回所有响应，然后批量解析为 rsp 数组
|
||
*
|
||
* 返回：成功解析出的响应条数（期望是 b->cnt）
|
||
*/
|
||
static int kvs_batch_recv_parse(int fd,
|
||
const kvs_batch_t *b,
|
||
kvs_response_t *rsps, // 输出数组，长度 >= b->cnt
|
||
uint8_t *recvbuf,
|
||
int recvbuf_cap)
|
||
{
|
||
int parsed = 0;
|
||
int used = 0;
|
||
|
||
while(parsed < b->cnt){
|
||
int nrecv = (int)recv(fd, recvbuf+used, recvbuf_cap, 0);
|
||
if (nrecv <= 0) return -1;
|
||
|
||
int off = 0;
|
||
|
||
while (parsed < b->cnt) {
|
||
int consumed = parse_response(recvbuf + used, nrecv - off, &rsps[parsed]);
|
||
if (consumed <= 0) break; // 不够解析/失败，简单处理：直接退出
|
||
|
||
off += consumed;
|
||
used+= consumed;
|
||
parsed++;
|
||
}
|
||
}
|
||
return parsed;
|
||
} |