Merge pull request #189 from aldenml/dht-storage-api

Add set_dht_storage to session API

include/libtorrent/aux_/session_impl.hpp

@@ -292,6 +292,7 @@ namespace libtorrent
 			void add_dht_router(std::pair<std::string, int> const& node);
 			void set_dht_settings(dht_settings const& s);
 			dht_settings const& get_dht_settings() const { return m_dht_settings; }
+			void set_dht_storage(dht::dht_storage_constructor_type sc);
 			void start_dht();
 			void stop_dht();
 			void start_dht(entry const& startup_state);
@@ -994,6 +995,7 @@ namespace libtorrent
 #ifndef TORRENT_DISABLE_DHT
 			boost::shared_ptr<dht::dht_tracker> m_dht;
 			dht_settings m_dht_settings;
+			dht::dht_storage_constructor_type m_dht_storage_constructor;
 
 			// these are used when starting the DHT
 			// (and bootstrapping it), and then erased

include/libtorrent/kademlia/dht_storage.hpp

@@ -30,22 +30,20 @@ POSSIBILITY OF SUCH DAMAGE.
 
 */
 
-#ifndef DHT_STORAGE_HPP
-#define DHT_STORAGE_HPP
+#ifndef TORRENT_DHT_STORAGE_HPP
+#define TORRENT_DHT_STORAGE_HPP
 
-#include <algorithm>
-#include <map>
-#include <set>
-#include <string>
+#include <boost/function.hpp>
 
-#include <libtorrent/config.hpp>
-#include <libtorrent/time.hpp>
 #include <libtorrent/socket.hpp>
 #include <libtorrent/sha1_hash.hpp>
 #include <libtorrent/address.hpp>
-#include <libtorrent/session_settings.hpp>
-#include <libtorrent/performance_counters.hpp>
-#include <libtorrent/kademlia/item.hpp>
+
+namespace libtorrent
+{
+	struct dht_settings;
+	class entry;
+}
 
 namespace libtorrent {
 namespace dht
@@ -200,9 +198,12 @@ namespace dht
 		virtual ~dht_storage_interface() {}
 	};
 
+	typedef boost::function<dht_storage_interface*(sha1_hash const& id
+		, dht_settings const& settings)> dht_storage_constructor_type;
+
 	TORRENT_EXPORT dht_storage_interface* dht_default_storage_constructor(sha1_hash const& id
 		, dht_settings const& settings);
 
 } } // namespace libtorrent::dht
 
-#endif //DHT_STORAGE_HPP
+#endif //TORRENT_DHT_STORAGE_HPP

include/libtorrent/kademlia/dht_tracker.hpp

@@ -72,7 +72,9 @@ namespace libtorrent { namespace dht
 		, boost::enable_shared_from_this<dht_tracker>
 	{
 		dht_tracker(dht_observer* observer, rate_limited_udp_socket& sock
-			, dht_settings const& settings, counters& cnt, entry const* state = 0);
+			, dht_settings const& settings, counters& cnt
+			, dht_storage_constructor_type storage_constructor
+			, entry const* state = 0);
 		virtual ~dht_tracker();
 
 		void start(entry const& bootstrap

include/libtorrent/kademlia/node.hpp

@@ -129,7 +129,8 @@ class TORRENT_EXTRA_EXPORT node : boost::noncopyable
 public:
 	node(udp_socket_interface* sock
 		, libtorrent::dht_settings const& settings, node_id nid
-		, dht_observer* observer, counters& cnt);
+		, dht_observer* observer, counters& cnt
+		, dht_storage_constructor_type storage_constructor = dht_default_storage_constructor);
 
 	virtual ~node();
 

include/libtorrent/session_handle.hpp

@@ -43,6 +43,8 @@ POSSIBILITY OF SUCH DAMAGE.
 #include "libtorrent/peer_class_type_filter.hpp"
 #include "libtorrent/session_settings.hpp"
 
+#include "libtorrent/kademlia/dht_storage.hpp"
+
 #ifndef TORRENT_NO_DEPRECATE
 #include "libtorrent/rss.hpp"
 #endif
@@ -338,7 +340,7 @@ namespace libtorrent
 		void stop_dht();
 #endif
 
-		// ``set_dht_settings`` sets some parameters availavle to the dht node.
+		// ``set_dht_settings`` sets some parameters available to the dht node.
 		// See dht_settings for more information.
 		//
 		// ``is_dht_running()`` returns true if the DHT support has been started
@@ -350,6 +352,20 @@ namespace libtorrent
 		bool is_dht_running() const;
 		dht_settings get_dht_settings() const;
 
+		// ``set_dht_storage`` set a dht custom storage constructor function
+		// to be used internally when the dht is created.
+		//
+		// Since the dht storage is a critical component for the dht behavior,
+		// this function will only be effective the next time the dht is started.
+		// If you never touch this feature, a default map-memory based storage
+		// is used.
+		//
+		// If you want to make sure the dht is initially created with your
+		// custom storage, create a session with the setting
+		// ``settings_pack::enable_dht`` to false, set your constructor function
+		// and call ``apply_settings`` with ``settings_pack::enable_dht`` to true.
+		void set_dht_storage(dht::dht_storage_constructor_type sc);
+
 		// ``add_dht_node`` takes a host name and port pair. That endpoint will be
 		// pinged, and if a valid DHT reply is received, the node will be added to
 		// the routing table.

src/kademlia/dht_tracker.cpp

@@ -98,9 +98,10 @@ namespace libtorrent { namespace dht
 		, rate_limited_udp_socket& sock
 		, dht_settings const& settings
 		, counters& cnt
+		, dht_storage_constructor_type storage_constructor
 		, entry const* state)
 		: m_counters(cnt)
-		, m_dht(this, settings, extract_node_id(state), observer, cnt)
+		, m_dht(this, settings, extract_node_id(state), observer, cnt, storage_constructor)
 		, m_sock(sock)
 		, m_log(observer)
 		, m_last_new_key(clock_type::now() - minutes(int(key_refresh)))

src/kademlia/node.cpp

@@ -88,7 +88,8 @@ node_id calculate_node_id(node_id const& nid, dht_observer* observer)
 node::node(udp_socket_interface* sock
 	, dht_settings const& settings, node_id nid
 	, dht_observer* observer
-	, struct counters& cnt)
+	, struct counters& cnt
+	, dht_storage_constructor_type storage_constructor)
 	: m_settings(settings)
 	, m_id(calculate_node_id(nid, observer))
 	, m_table(m_id, 8, settings, observer)
