peermanager.cpp

这是一个嵌入式linux系统下的BT下载工具包

/***************************************************************************
 *   Copyright (C) 2005-2006 Gao Xianchao                                  *
 *                 2007 Gao Xianchao gnap_an linux_lyb ahlongxp            *
 *                                                                         *
 *   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.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

/*
 * Author:	gxc
 * Create data:	2005-10-15 10:09
 */
 
#include "PeerManager.h"
#include "log.h"
#include "utils.h"
#include "PeerLink.h"

CPeerManager::CPeerManager()
: _task(NULL)
{
}

CPeerManager::~CPeerManager()
{
}

void CPeerManager::setBTTask(IBTTask* task)
{
	_task = task;
}

IBTTask* CPeerManager::getBTTask()
{
	return _task;
}

bool CPeerManager::start()
{
	LOG_INFO("PeerManager starting");
	_connectTimerID = _task->getSocketReactor()->addTimer(this, 2000, false);
	LOG_DEBUG("_connectTimerID = "<<_connectTimerID);
	
	_chokeTimerID = _task->getSocketReactor()->addTimer(this, 20000, false);
	
	LOG_INFO("PeerManager started");	
	return true;
}

void CPeerManager::stop()
{
	LOG_INFO("PeerManager stopping");
	
	_task->getSocketReactor()->removeTimer(_connectTimerID);
	_task->getSocketReactor()->removeTimer(_chokeTimerID);

	//free PeerLinks
	TPeerInfoMap::iterator iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end(); ++iter)
	{
		if(iter->second.peerLink != NULL)
		{
			iter->second.peerLink->closeLink();
			delete iter->second.peerLink;
			iter->second.peerLink = NULL;
		}
	}	
	
	iter = _connectingPeerList.begin();
	for(;iter != _connectingPeerList.end(); ++iter)
	{
		if(iter->second.peerLink != NULL)
		{
			iter->second.peerLink->closeLink();
			delete iter->second.peerLink;
			iter->second.peerLink = NULL;
		}
	}			
	
	_unusedPeerList.clear();
	_connectingPeerList.clear();
	_connectedPeerList.clear();
	_banedPeerList.clear();
	
	LOG_INFO("PeerManager stopped");
}

bool CPeerManager::addAcceptedPeer(int handle, const char* ip, unsigned short port)
{
	if(_connectedPeerList.size() >= _task->getPeerLinkMax())
	{
		return false;
	}
	
	std::string id = getPeerStr(ip, port);
	
	TPeerInfo peerInfo;
	peerInfo.linkID = id;
	peerInfo.ip = ip;
	peerInfo.port = port;
	peerInfo.connectFailedCount = 0;
	peerInfo.peerLink = new CPeerLink();
	peerInfo.peerLink->attach(handle, ip, port,  this);		
	
	_connectedPeerList[id]=peerInfo;
	
	return true;
}

void CPeerManager::addPeerInfoWithoutID(const char* ip, unsigned short port)
{
	std::string id = getPeerStr(ip, port);
	
//	LOG_DEBUG("addPeerInfoWithoutID, id="<<id);
	
	if(peerInUnusedList(id)
		|| peerInConnectingList(id)
		|| peerInConnectedList(id)
		|| peerInBanedList(id))
	{
		return;
	}	
	
	TPeerInfo peerInfo;
	peerInfo.linkID = id;
	peerInfo.ip = ip;
	peerInfo.port = port;
	peerInfo.connectFailedCount = 0;
	peerInfo.peerLink = NULL;
	
	_unusedPeerListMutex.lock();
	_unusedPeerList[id]=peerInfo;
	_unusedPeerListMutex.unlock();
}

void CPeerManager::onTimer(unsigned int id)
{
	//LOG_DEBUG("CPeerManager::onTimer connected="<<_connectedPeerList.size());
	if(id == _connectTimerID)
	{
		checkPeerConnection();
	}
	
	if(id == _chokeTimerID)
	{
		checkPeerChoke();
	}
}

void CPeerManager::countPeersSpeed()
{
	TPeerInfoMap::iterator iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end(); ++iter)
	{
		IPeerLink* peerLink = iter->second.peerLink;
		if(peerLink != NULL)
		{
			peerLink->countSpeed();
		}			
	}
}

