C++协程(2):使用ucontext实现Linux下的协程池

/*====================================
* file: ws_thread_proto.h 文件
* anchor: wensheng
* date: 2016-05-16
* info: 协程
* log: 1. 2016-05-16 wensheng create
======================================*/


#ifndef _WS_THREAD_PROTO_H_
#define _WS_THREAD_PROTO_H_


#include "../h/types.h"
#include "../ws_output_log_d/ws_log_output.h"
#include "ws_thread_pool.h"
using namespace ws_log_output;
using namespace std;




#ifdef WS_WINDOWS // Windows 使用swatch 奇技淫巧实现


#else // linux 使用ucontext库实现
#include "ucontext.h" 


// linux的线程形式
namespace ws_thread
{




#define WS_THREAD_PROTO_TASK_MAX (9999999) // 同时允许存在的任务数为9999999个
#define WS_DEFAULT_STACK_SZIE (1024*256*2) // 512k (注意:用户函数内的局部变量申请空间总和不能大于等于该值(或者要小于该值4000个字节以下)，否则会栈溢出),或者上调该值
#define WS_THREAD_PROTO_MAX_SIZE (160) // 一个线程允许的最大协程个数 消耗的空间为 WS_THREAD_PROTO_MAX_SIZE*WS_DEFAULT_STACK_SZIE


// 协程状态
enum ThreadPotoState
{
WS_THREAD_POTO_FREE = 0, // 空闲状态
WS_THREAD_POTO_RUNNABLE = 1, // 就绪状态
WS_THREAD_POTO_RUNING = 2, // 执行状态
WS_THREAD_POTO_SUSPEND = 3 // 挂起状态
};




// 任务类型
enum ThreadPotoTaskType
{
WS_THREAD_PROTO_TASK_TYPE_LOGIN = 0, // 正常业务逻辑
WS_THREAD_PROTO_TASK_TYPE_SUB = 1, // 控制逻辑,消减协程数量
};




// 从外部设置线程执行的函数
typedef void* (*WS_THREAD_PROTO_FUNCTION)(void* pObj);


// 协程信息结构体
typedef struct ws_thread_proto_t
{
ucontext_t m_ctx; // 线程上下文对象
WS_THREAD_PROTO_FUNCTION m_func; // 任务函数
void *m_arg; // 任务参数
ThreadPotoState m_state; // 协程状态
char m_stack[WS_DEFAULT_STACK_SZIE]; // 协程独立栈空间,linux线程默认栈空间大小是8M,所以一个linux默认线程允许同时最大协程个数是40个
}THREAD_PROTO_T;


//线程控制信息
typedef struct ws_thread_proto_control_t
{
ucontext_t m_main; // 主线程上下文(调度协程)
int32_t m_runningId; // 当前运行的协程id
vector m_protos; // 协程集合
ws_thread_proto_control_t() :m_runningId(-1) { m_protos.clear(); }
}WS_T_P_CONTROL_T;






// Task(任务)结构体
struct ProtoTask
{
uint32_t m_type; // 任务类型(0表示业务任务,1表示控制人物--消减协程)
WS_THREAD_PROTO_FUNCTION m_fun; // 函数
void* m_para; // 参数
ProtoTask(WS_THREAD_PROTO_FUNCTION afun = NULL, void *pPara = NULL):m_type(WS_THREAD_PROTO_TASK_TYPE_LOGIN), m_fun(afun), m_para(pPara) {}
};




// 协程池(在线程池的基础上实现)
// WS_DEFAULT_STACK_SZIE(注意:用户函数内的局部变量申请空间总和不能大于等于该值(或者要小于该值4000个字节以下)，否则会栈溢出),或者上调该值
class ThreadProtoPool
{
public:
//dwNum 线程池规模
ThreadProtoPool(uint32_t thread_num = 1); // thread_num 线程个数,都允许动态调整
virtual ~ThreadProtoPool();


// 添加协程任务: (注:不承诺销毁pObj，需要用户自己去销毁)
bool addProtoTask(WS_THREAD_PROTO_FUNCTION pFun, void* pObj);


//调整线程规模
uint32_t updateThreadNum(int32_t aNum);


// 所有异步I/O的地方都需要调用该函数
bool yieldIO();


// 结束协程池
bool endPool();


public:
std::map m_thread_info; // 所有线程控制信息<线程id, 控制信息>
std::queue m_TaskQueue; // 任务队列
WS_LOCK_CRITICAL m_TaskLock;   // 任务队列锁
WS_LOCK_CRITICAL m_PoolLock;   // 线程队列锁
bool m_ustopPool; // false表示结束协程池
uint32_t m_stopNum; // 关闭线程(软关闭)(可能不生效)
};


}


