使用iocp实现简易的线程池

1：定义线程池接口

thread_pool.h

//线程任务回调
typedef void(*thread_jb_cb)(uint64_t job_id, void* udata);

class thread_pool
{
public:
	//开启线程池，最大是16，最小是1
	virtual bool start(uint8_t thread_count) = 0;
	//停止
	virtual void stop() = 0;
	//添加任务,udata:自定义数据，callback：回调，thread_index：线程序号
	virtual uint64_t add_job(void* udata, thread_jb_cb callback, int8_t thread_index = -1) = 0;
};

2：实现线程锁功能

autolock.h

#pragma once
#include <winsock2.h>
#include <windows.h>

class CCriticalSection
{
public:
	void Lock()
	{
		EnterCriticalSection(&m_lock);
	}
	void UnLock()
	{
		LeaveCriticalSection(&m_lock);
	}

	CCriticalSection()
	{
		InitializeCriticalSection(&m_lock);
	}
	~CCriticalSection()
	{
		DeleteCriticalSection(&m_lock);
	}

private:
	CRITICAL_SECTION m_lock;
};


class CAutoLock
{
	CCriticalSection &m_cs;
	CAutoLock();
	const CAutoLock& operator= (const CAutoLock&);
public:
	explicit CAutoLock(CCriticalSection &lock) : m_cs(lock)
	{
		lock.Lock();
	}
	~CAutoLock()
	{
		m_cs.UnLock();
	}
};

3：线程池的实现

3.1：thread_pool_impl.h

#pragma once
#include "thread_pool.h"
#include "autolock.h"

struct iocp_overlapped :public OVERLAPPED
{
	int type;
	void* udata;
	thread_jb_cb callback;
	uint64_t job_id;

	iocp_overlapped(){
		memset(this, 0, sizeof(*this));
	}
};

class thread_pool_impl : public thread_pool{

public:
	thread_pool_impl();
	~thread_pool_impl();

	enum JobType
	{
		Stop,		//没有处理，也就是关闭
		DealMsg,	//处理消息
	};

public:
	//开启线程池，最大是16，最小是1
	virtual bool start(uint8_t thread_count);
	//停止
	virtual void stop();
	//添加任务,
	virtual uint64_t add_job(void* udata, thread_jb_cb callback, int8_t thread_index = -1);

public:
	//处理线程
	void deal_thread(int index);

private:
	uint64_t post_msg_to_iocp(int type, void* udata, thread_jb_cb callback, int8_t thread_index = -1);
	uint64_t get_new_job_id();
	uint8_t get_new_thread_index();

private:
	uint64_t m_job_id;						//jobid
	CCriticalSection m_mutex;				//锁
	uint8_t m_thread_count;					//线程个数
	HANDLE m_iocp_list[MAX_THREAD_COUNT];	//iocp列表
	uint32_t m_thread_id[MAX_THREAD_COUNT];	//线程id
	uint8_t m_thread_index;					//当前的线程序号
	DWORD m_run_thread_count;				//当前运行的线程数量
};

3.2：thread_pool_impl.cpp

#pragma once
#include "thread_pool_impl.h"
#include <process.h>

static thread_pool* s_thread_pool = nullptr;
static thread_pool* get_thread_pool()
{
	if (!s_thread_pool){
		s_thread_pool = new thread_pool_impl();
	}
	return s_thread_pool;
}

static thread_pool* create_thread_pool(){
	return new thread_pool_impl();
}

static void destory_thread_pool(thread_pool* data){
	delete data;
}

struct thread_pool_unit
{
	thread_pool_impl* pool;
	int index;
};

thread_pool_impl::thread_pool_impl()
: m_job_id(0)
, m_thread_index(0)
, m_run_thread_count(0)
, m_thread_count(0)
{
	memset(m_iocp_list, -1, sizeof(m_iocp_list));
}

thread_pool_impl::~thread_pool_impl()
{

}

