added support for rate limiting on http_connection

include/libtorrent/http_connection.hpp

@@ -66,10 +66,19 @@ struct http_connection : boost::enable_shared_from_this<http_connection>, boost:
 		, m_last_receive(boost::posix_time::second_clock::universal_time())
 		, m_bottled(bottled)
 		, m_called(false)
+		, m_rate_limit(0)
+		, m_download_quota(0)
+		, m_limiter_timer_active(false)
+		, m_limiter_timer(ios)
 	{
 		assert(!m_handler.empty());
 	}
 
+	void rate_limit(int limit);
+
+	int rate_limit() const
+	{ return m_rate_limit; }
+
 	std::string sendbuffer;
 
 	void get(std::string const& url, boost::posix_time::time_duration timeout
@@ -89,6 +98,7 @@ private:
 	void on_read(asio::error_code const& e, std::size_t bytes_transferred);
 	static void on_timeout(boost::weak_ptr<http_connection> p
 		, asio::error_code const& e);
+	void on_assign_bandwidth(asio::error_code const& e);
 
 	std::vector<char> m_recvbuffer;
 	tcp::socket m_sock;
@@ -108,6 +118,22 @@ private:
 	bool m_called;
 	std::string m_hostname;
 	std::string m_port;
+
+	// the current download limit, in bytes per second
+	// 0 is unlimited.
+	int m_rate_limit;
+
+	// the number of bytes we are allowed to receive
+	int m_download_quota;
+
+	// only hand out new quota 4 times a second if the
+	// quota is 0. If it isn't 0 wait for it to reach
+	// 0 and continue to hand out quota at that time.
+	bool m_limiter_timer_active;
+
+	// the timer fires every 250 millisecond as long
+	// as all the quota was used.
+	deadline_timer m_limiter_timer;
 };
 
 }

src/bandwidth_manager.cpp

@@ -242,3 +242,4 @@ namespace libtorrent
 	{ assert(false); };
 
 }
+

src/http_connection.cpp

@@ -31,9 +31,13 @@ POSSIBILITY OF SUCH DAMAGE.
 */
 
 #include "http_connection.hpp"
+#include <boost/date_time/posix_time/posix_time_duration.hpp>
+#include <boost/bind.hpp>
 #include <asio/ip/tcp.hpp>
 
 using boost::posix_time::second_clock;
+using boost::posix_time::millisec;
+using boost::bind;
 
 namespace libtorrent
 {
@@ -160,14 +164,33 @@ void http_connection::on_write(asio::error_code const& e)
 
 	std::string().swap(sendbuffer);
 	m_recvbuffer.resize(4096);
+
+	int amount_to_read = m_recvbuffer.size() - m_read_pos;
+	if (m_rate_limit > 0 && amount_to_read > m_download_quota)
+	{
+		amount_to_read = m_download_quota;
+		if (m_download_quota == 0)
+		{
+			if (!m_limiter_timer_active)
+				on_assign_bandwidth(asio::error_code());
+			return;
+		}
+	}
 	m_sock.async_read_some(asio::buffer(&m_recvbuffer[0] + m_read_pos
-		, m_recvbuffer.size() - m_read_pos)
+		, amount_to_read)
-		, bind(&http_connection::on_read, shared_from_this(), _1, _2));
+		, bind(&http_connection::on_read
+		, shared_from_this(), _1, _2));
 }
 
 void http_connection::on_read(asio::error_code const& e
 	, std::size_t bytes_transferred)
 {
+	if (m_rate_limit)
+	{
+		m_download_quota -= bytes_transferred;
+		assert(m_download_quota >= 0);
+	}
+
 	if (e == asio::error::eof)
 	{
 		close();
@@ -228,11 +251,58 @@ void http_connection::on_read(asio::error_code const& e
 		m_handler(asio::error::eof, m_parser, 0, 0);
 		return;
 	}
+	int amount_to_read = m_recvbuffer.size() - m_read_pos;
+	if (m_rate_limit > 0 && amount_to_read > m_download_quota)
+	{
+		amount_to_read = m_download_quota;
+		if (m_download_quota == 0)
+		{
+			if (!m_limiter_timer_active)
+				on_assign_bandwidth(asio::error_code());
+			return;
+		}
+	}
 	m_sock.async_read_some(asio::buffer(&m_recvbuffer[0] + m_read_pos
-		, m_recvbuffer.size() - m_read_pos)
+		, amount_to_read)
 		, bind(&http_connection::on_read
 		, shared_from_this(), _1, _2));
 }
 
+void http_connection::on_assign_bandwidth(asio::error_code const& e)
+{
+	m_limiter_timer_active = false;
+	if (e) return;
+
+	if (m_download_quota > 0) return;
+
+	m_download_quota = m_rate_limit / 4;
+
+	int amount_to_read = m_recvbuffer.size() - m_read_pos;
+	if (amount_to_read > m_download_quota)
+		amount_to_read = m_download_quota;
+
+	m_sock.async_read_some(asio::buffer(&m_recvbuffer[0] + m_read_pos
+		, amount_to_read)
+		, bind(&http_connection::on_read
+		, shared_from_this(), _1, _2));
+
+	m_limiter_timer_active = true;
+	m_limiter_timer.expires_from_now(millisec(250));
+	m_limiter_timer.async_wait(bind(&http_connection::on_assign_bandwidth
+		, shared_from_this(), _1));
+}
+
+void http_connection::rate_limit(int limit)
+{
+	if (!m_limiter_timer_active)
+	{
+		m_limiter_timer_active = true;
+		m_limiter_timer.expires_from_now(millisec(250));
+		m_limiter_timer.async_wait(bind(&http_connection::on_assign_bandwidth
+			, shared_from_this(), _1));
+	}
+	m_rate_limit = limit;
+}
+
 }