int main(int argc, char *argv[])
{
void * handle;
if (argc < 2)
{
fprintf(stderr, "Invalid arguments\n");
exit(1);
}
if (msd_load_so(&handle, syms, argv[1]) < 0)
{
fprintf(stderr, "load so file failed\n");
exit(1);
}
if (dll.handle_init)
{
dll.handle_init("handle_init", 0);
}
dll.handle_task("handle_task");
if (dll.handle_fini)
{
dll.handle_fini("handle_fnit", 1);
}
msd_unload_so(&handle);
exit(0);
}
/**
* @brief 底层关闭连接时通知用户
* @param p_net_io_server
* @param connection_id 指示关闭的是哪个连接
* @return 0-success -1-failed
*/
int c_newbench_net_io_event::on_connection_closed(void * /* p_net_io_server */,
const int connection_id,
const int connection_fd,
const char *ip,
const int port,
union net_io_storage *p_storage)
{
TRACE_LOG("Close connection on conn id: %d, sock fd: %d", connection_id, connection_fd);
if (dll.handle_close)
{
m_push_mb->skinfo.sockfd = connection_fd;
m_push_mb->skinfo.remote_ip = inet_addr(ip);
m_push_mb->skinfo.remote_port = htons(port);
m_push_mb->skinfo.storage = *p_storage;
m_push_mb->skinfo.connection_id = connection_id;
if (dll.handle_close(&m_push_mb->skinfo) != 0)
{
ERROR_LOG("dll handle_close failed!");
return -1;
}
}
return 0;
}
int c_newbench_net_io_event::on_new_connection(void *p_net_io_server,
const int connection_id,
const int connection_fd,
const char *peer_ip,
const int peer_port,
union net_io_storage *p_storage)
{
if (dll.handle_open)
{
if (p_net_io_server == NULL)
{
return -1;
}
m_push_mb->skinfo.sockfd = connection_fd;
m_push_mb->skinfo.remote_ip = inet_addr(peer_ip);
m_push_mb->skinfo.remote_port = htons(peer_port);
// default for m_push_mb->skinfo.storage
// ...
m_push_mb->skinfo.ptr_lookout = p_storage;
m_push_mb->skinfo.connection_id = connection_id;
if (dll.handle_open(NULL, NULL, &m_push_mb->skinfo) != 0)
{
return -1;
}
*p_storage = m_push_mb->skinfo.storage;
}
return 0;
}
/**
* @brief 重启已停止work进程
* @param
* @return
*/
int monitor_work_proc(int argc, char* argv[], pid_t *wpids, int pid_count, bool /* use_barrier */)
{
int result;
pid_t old_pid, new_pid;
int sec = 10;
while (!g_stop) {
struct timeval tv = {sec, 0};
int sleep_ok = select(0, NULL, NULL, NULL, &tv);
for (int i = 0; i < pid_count && sleep_ok == 0; ++i) {
if(wpids[i] != 0) {
result = kill(wpids[i], 0);
if (0 == result || errno != ESRCH) {
continue;
}
}
if (!g_stop) {
old_pid = wpids[i];
/// WORK进程异常退出时，需要清理过滤器中这个进程的资源
if (g_bench_conf.use_barrier) {
AFUTEX_LOCK_ERR_CHK(g_rcv_rq_lock);
g_prcv_brq->clean_4pid(old_pid);
AFUTEX_UNLOCK_ERR_CHK(g_rcv_rq_lock);
}
new_pid = fork();
if(new_pid < 0) {
wpids[i] = 0;
ERROR_LOG("Work[%u] exited, and fork failed: %s", new_pid, strerror(errno));
} else if(new_pid > 0) {///parent process
wpids[i] = new_pid;
ERROR_LOG("Work[%u] exited, fork[%u] to replace.", old_pid, new_pid);
} else {
run_work_proc(argc, argv, g_bench_conf.use_barrier);
}
}
}
if (dll.handle_timer) {
dll.handle_timer(&sec);
}
}
if (dll.handle_fini) {
dll.handle_fini(PROC_MAIN);
}
///wait for all child processes end
for (int i = 0; i < pid_count; ) {
if (wpids[i] != 0) {
result = kill(wpids[i], 0);
if (result != 0 && errno == ESRCH) {
/// 等待下一个进程
++i;
} else {
/// 继续等待 50 毫秒
usleep(50000);
continue;
}
}
}
return 0;
}
/**
* @brief 启动work进程
* @param
* @return
*/
int run_work_proc(int argc, char *argv[], bool /* use_barrier */)
{
setproctitle("%s:[WORK]", g_bench_conf.prog_name);
if (dll.handle_init != NULL && dll.handle_init(argc, argv, PROC_WORK) != 0) {
print_prompt(false, "Worker process handle_init failed.");
return -1;
}
/// 由于环形队列中有头 shm_block_t，所以需要给头预留空间
int prcv_data_len = g_bench_conf.max_pkg_len + sizeof(shm_block_t);
char *prcv_data = (char *)malloc(prcv_data_len);
if (prcv_data == NULL) {
print_prompt(false, "Malloc work proc pop buffer fail, size: %d B", prcv_data_len);
return -1;
}
shm_block_t *mb = (shm_block_t *)prcv_data;
int sndlen = 0;
int ret = 0;
char *puser_data = NULL;
if (g_bench_conf.use_barrier == true) {
if (g_prcv_brq == NULL) {
print_prompt(false, "Use barrier mode, but i_barrier_ring_queue is null.");
goto work_proc_end;
}
while (! g_stop) {
if (dll.handle_schedule) {
dll.handle_schedule();
}
if (!g_prcv_brq->is_able_pop()) {
/// 等待数据
struct timespec tv = {0, 200000};
nanosleep(&tv, NULL);
continue;
}
AFUTEX_LOCK_ERR_CHK(g_rcv_rq_lock);
ret = g_prcv_brq->pop(prcv_data, prcv_data_len);
AFUTEX_UNLOCK_ERR_CHK(g_rcv_rq_lock);
if (ret < 0) {
ERROR_LOG("Barrier pop error: %d", g_prcv_brq->get_last_errno());
sleep(1); /// 防止出错以后打印日志速度太快
continue;
} else if (ret == 0) {
/// 没有取出任何数据
continue;
} else {
if(ret != mb->length || (u_int)mb->length <= sizeof(shm_block_t) ||
mb->length > prcv_data_len) {
/// 数据包包长有误，关闭这个连接
ERROR_LOG("Pop data len(%d) err, pkg len(%d), prcv_data_len(%d), shm_block_t(%zd), close conn: %d.",
ret, mb->length, prcv_data_len, sizeof(shm_block_t), mb->blk_id);
work_proc_close_conn(mb);
continue;
}
}
TRACE_LOG("Pop barrier ring queue len: %d", mb->length);
sndlen = 0;
ret = dll.handle_process(mb->data, mb->length - sizeof(shm_block_t), &puser_data, &sndlen, &mb->skinfo);
if (ret != 0) {
work_proc_close_conn(mb);
TRACE_LOG("Handle_process fail, close conn: %d", mb->blk_id);
} else if(sndlen + sizeof(shm_block_t) > (u_int)prcv_data_len) {
work_proc_close_conn(mb);
ERROR_LOG("Handle_process sndlen(%d) > buf len, close conn: %d", sndlen, mb->blk_id);
} else {
if (sndlen > 0) {
mb->length = sndlen + sizeof(shm_block_t);
mb->type = DAT_BLOCK;
memcpy(mb->data, puser_data, sndlen);
AFUTEX_LOCK_ERR_CHK(g_snd_rq_lock);
ret = g_psnd_rq->push_data(prcv_data, mb->length, true);
AFUTEX_UNLOCK_ERR_CHK(g_snd_rq_lock);
if (ret != mb->length) {
ERROR_LOG("Push data(len: %d) error (ret: %d, data: %d, empty: %d)", mb->length,
g_psnd_rq->get_last_errno(), g_psnd_rq->get_data_len(), g_psnd_rq->get_empty_buffer_len());
} else {
TRACE_LOG("ring queue push data len: %d", mb->length);
g_p_net_io_notifier->popup();
}
}
}
}
} else {
if (g_prcv_rq == NULL) {
print_prompt(false, "Use nonbarrier mode, but i_ring_queue object is null.");
goto work_proc_end;
}
while (!g_stop) {
if (dll.handle_schedule != NULL) {
dll.handle_schedule();
}
if (g_prcv_rq->get_data_len() <= 0) {
struct timespec tv = {0, 200000};
nanosleep (&tv, NULL);
continue;
}
AFUTEX_LOCK_ERR_CHK(g_rcv_rq_lock);
ret = g_prcv_rq->pop_data(prcv_data, prcv_data_len, 0);
AFUTEX_UNLOCK_ERR_CHK(g_rcv_rq_lock);
if (ret < 0) {
ERROR_LOG("Common ring queue pop fail(ret: %d): %s", ret, g_prcv_rq->get_last_errstr());
sleep(1); /// 防止日志打印过快
continue;
} else if (ret == 0) {
/// 数据还没准备好
continue;
} else {
if(ret != mb->length || (u_int)mb->length <= sizeof(shm_block_t) || mb->length > prcv_data_len) {
/// 数据包包长有误，关闭这个连接
ERROR_LOG("Pop data len(%d) err, pkg len(%d), prcv_data_len(%d), shm_block_t(%zd), close conn: %d.",
ret, mb->length, prcv_data_len, sizeof(shm_block_t), mb->blk_id);
work_proc_close_conn(mb);
continue;
}
}
TRACE_LOG("Common ring queue pop len: %d", mb->length);
sndlen = 0;
ret = dll.handle_process(mb->data, mb->length - sizeof(shm_block_t), &puser_data, &sndlen, &mb->skinfo);
if (ret != 0) {
work_proc_close_conn(mb);
TRACE_LOG("Handle process fail, close conn: %d", mb->blk_id);
} else if(sndlen + sizeof(shm_block_t) > (u_int)prcv_data_len) {
work_proc_close_conn(mb);
ERROR_LOG("Handle_process sndlen(%zd) > buf len %d, close conn: %d",
sndlen + sizeof(shm_block_t), prcv_data_len, mb->blk_id);
} else {
if (sndlen > 0) {
mb->length = sndlen + sizeof(shm_block_t);
mb->type = DAT_BLOCK;
memcpy(mb->data, puser_data, sndlen);
AFUTEX_LOCK_ERR_CHK(g_snd_rq_lock);
ret = g_psnd_rq->push_data(prcv_data, mb->length, true);
AFUTEX_UNLOCK_ERR_CHK(g_snd_rq_lock);
if (ret != mb->length) {
ERROR_LOG("Push data(len: %d) error (ret: %d, data: %d, empty: %d)", mb->length,
g_psnd_rq->get_last_errno(), g_psnd_rq->get_data_len(), g_psnd_rq->get_empty_buffer_len());
} else {
TRACE_LOG("Common ring queue push len: %d", mb->length);
g_p_net_io_notifier->popup();
}
}
}
}
}
work_proc_end:
free(prcv_data);
if (g_p_send_buff != NULL) {
free(g_p_send_buff);
g_p_send_buff = NULL;
g_send_buff_len = 0;
}
if (dll.handle_fini != NULL) {
dll.handle_fini(PROC_WORK);
}
return 0;
}
int main(int argc, char* argv[])
{
memset(&g_bench_conf, 0, sizeof(bench_config_t));
if (parse_args(argc, argv) != 0)
{
return -1;
}
if (nb_init() != 0)
{
return -1;
}
if (dll.handle_init)
{
NB_BOOT_LOG(dll.handle_init(argc, argv, PROC_MAIN), "Call main process handle_init");
}
daemon_start(argc, argv);
initproctitle(argc, argv);
setproctitle("%s:[MAIN]", g_bench_conf.prog_name);
if (g_bench_conf.use_barrier && dll.handle_dispatch == NULL)
{
ERROR_LOG("Warning: handle_dispatch not found, pkg filter disabled.");
}
pid_t pid;
pid = fork();
if (pid < 0)
{
NB_BOOT_LOG(-1, "Fork child process failed: %s", strerror(errno));
return -1;
}
else if(pid > 0)
{///父进程
///查看CONN进程是否已启动，等待 1 秒
usleep (100000);
if (kill(pid, 0) == 0 || errno != ESRCH)
{
pid_t *wpids = (pid_t*) calloc(g_bench_conf.worker_num, sizeof(pid_t));
if (spawn_work_proc(argc, g_bench_conf.saved_argv, wpids,
g_bench_conf.worker_num, g_bench_conf.use_barrier) != 1)
{
///父进程返回
monitor_work_proc(argc, g_bench_conf.saved_argv, wpids,
g_bench_conf.worker_num, g_bench_conf.use_barrier);
daemon_stop();
}
else
{
///子进程返回
}
free(wpids);
}
}
else
{///子进程
conn_proc(argc, g_bench_conf.saved_argv);
}
nb_uninit();
///只能增加资源释放相关代码
return 0;
}
/**
* @brief 底层接收到数据时通知用户:将数据放入环形队列,并通知work进程去取数据
* @param p_net_io_server
* @param connection_id 指示哪个连接ID上接收到数据
* @param p_data 指向接收到的数据
* @param data_len 接收到的数据长度
* @return 0-success -1-failed
*/
int c_newbench_net_io_event::on_recv_data(void *p_net_io_server,
const int connection_id,
const int connection_fd,
const char *ip,
const int port,
char *p_data,
int data_len,
union net_io_storage *p_storage)
{
if(p_net_io_server == NULL || p_data == NULL)
{
ERROR_LOG("Parameter error on receive data!");
return -1;
}
TRACE_LOG("Recv from conn id: %d, sock fd: %d, data length: %d, sto: %lld",
connection_id,
connection_fd,
data_len,
static_cast<long long>(p_storage->u64));
m_push_mb->blk_id = connection_id;
m_push_mb->skinfo.sockfd = connection_fd;
m_push_mb->skinfo.remote_ip = inet_addr(ip);
m_push_mb->skinfo.remote_port = htons(port);
m_push_mb->skinfo.storage = *p_storage;
m_push_mb->skinfo.ptr_lookout = p_storage;
m_push_mb->skinfo.connection_id = connection_id;
int pkg_len = 0;
int mimi_id = 0;
int barr_flag = 0;
char *p_push_data = p_data;
do
{
pkg_len = dll.handle_input(p_push_data, data_len, &m_push_mb->skinfo);
if (pkg_len < 0)
{
/// handle_input 返回值小于0时关闭连接
ERROR_LOG("dll handle_input(...) return %d", pkg_len);
return -1;
}
else if ((pkg_len == 0) || (pkg_len > data_len))
{
*p_storage = m_push_mb->skinfo.storage;
/// 数据长度不够一个包，继续接收
memmove(p_data, p_push_data, data_len);
return data_len;
}
else
{
/// pkg_len > 0 && pkg_len < data_len, 接收到了完整的包，处理这个包
*p_storage = m_push_mb->skinfo.storage;
}
m_push_mb->length = sizeof(shm_block_t) + pkg_len;
if (m_push_mb->length > m_push_buf_len)
{
/// 数据包溢出，说明平台中的缓冲区与 net_io_server 的缓冲区不匹配
/// 可能是最大包长配置有问题，直接关闭这个连接，可以起到一定的警告使用者的目的
ERROR_LOG("IP:Port[%s:%d] recv data len %d > max size %d!",
ip, port, pkg_len, g_bench_conf.max_pkg_len);
return -1;
}
memcpy(m_push_mb->data, p_push_data, pkg_len);
if (g_bench_conf.use_barrier)
{
mimi_id = 0;
barr_flag = 0;
if (dll.handle_dispatch == NULL)
{
ERROR_LOG("Warning: handle_dispatch not found, pkg filter disabled.");
}
else
{
if ((barr_flag = dll.handle_dispatch(p_push_data, pkg_len, 0, &mimi_id)) < 0)
{
ERROR_LOG("dll handle_dispatch failed!");
return -1;
}
}
AFUTEX_LOCK_ERR_CHK(g_rcv_rq_lock);
int ret = g_prcv_brq->push((char*)m_push_mb, m_push_mb->length, mimi_id, barr_flag);
AFUTEX_UNLOCK_ERR_CHK(g_rcv_rq_lock);
if (ret != 0)
{
/// 如果 push 失败，则丢弃这个包
ERROR_LOG("Barrier push data failed, discard pkg (len: %d): %s.", pkg_len, g_prcv_brq->get_last_errstr());
}
else
{
TRACE_LOG("Barrier ring queue push data length: %d, flag: %d, key: %d", m_push_mb->length, barr_flag, mimi_id);
}
}
else
{
if (g_prcv_rq->push_data(m_push_buf, m_push_mb->length, true) != m_push_mb->length)
{
/// 如果 push 失败，则丢弃这个包
ERROR_LOG("Common ring queue push failed(ret: %d, data: %d, empty: %d), discard pkg (len: %d)",
g_prcv_rq->get_last_errno(), g_prcv_rq->get_data_len(), g_prcv_rq->get_empty_buffer_len(), pkg_len);
}
else
{
TRACE_LOG("Common ring queue push data length: %d", m_push_mb->length);
}
}
p_push_data += pkg_len;
data_len -= pkg_len;
} while (data_len > 0);
return 0;
}
/**
* @brief 创建conn进程
* @param
* @return 0:success, -1:fail
*/
int conn_proc(int argc, char **argv)
{
//建立TCP连接
if (g_bench_conf.bind_list == NULL)
{
NB_BOOT_LOG(-1, "Bind info is null.");
return -1;
}
//暂时仅支持一个端口
struct bind_config *bc = g_bench_conf.bind_list;
if (bc->type != SOCK_STREAM)
{
NB_BOOT_LOG(-1, "Warning: cannot support UDP now.");
return -1;
}
c_newbench_net_io_event nb_event;
if (nb_event.init() != 0)
{
NB_BOOT_LOG(-1, "Init c_newbench_net_io_event object failed.");
return -1;
}
if (pnb_io_server->init(bc->bind_ip,
bc->bind_port,
&nb_event,
g_p_net_io_notifier,
1,
1) != 0)
{
NB_BOOT_LOG(-1, "Bind [%s:%u]: %s", bc->bind_ip,
bc->bind_port, pnb_io_server->get_last_errstr());
return -1;
}
NB_BOOT_LOG(0, "Listen on %s:%u", bc->bind_ip, bc->bind_port);
int psnd_data_len = g_bench_conf.max_pkg_len + sizeof(shm_block_t);
char *psnd_data = (char *)malloc(psnd_data_len);
if (psnd_data == NULL)
{
NB_BOOT_LOG(-1, "Malloc package send buffer fail, size: %d B", g_bench_conf.max_pkg_len);
return -1;
}
if (dll.handle_init && dll.handle_init(argc, argv, PROC_CONN) != 0)
{
NB_BOOT_LOG(-1, "Conn handle init fail.");
return -1;
}
setproctitle("%s:[CONN]", g_bench_conf.prog_name);
shm_block_t* mb = (shm_block_t*)psnd_data;
while(!g_stop)
{
// 每次阻塞1毫秒, 因后面要发送数据
// TODO:时间应该是可配的
if (pnb_io_server->do_io(4,
NET_IO_SERVER_CMD_ACCEPT
| NET_IO_SERVER_CMD_READ
| NET_IO_SERVER_CMD_WRITE) != 0)
{
ERROR_LOG("do_io fail: %s", pnb_io_server->get_last_errstr());
}
// 发送发送环形队列中的数据包, 发干净为止
while (!g_stop)
{
// 取出数据包
int ret = g_psnd_rq->pop_data(psnd_data, psnd_data_len, 0);
if (ret == 0)
{
break;
}
else if (ret < 0)
{
// 发生错误
ERROR_LOG("Pop data error, ret: %d, err: %s", ret, g_psnd_rq->get_last_errstr());
break;
}
else
{
if (ret != mb->length || mb->length > psnd_data_len)
{
// 出错, 打印日志以后退出发送
ERROR_LOG("pop data len %d != pkg len %d or > psnd_data_len", mb->length, psnd_data_len);
break;
}
}
if (mb->type == FIN_BLOCK)
{
pnb_io_server->close_connection(mb->blk_id, true);
continue;
}
TRACE_LOG("common ring queue pop data length: %d", mb->length);
// 发送数据
int bytes_sent = 0;
int data_len = mb->length - sizeof(shm_block_t);
if (data_len)
{
if (mb->type == BROADCAST_BLOCK)
{
ret = pnb_io_server->broadcast(
mb->blk_id,
(char*)mb->data + bytes_sent,
data_len - bytes_sent,
1);
if (ret)
{
// 广播报文为全部成功
ERROR_LOG("broadcast not success at all, %d failed",
ret);
}
}
else
{
ret = pnb_io_server->send_data_atomic(
mb->blk_id,
(char*)mb->data + bytes_sent,
data_len - bytes_sent);
if (ret != data_len - bytes_sent)
{
// 发送数据失败时丢弃这个包
ERROR_LOG("Send to conn(%d) fail, len: %d, discard this pkg.", mb->blk_id, mb->length);
}
else
{
if (mb->skinfo.ptr_lookout)
{
// 传回storage
reinterpret_cast<union net_io_storage *>(mb->skinfo.ptr_lookout)->u64
= mb->skinfo.storage.u64;
}
}
}
}
}
}
free(psnd_data);
if (dll.handle_fini)
{
dll.handle_fini(PROC_CONN);
}
return 0;
}