unsigned __stdcall iocp_thread(void *p){
	thread_pool_unit* unit = (thread_pool_unit*)p;
	unit->pool->deal_thread(unit->index);
	delete unit;
	return 0;
}

bool thread_pool_impl::start(uint8_t thread_count)
{
	if (m_run_thread_count != 0){	//当前有线程在运行
		return false;
	}

	if (thread_count < 1)
		thread_count = 1;
	else if (thread_count > MAX_THREAD_COUNT)
		thread_count = MAX_THREAD_COUNT;
	m_thread_count = thread_count;

	for (int i = 0; i < m_thread_count; ++i)
	{
		m_iocp_list[i] = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, NULL, 1);
	}

	for (int i = 0; i < m_thread_count; ++i)
	{
		thread_pool_unit* unit = new thread_pool_unit();
		unit->pool = this;
		unit->index = i;
		_beginthreadex(NULL, 0, iocp_thread, unit, NULL, &m_thread_id[i]);
	}

	return true;
}

void thread_pool_impl::stop()
{
	for (int i = 0; i < m_thread_count; ++i){
		post_msg_to_iocp(JobType::Stop, nullptr, nullptr, i);
	}
}

uint64_t thread_pool_impl::add_job(void* udata, thread_jb_cb callback, int8_t thread_index /*= -1*/)
{
	return post_msg_to_iocp(JobType::DealMsg, udata, callback, thread_index);;
}

void thread_pool_impl::deal_thread(int index)
{
	CoInitialize(NULL);

	InterlockedIncrement(&m_run_thread_count);

	DWORD nIOBytes = 0;
	ULONG_PTR pKey;
	iocp_overlapped* pOverlap = nullptr;

	while (true)
	{
		BOOL ret = ::GetQueuedCompletionStatus(m_iocp_list[index], &nIOBytes, &pKey, (LPOVERLAPPED*)&pOverlap, INFINITE);
		if (!pOverlap || !ret){
			break;
		}

		if (pOverlap->type <= 0){	//关闭消息
			delete pOverlap;
			pOverlap = nullptr;
			break;
		}

		if (pOverlap->callback){
			pOverlap->callback(pOverlap->job_id, pOverlap->udata);
		}

		delete pOverlap;
		pOverlap = nullptr;
	}

	InterlockedDecrement(&m_run_thread_count);
}

uint64_t thread_pool_impl::post_msg_to_iocp(int type, void* udata, thread_jb_cb callback, int8_t thread_index /*= -1*/)
{
	if (m_thread_count <= 0 || m_run_thread_count <= 0){//线程数量或者正在运行的线程数量都小于0，则关闭
		return 0;
	}

	if (thread_index < 0){
		thread_index = get_new_thread_index();
	}
	if (thread_index >= m_thread_count){
		thread_index = 0;
	}

	uint64_t new_job_id = get_new_job_id();

	iocp_overlapped* overlap = new iocp_overlapped();
	overlap->type = type;
	overlap->udata = udata;
	overlap->callback = callback;
	overlap->job_id = new_job_id;
	PostQueuedCompletionStatus(m_iocp_list[thread_index], 0, (ULONG_PTR)overlap, overlap);

	return new_job_id;
}

uint64_t thread_pool_impl::get_new_job_id()
{
	CAutoLock lock(m_mutex);
	++m_job_id;
	return m_job_id;
}

uint8_t thread_pool_impl::get_new_thread_index()
{
	CAutoLock lock(m_mutex);
	++m_thread_index;
	return m_thread_index;
}

4：线程池的使用

	//创建线程池
	thread_pool* m_thread_pool = new thread_pool_impl();

	//开启
	m_thread_pool->start(16);
	
	//添加任务
	int udata = 1001;
	m_thread_pool->add_job(&udata, [](uint64_t job_id, void* udata){
		int num = *((int*)udata);
		printf("job_id:%lld,num:%d", job_id, num);
	});