* Better handling of resume after crash/reboot. Sponsored by anonymous source.

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

src/core/download_factory.cc

@@ -42,6 +42,7 @@
 #include <stdexcept>
 #include <rak/path.h>
 #include <torrent/object.h>
+#include <torrent/object_stream.h>
 #include <torrent/exceptions.h>
 #include <torrent/rate.h>
 #include <torrent/resume.h>
@@ -161,6 +162,23 @@ DownloadFactory::receive_commit() {
     receive_success();
 }
 
+static bool
+download_factory_add_stream(torrent::Object* root, const char* key, const char* filename) {
+  std::fstream stream(filename, std::ios::in | std::ios::binary);
+    
+  if (!stream.is_open())
+    return false;
+
+  torrent::Object obj;
+  stream >> obj;
+
+  if (!stream.good())
+    return false;
+
+  root->insert_key_swap(key, obj);
+  return true;
+}
+
 void
 DownloadFactory::receive_success() {
   if (m_stream == NULL)
@@ -177,15 +195,19 @@ DownloadFactory::receive_success() {
 
   torrent::Object* root = download->bencode();
 
-  if (!m_session) {
+  if (m_session) {
+    download_factory_add_stream(root, "rtorrent", (rak::path_expand(m_uri) + ".rtorrent").c_str());
+    download_factory_add_stream(root, "libtorrent_resume", (rak::path_expand(m_uri) + ".libtorrent_resume").c_str());
+    
+  } else {
     // We only allow session torrents to keep their
     // 'rtorrent/libtorrent' sections. The "fast_resume" section
     // should be safe to keep.
     root->erase_key("rtorrent");
-    root->erase_key("libtorrent");
   }
 
   torrent::Object* rtorrent = &root->insert_preserve_copy("rtorrent", torrent::Object::create_map()).first->second;
+  torrent::Object& resumeObject = root->insert_preserve_copy("libtorrent_resume", torrent::Object::create_map()).first->second;
 
   initialize_rtorrent(download, rtorrent);
 
@@ -233,7 +255,6 @@ DownloadFactory::receive_success() {
 
   rpc::call_command("d.set_peer_exchange", rpc::call_command_value("get_peer_exchange"), rpc::make_target(download));
 
-  torrent::Object& resumeObject = root->insert_preserve_copy("libtorrent_resume", torrent::Object::create_map()).first->second;
   torrent::resume_load_addresses(*download->download(), resumeObject);
   torrent::resume_load_file_priorities(*download->download(), resumeObject);
   torrent::resume_load_tracker_settings(*download->download(), resumeObject);

src/core/download_list.cc

@@ -82,7 +82,7 @@ DownloadList::clear() {
 
 void
 DownloadList::session_save() {
-  unsigned int c = std::count_if(begin(), end(), std::bind1st(std::mem_fun(&DownloadStore::save), control->core()->download_store()));
+  unsigned int c = std::count_if(begin(), end(), std::bind1st(std::mem_fun(&DownloadStore::save_resume), control->core()->download_store()));
 
   if (c != size())
     control->core()->push_log("Failed to save session torrents.");
@@ -342,7 +342,9 @@ DownloadList::resume(Download* download, int flags) {
       // on non-complete downloads after a crash. This shouldn't be
       // needed, but for some reason linux 2.6 is very lazy about
       // updating mtime.
-      torrent::resume_save_progress(*download->download(), download->download()->bencode()->get_key("libtorrent_resume"), true);
+      //
+      // Disabling this due to the new resume code.
+      //      torrent::resume_save_progress(*download->download(), download->download()->bencode()->get_key("libtorrent_resume"), true);
     }
 
     // If the DHT server is set to auto, start it now.
@@ -559,8 +561,8 @@ DownloadList::confirm_finished(Download* download) {
   //
   // Obsolete.
   if (!download->is_active() && rpc::call_command_value("get_session_on_completion") != 0) {
-    torrent::resume_save_progress(*download->download(), download->download()->bencode()->get_key("libtorrent_resume"));
-    control->core()->download_store()->save(download);
+    //    torrent::resume_save_progress(*download->download(), download->download()->bencode()->get_key("libtorrent_resume"));
+    control->core()->download_store()->save_resume(download);
   }
 
   // Send the completed request before resuming so we don't reset the

src/core/download_store.cc

@@ -99,56 +99,82 @@ DownloadStore::set_path(const std::string& path) {
 }
 
 bool
-DownloadStore::save(Download* d) {
-  if (!is_enabled())
-    return true;
-
-  std::fstream f((create_filename(d) + ".new").c_str(), std::ios::out | std::ios::trunc);
-
-  if (!f.is_open())
-    return false;
-
-  // Move this somewhere else?
-  d->bencode()->get_key("rtorrent").insert_key("total_uploaded", d->download()->up_rate()->total());
-  d->bencode()->get_key("rtorrent").insert_key("chunks_done", d->download()->file_list()->completed_chunks());
-
+DownloadStore::write_bencode(const std::string& filename, const torrent::Object& obj, uint32_t skip_mask) {
   torrent::Object tmp;
-  torrent::Object& resumeObject = d->download()->bencode()->get_key("libtorrent_resume");
+  std::fstream output(filename.c_str(), std::ios::out | std::ios::trunc);
 
-  torrent::resume_save_addresses(*d->download(), resumeObject);
-  torrent::resume_save_file_priorities(*d->download(), resumeObject);
-  torrent::resume_save_tracker_settings(*d->download(), resumeObject);
-
-  f << *d->bencode();
-
-  if (!f.good())
+  if (!output.is_open())
     goto download_store_save_error;
 
-  f.close();
+  torrent::object_write_bencode(&output, &obj, skip_mask);
+
+  if (!output.good())
+    goto download_store_save_error;
+
+  output.close();
 
   // Test the new file, to ensure it is a valid bencode string.
-  f.open((create_filename(d) + ".new").c_str(), std::ios::in);
+  output.open(filename.c_str(), std::ios::in);
+  output >> tmp;
 
-  f >> tmp;
-
-  if (!f.good())
+  if (!output.good())
     goto download_store_save_error;
 
-  f.close();
-
-  ::rename((create_filename(d) + ".new").c_str(), create_filename(d).c_str());
+  output.close();
   return true;
 
 download_store_save_error:
-  f.close();
+  output.close();
   return false;
 }
 
+bool
+DownloadStore::save(Download* d, int flags) {
+  if (!is_enabled())
+    return true;
+
+  torrent::Object* resume_base   = &d->download()->bencode()->get_key("libtorrent_resume");
+  torrent::Object* rtorrent_base = &d->download()->bencode()->get_key("rtorrent");
+
+  // Move this somewhere else?
+  rtorrent_base->insert_key("chunks_done", d->download()->file_list()->completed_chunks());
+  rtorrent_base->insert_key("total_uploaded", d->download()->up_rate()->total());
+
+  // Don't save for completed torrents when we've cleared the uncertain_pieces.
+  torrent::resume_save_progress(*d->download(), *resume_base);
+  torrent::resume_save_uncertain_pieces(*d->download(), *resume_base);
+
+  torrent::resume_save_addresses(*d->download(), *resume_base);
+  torrent::resume_save_file_priorities(*d->download(), *resume_base);
+  torrent::resume_save_tracker_settings(*d->download(), *resume_base);
+
+  // Temp fixing of all flags, move to a better place:
+  resume_base->set_flags(torrent::Object::flag_session_data);
+  rtorrent_base->set_flags(torrent::Object::flag_session_data);
+
+  std::string base_filename = create_filename(d);
+
+  if (!write_bencode(base_filename + ".libtorrent_resume.new", *resume_base, 0) ||
+      !write_bencode(base_filename + ".rtorrent.new", *rtorrent_base, 0))
+    return false;
+
+  ::rename((base_filename + ".libtorrent_resume.new").c_str(), (base_filename + ".libtorrent_resume").c_str());
+  ::rename((base_filename + ".rtorrent.new").c_str(), (base_filename + ".rtorrent").c_str());
+  
+  if (!(flags & flag_skip_static) &&
+      write_bencode(base_filename + ".new", *d->bencode(), torrent::Object::flag_session_data))
+    ::rename((base_filename + ".new").c_str(), base_filename.c_str());
+
+  return true;
+}
+
 void
 DownloadStore::remove(Download* d) {
   if (!is_enabled())
     return;
 
+  ::unlink((create_filename(d) + ".libtorrent_resume").c_str());
+  ::unlink((create_filename(d) + ".rtorrent").c_str());
   ::unlink(create_filename(d).c_str());
 }
 

src/core/download_store.h

@@ -51,6 +51,7 @@ class Download;
 
 class DownloadStore {
 public:
+  static const int flag_skip_static = 0x1;
 
   bool                is_enabled()                            { return m_lockfile.is_locked(); }
 
@@ -60,7 +61,9 @@ public:
   const std::string&  path() const                            { return m_path; }
   void                set_path(const std::string& path);
 
-  bool                save(Download* d);
+  bool                save(Download* d, int flags);
+  bool                save_full(Download* d)   { return save(d, 0); }
+  bool                save_resume(Download* d) { return save(d, flag_skip_static); }
   void                remove(Download* d);
 
   // Currently shows all entries in the correct format.
@@ -71,6 +74,8 @@ public:
 private:
   std::string         create_filename(Download* d);
 
+  bool                write_bencode(const std::string& filename, const torrent::Object& obj, uint32_t skip_mask);
+
   std::string         m_path;
   utils::Lockfile     m_lockfile;
 };