#endif

#endif /* _WS_THREAD_PROTO_H_ */

//===================================================分割线============================================//

/*====================================
* file: ws_thread_proto.cpp 文件
* anchor: wensheng
* date: 2016-05-16
* info: 协程
* log: 1. 2016-05-16 wensheng create
======================================*/
#include "ws_thread_proto.h"
#include "../ws_output_log_d/ws_log_output.h"
#include
using namespace ws_log_output;
using namespace ws_thread;






// 协程处理函数
static void* thread_proto_runing(void* a_pObj)
{
// 获取协程管理对象
WS_T_P_CONTROL_T* pObj = (WS_T_P_CONTROL_T*)a_pObj;
if (NULL == pObj) { return NULL; }
if ((pObj->m_runningId < 0) || (pObj->m_runningId >= pObj->m_protos.size()))
{
pObj->m_runningId++;
}
// 获取当前正在执行的协程
THREAD_PROTO_T* pProto = pObj->m_protos[pObj->m_runningId];
if (NULL == pProto) { return NULL; }
if (NULL != pProto->m_func)
{
WS_LOG(Info, "thread_proto_runing !thread_id:%u, proto_num:%u", pthread_self(), pObj->m_runningId);
// 执行协程函数
pProto->m_func(pProto->m_arg);
}
// 执行结束,清理栈空间，重置状态
pProto->m_func = NULL; // 处理函数
pProto->m_arg = NULL; // 参数
pProto->m_state = WS_THREAD_POTO_FREE; // 起始空闲状态
//memset(&(pProto->m_stack), 0, WS_DEFAULT_STACK_SZIE - 2); // 这个地方不能清理，保存的有寄存器信息
++pObj->m_runningId;
// 切回到主协程
swapcontext(&(pProto->m_ctx), &(pObj->m_main));


return NULL;
}






