├── epoller ├── epoller.h └── epoller.cpp ├── README.md ├── hook ├── socket_f.h └── socket_f.cpp ├── Makefile ├── main.cpp ├── coroutine ├── coroutine.cpp └── coroutine.h └── schedule └── schedule.cpp /epoller/epoller.h: -------------------------------------------------------------------------------- 1 | #include "../coroutine/coroutine.h" 2 | 3 | class Epoller { 4 | public: 5 | static int epoll_init(); 6 | static int epoll_ev_register(); 7 | static int epoll_wait_for(uint64_t interval); 8 | 9 | }; 10 | 11 | 12 | 13 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # TinyCo 2 | Linux下一个简易的基于x86_64的协程框架，使用C++实现 3 | + 协程可绑定**变参入口函数** 4 | + 使用**C++11的chorno**调整定时器 5 | + 协程的状态有：就绪、等待、睡眠、退出 6 | + 协程对系统API进行了hook 7 | 8 | # 编译 9 | ``` 10 | make 11 | ``` 12 | 13 | 14 | # 运行 15 | ``` 16 | ./exes/test 17 | ``` 18 | 19 | # 待优化的地方 20 | + 协程的调度策略 21 | + 加入协程的优先级 22 | + 协程性能测试以及应用 23 | -------------------------------------------------------------------------------- /epoller/epoller.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "epoller.h" 3 | #include 4 | 5 | int Epoller::epoll_init() { 6 | // 1024 just a hint 7 | return epoll_create(1024); 8 | } 9 | 10 | int Epoller::epoll_ev_register() { 11 | Schedule *sched = Schedule::get_schedule(); 12 | 13 | if (sched->eventfd == 0) { 14 | sched->eventfd = eventfd(0, EFD_NONBLOCK); //创建一个用于事件通知的 fd 15 | assert(sched->eventfd != -1); 16 | } 17 | 18 | epoll_event ev; 19 | ev.events = EPOLLIN; 20 | ev.data.fd = sched->eventfd; 21 | int ret = epoll_ctl(sched->poller_fd, EPOLL_CTL_ADD, sched->eventfd, &ev); 22 | 23 | 24 | return ret; 25 | } 26 | 27 | int Epoller::epoll_wait_for(uint64_t interval) { 28 | Schedule *sched = Schedule::get_schedule(); 29 | return epoll_wait(sched->poller_fd, sched->evs, MAX_EVENTS, interval); 30 | } -------------------------------------------------------------------------------- /hook/socket_f.h: -------------------------------------------------------------------------------- 1 | #define COROUTINE_HOOK 2 | #ifdef COROUTINE_HOOK 3 | 4 | #include "../coroutine/coroutine.h" 5 | #include 6 | 7 | typedef int (*socket_t)(int domain, int type, int protocol); 8 | typedef int(*connect_t)(int, const struct sockaddr *, socklen_t); 9 | typedef ssize_t(*read_t)(int, void *, size_t); 10 | typedef ssize_t(*recv_t)(int sockfd, void *buf, size_t len, int flags); 11 | typedef ssize_t(*recvfrom_t)(int sockfd, void *buf, size_t len, int flags, 12 | struct sockaddr *src_addr, socklen_t *addrlen); 13 | typedef ssize_t(*write_t)(int, const void *, size_t); 14 | typedef ssize_t(*send_t)(int sockfd, const void *buf, size_t len, int flags); 15 | typedef ssize_t(*sendto_t)(int sockfd, const void *buf, size_t len, int flags, 16 | const struct sockaddr *dest_addr, socklen_t addrlen); 17 | typedef int(*accept_t)(int sockfd, struct sockaddr *addr, socklen_t *addrlen); 18 | typedef int(*close_t)(int); 19 | 20 | extern socket_t socket_f; 21 | extern connect_t connect_f; 22 | extern read_t read_f; 23 | extern recv_t recv_f; 24 | extern recvfrom_t recvfrom_f; 25 | extern write_t write_f; 26 | extern send_t send_f; 27 | extern sendto_t sendto_f; 28 | extern accept_t accept_f; 29 | extern close_t close_f; 30 | 31 | 32 | int init_hook(void); 33 | 34 | 35 | #endif // -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: all clean 2 | MKDIR = mkdir 3 | RM = rm 4 | FLAGS = -lpthread -l dl -g 5 | RMFLAGS = -rf 6 | CC = g++ 7 | 8 | DIR_OBJS = objs 9 | DIR_EXES = exes 10 | 11 | 12 | DIR_COROUTINE = coroutine 13 | DIR_SCHEDULE = schedule 14 | DIR_EPOLLER = epoller 15 | DIR_HOOK = hook 16 | 17 | DIRS = $(DIR_OBJS) $(DIR_EXES) $(DIR_OBJS)/$(DIR_COROUTINE) $(DIR_OBJS)/$(DIR_SCHEDULE) $(DIR_OBJS)/$(DIR_EPOLLER) $(DIR_OBJS)/$(DIR_HOOK) 18 | 19 | EXE = test 20 | EXE := $(addprefix $(DIR_EXES)/, $(EXE)) 21 | SRCCUR = $(wildcard *.cpp) 22 | SRC_COR = $(wildcard $(DIR_COROUTINE)/*.cpp) 23 | SRC_SCHE = $(wildcard $(DIR_SCHEDULE)/*.cpp) 24 | SRC_EPOLL = $(wildcard $(DIR_EPOLLER)/*.cpp) 25 | SRC_HOOK = $(wildcard $(DIR_HOOK)/*.cpp) 26 | 27 | OBJ_CUR = $(patsubst %.cpp,%.o,$(SRCCUR)) 28 | OBJ_COR = $(patsubst %.cpp,%.o,$(SRC_COR)) 29 | OBJ_SCHE = $(patsubst %.cpp,%.o,$(SRC_SCHE)) 30 | OBJ_EPOLL = $(patsubst %.cpp,%.o,$(SRC_EPOLL)) 31 | OBJ_HOOK = $(patsubst %.cpp,%.o,$(SRC_HOOK)) 32 | 33 | OBJ_CUR := $(addprefix $(DIR_OBJS)/, $(OBJ_CUR)) 34 | OBJ_COR := $(addprefix $(DIR_OBJS)/, $(OBJ_COR)) 35 | OBJ_SCHE := $(addprefix $(DIR_OBJS)/, $(OBJ_SCHE)) 36 | OBJ_EPOLL := $(addprefix $(DIR_OBJS)/, $(OBJ_EPOLL)) 37 | OBJ_HOOK := $(addprefix $(DIR_OBJS)/, $(OBJ_HOOK)) 38 | 39 | OBJS = $(OBJ_COR) $(OBJ_CUR) $(OBJ_SCHE) $(OBJ_EPOLL) $(OBJ_HOOK) $(OBJ_CUR) 40 | 41 | 42 | 43 | all: $(DIRS) $(EXE) 44 | 45 | $(DIRS): 46 | $(MKDIR) $@ 47 | 48 | 49 | $(EXE): $(OBJS) 50 | $(CC) $^ -o $@ $(FLAGS) 51 | $(OBJ_COR) : $(SRC_COR) 52 | $(CC) -c $^ -o $@ 53 | $(OBJ_SCHE) : $(SRC_SCHE) 54 | $(CC) -c $^ -o $@ 55 | $(OBJ_EPOLL) : $(SRC_EPOLL) 56 | $(CC) -c $^ -o $@ 57 | $(OBJ_HOOK) : $(SRC_HOOK) 58 | $(CC) -c $^ -o $@ 59 | $(OBJ_CUR) : $(SRCCUR) 60 | $(CC) -c $^ -o $@ 61 | 62 | 63 | clean: 64 | $(RM) $(RMFLAGS) $(DIRS) 65 | -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- 1 | #include "coroutine/coroutine.h" 2 | #include "epoller/epoller.h" 3 | #include "hook/socket_f.h" 4 | 5 | #include 6 | #include 7 | 8 | #define MAX_CLIENT_NUM 100 9 | 10 | 11 | void read_cb(void *arg) { 12 | int fd = *(int *)arg; 13 | int ret = 0; 14 | 15 | 16 | struct pollfd fds; 17 | fds.fd = fd; 18 | fds.events = POLLIN; 19 | 20 | while (1) { 21 | 22 | char buf[1024] = {0}; 23 | ret = recv(fd, buf, 1024, 0); 24 | if (ret > 0) { 25 | if(fd > MAX_CLIENT_NUM) 26 | printf("read from server: %.*s\n", ret, buf); 27 | 28 | ret = send(fd, buf, strlen(buf), 0); 29 | if (ret == -1) { 30 | close(fd); 31 | break; 32 | } 33 | } else if (ret == 0) { 34 | close(fd); 35 | break; 36 | } 37 | 38 | } 39 | } 40 | 41 | 42 | void server(void *arg) { 43 | 44 | unsigned short port = *(unsigned short *)arg; 45 | free(arg); 46 | 47 | int fd = socket(AF_INET, SOCK_STREAM, 0); 48 | if (fd < 0) return ; 49 | 50 | struct sockaddr_in local, remote; 51 | local.sin_family = AF_INET; 52 | local.sin_port = htons(port); 53 | local.sin_addr.s_addr = INADDR_ANY; 54 | bind(fd, (struct sockaddr*)&local, sizeof(struct sockaddr_in)); 55 | 56 | listen(fd, 20); 57 | printf("listen port : %d\n", port); 58 | 59 | steady_clock::time_point t_begin = steady_clock::now(); 60 | 61 | while (1) { 62 | socklen_t len = sizeof(struct sockaddr_in); 63 | 64 | int client_fd = accept(fd, (struct sockaddr*)&remote, &len); 65 | 66 | if (client_fd % 1000 == 999) { 67 | 68 | steady_clock::time_point t_now = steady_clock::now(); 69 | milliseconds time_used = duration_cast(t_now.time_since_epoch() - t_begin.time_since_epoch()); 70 | 71 | printf("client fd : %d, time_used: %ld\n", client_fd, time_used.count()); 72 | } 73 | 74 | coroutine_create(NULL, read_cb, &client_fd); 75 | } 76 | 77 | } 78 | 79 | 80 | int main(int argc, char *argv[]) { 81 | Coroutine *co = NULL; 82 | 83 | int i = 0; 84 | unsigned short base_port = 8888; 85 | for (i = 0;i < MAX_CLIENT_NUM;i ++) { 86 | unsigned short *port = (unsigned short *)calloc(1, sizeof(unsigned short)); 87 | *port = base_port + i; 88 | coroutine_create(NULL, server, port); // not run 89 | } 90 | 91 | Schedule *sched = Schedule::get_schedule(); 92 | sched->run(); // run 93 | schedule_free(sched); 94 | 95 | return 0; 96 | } 97 | -------------------------------------------------------------------------------- /coroutine/coroutine.cpp: -------------------------------------------------------------------------------- 1 | #include "coroutine.h" 2 | 3 | pthread_key_t global_sched_key; 4 | 5 | pthread_once_t sched_key_once = PTHREAD_ONCE_INIT; // 在线程中只执行一次 6 | 7 | #define SCHEDULE_PAGE_SIZE (1024*4) 8 | 9 | 10 | extern "C" 11 | { 12 | 13 | int _switch(cpu_ctx *new_ctx, cpu_ctx *cur_ctx); 14 | // 内嵌入 _switch的定义 15 | 16 | 17 | #ifdef __x86_64__ 18 | __asm__ ( 19 | " .text \n" 20 | " .p2align 4,,15 \n" 21 | ".globl _switch \n" 22 | "_switch: \n" 23 | " movq %rsp, 0(%rsi) # save stack_pointer \n" 24 | " movq %rbp, 8(%rsi) # save frame_pointer \n" 25 | " movq (%rsp), %rax # save insn_pointer \n" 26 | " movq %rax, 16(%rsi) \n" 27 | " movq %rbx, 24(%rsi) # save rbx,r12-r15 \n" 28 | " movq %r12, 32(%rsi) \n" 29 | " movq %r13, 40(%rsi) \n" 30 | " movq %r14, 48(%rsi) \n" 31 | " movq %r15, 56(%rsi) \n" 32 | " movq 56(%rdi), %r15 \n" 33 | " movq 48(%rdi), %r14 \n" 34 | " movq 40(%rdi), %r13 # restore rbx,r12-r15 \n" 35 | " movq 32(%rdi), %r12 \n" 36 | " movq 24(%rdi), %rbx \n" 37 | " movq 8(%rdi), %rbp # restore frame_pointer \n" 38 | " movq 0(%rdi), %rsp # restore stack_pointer \n" 39 | " movq 16(%rdi), %rax # restore insn_pointer \n" 40 | " movq %rax, (%rsp) \n" 41 | " ret \n" 42 | ); 43 | } 44 | #endif 45 | 46 | 47 | void _exec(void *cor) { 48 | Coroutine *co = (Coroutine *)cor; 49 | co->task._func(); 50 | co->status = co->status | (BIT(STATUS_EXIT)); 51 | co->sche->exitQueue.push(co); 52 | co->yield(); 53 | } 54 | 55 | void coroutine_init(Coroutine *co) { 56 | // 一级指针强转为二级指针，也就是将一级指针所指内存中的值强转为 void* 57 | void **stack = (void **)((char *)co->stack_addr + co->stack_size); 58 | 59 | // 从 stack[-1] 到 stack[-4] 是用来保存栈指针、某些数据、协程信息等的预留空间 60 | stack[-3] = NULL; 61 | stack[-2] = (void *)co; // 协程信息保存 stack[-2] 62 | 63 | co->ctx.rsp = (char*)stack - (4 * sizeof(char*)); // 栈顶指针保存 stack[-4] 64 | co->ctx.rbp = (char*)stack - (3 * sizeof(char*)); // 栈基指针保存 stack[-3] 65 | co->ctx.rip = (void*)_exec; //下一条指令是协程的入口函数 66 | co->status &= CLEARBIT(STATUS_NEW); 67 | co->status |= BIT(STATUS_READY); 68 | } 69 | 70 | 71 | 72 | /* 73 | Coroutine::Coroutine(Coroutine **new_co, coroutine_entry funcp, void *argp) { 74 | int ret = pthread_once(&sched_key_once, Schedule::schedule_init_key); 75 | if (ret != 0) { 76 | perror("pthread_once"); 77 | } 78 | Schedule *sched = Schedule::get_schedule(); 79 | if (sched == NULL) { 80 | sched = schedule_create(); 81 | if (sched == NULL) { 82 | printf("Failed to create scheduler\n"); 83 | return ; 84 | } 85 | } 86 | 87 | // 一个栈分配 4k 的整数倍，一般来说就是一页 88 | 89 | ret = posix_memalign(&stack_addr, getpagesize(), sched->stack_size); 90 | 91 | sche = sched; 92 | stack_size = sched->stack_size; 93 | status = BIT(STATUS_NEW); 94 | func = funcp; 95 | arg = argp; 96 | id = sched->spawned_coroutines++; 97 | 98 | fd = -1; 99 | events = EPOLLIN; 100 | 101 | if (new_co) 102 | *new_co = this; // 获取当前协程 103 | 104 | bzero(&ctx, sizeof(cpu_ctx)); 105 | sche->ready.push(this); 106 | } 107 | */ 108 | /* void coroutine_create(Coroutine **new_co, coroutine_entry func, void *arg) { 109 | new Coroutine(new_co, func, arg); 110 | } 111 | */ 112 | 113 | int Coroutine::yield() { 114 | precedence = 0; 115 | _switch(&sche->ctx, &ctx); 116 | } 117 | 118 | 119 | 120 | int Coroutine::resume() { 121 | if (status & BIT(STATUS_NEW)) { 122 | coroutine_init(this); 123 | 124 | } 125 | Schedule *sched = Schedule::get_schedule(); 126 | sched->cur_co = this; 127 | _switch(&ctx, &sched->ctx); 128 | // 由于切换了到了协程，因此调度器的当前协程为NULL 129 | 130 | sched->cur_co = NULL; 131 | 132 | } 133 | 134 | void Coroutine::sleep_for(uint64_t msec) { 135 | if (msec == 0) { 136 | sche->ready.push(this); 137 | } else { 138 | Schedule::get_schedule()->add_sleep(this, msec); 139 | } 140 | } 141 | 142 | void Coroutine::detach() { 143 | status |= STATUS_DETACH; 144 | } 145 | 146 | Coroutine::~Coroutine() { 147 | if (stack_addr) { 148 | free(stack_addr); 149 | stack_addr = NULL; 150 | } 151 | } 152 | 153 | void coroutine_free(Coroutine *co){ 154 | delete co; 155 | } 156 | 157 | 158 | 159 | 160 | 161 | -------------------------------------------------------------------------------- /coroutine/coroutine.h: -------------------------------------------------------------------------------- 1 | #ifndef _COROUTINE_H_ 2 | #define _COROUTINE_H_ 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include 20 | #include 21 | #include 22 | using namespace std::chrono; 23 | 24 | 25 | #define MAX_EVENTS 1024 26 | #define MAX_STACKSIZE (4*1024) // 最大的栈大小 27 | 28 | extern pthread_key_t global_sched_key; 29 | extern pthread_once_t sched_key_once; // 在线程中只执行一次 30 | 31 | #define BIT(x) (1 << (x)) 32 | #define CLEARBIT(x) ~(1 << (x)) 33 | 34 | struct cpu_ctx { 35 | void *rsp; // 栈顶指针 36 | void *rbp; // 被调用者保存，用作数据存储 37 | void *rip; // 下一条指令 38 | void *rbx; // 被调用者保存 39 | void *r12; // 被调用者保存 40 | void *r13; // 被调用者保存 41 | void *r14; // 被调用者保存 42 | void *r15; // 被调用者保存 43 | }; 44 | 45 | enum CorStatus { 46 | STATUS_READY = 0, 47 | STATUS_WAIT_READ, 48 | STATUS_WAIT_WRITE, 49 | STATUS_SLEEPING, 50 | STATUS_EXIT, 51 | STATUS_DETACH, 52 | STATUS_NEW, 53 | STATUS_FDEOF 54 | 55 | }; 56 | 57 | 58 | struct Schedule; 59 | extern Schedule* schedule_create(); 60 | 61 | 62 | class Coroutine { 63 | public: 64 | Coroutine() : stack_addr(NULL) {}; 65 | 66 | struct TaskFunc 67 | { 68 | TaskFunc(std::function func = NULL) : _func(func) 69 | { } 70 | 71 | std::function _func; 72 | }; 73 | //Coroutine(Coroutine **new_co, coroutine_entry func, void *arg); 74 | template 75 | Coroutine(Coroutine **new_co, Entry &&f, Args &&... args); 76 | 77 | template 78 | friend void coroutine_create(Coroutine **new_co, Entry &&f, Args &&... args); 79 | 80 | ~Coroutine(); 81 | 82 | int yield(); 83 | int resume(); 84 | void sleep_for(uint64_t msec); 85 | void detach(); 86 | 87 | // 允许以下的函数和类对协程进行修改 88 | friend void coroutine_free(Coroutine *co); 89 | friend void _exec(void *cor); 90 | friend void coroutine_init(Coroutine *co); 91 | friend int poll_inner(pollfd *fds, nfds_t num_fd, int timeout); 92 | friend struct SleepCmp; 93 | friend struct WaitCmp; 94 | friend struct Schedule; 95 | 96 | 97 | private: 98 | cpu_ctx ctx; // 用来 save 上下文 99 | TaskFunc task; 100 | void *stack_addr; 101 | size_t stack_size; 102 | size_t last_stack_size; //剩余栈空间 103 | u_int64_t id; 104 | 105 | // 一个 fd，对应一个事件，在一个协程中运行 106 | int fd; 107 | unsigned short events; 108 | int status; 109 | 110 | // 协程所属的调度器 111 | Schedule *sche; 112 | 113 | int precedence; // 运行优先级 114 | uint64_t sleep_usecs; //睡眠时间 115 | 116 | }; 117 | 118 | struct SleepCmp { 119 | bool operator()(const Coroutine *co1, const Coroutine *co2) const { 120 | return co1->sleep_usecs < co2->sleep_usecs; 121 | } 122 | }; 123 | 124 | struct WaitCmp { 125 | bool operator()(const Coroutine *co1, const Coroutine *co2) const { 126 | return co1->fd < co2->fd; 127 | } 128 | }; 129 | 130 | 131 | // 调度器的定义 132 | class Schedule { 133 | public: 134 | Schedule(int stack_size); 135 | ~Schedule(); 136 | 137 | static Schedule *get_schedule(void); 138 | static void sched_key_destructor(void *data); 139 | static int schedule_epoll(Schedule *sched); 140 | static void schedule_init_key(); 141 | // 获取sleep 超时的 coroutine 142 | static Coroutine *get_expired_coroutine(Schedule *sched); 143 | // 获取 sleep 可能最近超时的时间间隔 144 | static uint64_t get_min_timeout(Schedule *sched); 145 | 146 | void del_sleep(Coroutine *co); 147 | void add_sleep(Coroutine *co, uint64_t msecs); 148 | 149 | Coroutine* search_wait(int fd); 150 | Coroutine* del_wait(int fd); 151 | void add_wait(Coroutine *co, int fd, unsigned short events, uint64_t timeout); 152 | 153 | void run(); 154 | friend void _exec(void *cor); 155 | friend void schedule_free(Schedule *sched); 156 | friend Schedule* schedule_create(); 157 | friend int poll_inner(pollfd *fds, nfds_t num_fd, int timeout); 158 | friend class Coroutine; 159 | friend class Epoller; 160 | 161 | 162 | private: 163 | cpu_ctx ctx; 164 | 165 | uint64_t birth; 166 | size_t stack_size; 167 | uint64_t default_timeout; 168 | int spawned_coroutines; 169 | 170 | Coroutine *cur_co; //正在运行的协程 171 | int poller_fd; // epfd -- IO操作之前用来检测是否有事件 172 | int eventfd; //事件通知的 fd，用于和其他进程通信 173 | 174 | int num_new_events; // poll 检测的新的IO数量 175 | epoll_event evs[MAX_EVENTS]; 176 | 177 | int pagesize; //页大小，用于内存优化 178 | 179 | std::queue ready; 180 | std::queue exitQueue; 181 | std::set sleepSet; 182 | std::set wait; 183 | 184 | }; 185 | 186 | 187 | 188 | static uint64_t getTimeNow() { 189 | steady_clock::time_point t1 = steady_clock::now(); 190 | microseconds t2 = duration_cast (t1.time_since_epoch()); 191 | return t2.count(); 192 | } 193 | 194 | 195 | 196 | template 197 | void coroutine_create(Coroutine **new_co, Entry &&f, Args &&... args) { 198 | new Coroutine(new_co, f, args...); 199 | } 200 | 201 | 202 | template 203 | Coroutine::Coroutine(Coroutine **new_co, Entry &&f, Args &&... args) { 204 | 205 | int ret = pthread_once(&sched_key_once, Schedule::schedule_init_key); 206 | if (ret != 0) { 207 | perror("pthread_once"); 208 | } 209 | Schedule *sched = Schedule::get_schedule(); 210 | if (sched == NULL) { 211 | sched = schedule_create(); 212 | if (sched == NULL) { 213 | printf("Failed to create scheduler\n"); 214 | return ; 215 | } 216 | } 217 | 218 | // 一个栈分配 4k 的整数倍，一般来说就是一页 219 | 220 | ret = posix_memalign(&stack_addr, getpagesize(), sched->stack_size); 221 | 222 | sche = sched; 223 | stack_size = sched->stack_size; 224 | status = BIT(STATUS_NEW); 225 | task._func = std::bind(std::forward(f), std::forward(args)...); 226 | id = sched->spawned_coroutines++; 227 | 228 | fd = -1; 229 | events = EPOLLIN; 230 | 231 | if (new_co) 232 | *new_co = this; // 获取当前协程 233 | 234 | bzero(&ctx, sizeof(cpu_ctx)); 235 | sche->ready.push(this); 236 | } 237 | 238 | 239 | 240 | 241 | #endif // !_COROUTINE_H_ -------------------------------------------------------------------------------- /schedule/schedule.cpp: -------------------------------------------------------------------------------- 1 | #include "../coroutine/coroutine.h" 2 | #include "../epoller/epoller.h" 3 | 4 | void Schedule::schedule_init_key() { 5 | pthread_key_create(&global_sched_key, Schedule::sched_key_destructor); 6 | pthread_setspecific(global_sched_key, NULL); 7 | } 8 | 9 | Schedule::~Schedule() { 10 | schedule_free(this); 11 | } 12 | 13 | Schedule* Schedule::get_schedule(void) { // 取出线程特有的值 14 | return (Schedule *)pthread_getspecific(global_sched_key); 15 | } 16 | 17 | void Schedule::sched_key_destructor(void *data) { 18 | free(data); 19 | } 20 | 21 | 22 | Schedule::Schedule(int size) { 23 | eventfd = 0; 24 | int sched_stack_size = size > 0 ? size : MAX_STACKSIZE; 25 | pthread_setspecific(global_sched_key, this); 26 | poller_fd = Epoller::epoll_init(); 27 | if (poller_fd == -1) { 28 | printf("Failed to initialize epoller\n"); 29 | return ; 30 | } 31 | 32 | Epoller::epoll_ev_register(); 33 | 34 | stack_size = sched_stack_size; 35 | pagesize = getpagesize(); 36 | 37 | spawned_coroutines = 0; 38 | default_timeout = 3000000u; 39 | 40 | birth = getTimeNow(); 41 | 42 | bzero(&ctx, sizeof(cpu_ctx)); 43 | } 44 | 45 | 46 | 47 | void Schedule::del_sleep(Coroutine *co) { 48 | if (co->status & BIT(STATUS_SLEEPING)) { 49 | co->sche->sleepSet.erase(co); 50 | co->status &= CLEARBIT(STATUS_SLEEPING); 51 | co->status |= BIT(STATUS_READY); 52 | } 53 | } 54 | 55 | 56 | 57 | void Schedule::add_sleep(Coroutine *co, uint64_t msecs) { 58 | uint64_t usecs = msecs * 1000u; 59 | co->sleep_usecs = getTimeNow() - co->sche->birth + usecs; 60 | 61 | auto co_tmp = co->sche->sleepSet.find(co); 62 | if (co_tmp != co->sche->sleepSet.end()) { // 如果找到了，需要删除后重新设置插入新的协程 63 | co->sche->sleepSet.erase(co_tmp); 64 | co->sche->sleepSet.insert(co); 65 | } else { //如果没找到，设置状态，然后插入协程 66 | co->status |= BIT(STATUS_SLEEPING); 67 | co->sche->sleepSet.insert(co); 68 | } 69 | } 70 | 71 | Coroutine* Schedule::search_wait(int fd) { 72 | Coroutine target; 73 | target.fd = fd; 74 | Schedule *sched = get_schedule(); 75 | auto it = sched->wait.find(&target); 76 | if (it == wait.end()) { 77 | return NULL; 78 | } 79 | return *it; 80 | } 81 | 82 | 83 | Coroutine* Schedule::del_wait(int fd) { 84 | Coroutine target; 85 | target.fd = fd; 86 | 87 | Schedule *sched = get_schedule(); 88 | std::set::iterator it = sched->wait.find(&target); 89 | del_sleep(*it); 90 | if (it != wait.end()) 91 | sched->wait.erase(it); 92 | 93 | (*it)->status = BIT(STATUS_READY); 94 | return *it; 95 | } 96 | 97 | void Schedule::add_wait(Coroutine *co, int fd, unsigned short events, uint64_t timeout) { 98 | if (events & POLLIN) { 99 | co->status |= STATUS_WAIT_READ; 100 | } else if (events & POLLOUT) { 101 | co->status |= STATUS_WAIT_WRITE; 102 | } else { 103 | printf("events : %d\n", events); 104 | assert(0); 105 | } 106 | co->fd = fd; 107 | co->events = events; 108 | 109 | add_sleep(co, timeout); 110 | co->sche->wait.insert(co); 111 | } 112 | 113 | void Schedule::run() { 114 | while (!sleepSet.empty() || !ready.empty() || !wait.empty() || !exitQueue.empty()) { 115 | // sleep rbtree-- expired 116 | while (!sleepSet.empty()) { 117 | Coroutine *expired_co = get_expired_coroutine(this); 118 | if (expired_co == NULL) { 119 | break; 120 | } 121 | expired_co->resume();// 让sleep超时的线程先执行 122 | } 123 | // ready queue 124 | while (!ready.empty()) { 125 | Coroutine *co = ready.front(); 126 | ready.pop(); 127 | 128 | co->resume(); 129 | } 130 | // wait rbtree 131 | schedule_epoll(this); 132 | while (num_new_events > 0) { 133 | int idx = --num_new_events; 134 | epoll_event *ev = evs + idx; //第i个事件 135 | 136 | int fd = ev->data.fd; 137 | int is_eof = ev->events & EPOLLHUP; 138 | if (is_eof) errno = ECONNRESET; // 如果对端关闭，设置errno 139 | 140 | 141 | Coroutine* co = search_wait(fd); 142 | if (co != NULL) { 143 | if (is_eof) { // 如果对端关闭，设置协程状态 144 | co->status |= BIT(STATUS_FDEOF); 145 | } 146 | co->resume(); 147 | } 148 | } 149 | // exit queue 150 | while (!exitQueue.empty()) { 151 | Coroutine *co = exitQueue.front(); 152 | exitQueue.pop(); 153 | coroutine_free(co); 154 | } 155 | } 156 | 157 | } 158 | 159 | Schedule* schedule_create() { 160 | Schedule *sched = new Schedule(0); 161 | return sched; 162 | } 163 | 164 | void schedule_free(Schedule *sched) { 165 | if (sched->poller_fd > 0) { 166 | close(sched->poller_fd); 167 | } 168 | if (sched->eventfd > 0) { 169 | close(sched->eventfd); 170 | } 171 | 172 | free(sched); 173 | assert(pthread_setspecific(global_sched_key, NULL) == 0); 174 | } 175 | 176 | Coroutine* Schedule::get_expired_coroutine(Schedule *sched) { 177 | uint64_t diff_usecs = getTimeNow() - sched->birth; 178 | auto it = sched->sleepSet.begin(); // 获取最小的 179 | //printf("get_expired_coroutine avaiable: [co->sleep_usecs: (%ld)]\n", it->sleep_usecs); 180 | if (it == sched->sleepSet.end()) { 181 | return NULL; 182 | } 183 | 184 | if ((*it)->sleep_usecs <= diff_usecs) { //如果有超时的协程 185 | sched->sleepSet.erase(it); 186 | return *it; 187 | } 188 | //没有超时的协程 189 | return NULL; 190 | } 191 | 192 | uint64_t Schedule::get_min_timeout(Schedule *sched) { 193 | uint64_t diff_usecs = getTimeNow() - sched->birth; 194 | uint64_t min = sched->default_timeout; 195 | 196 | auto it = sched->sleepSet.begin(); 197 | if (it == sched->sleepSet.end()) { //如果没有找到，就是默认超时时间 198 | return min; 199 | } 200 | 201 | min = (*it)->sleep_usecs; 202 | if (min > diff_usecs) { 203 | return min - diff_usecs; // 如果协程的睡眠时间还没有超时，返回差值 204 | } 205 | return 0; //否则就判定为超时，返回0 206 | 207 | } 208 | 209 | // 在 ready queue 处理完后进行 schedule_epoll 210 | int Schedule::schedule_epoll(Schedule *sched) { 211 | sched->num_new_events = 0; 212 | uint64_t usecs = get_min_timeout(sched); 213 | if (usecs <= 0 || !sched->ready.empty()) { // 如果超时时间大于0，并且就绪队列为空,就继续执行 214 | return 0; 215 | } 216 | 217 | int nready = 0; 218 | // 使用循环主要是因为可能会被信号打断而没有 epoll_wait成功 219 | while (1) { 220 | nready = Epoller::epoll_wait_for(usecs); 221 | if (nready == -1) { 222 | if (errno == EINTR) continue; 223 | else assert(0); 224 | } 225 | break; 226 | } 227 | //sched->nevents = 0; 228 | sched->num_new_events = nready; 229 | } 230 | 231 | -------------------------------------------------------------------------------- /hook/socket_f.cpp: -------------------------------------------------------------------------------- 1 | #include "socket_f.h" 2 | 3 | 4 | // 如果已经超时，需要先使用poll检测，否则使用epoll检测 5 | 6 | //将 poll 事件转换为 epoll 7 | static uint32_t ev_poll_2_epoll( short events ) 8 | { 9 | uint32_t e = 0; 10 | if( events & POLLIN ) e |= EPOLLIN; 11 | if( events & POLLOUT ) e |= EPOLLOUT; 12 | if( events & POLLHUP ) e |= EPOLLHUP; 13 | if( events & POLLERR ) e |= EPOLLERR; 14 | if( events & POLLRDNORM ) e |= EPOLLRDNORM; 15 | if( events & POLLWRNORM ) e |= EPOLLWRNORM; 16 | return e; 17 | } 18 | 19 | 20 | int poll_inner(pollfd *fds, nfds_t num_fd, int timeout) { 21 | if (timeout == 0) { 22 | return poll(fds, num_fd, 0); //如果超时时间为0，那么直接使用 poll检测 23 | } 24 | if (timeout < 0) { 25 | timeout = INT_MAX; 26 | } 27 | 28 | Schedule *sched = Schedule::get_schedule(); 29 | Coroutine *co = sched->cur_co; 30 | 31 | for (int i = 0; i < num_fd; ++i) { 32 | epoll_event ev; 33 | ev.events = ev_poll_2_epoll(fds[i].events); 34 | ev.data.fd = fds[i].fd; 35 | 36 | epoll_ctl(sched->poller_fd, EPOLL_CTL_ADD, fds[i].fd, &ev); 37 | co->events = fds[i].events; 38 | sched->add_wait(co, fds[i].fd, fds[i].events, timeout); 39 | } 40 | co->yield(); 41 | 42 | //注册完事件加入等待队列后，删除注册的事件 43 | for (int i = 0; i < num_fd; ++i) { 44 | 45 | struct epoll_event ev; 46 | ev.events = ev_poll_2_epoll(fds[i].events); 47 | ev.data.fd = fds[i].fd; 48 | epoll_ctl(sched->poller_fd, EPOLL_CTL_DEL, fds[i].fd, &ev); 49 | 50 | sched->del_wait(fds[i].fd); 51 | } 52 | 53 | return num_fd; 54 | } 55 | 56 | int init_socket(int domain, int type, int protocol) { 57 | 58 | int fd = socket(domain, type, protocol); 59 | if (fd == -1) { 60 | printf("Failed to create a new socket\n"); 61 | return -1; 62 | } 63 | int ret = fcntl(fd, F_SETFL, O_NONBLOCK); 64 | if (ret == -1) { 65 | close(ret); 66 | return -1; 67 | } 68 | int reuse = 1; 69 | setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (char *)&reuse, sizeof(reuse)); 70 | return fd; 71 | } 72 | 73 | int cor_accept(int fd, struct sockaddr *addr, socklen_t *len) { 74 | int sockfd = -1; 75 | int timeout = 1; 76 | 77 | // 使用循环是防止被信号打断 78 | while (1) { 79 | pollfd fds; 80 | fds.fd = fd; 81 | fds.events = POLLIN | POLLERR | POLLHUP; 82 | poll_inner(&fds, 1, timeout); 83 | 84 | sockfd = accept(fd, addr, len); 85 | if (sockfd < 0) { 86 | if (errno == EAGAIN) { 87 | continue; 88 | } else if (errno == ECONNABORTED) { 89 | printf("accept : ECONNABORTED\n"); 90 | 91 | } else if (errno == EMFILE || errno == ENFILE) { 92 | printf("accept : EMFILE || ENFILE\n"); 93 | } 94 | return -1; 95 | } else { 96 | break; 97 | } 98 | } 99 | 100 | int ret = fcntl(sockfd, F_SETFL, O_NONBLOCK); 101 | if (ret == -1) { 102 | close(sockfd); 103 | return -1; 104 | } 105 | int reuse = 1; 106 | setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (char *)&reuse, sizeof(reuse)); 107 | 108 | return sockfd; 109 | } 110 | 111 | int cor_connect(int fd, struct sockaddr *name, socklen_t namelen) { 112 | 113 | int ret = 0; 114 | 115 | while (1) { 116 | 117 | struct pollfd fds; 118 | fds.fd = fd; 119 | fds.events = POLLOUT | POLLERR | POLLHUP; 120 | poll_inner(&fds, 1, 1); 121 | 122 | ret = connect(fd, name, namelen); 123 | if (ret == 0) break; 124 | 125 | if (ret == -1 && (errno == EAGAIN || 126 | errno == EWOULDBLOCK || 127 | errno == EINPROGRESS)) { 128 | continue; 129 | } else { 130 | break; 131 | } 132 | } 133 | 134 | return ret; 135 | } 136 | 137 | ssize_t cor_recv(int fd, void *buf, size_t len, int flags) { 138 | 139 | struct pollfd fds; 140 | fds.fd = fd; 141 | fds.events = POLLIN | POLLERR | POLLHUP; 142 | 143 | poll_inner(&fds, 1, 1); 144 | 145 | int ret = recv(fd, buf, len, flags); 146 | if (ret < 0) { 147 | //if (errno == EAGAIN) return ret; 148 | if (errno == ECONNRESET) return -1; 149 | //printf("recv error : %d, ret : %d\n", errno, ret); 150 | 151 | } 152 | return ret; 153 | } 154 | 155 | ssize_t cor_send(int fd, const void *buf, size_t len, int flags) { 156 | 157 | int sent = 0, ret = 0; 158 | 159 | while (sent < len) { 160 | struct pollfd fds; 161 | fds.fd = fd; 162 | fds.events = POLLOUT | POLLERR | POLLHUP; 163 | 164 | poll_inner(&fds, 1, 1); 165 | ret = send(fd, ((char*)buf)+sent, len-sent, flags); 166 | //printf("send --> len : %d\n", ret); 167 | if (ret <= 0) { 168 | break; 169 | } 170 | sent += ret; 171 | } 172 | 173 | if (ret <= 0 && sent == 0) return ret; 174 | 175 | return sent; 176 | } 177 | 178 | int cor_close(int fd) { 179 | 180 | return close(fd); 181 | } 182 | 183 | #ifdef COROUTINE_HOOK 184 | 185 | socket_t socket_f = NULL; 186 | 187 | read_t read_f = NULL; 188 | recv_t recv_f = NULL; 189 | recvfrom_t recvfrom_f = NULL; 190 | 191 | write_t write_f = NULL; 192 | send_t send_f = NULL; 193 | sendto_t sendto_f = NULL; 194 | 195 | accept_t accept_f = NULL; 196 | close_t close_f = NULL; 197 | connect_t connect_f = NULL; 198 | 199 | 200 | int init_hook(void) { 201 | 202 | socket_f = (socket_t)dlsym(RTLD_NEXT, "socket"); 203 | 204 | recv_f = (recv_t)dlsym(RTLD_NEXT, "recv"); 205 | recvfrom_f = (recvfrom_t)dlsym(RTLD_NEXT, "recvfrom"); 206 | 207 | send_f = (send_t)dlsym(RTLD_NEXT, "send"); 208 | sendto_f = (sendto_t)dlsym(RTLD_NEXT, "sendto"); 209 | 210 | accept_f = (accept_t)dlsym(RTLD_NEXT, "accept"); 211 | close_f = (close_t)dlsym(RTLD_NEXT, "close"); 212 | connect_f = (connect_t)dlsym(RTLD_NEXT, "connect"); 213 | 214 | } 215 | 216 | int socket(int domain, int type, int protocol) { 217 | 218 | if (!socket_f) init_hook(); 219 | int fd = socket_f(domain, type, protocol); 220 | if (fd == -1) { 221 | printf("Failed to create a new socket\n"); 222 | return -1; 223 | } 224 | int ret = fcntl(fd, F_SETFL, O_NONBLOCK); 225 | if (ret == -1) { 226 | close(ret); 227 | return -1; 228 | } 229 | int reuse = 1; 230 | setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (char *)&reuse, sizeof(reuse)); 231 | return fd; 232 | } 233 | 234 | int accept(int fd, struct sockaddr *addr, socklen_t *len) { 235 | if (!accept_f) init_hook(); 236 | int sockfd = -1; 237 | int timeout = 1; 238 | 239 | // 使用循环是防止被信号打断 240 | while (1) { 241 | pollfd fds; 242 | fds.fd = fd; 243 | fds.events = POLLIN | POLLERR | POLLHUP; 244 | poll_inner(&fds, 1, timeout); 245 | 246 | sockfd = accept_f(fd, addr, len); 247 | if (sockfd < 0) { 248 | if (errno == EAGAIN) { 249 | continue; 250 | } else if (errno == ECONNABORTED) { 251 | printf("accept : ECONNABORTED\n"); 252 | 253 | } else if (errno == EMFILE || errno == ENFILE) { 254 | printf("accept : EMFILE || ENFILE\n"); 255 | } 256 | return -1; 257 | } else { 258 | break; 259 | } 260 | } 261 | 262 | int ret = fcntl(sockfd, F_SETFL, O_NONBLOCK); 263 | if (ret == -1) { 264 | close(sockfd); 265 | return -1; 266 | } 267 | int reuse = 1; 268 | setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (char *)&reuse, sizeof(reuse)); 269 | 270 | return sockfd; 271 | } 272 | 273 | int connect(int fd, struct sockaddr *name, socklen_t namelen) { 274 | if (!connect_f) init_hook(); 275 | int ret = 0; 276 | 277 | while (1) { 278 | 279 | struct pollfd fds; 280 | fds.fd = fd; 281 | fds.events = POLLOUT | POLLERR | POLLHUP; 282 | poll_inner(&fds, 1, 1); 283 | 284 | ret = connect_f(fd, name, namelen); 285 | if (ret == 0) break; 286 | 287 | if (ret == -1 && (errno == EAGAIN || 288 | errno == EWOULDBLOCK || 289 | errno == EINPROGRESS)) { 290 | continue; 291 | } else { 292 | break; 293 | } 294 | } 295 | 296 | return ret; 297 | } 298 | 299 | ssize_t recv(int fd, void *buf, size_t len, int flags) { 300 | 301 | if (!recv_f) init_hook(); 302 | struct pollfd fds; 303 | fds.fd = fd; 304 | fds.events = POLLIN | POLLERR | POLLHUP; 305 | 306 | poll_inner(&fds, 1, 1); 307 | 308 | int ret = recv_f(fd, buf, len, flags); 309 | if (ret < 0) { 310 | //if (errno == EAGAIN) return ret; 311 | if (errno == ECONNRESET) return -1; 312 | //printf("recv error : %d, ret : %d\n", errno, ret); 313 | 314 | } 315 | return ret; 316 | } 317 | 318 | ssize_t send(int fd, const void *buf, size_t len, int flags) { 319 | if (!send_f) init_hook(); 320 | int sent = 0, ret = 0; 321 | 322 | while (sent < len) { 323 | struct pollfd fds; 324 | fds.fd = fd; 325 | fds.events = POLLOUT | POLLERR | POLLHUP; 326 | 327 | poll_inner(&fds, 1, 1); 328 | ret = send_f(fd, ((char*)buf)+sent, len-sent, flags); 329 | //printf("send --> len : %d\n", ret); 330 | if (ret <= 0) { 331 | break; 332 | } 333 | sent += ret; 334 | } 335 | 336 | if (ret <= 0 && sent == 0) return ret; 337 | 338 | return sent; 339 | } 340 | 341 | int close(int fd) { 342 | if (!close_f) init_hook(); 343 | return close_f(fd); 344 | } 345 | 346 | #endif --------------------------------------------------------------------------------