socksvr.cpp

一套linux下的C++开发库

/***************************************************************************
                          socksvr.cpp  -  description
                             -------------------
    begin                : Fri Jul 20 2001
    copyright            : (C) 2001 by Mark
    email                : [e-mail]
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/


#include <errno.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <setjmp.h>

#include "exception.h"
#include "socksvr.h"
#include "etc.h"


CNBTcpSvr::CNBTcpSvr(const CInetAddress& stInetAddr, tcport_t tPort, int iBufSize,
		int iMaxClients, int iTimeout, int iWaitTime)
		: CTcpSocket(stInetAddr, tPort)
{
	m_iBufSize = iBufSize;
	m_iMaxClients = iMaxClients;
	m_iTimeout = iTimeout;
	m_iWaitTime = iWaitTime;
	m_bShutdown = false;
	
	SetBlock(false);
	Open();
}


CNBTcpSvr::CNBTcpSvr(const CUnixAddress& stUnixAddr, int iBufSize, int iMaxClients,
		int iTimeout, int iWaitTime)
		: CTcpSocket(stUnixAddr)
{
	m_iBufSize = iBufSize;
	m_iMaxClients = iMaxClients;
	m_iTimeout = iTimeout;
	m_iWaitTime = iWaitTime;
	m_bShutdown = false;
	
	SetBlock(false);
	Open();
}


CNBTcpSvr::~CNBTcpSvr()
{
	CConnList::iterator i = m_stConnList.begin();
	while (i != m_stConnList.end())
		delete *i++;
}


void CNBTcpSvr::Run()
{
	while (true)
	{
		try
		{
			if (m_bShutdown)
				break;
				
			CheckConnTimeout();
			CheckEvent();
			OnCheck();
		}
		catch (CException& e)
		{
			if (!OnException(e))
				break;
		}
	}
}


void CNBTcpSvr::CheckConnTimeout()
{
	if (m_iTimeout == 0)	//no timeout
		return;
		
	time_t tCurrTime = time(NULL);
	CConnList::iterator i = m_stConnList.begin();
	while (i != m_stConnList.end())
	{
		if (tCurrTime - (*i)->m_iLastAccessTime > m_iTimeout)	//timeout
		{
			delete (*i);
			i = m_stConnList.erase(i);
			continue;
		}
		
		i++;
	}
}


void CNBTcpSvr::PrepareSelect(int& iMaxFd, fd_set& rdset, fd_set& wrset)
{
	if (m_iSocket > iMaxFd)
		iMaxFd = m_iSocket;
	FD_SET(m_iSocket, &rdset);
	CConnList::iterator i = m_stConnList.begin();
	while (i != m_stConnList.end())
	{
		if ((*i)->m_pstTcpStream->SocketFd() > iMaxFd)
			iMaxFd = (*i)->m_pstTcpStream->SocketFd();
			
		if ((*i)->m_pstTcpStream->rdbuf()->out_waiting())
			FD_SET((*i)->m_pstTcpStream->SocketFd(), &wrset);
			
		FD_SET((*i)->m_pstTcpStream->SocketFd(), &rdset);
		i++;
	}
}


void CNBTcpSvr::CheckPendingConn(const fd_set& rdset)
{
	if (FD_ISSET(m_iSocket, &rdset))
	{
		CTcpStream* p = new CTcpStream(m_iBufSize);
		if (Accept(p) && m_stConnList.size() < m_iMaxClients)
		{
			CConn* pstConn = new CConn;
			pstConn->m_iLastAccessTime = (int)time(NULL);
			pstConn->m_pstTcpStream = p;
			m_stConnList.push_back(pstConn);
		}
		else
		{
			delete p;
		}
	}
}


void CNBTcpSvr::CheckPendingRead(const fd_set& rdset)
{
	CConnList::iterator i = m_stConnList.begin();
	while (i != m_stConnList.end())
	{
		if (FD_ISSET((*i)->m_pstTcpStream->SocketFd(), &rdset))
		{
			(*i)->m_pstTcpStream->underflow();
			if (!(*i)->m_pstTcpStream->good())
			{
				delete (*i);
				i = m_stConnList.erase(i);
				continue;
			}

			(*i)->m_iLastAccessTime = (int)time(NULL);
			if (!OnRead(*((*i)->m_pstTcpStream)))
			{
				delete (*i);
				i = m_stConnList.erase(i);
				continue;			
			}
		}
		
		i++;
	}	
}


void CNBTcpSvr::CheckPendingWrite(const fd_set& wrset)
{
	CConnList::iterator i = m_stConnList.begin();
	while (i != m_stConnList.end())
	{
		if (FD_ISSET((*i)->m_pstTcpStream->SocketFd(), &wrset))
		{
			(*i)->m_pstTcpStream->flush();
			if (!(*i)->m_pstTcpStream->good())
			{
				delete (*i);
				i = m_stConnList.erase(i);
				continue;
			}

			(*i)->m_iLastAccessTime = (int)time(NULL);
			if (!OnWrite(*((*i)->m_pstTcpStream)))
			{
				delete (*i);
				i = m_stConnList.erase(i);
				continue;			
			}
		}
		
		i++;
	}	
}


void CNBTcpSvr::CheckEvent()
{
	struct timeval stTimeVal;
	stTimeVal.tv_sec = m_iWaitTime;
	stTimeVal.tv_usec = 0;

	fd_set rdset, wrset;
	FD_ZERO(&rdset);
	FD_ZERO(&wrset);
	int iMaxFd = -1;
	PrepareSelect(iMaxFd, rdset, wrset);
	
	int iRet = select(iMaxFd+1, &rdset, &wrset, NULL, &stTimeVal);
	if (iRet == 0)	//timeout
		return;
		
	if (iRet < 0)
	{
		if (errno == EINTR)
			return;
		else
			throw CNBTcpSvrException(strerror(errno), __FILE__, __LINE__);
	}
	
	CheckPendingConn(rdset);
	CheckPendingRead(rdset);
	CheckPendingWrite(wrset);
}



static sigjmp_buf s_ParentProcEnv;
static sigjmp_buf s_ChildProcEnv;

int CPreforkSvr::s_iMaxProcs;
pid_t CPreforkSvr::s_atPId[CPreforkSvr::PROCESS_COUNT_HARD_LIMIT];

static void OnParentProcBeKilled(int iSigNo)
{
	siglongjmp(s_ParentProcEnv, 1);
}

static void OnChildProcEnd(int iSigNo)
{
	pid_t tPId;
	while ((tPId = waitpid(-1, NULL, WNOHANG)) > 0)
	{
		for (int iIndex = 0; iIndex < CPreforkSvr::s_iMaxProcs; iIndex++)
		{
			if (CPreforkSvr::s_atPId[iIndex] == tPId)
				CPreforkSvr::s_atPId[iIndex] = -1;
		}
	}
}

static void OnChildProcBeKilled(int iSigNo)
{
	siglongjmp(s_ChildProcEnv, 1);
}


CPreforkSvr::CPreforkSvr(key_t tKey, const char* sFileLockPath, int iMaxProcs, int iMaxIdleProcs,
		int iMinIdleProcs, int iStartProcs, unsigned int uiCheckProcTimeval)
		: m_stScoreBoard(tKey, PROCESS_COUNT_HARD_LIMIT),
		  m_stLock(sFileLockPath)
{
	if (iMaxProcs > PROCESS_COUNT_HARD_LIMIT)
		throw CPreforkSvrException("\"MaxProcs\" cannot be greater than PROCESS_COUNT_HARD_LIMIT", __FILE__, __LINE__);
		
	if (iMaxIdleProcs < iMinIdleProcs)
		throw CPreforkSvrException("\"MinIdleProcs\" cannot be greater than \"MaxIdleProcs\"", __FILE__, __LINE__);
		
	s_iMaxProcs = iMaxProcs;
	m_iMaxIdleProcs = iMaxIdleProcs;
	m_iMinIdleProcs = iMinIdleProcs;
	m_iStartProcs = iStartProcs;
	m_uiCheckProcTimeval = uiCheckProcTimeval;
	for (int i = 0; i < PROCESS_COUNT_HARD_LIMIT; i++)
		s_atPId[i] = (pid_t)-1;
		
	m_iProcIndex = -1;
}


void CPreforkSvr::Run()
{
	Signal(SIGTERM, OnParentProcBeKilled, false);
	Signal(SIGCHLD, OnChildProcEnd, true);
	if (sigsetjmp(s_ParentProcEnv, 0))
	{
		AllWriteLockW();
		KillAllChildSvrs();
		AllUnlock();
		return;
	}
	
	sigset_t tBlockSet, tOldSet;
	sigemptyset(&tBlockSet);
	sigaddset(&tBlockSet, SIGTERM);
	sigaddset(&tBlockSet, SIGCHLD);	
	
	Sigprocmask(SIG_BLOCK, &tBlockSet, &tOldSet);	
	AllWriteLockW();
	ForkChildSvrs(m_iStartProcs);
	AllUnlock();
	Sigprocmask(SIG_SETMASK, &tOldSet, NULL);
	
	while (true)
	{
		unsigned int uiSleep = m_uiCheckProcTimeval;
		while ((uiSleep = sleep(uiSleep)));

		Sigprocmask(SIG_BLOCK, &tBlockSet, &tOldSet);
		AllWriteLockW();
		int iIdleSvrs = 0;
		int iEmptySvrs = 0;
		for (int iIndex = 0; iIndex < s_iMaxProcs; iIndex++)
		{
			if (m_stScoreBoard[iIndex] == CScoreBoard::S_IDLE)
				iIdleSvrs++;
				
			if (m_stScoreBoard[iIndex] == CScoreBoard::S_EMPTY && s_atPId[iIndex] < 0)
				iEmptySvrs++;
		}
		
		if (iIdleSvrs > m_iMaxIdleProcs)
			KillChildSvrs(iIdleSvrs - m_iMaxIdleProcs);
			
		if (iIdleSvrs < m_iMinIdleProcs && iEmptySvrs > 0)
			ForkChildSvrs(m_iMinIdleProcs - iIdleSvrs < iEmptySvrs ? m_iMinIdleProcs - iIdleSvrs : iEmptySvrs);
		
		AllUnlock();
		Sigprocmask(SIG_SETMASK, &tOldSet, NULL);
	}
}


int CPreforkSvr::ForkChildSvrs(int iSvrs)
{
	int iForkedSvrs = 0;
	try
	{
		for (int iIndex = 0; iIndex < s_iMaxProcs; iIndex++)
		{
			if (iForkedSvrs >= iSvrs)
				break;
			
			if (m_stScoreBoard[iIndex] == CScoreBoard::S_EMPTY && s_atPId[iIndex] < 0)
			{
				pid_t tPId;
				tPId = Fork();
				if (tPId == 0)
				{
					try
					{
						m_iProcIndex = iIndex;
					
						sigset_t tUnblockSet;
						sigemptyset(&tUnblockSet);
						sigaddset(&tUnblockSet, SIGTERM);
						sigaddset(&tUnblockSet, SIGCHLD);
						Sigprocmask(SIG_UNBLOCK, &tUnblockSet, NULL);
				
						Signal(SIGTERM, OnChildProcBeKilled, false);
						if (sigsetjmp(s_ChildProcEnv, 0))
						{
							ChildEnd();
							exit(0);
						}

						SetChildIdle();
						ChildRun();
					}
					catch (CException& e)
					{
						OnChildException(e);				
					}
				
					ChildEnd();
					exit(0);
				}

				iForkedSvrs++;
				s_atPId[iIndex] = tPId;
				m_stScoreBoard[iIndex] = CScoreBoard::S_BUSY;	//because of the lock this would run before child SetIdle()
			}
		}
	}
	catch (CException& e)
	{
		OnNoticeException(e);
	}			
	
	return iForkedSvrs;
}


void CPreforkSvr::KillChildSvrs(int iSvrs)
{
	int iKilledSvrs = 0;
	for (int iIndex = 0; iIndex < s_iMaxProcs; iIndex++)
	{
		if (iKilledSvrs >= iSvrs)
			break;
			
		if (m_stScoreBoard[iIndex] == CScoreBoard::S_IDLE)
		{
			try
			{
				if (s_atPId[iIndex] > 0)
				{
					Kill(s_atPId[iIndex], SIGTERM);
					iKilledSvrs++;
				}
				else
				{
					throw CPreforkSvrException("one child's state is IDLE but pid < 0", __FILE__, __LINE__);
				}
			}
			catch (CException& e)
			{
				OnNoticeException(e);
			}		
		}
	}
}


void CPreforkSvr::KillAllChildSvrs()
{
	for (int iIndex = 0; iIndex < s_iMaxProcs; iIndex++)
	{
		if (m_stScoreBoard[iIndex] != CScoreBoard::S_EMPTY)
		{
			try
			{
				if (s_atPId[iIndex] > 0)
					Kill(s_atPId[iIndex], SIGTERM);
				else
					throw CPreforkSvrException("one child's state is NOT EMPTY but pid < 0", __FILE__, __LINE__);
			}
			catch (CException& e)
			{
				OnNoticeException(e);
			}		
		}
	}
}