// 主协程去取任务，然后添加到分协程中
static void* main_thread_proto_runing(void* pObj)
{
// 获取对象
ThreadProtoPool* pProtoPool = static_cast(pObj);
assert(NULL != pProtoPool);
// 获取线程id
uint32_t pid = pthread_self();
// 创建协程管理对象
WS_T_P_CONTROL_T* proto_control = new WS_T_P_CONTROL_T();
assert(NULL != proto_control);
// 获取当前的主协程上下文
getcontext(&(proto_control->m_main)); // 获取当前上下文
    // 加入管理池中
pthread_mutex_lock(&pProtoPool->m_PoolLock);
pProtoPool->m_thread_info.insert(make_pair(pid, proto_control));
pthread_mutex_unlock(&pProtoPool->m_PoolLock);


ProtoTask* pJob = NULL;
// 是否结束
while (pProtoPool->m_ustopPool)
{
usleep(10000);
// 检查是不是关闭线程
if (pProtoPool->m_stopNum > 0)
{
// 不再接收新任务
bool can_stop = true;
for (uint32_t index = 0; index < proto_control->m_protos.size(); ++index)
{
if (WS_THREAD_POTO_FREE != proto_control->m_protos[index]->m_state)
{
can_stop = false;
break; // 跳出第一循环执行剩余的协程任务
}
}
if (can_stop)
{
// 删除数量
__sync_sub_and_fetch(&pProtoPool->m_stopNum, 1);
break; // 直接跳出线程循环
}
}
else
{
pJob = NULL;
// 锁, 获取任务
pthread_mutex_lock(&pProtoPool->m_TaskLock);
if (!pProtoPool->m_TaskQueue.empty())
{
pJob = pProtoPool->m_TaskQueue.front(); // 获取任务
pProtoPool->m_TaskQueue.pop();
}
pthread_mutex_unlock(&pProtoPool->m_TaskLock);
if (NULL != pJob)
{
uint32_t poroSize = proto_control->m_protos.size();
// 检查任务类型(若是协程控制任务，就执行控制属性)
switch (pJob->m_type)
{
case WS_THREAD_PROTO_TASK_TYPE_SUB:  //控制逻辑,消减协程数量
{
vector::iterator it = proto_control->m_protos.begin();
for (; it != proto_control->m_protos.end(); ++it)
{
if ((*it)->m_state == WS_THREAD_POTO_FREE)// 空闲状态才能删除
{
THREAD_PROTO_T* pProto = *it;
proto_control->m_protos.erase(it++);
delete[] pProto; // 销毁协程对象
break; // 跳出第一层循环
}
}
break;
}
case WS_THREAD_PROTO_TASK_TYPE_LOGIN: // 正常业务逻辑
{
// 先查找是否有空闲协程
THREAD_PROTO_T* pProto = NULL;
uint32_t index = 0;
for (; index < poroSize; ++index)
{
assert(NULL != proto_control->m_protos[index]);
if (WS_THREAD_POTO_FREE == proto_control->m_protos[index]->m_state)
{
pProto = proto_control->m_protos[index];
break;
}
}
// 如果没找到就创建一个协程
if (NULL == pProto)
{
if (WS_THREAD_PROTO_MAX_SIZE > poroSize) // 协程个数不能太大
{
THREAD_PROTO_T* newProto = new THREAD_PROTO_T(); // 创建一个新协程
newProto->m_func = NULL; // 处理函数
newProto->m_arg = NULL; // 参数
newProto->m_state = WS_THREAD_POTO_FREE; // 起始空闲状态
memset(&newProto->m_stack, 0, WS_DEFAULT_STACK_SZIE); // 清理栈空间
proto_control->m_protos.push_back(newProto);  // 产生了拷贝
pProto = proto_control->m_protos[index];
}
}
// 添加任务
if (NULL != pProto)
{
pProto->m_func = pJob->m_fun; // 处理函数
pProto->m_arg = pJob->m_para; // 参数
pProto->m_state = WS_THREAD_POTO_RUNNABLE; // 就绪状态
delete pJob; // 删除任务原语
}
else // 任务没有被执行，回插入任务队列
{
pthread_mutex_lock(&pProtoPool->m_TaskLock);
pProtoPool->m_TaskQueue.push(pJob); // 插入任务
pthread_mutex_unlock(&pProtoPool->m_TaskLock);
}
break;
}
default:
{
break;
}
}
}

}




// 执行协程
{
if ((proto_control->m_runningId < 0) || (proto_control->m_runningId >= proto_control->m_protos.size()))
{
proto_control->m_runningId = 0;
}


// 获取协程对象
THREAD_PROTO_T *pProto = NULL;
if (proto_control->m_protos.size() > proto_control->m_runningId)
{
pProto = proto_control->m_protos[proto_control->m_runningId];


// 检查执行状态
switch (pProto->m_state)
{
case WS_THREAD_POTO_RUNNABLE: // 就绪状态
{
getcontext(&(pProto->m_ctx)); // 获取当前上下文
pProto->m_ctx.uc_stack.ss_sp = pProto->m_stack; // 指定运行栈空间
pProto->m_ctx.uc_stack.ss_size = WS_DEFAULT_STACK_SZIE - 2; // 指定运行栈空间大小
pProto->m_ctx.uc_stack.ss_flags = 0;
pProto->m_ctx.uc_link = &proto_control->m_main;//设置后继上下文
pProto->m_state = WS_THREAD_POTO_RUNING; // 执行状态
//proto_control->m_runningId = index; // 当前只在执行的线程id
//设置上下文执行函数
makecontext(&(pProto->m_ctx), (void(*)(void))(thread_proto_runing), 1, proto_control);
// 这里不需要break
}
case WS_THREAD_POTO_SUSPEND:  // 挂起状态(就是执行了一部分)
// 跳转到指定协程，保留当前主协程信息
swapcontext(&(proto_control->m_main), &(pProto->m_ctx));
break;
default:
break;
}
}


}


}


// 关闭协程池
proto_control->m_protos.clear();

// 删除记录
pthread_mutex_lock(&pProtoPool->m_PoolLock);
pProtoPool->m_thread_info.erase(pProtoPool->m_thread_info.find(pid));
pthread_mutex_unlock(&pProtoPool->m_PoolLock);
return NULL;
}












//dwNum 线程池规模
//thread_num 线程个数
ThreadProtoPool::ThreadProtoPool(uint32_t thread_num)
{
assert(0 == pthread_mutex_init(&m_TaskLock, NULL)); 
assert(0 == pthread_mutex_init(&m_PoolLock, NULL)); 
m_ustopPool = true;
m_stopNum = 0;
// 创建线程
updateThreadNum(thread_num);
}