@@ -98,7 +99,7 @@ node::node(udp_socket_interface* sock
 	, m_last_self_refresh(min_time())
 	, m_sock(sock)
 	, m_counters(cnt)
-	, m_storage(dht_default_storage_constructor(m_id, m_settings))
+	, m_storage(storage_constructor(m_id, m_settings))
 {
 	m_secret[0] = random();
 	m_secret[1] = random();

src/session_handle.cpp

@@ -351,6 +351,15 @@ namespace libtorrent
 #endif
 	}
 
+	void session_handle::set_dht_storage(dht::dht_storage_constructor_type sc)
+	{
+#ifndef TORRENT_DISABLE_DHT
+		TORRENT_ASYNC_CALL1(set_dht_storage, sc);
+#else
+		TORRENT_UNUSED(sc);
+#endif
+	}
+
 	void session_handle::add_dht_node(std::pair<std::string, int> const& node)
 	{
 #ifndef TORRENT_DISABLE_DHT

src/session_impl.cpp

@@ -403,6 +403,7 @@ namespace aux {
 		, m_last_auto_manage(m_created)
 		, m_next_port(0)
 #ifndef TORRENT_DISABLE_DHT
+		, m_dht_storage_constructor(dht::dht_default_storage_constructor)
 		, m_dht_announce_timer(m_io_service)
 		, m_dht_interval_update_torrents(0)
 		, m_outstanding_router_lookups(0)
@@ -5483,7 +5484,9 @@ retry:
 
 		m_dht = boost::make_shared<dht::dht_tracker>(static_cast<dht_observer*>(this)
 			, boost::ref(m_udp_socket), boost::cref(m_dht_settings)
-			, boost::ref(m_stats_counters), &startup_state);
+			, boost::ref(m_stats_counters)
+			, m_dht_storage_constructor
+			, &startup_state);
 
 		for (std::vector<udp::endpoint>::iterator i = m_dht_router_nodes.begin()
 			, end(m_dht_router_nodes.end()); i != end; ++i)
@@ -5516,6 +5519,11 @@ retry:
 		m_dht_settings = settings;
 	}
 
+	void session_impl::set_dht_storage(dht::dht_storage_constructor_type sc)
+	{
+		m_dht_storage_constructor = sc;
+	}
+
 #ifndef TORRENT_NO_DEPRECATE
 	entry session_impl::dht_state() const
 	{

test/test_dht_storage.cpp

@@ -55,6 +55,7 @@ POSSIBILITY OF SUCH DAMAGE.
 
 using namespace libtorrent;
 using namespace libtorrent::dht;
+namespace lt = libtorrent;
 
 namespace
 {
@@ -71,6 +72,15 @@ namespace
 		from_hex(s, 40, (char *) &ret[0]);
 		return ret;
 	}
+
+	bool g_storage_constructor_invoked = false;
+
+	dht_storage_interface* dht_custom_storage_constructor(sha1_hash const& id
+		, dht_settings const& settings)
+	{
+		g_storage_constructor_invoked = true;
+		return dht_default_storage_constructor(id, settings);
+	}
 }
 
 TORRENT_TEST(dht_storage)
@@ -177,4 +187,50 @@ TORRENT_TEST(dht_storage_counters)
 	TEST_EQUAL(s->counters().mutable_data, 1);
 }
 
+TORRENT_TEST(dht_storage_set_custom)
+{
+	g_storage_constructor_invoked = false;
+	settings_pack p;
+	p.set_bool(settings_pack::enable_dht, false);
+	lt::session ses(p);
+
+	bool r = ses.is_dht_running();
+	TEST_CHECK(!r);
+
+	ses.set_dht_storage(dht_custom_storage_constructor);
+
+	p.set_bool(settings_pack::enable_dht, true);
+	ses.apply_settings(p); // async with dispatch
+	r = ses.is_dht_running();
+	TEST_CHECK(r);
+	TEST_EQUAL(g_storage_constructor_invoked, true);
+}
+
+TORRENT_TEST(dht_storage_default_set_custom)
+{
+	g_storage_constructor_invoked = false;
+	settings_pack p;
+	p.set_bool(settings_pack::enable_dht, true);
+	lt::session ses(p);
+
+	bool r = ses.is_dht_running();
+	TEST_CHECK(r);
+
+	ses.set_dht_storage(dht_custom_storage_constructor);
+
+	p.set_bool(settings_pack::enable_dht, false);
+	ses.apply_settings(p); // async with dispatch
+	r = ses.is_dht_running();
+	TEST_CHECK(!r);
+	TEST_EQUAL(g_storage_constructor_invoked, false);
+
+	ses.set_dht_storage(dht_custom_storage_constructor);
+
+	p.set_bool(settings_pack::enable_dht, true);
+	ses.apply_settings(p); // async with dispatch
+	r = ses.is_dht_running();
+	TEST_CHECK(r);
+	TEST_EQUAL(g_storage_constructor_invoked, true);
+}
+
 #endif