* Enable custom throttles, both per-download or per-IP. See http://libtorrent.rakshasa.no/ticket/20 for info and instructions. Patch by Josef Drexler.

git-svn-id: svn://rakshasa.no/libtorrent/trunk/rtorrent@1091 e378c898-3ddf-0310-93e7-cc216c733640

src/core/Makefile.am

@@ -32,6 +32,7 @@ libsub_core_a_SOURCES = \
 	poll_manager_kqueue.h \
 	poll_manager_select.cc \
 	poll_manager_select.h \
+	range_map.h \
 	view.cc \
 	view.h \
 	view_manager.cc \

src/core/dht_manager.cc

@@ -95,6 +95,10 @@ DhtManager::start_dht() {
   if (!torrent::dht_manager()->is_valid() || torrent::dht_manager()->is_active())
     return;
 
+  torrent::ThrottlePair throttles = control->core()->get_throttle(m_throttleName);
+  torrent::dht_manager()->set_upload_throttle(throttles.first);
+  torrent::dht_manager()->set_download_throttle(throttles.second);
+
   int port = rpc::call_command_value("get_dht_port");
   if (port <= 0)
     return;
@@ -283,6 +287,7 @@ DhtManager::dht_statistics() {
 
   dhtStats.insert_key("dht",              dht_settings[m_start]);
   dhtStats.insert_key("active",           torrent::dht_manager()->is_active());
+  dhtStats.insert_key("throttle",         m_throttleName);
 
   if (torrent::dht_manager()->is_active()) {
     torrent::DhtManager::statistics_type stats = torrent::dht_manager()->get_statistics();
@@ -303,4 +308,12 @@ DhtManager::dht_statistics() {
   return dhtStats;
 }
 
+void
+DhtManager::set_throttle_name(const std::string& throttleName) {
+  if (torrent::dht_manager()->is_active())
+    throw torrent::input_error("Cannot set DHT throttle while active.");
+
+  m_throttleName = throttleName;
+}
+
 }

src/core/dht_manager.h

@@ -58,6 +58,9 @@ public:
 
   void                set_start(const std::string& arg);
 
+  void                set_throttle_name(const std::string& throttleName);
+  const std::string&  throttle_name() const        { return m_throttleName; }
+
 private:
   static const int    dht_disable = 0;
   static const int    dht_off     = 1;
@@ -82,6 +85,7 @@ private:
   bool                m_warned;
 
   int                 m_start;
+  std::string         m_throttleName;
 };
 
 }

src/core/download.cc

@@ -155,6 +155,18 @@ Download::receive_chunk_failed(__UNUSED uint32_t idx) {
   m_chunksFailed++;
 }
 