ThreadProtoPool::~ThreadProtoPool()
{
endPool();
// 删除剩余任务
for (uint32_t i = 0; i < m_TaskQueue.size(); ++i)
{
ProtoTask* pJob = m_TaskQueue.front(); // 获取任务
m_TaskQueue.pop();
delete pJob;
}
pthread_mutex_destroy(&m_TaskLock); 
pthread_mutex_destroy(&m_PoolLock); 
}


// 添加协程任务
bool ThreadProtoPool::addProtoTask(WS_THREAD_PROTO_FUNCTION pFun, void* pObj)
{
assert(pFun);
if (m_TaskQueue.size() <= WS_THREAD_PROTO_TASK_MAX)
{
ProtoTask* task = new ProtoTask(pFun, pObj);
pthread_mutex_lock(&m_TaskLock);
m_TaskQueue.push(task);
pthread_mutex_unlock(&m_TaskLock);
//WS_LOG(Info, "addProtoTask size:%u", m_TaskQueue.size());
return true;
}


WS_LOG(Error, "addProtoTask have too match! size:%u", m_TaskQueue.size());
return false;
}




//调整线程规模
uint32_t ThreadProtoPool::updateThreadNum(int32_t aNum)
{
if (aNum > 0)
{
// 创建线程
for (uint32_t i = 0; i < aNum; ++i)
{
// 创建线程
Thread* pNewThread = new Thread;
uint32_t pid = (uint32_t)pNewThread->CreateThread((void*)this, &main_thread_proto_runing);
assert(pid != 0);
}
}
else
{
// 不需要锁
m_stopNum = -1 * aNum;
}
}




// 所有异步I/O的地方都需要调用该函数
bool ThreadProtoPool::yieldIO()
{
uint32_t pid = pthread_self();
std::map::iterator it = m_thread_info.find(pid);
if (it != m_thread_info.end())
{
WS_T_P_CONTROL_T* pControl = it->second;
assert(NULL != pControl);
THREAD_PROTO_T *pProto = pControl->m_protos[pControl->m_runningId];
pProto->m_state = WS_THREAD_POTO_SUSPEND; // 挂起
++pControl->m_runningId; // 下一个协程
// 切回到主协程
swapcontext(&(pProto->m_ctx), &(pControl->m_main));
return true;
}


return false;
}




// 结束协程池
bool ThreadProtoPool::endPool()
{
m_ustopPool = false;
// 等待销毁
while (m_thread_info.size() > 0)
{
sleep(1);
}
}

//============================================分割线 测试test.cpp==================================//

// thread_test.cpp : 定义控制台应用程序的入口点。
//




#include "ws_thread_pool.h"
#include "iostream"
#include "ws_thread_proto.h"
#include
using namespace ws_thread;
using namespace std;


struct someThin
{
public:
uint32_t i;
};


ThreadProtoPool one_proto_pool(2);


#ifdef WS_WINDOWS
static unsigned int __stdcall RUN(void* pObj)
{
for (uint32_t i = 0; i < 10; ++i)
{
someThin* pST = ((someThin*)pObj);
pST->i = pST->i + 1;
WS_LOG(Info, "线程值 :%d thread_id:%d", pST->i, GetCurrentThreadId());
}



return NULL;
}
#else


static void* RUN(void* pObj)
{
someThin* pST = ((someThin*)pObj);
int* pone = new int[65535];
//int one[65536] = {0};
memset(pone, 0, 65535*sizeof(int));
pST->i = pST->i + 1;
WS_LOG(Info, "1: num :%d thread_id:%d %d", pST->i, pthread_self(), pone[0]);
one_proto_pool.yieldIO();
pST->i = pST->i + 1;
pone[0]++;
WS_LOG(Info, "2: num :%d thread_id:%d %d", pST->i, pthread_self(), pone[0]);
delete[] pone;
return NULL;
}


#endif




int main()
{
int i = 0;
//ThreadPool one_pool(40);
someThin one;



one.i = 0;
for (uint32_t j = 0; j < 100000; ++j)
{
//one.i = j;
WS_LOG(Info, "yuanshizhi :%d ", j);
//one_pool.addTheadTask(&RUN, &one);
one_proto_pool.addProtoTask(&RUN, &one);
usleep(10000);
}




cin >> i;
    return 0;
}