void CPeerManager::checkPeerChoke()
{
	countPeersSpeed();

/*
{	
	LOG_DEBUG("");
	LOG_DEBUG("");
	//log
	TPeerInfoMap::iterator iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end(); ++iter)
	{
		IPeerLink* peerLink = iter->second.peerLink;
		if(peerLink == NULL)
		{
			continue;
		}
		
		int i = 0;
		if(peerLink->peerInterested())
		{
			i = 1;
		}
		
		int c = 0;
		if(peerLink->peerChoked())
		{
			c = 1;
		}
		
		LOG_DEBUG("ip="<<peerLink->getIP()<<"     peerInterest="<<i<<"      peerIsChoked="<<c<< "   downloadSpeed="<<peerLink->getDownloadSpeed());
	}	
	LOG_DEBUG("");
	LOG_DEBUG("");
}	
*/
	//达到最大上传数
	unsigned int downloaderCount = 0;
	TPeerInfoMap::iterator iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end(); ++iter)
	{
		IPeerLink* peerLink = iter->second.peerLink;
		if(peerLink == NULL)
		{
			continue;
		}
		
		if(!peerLink->peerInterested())
		{
			peerLink->chokePeer(true);
		}
		if(!peerLink->peerChoked())
		{
			downloaderCount++;
		}
	}
	
	for(;downloaderCount<_task->getUploadPeerLinkMax();)
	{
		iter = _connectedPeerList.begin();
		TPeerInfoMap::iterator iter2 = _connectedPeerList.end();
		for(;iter != _connectedPeerList.end(); ++iter)
		{
			IPeerLink* peerLink = iter->second.peerLink;
			if(peerLink == NULL)
			{
				continue;
			}
			
			if(peerLink->peerChoked()
				&& peerLink->peerInterested())
			{
				if(iter2 == _connectedPeerList.end())
				{
					iter2 = iter;
				}
				else if(peerLink->getDownloadSpeed()>iter2->second.peerLink->getDownloadSpeed())
				{
					iter2 = iter;
				}
			}
		}
		
		if(iter2 != _connectedPeerList.end())
		{
			iter2->second.peerLink->chokePeer(false);
			downloaderCount++;	
		}
		else
		{
			break;
		}
	}
	
	if(downloaderCount >= _task->getUploadPeerLinkMax())
	{
		//得到当前上传peer中下载速度最慢的一个
		iter = _connectedPeerList.begin();
		TPeerInfoMap::iterator worstPeerIter = _connectedPeerList.end();
		for(;iter != _connectedPeerList.end(); ++iter)
		{
			IPeerLink* peerLink = iter->second.peerLink;
			if(peerLink == NULL)
			{
				continue;
			}
			
			if(!peerLink->peerChoked()
				&& peerLink->peerInterested())
			{
				if(worstPeerIter == _connectedPeerList.end())
				{
					worstPeerIter = iter;
				}
				else if(peerLink->getDownloadSpeed()<worstPeerIter->second.peerLink->getDownloadSpeed())
				{
					worstPeerIter = iter;
				}
			}
		}	
		
		//得到非上传peer中下载速度最快的一个				
		iter = _connectedPeerList.begin();
		TPeerInfoMap::iterator bestPeerIter = _connectedPeerList.end();
		for(;iter != _connectedPeerList.end(); ++iter)
		{
			IPeerLink* peerLink = iter->second.peerLink;
			if(peerLink == NULL)
			{
				continue;
			}
			
			if(peerLink->peerChoked()
				&& peerLink->peerInterested())
			{
				if(bestPeerIter == _connectedPeerList.end())
				{
					bestPeerIter = iter;
				}
				else if( peerLink->getDownloadSpeed()>bestPeerIter->second.peerLink->getDownloadSpeed())
				{
					bestPeerIter = iter;
				}
			}
		}
		
		if(worstPeerIter != _connectedPeerList.end()
			&& bestPeerIter != _connectedPeerList.end()
			&& worstPeerIter->second.peerLink->getDownloadSpeed()<bestPeerIter->second.peerLink->getDownloadSpeed())
		{
			//阻塞下载最慢的downloader
			worstPeerIter->second.peerLink->chokePeer(true);
			//给下载最快的非downloader机会
			bestPeerIter->second.peerLink->chokePeer(false);
		}
	}
}

