/TGame/TTools/MysqlStressTest/MysqlStressTestManager/MysqlStressTestManager/Source/ConnecterManager.cpp

  1//////////////////////////////////////////////////////////////////////////
  2// ConnecterManager.cpp													//
  3//																		//
  4// WYJ																	//
  5// 2010.4																//
  6//																		//
  7//////////////////////////////////////////////////////////////////////////
  8
  9#include "stdafx.h"
 10#include "ConnecterManager.h"
 11#include <assert.h>
 12#include "Connecter.h"
 13#include "MessageHandler.h"
 14
 15#define PORT			5333
 16#define MAX_CONNECTER	2000
 17#define MSGSIZE			1024
 18#define WM_CONNECTERNOTIFY WM_USER + 101
 19
 20DWORD WINAPI ReceiveThread( LPVOID lpParam )
 21{
 22	ConnecterManager* pConnecterManager = (ConnecterManager *)lpParam;
 23
 24	int            i;
 25	fd_set         fdread;
 26	int            ret;
 27	struct timeval tv = {1, 0};
 28	char           szMessage[MSGSIZE];
 29
 30	while (TRUE)
 31	{
 32		FD_ZERO(&fdread);
 33		for (i = 0; i < pConnecterManager->GetConnecterNum(); i++)
 34		{
 35			FD_SET( pConnecterManager->GetConnecter( i )->sock, &fdread );
 36		}
 37
 38		// We only care read event
 39		ret = select( 0, &fdread, NULL, NULL, &tv );
 40
 41		if (ret == 0)
 42		{
 43			// Time expired
 44			continue;
 45		}
 46
 47		for ( i = 0; i < pConnecterManager->GetConnecterNum(); i++ )
 48		{
 49			if ( FD_ISSET( pConnecterManager->GetConnecter( i )->sock, &fdread ) )
 50			{
 51				// A read event happened on g_CliSocketArr[i]
 52				ret = recv( pConnecterManager->GetConnecter( i )->sock, szMessage, MSGSIZE, 0 );
 53				if ( ret == 0 || ( ret == SOCKET_ERROR && WSAGetLastError() == WSAECONNRESET ) )
 54				{
 55					// Client socket closed
 56					pConnecterManager->RemoveConnecter( i );
 57				}
 58				else
 59				{
 60					// We received a message from client
 61					pConnecterManager->DispatchMessage( i, szMessage );
 62				}
 63			}
 64		}
 65	}
 66
 67	return 0;
 68}
 69
 70DWORD WINAPI ListenThread( LPVOID lpParam )
 71{
 72	ConnecterManager* pConnecterManager = (ConnecterManager *)lpParam;
 73
 74	WSADATA		wsaData;
 75	SOCKET		sListen;
 76	SOCKET		sConnecter;
 77	SOCKADDR_IN	connecter;
 78	int         iaddrSize = sizeof(SOCKADDR_IN);
 79	DWORD		dwReceiveThreadId;
 80
 81	if ( WSAStartup( 0x0202, &wsaData ) )
 82	{
 83		
 84	}
 85	sListen = socket( AF_INET, SOCK_STREAM, IPPROTO_TCP );
 86	pConnecterManager->SetListenSocket( sListen );
 87	bind(sListen, (struct sockaddr *)pConnecterManager->GetListenAddr(), sizeof(SOCKADDR_IN));
 88
 89	listen( sListen, 3 );
 90
 91	CreateThread(NULL, 0, ReceiveThread, pConnecterManager, 0, &dwReceiveThreadId); 
 92	pConnecterManager->SetReceiveThreadID( dwReceiveThreadId );
 93
 94	while (TRUE)
 95	{
 96		// Accept a connection
 97		sConnecter = accept( sListen, (struct sockaddr *)&connecter, &iaddrSize );
 98		pConnecterManager->AddConnecter( sConnecter, &connecter );
 99	}
100
101	return 0;
102}
103
104///////////////////////////////////////////////////////////////////////////////////
105
106ConnecterManager::ConnecterManager( void )
107: m_ConnecterNum( 0 ), m_Connecter( NULL ), m_AddrMapNum( 0 ), m_ListenSocket( NULL ),
108  m_ListenThreadID( 0 ), m_ReceiveThreadID( 0 ), m_pMessageHandler( NULL ), m_hWnd( NULL )
109{
110	m_ListenAddr.sin_addr.S_un.S_addr = htonl( INADDR_ANY );
111	m_ListenAddr.sin_family = AF_INET;
112	m_ListenAddr.sin_port = htons( PORT );
113}
114
115ConnecterManager::~ConnecterManager( void )
116{
117	m_Mutex.Lock();
118	for ( AddrMapIter iter = m_AddrMap.begin(); iter != m_AddrMap.end(); iter++ )
119	{
120		for ( ThreadMapIter i = iter->second.begin(); i != iter->second.end(); i++ )
121		{
122			if ( i->second->status != e_Close )
123			{
124				closesocket( i->second->sock );
125			}
126			delete i->second;
127		}
128	}
129	delete [] m_Connecter;
130	m_Mutex.Unlock();
131
132	if ( m_pMessageHandler )
133	{
134		delete m_pMessageHandler;
135	}
136}
137bool ConnecterManager::Initialize( void )
138{
139	assert( m_pMessageHandler == NULL );
140	m_pMessageHandler = new MessageHandler( this );
141
142	assert( m_Connecter == NULL );
143	m_Connecter = new Connecter*[MAX_CONNECTER];
144	if ( !m_Connecter )
145	{
146		char buf[1024];
147		sprintf_s( buf, 1024, "Initialize connecter manager failed." );
148		MessageBox( NULL, buf, 0, 0 );
149		return false;
150	}
151
152	CreateThread( NULL, 0, ListenThread, this, 0, &m_ListenThreadID );
153	return true;
154}
155
156void ConnecterManager::SetListenSocket( SOCKET sock )
157{
158	m_ListenSocket = sock;
159}
160
161const SOCKADDR_IN* ConnecterManager::GetListenAddr( void )
162{
163	return &m_ListenAddr;
164}
165
166void ConnecterManager::AddConnecter( SOCKET sock, const SOCKADDR_IN* addr )
167{
168	if ( m_ConnecterNum >= MAX_CONNECTER )
169	{
170		char buf[1024];
171		sprintf_s( buf, 1024, "Add a new connecter failed for over the max connecter number %d", MAX_CONNECTER );
172		MessageBox( NULL, buf, 0, 0 );
173	}
174
175	m_Mutex.Lock();
176	m_Connecter[m_ConnecterNum] = new Connecter;
177	m_Connecter[m_ConnecterNum]->sock = sock;
178	memcpy( m_Connecter[m_ConnecterNum]->ip, inet_ntoa( addr->sin_addr ), 16 );
179	m_Connecter[m_ConnecterNum]->thread_id = 0;
180	m_Connecter[m_ConnecterNum]->status = e_None;
181	m_ConnecterNum++;
182	m_Mutex.Unlock();
183}
184
185void ConnecterManager::SetReceiveThreadID( DWORD threadID )
186{
187	m_ReceiveThreadID = threadID;
188}
189
190int	ConnecterManager::GetConnecterNum( void )
191{
192	return m_ConnecterNum;
193}
194
195Connecter* ConnecterManager::GetConnecter( int index )
196{
197	return index >= m_ConnecterNum ? NULL : m_Connecter[index];
198}
199
200bool ConnecterManager::RemoveConnecter( int index )
201{
202	if ( index >= m_ConnecterNum )
203	{
204		char buf[1024];
205		sprintf_s( buf, 1024, "Remove a not exist connect index %d", index );
206		MessageBox( NULL, buf, 0, 0 );
207	}
208	
209	m_Mutex.Lock();
210	if ( m_Connecter[index]->status == e_Close )
211	{
212		char buf[1024];
213		sprintf_s( buf, 1024, "Remove a closed connect index %d", index );
214		MessageBox( NULL, buf, 0, 0 );
215		m_Mutex.Unlock();
216		return false;
217	}
218	closesocket( m_Connecter[index]->sock );
219	m_Connecter[index]->status = e_Close;
220	Notify( m_Connecter[index] );
221	m_Connecter[index] = m_Connecter[--m_ConnecterNum];
222	m_Connecter[m_ConnecterNum] = NULL;
223	m_Mutex.Unlock();
224
225	return true;
226}
227
228void ConnecterManager::DispatchMessage( int index, const char * msg )
229{
230	m_pMessageHandler->HandleMessage( index, msg );
231}
232
233bool ConnecterManager::InitConnecter( int index, DWORD threadID )
234{
235	if ( index >= m_ConnecterNum )
236	{
237		char buf[1024];
238		sprintf_s( buf, 1024, "Initialize a not exist connect index %d", index );
239		MessageBox( NULL, buf, 0, 0 );
240	} 
241
242	m_Mutex.Lock();
243	m_Connecter[index]->thread_id = threadID;
244	m_Connecter[index]->status = e_Ready;
245	AddrMapIter iter = m_AddrMap.find( inet_addr( m_Connecter[index]->ip ) );
246	if ( iter == m_AddrMap.end() )
247	{
248		m_AddrMap.insert( AddrMap::value_type( inet_addr( m_Connecter[index]->ip ), ThreadMap() ) );
249		iter = m_AddrMap.find( inet_addr( m_Connecter[index]->ip ) );
250	}
251	if ( iter != m_AddrMap.end() )
252	{
253		ThreadMapIter i = iter->second.find( threadID );
254		if ( i != iter->second.end() )
255		{
256			delete i->second;
257			i->second = m_Connecter[index];
258		}
259		else
260		{
261			iter->second.insert( ThreadMap::value_type( threadID, m_Connecter[index] ) );
262		}
263		Notify( m_Connecter[index] );
264	}
265	m_Mutex.Unlock();
266
267	return true;
268}
269
270void ConnecterManager::SetHwnd( HWND hWnd )
271{
272	m_hWnd = hWnd;
273}
274
275Connecter* ConnecterManager::GetConnecter( unsigned long addr, DWORD threadID )
276{
277	AddrMapIter iter = m_AddrMap.find( addr );
278	if ( iter != m_AddrMap.end() )
279	{
280		ThreadMapIter i = iter->second.find( threadID );
281		if ( i != iter->second.end() )
282			return i->second;
283	}
284	return NULL;
285}
286
287void ConnecterManager::Notify( Connecter* pConnecter )
288{
289	SendMessage( m_hWnd, WM_CONNECTERNOTIFY, (WPARAM)0, (LPARAM)pConnecter );
290}