+void
+Download::set_throttle_name(const std::string& throttleName) {
+  if (m_download.is_active())
+    throw torrent::input_error("Cannot set throttle on active download.");
+
+  torrent::ThrottlePair throttles = control->core()->get_throttle(throttleName);
+  m_download.set_upload_throttle(throttles.first);
+  m_download.set_download_throttle(throttles.second);
+
+  m_download.bencode()->get_key("rtorrent").insert_key("throttle_name", throttleName);
+}
+
 void
 Download::set_root_directory(const std::string& path) {
   torrent::FileList* fileList = m_download.file_list();

src/core/download.h

@@ -117,6 +117,8 @@ public:
 
   void                set_root_directory(const std::string& path);
 
+  void                set_throttle_name(const std::string& throttleName);
+
   bool                operator == (const std::string& str) const;
 
   float               distributed_copies() const;

src/core/download_factory.cc

@@ -336,6 +336,10 @@ DownloadFactory::initialize_rtorrent(Download* download, torrent::Object* rtorre
     download->download()->set_chunks_done(std::min<uint32_t>(rtorrent->get_key_value("chunks_done"),
                                                              download->download()->file_list()->size_chunks()));
 
+  download->set_throttle_name(rtorrent->has_key_string("throttle_name")
+                              ? rtorrent->get_key_string("throttle_name")
+                              : std::string());
+
   rtorrent->insert_preserve_copy("ignore_commands", (int64_t)0);
   rtorrent->insert_preserve_copy("views", torrent::Object::create_list());
 

src/core/manager.cc

@@ -54,6 +54,7 @@
 #include <torrent/exceptions.h>
 #include <torrent/resume.h>
 #include <torrent/tracker_list.h>
+#include <torrent/throttle.h>
 
 #include "rpc/parse_commands.h"
 #include "utils/directory.h"
@@ -166,9 +167,14 @@ Manager::Manager() :
   m_fileStatusCache = new FileStatusCache();
   m_httpQueue       = new HttpQueue();
   m_httpStack       = new CurlStack();
+
+  torrent::Throttle* unthrottled = torrent::Throttle::create_throttle();
+  unthrottled->set_max_rate(0);
+  m_throttles["NULL"] = std::make_pair(unthrottled, unthrottled);
 }
 
 Manager::~Manager() {
+  torrent::Throttle::destroy_throttle(m_throttles["NULL"].first);
   delete m_downloadList;
 
   delete m_downloadStore;
@@ -185,6 +191,31 @@ Manager::set_hashing_view(View* v) {
   v->signal_changed().connect(sigc::mem_fun(this, &Manager::receive_hashing_changed));
 }
 
+torrent::ThrottlePair
+Manager::get_throttle(const std::string& name) {
+  ThrottleMap::const_iterator itr = m_throttles.find(name);
+  torrent::ThrottlePair throttles = (itr == m_throttles.end() ? torrent::ThrottlePair(NULL, NULL) : itr->second);
+
+  if (throttles.first == NULL)
+    throttles.first = torrent::up_throttle_global();
+
+  if (throttles.second == NULL)
+    throttles.second = torrent::down_throttle_global();
+
+  return throttles;
+}
+
+void
+Manager::set_address_throttle(uint32_t begin, uint32_t end, torrent::ThrottlePair throttles) {
+  m_addressThrottles.set_merge(begin, end, throttles);
+  torrent::connection_manager()->set_address_throttle(sigc::mem_fun(control->core(), &core::Manager::get_address_throttle));
+}
+
+torrent::ThrottlePair
+Manager::get_address_throttle(const sockaddr* addr) {
+  return m_addressThrottles.get(rak::socket_address::cast_from(addr)->sa_inet()->address_h(), torrent::ThrottlePair(NULL, NULL));
+}
+
 void
 Manager::initialize_first() {
   const char* poll = getenv("RTORRENT_POLL");

src/core/manager.h

@@ -40,8 +40,11 @@
 #include <iosfwd>
 #include <vector>
 
+#include <torrent/connection_manager.h>
+
 #include "download_list.h"
 #include "poll_manager.h"
+#include "range_map.h"
 #include "log.h"
 
 namespace torrent {
@@ -57,6 +60,8 @@ namespace core {
 class DownloadStore;
 class HttpQueue;
 
+typedef std::map<std::string, torrent::ThrottlePair> ThrottleMap;
+
 class View;
 
 class Manager {
@@ -84,6 +89,13 @@ public:
   Log&                get_log_important()                 { return m_logImportant; }
   Log&                get_log_complete()                  { return m_logComplete; }
 
+  ThrottleMap&          throttles()                       { return m_throttles; }
+  torrent::ThrottlePair get_throttle(const std::string& name);
+
+  // Use custom throttle for the given range of IP addresses.
+  void                  set_address_throttle(uint32_t begin, uint32_t end, torrent::ThrottlePair throttles);
+  torrent::ThrottlePair get_address_throttle(const sockaddr* addr);
+
   // Really should find a more descriptive name.
   void                initialize_first();
   void                initialize_second();
@@ -120,6 +132,8 @@ public:
   void                try_create_download_expand(const std::string& uri, int flags, command_list_type commands = command_list_type());
 
 private:
+  typedef RangeMap<uint32_t, torrent::ThrottlePair> AddressThrottleMap;
+
   void                create_http(const std::string& uri);
   void                create_final(std::istream* s);
 
@@ -136,6 +150,9 @@ private:
 
   View*               m_hashingView;
 
+  ThrottleMap         m_throttles;
+  AddressThrottleMap  m_addressThrottles;
+
   PollManager*        m_pollManager;
   Log                 m_logImportant;
   Log                 m_logComplete;

src/core/range_map.h

@@ -0,0 +1,194 @@
+// rTorrent - BitTorrent client
+// Copyright (C) 2005-2008, Jari Sundell
+//
+// 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
+//
+// In addition, as a special exception, the copyright holders give
+// permission to link the code of portions of this program with the
+// OpenSSL library under certain conditions as described in each
+// individual source file, and distribute linked combinations
+// including the two.
+//
+// You must obey the GNU General Public License in all respects for
+// all of the code used other than OpenSSL.  If you modify file(s)
+// with this exception, you may extend this exception to your version
+// of the file(s), but you are not obligated to do so.  If you do not
+// wish to do so, delete this exception statement from your version.
+// If you delete this exception statement from all source files in the
+// program, then also delete it here.
+//
+// Contact:  Jari Sundell <jaris@ifi.uio.no>
+//
+//           Skomakerveien 33
+//           3185 Skoppum, NORWAY
+
+#ifndef RTORRENT_CORE_RANGE_MAP_H
+#define RTORRENT_CORE_RANGE_MAP_H
+
+#include <map>
+#include <stdexcept>
+
+namespace core {
+
+// Associate values with a range of keys, and retrieve for any key in the range.
+
+// The template arguments have the same semantics as std::map.
+// Exception: if set_merge is used, the value type must have a defined operator ==.
+template<typename Key, typename T, typename Compare = std::less<Key>,
+         typename Alloc = std::allocator<std::pair<const Key, T> > >
+class RangeMap : private std::map<Key, std::pair<Key, T>, Compare,
+                                 typename Alloc::template rebind<std::pair<const Key, std::pair<Key, T> > >::other> {
+
+  typedef std::map<Key, std::pair<Key, T>, Compare,
+                   typename Alloc::template rebind<std::pair<const Key, std::pair<Key, T> > >::other> base_type;
+
+public:
+  RangeMap() {}
+  RangeMap(const Compare& c) : base_type(c) {}
+
+  using base_type::const_iterator;
+  using base_type::const_reverse_iterator;
+
+  using base_type::clear;
+  using base_type::swap;
+
+  using base_type::size;
+  using base_type::empty;
+
+  using base_type::begin;
+  using base_type::end;
+  using base_type::rbegin;
+  using base_type::rend;
+
+  using base_type::key_comp;
+  using base_type::value_comp;
+
+  // Store a value for the range [begin, end). Returns iterator for the range.
+  typename RangeMap::const_iterator   set_range(const Key& begin, const Key& end, const T& value);
+
+  // Same, but merge adjacent ranges having the same value. Returns iterator for the merged range.
+  typename RangeMap::const_iterator   set_merge(Key begin, const Key& end, const T& value);
+
+  // Find range containing the given key, or end().
+  typename RangeMap::const_iterator   find(const Key& key) const;
+
+  // Retrieve value for key in a range, throw std::out_of_range if range does not exist.
+  const T&                            get(const Key& key) const;
+
+  // Retrieve value for key in a range, return def if range does not exist.
+  T                                   get(const Key& key, T def) const;
+
+private:
+  typename RangeMap::iterator         crop_overlap(const Key& begin, const Key& end);
+};
+
+// Semantics of an entry:
+// .first          End of range (exclusive), map key.
+// .second.first   Beginning of range.
+// .second.second  Value.
+
+template<typename Key, typename T, typename C, typename A>
+inline typename RangeMap<Key,T,C,A>::iterator
+RangeMap<Key,T,C,A>::crop_overlap(const Key& _begin, const Key& _end) {
+  typename RangeMap::iterator itr = upper_bound(_begin);
+
+  while (itr != end() && key_comp()(itr->second.first, _end)) {
+    // There's a subrange before the new begin: need new entry (new range end means new key).
+    if (key_comp()(itr->second.first, _begin))
+      insert(itr, typename RangeMap::value_type(_begin, itr->second));
+
+    // Old end is within our range: erase entry.
+    if (!key_comp()(_end, itr->first)) {
+      erase(itr++);
+
+    // Otherwise simply set the new begin of the old range.
+    } else {
+      itr->second.first = _end;
+      ++itr;
+    }
+  }
+
+  return itr;
+}
+
+template<typename Key, typename T, typename C, typename A>
+inline typename RangeMap<Key,T,C,A>::const_iterator
+RangeMap<Key,T,C,A>::set_merge(Key _begin, const Key& _end, const T& value) {
+  if (!key_comp()(_begin, _end))
+    return end();
+
+  // Crop overlapping ranges and return iterator to first range after the one we're inserting.
+  typename RangeMap::iterator itr = crop_overlap(_begin, _end);
+
+  // Check if range before new one is adjacent and has same value: if so erase it and use its beginning.
+  if (itr != begin()) {
+    typename RangeMap::iterator prev = itr;
+    if (!key_comp()((--prev)->first, _begin) && prev->second.second == value) {
+      _begin = prev->second.first;
+      erase(prev);
+    }
+  }
+
+  // Range after new one is adjacent and has same value: set new beginning.
+  if (itr != end() && !key_comp()(_end, itr->second.first) && itr->second.second == value) {
+    itr->second.first = _begin;
+    return itr;
+  }
+
+  // Otherwise, this range isn't mergeable, make new entry.
+  return insert(itr, typename RangeMap::value_type(_end, typename RangeMap::mapped_type(_begin, value)));
+}
+
+template<typename Key, typename T, typename C, typename A>
+inline typename RangeMap<Key,T,C,A>::const_iterator
+RangeMap<Key,T,C,A>::set_range(const Key& _begin, const Key& _end, const T& value) {
+  if (!key_comp()(_begin, _end))
+    return end();
+
+  return insert(crop_overlap(_begin, _end), typename RangeMap::value_type(_end, typename RangeMap::mapped_type(_begin, value)));
+}
+
+template<typename Key, typename T, typename C, typename A>
+inline typename RangeMap<Key,T,C,A>::const_iterator
+RangeMap<Key,T,C,A>::find(const Key& key) const {
+  typename RangeMap::const_iterator itr = upper_bound(key);
+
+  if (itr != end() && key_comp()(key, itr->second.first))
+    itr = end();
+
+  return itr;
+}
+
+template<typename Key, typename T, typename C, typename A>
+inline const T&
+RangeMap<Key,T,C,A>::get(const Key& key) const {
+  typename RangeMap::const_iterator itr = find(key);
+
+  if (itr == end())
+    throw std::out_of_range("RangeMap::get");
+
+  return itr->second.second;
+}
+
+template<typename Key, typename T, typename C, typename A>
+inline T
+RangeMap<Key,T,C,A>::get(const Key& key, T def) const {
+  typename RangeMap::const_iterator itr = find(key);
+  return (itr == end() ? def : itr->second.second);
+}
+
+}
+
+#endif