void CPeerManager::checkPeerConnection()
{	
	TPeerInfoMap::iterator iter = _connectingPeerList.begin();
	for(; iter!=_connectingPeerList.end();)
	{
		if(iter->second.peerLink == NULL)
		{
			_connectingPeerList.erase(iter++);
			continue;
		}
		
		// move new connection to _connectedPeerList
		if(iter->second.peerLink->getState() == PS_ESTABLISHED)
		{
			iter->second.connectFailedCount = 0;
			_connectedPeerList[iter->second.linkID] = iter->second;
			_connectingPeerList.erase(iter++);		
			
			continue;
		}
		
		// move failed connection to _unusedPeerList & _banedPeerList		
		if(iter->second.peerLink->getState() == PS_CONNECTFAILED)
		{
			iter->second.connectFailedCount++;
			iter->second.peerLink->closeLink();
			delete iter->second.peerLink;
			iter->second.peerLink = NULL;			
			
			if(iter->second.connectFailedCount >= 3)
			{	
				_banedPeerList[iter->second.linkID] = iter->second;
			}
			else
			{
				_unusedPeerListMutex.lock();
				_unusedPeerList[iter->second.linkID] = iter->second;
				_unusedPeerListMutex.unlock();
			}
			
			_connectingPeerList.erase(iter++);
			continue;
		}	
		
		// move closed connection to _unusedPeerList	
		if(iter->second.peerLink->getState() == PS_CLOSED)
		{
			iter->second.peerLink->closeLink();
			delete iter->second.peerLink;
			iter->second.peerLink = NULL;
			
			_unusedPeerListMutex.lock();
			_unusedPeerList[iter->second.linkID] = iter->second;
			_unusedPeerListMutex.unlock();
			
			_connectingPeerList.erase(iter++);
			
			continue;
		}	
		
		++iter;	
	}
	
	//move closed connection to _banedPeerList;
	iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end();)
	{
		if(iter->second.peerLink->getState() == PS_CLOSED)
		{			
			bool accepted = iter->second.peerLink->isAccepted();
			iter->second.peerLink->closeLink();
			delete iter->second.peerLink;
			iter->second.peerLink = NULL;
			
			if(!accepted)
			{
				_unusedPeerListMutex.lock();
				_unusedPeerList[iter->second.linkID] = iter->second;
				_unusedPeerListMutex.unlock();
			}
			_connectedPeerList.erase(iter++);
			continue;
		}
		
		++iter;
	}
	
	_unusedPeerListMutex.lock();
	// launch new connection
	if(_connectedPeerList.size() < _task->getPeerLinkMax())
	{		
		for(;;)
		{
			if(_unusedPeerList.size() == 0)
			{
				break;
			}
			
			if(_connectedPeerList.size() >= _task->getPeerLinkMax())
			{
				break;
			}
			
			if(_connectingPeerList.size() < _task->getConnectingPeerLinkMax())
			{
				TPeerInfo peerInfo = _unusedPeerList.begin()->second;
				peerInfo.peerLink = new CPeerLink();
				peerInfo.peerLink->setPeerManager(this);
				//LOG_DEBUG("connect, "<<peerInfo.linkID);
				peerInfo.peerLink->connect(peerInfo.ip.c_str(), peerInfo.port);			
				
				_unusedPeerList.erase(_unusedPeerList.begin());
				_connectingPeerList[peerInfo.linkID]=peerInfo;
			}
			else
			{
				break;
			}
		}
	}
	_unusedPeerListMutex.unlock();

	// close some useless peers
	if(_connectedPeerList.size() > _task->getPeerLinkMax()-10)
	{
		int i=0;
		iter = _connectedPeerList.begin();
		for(;iter != _connectedPeerList.end();)
		{
			if(iter->second.peerLink->checkNeedClose())
			{
				iter->second.peerLink->closeLink();
				i++;
			}
			if(i>=5)
			{
				break;
			}
			++iter;
		}
	}	
}

bool CPeerManager::peerInUnusedList(std::string id)
{
	_unusedPeerListMutex.lock();
	bool result =  (_unusedPeerList.find(id) != _unusedPeerList.end());
	_unusedPeerListMutex.unlock();
	
	return result;
}

bool CPeerManager::peerInConnectingList(std::string id)
{
	return (_connectingPeerList.find(id) != _connectingPeerList.end());
}

bool CPeerManager::peerInConnectedList(std::string id)
{
	return (_connectedPeerList.find(id) != _connectedPeerList.end());
}

bool CPeerManager::peerInBanedList(std::string id)
{;
	return (_unusedPeerList.find(id) != _unusedPeerList.end());
}

void CPeerManager::broadcastHave(unsigned int pieceIndex)
{
	TPeerInfoMap::iterator iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end(); ++iter)
	{
		if(iter->second.peerLink != NULL)
		{
			iter->second.peerLink->notifyHavePiece(pieceIndex);
		}
	}	
}

void CPeerManager::cancelPieceRequest(unsigned int pieceIndex)
{
	TPeerInfoMap::iterator iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end(); ++iter)
	{
		if(iter->second.peerLink != NULL)
		{
			iter->second.peerLink->cancelPieceRequest(pieceIndex);
		}
	}		
}

unsigned int CPeerManager::getConnectedPeerCount()
{
	return _connectedPeerList.size();
}

void CPeerManager::onDownloadComplete()
{
	TPeerInfoMap::iterator iter = _connectedPeerList.begin();
	for(;iter != _connectedPeerList.end(); ++iter)
	{
		if(iter->second.peerLink != NULL)
		{
			iter->second.peerLink->onDownloadComplete();
		}
	}		
}