From fde376efbc745908f4cf090c8aef5079ad1ff393 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Tue, 28 May 2013 21:44:17 +0900
Subject: [PATCH] SessionSerializer: Save spent URIs as well as remaining ones

---
 src/FileEntry.h               |  6 +++++
 src/SessionSerializer.cc      | 45 +++++++++++++++++++++++++++++++----
 test/SessionSerializerTest.cc | 24 +++++++++++++++++++
 3 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/src/FileEntry.h b/src/FileEntry.h
index a1ed29de7..9ba57f005 100644
--- a/src/FileEntry.h
+++ b/src/FileEntry.h
@@ -136,6 +136,12 @@ public:
     return spentUris_;
   }
 
+  // Exposed for unitest
+  std::deque<std::string>& getSpentUris()
+  {
+    return spentUris_;
+  }
+
   size_t setUris(const std::vector<std::string>& uris);
 
   template<typename InputIterator>
diff --git a/src/SessionSerializer.cc b/src/SessionSerializer.cc
index 37dfe7b8b..ee5c4e61e 100644
--- a/src/SessionSerializer.cc
+++ b/src/SessionSerializer.cc
@@ -36,6 +36,7 @@
 
 #include <cstdio>
 #include <iterator>
+#include <set>
 
 #include "RequestGroupMan.h"
 #include "a2functional.h"
@@ -135,6 +136,27 @@ bool writeOption(IOFile& fp, const SharedHandle<Option>& op)
 }
 } // namespace
 
+namespace {
+bool writeUri(IOFile& fp, const std::string& uri)
+{
+  return fp.write(uri.c_str(), uri.size()) == uri.size() &&
+    fp.write("\t", 1) == 1;
+}
+} // namespace
+
+namespace {
+template<typename InputIterator>
+bool writeUri(IOFile& fp, InputIterator first, InputIterator last)
+{
+  for(; first != last; ++first) {
+    if(!writeUri(fp, *first)) {
+      return false;
+    }
+  }
+  return true;
+}
+} // namespace
+
 // The downloads whose followedBy() is empty is persisited with its
 // GID without no problem. For other cases, there are several patterns.
 //
@@ -173,17 +195,30 @@ bool writeDownloadResult
       return true;
     }
     const SharedHandle<FileEntry>& file = dr->fileEntries[0];
-    if(file->getRemainingUris().empty()) {
+    // Don't save download if there are no URIs.
+    if(file->getRemainingUris().empty() &&
+       file->getSpentUris().empty()) {
       return true;
     }
+    // Save spent URIs + remaining URIs. Remove URI in spent URI which
+    // also exists in remaining URIs.
+    std::set<std::string> uriSet(file->getRemainingUris().begin(),
+                                 file->getRemainingUris().end());
     for(std::deque<std::string>::const_iterator i =
-          file->getRemainingUris().begin(),
-          eoi = file->getRemainingUris().end(); i != eoi; ++i) {
-      if (fp.write((*i).c_str(), (*i).size()) != (*i).size() ||
-          fp.write("\t", 1) != 1) {
+          file->getSpentUris().begin(), eoi = file->getSpentUris().end();
+        i != eoi; ++i) {
+      if(uriSet.count(*i)) {
+        continue;
+      }
+      uriSet.insert(*i);
+      if(!writeUri(fp, *i)) {
         return false;
       }
     }
+    if(!writeUri(fp, file->getRemainingUris().begin(),
+                 file->getRemainingUris().end())) {
+      return false;
+    }
     if(fp.write("\n", 1) != 1) {
       return false;
     }
diff --git a/test/SessionSerializerTest.cc b/test/SessionSerializerTest.cc
index e9da45f68..b2dc5bd53 100644
--- a/test/SessionSerializerTest.cc
+++ b/test/SessionSerializerTest.cc
@@ -22,9 +22,11 @@ class SessionSerializerTest:public CppUnit::TestFixture {
 
   CPPUNIT_TEST_SUITE(SessionSerializerTest);
   CPPUNIT_TEST(testSave);
+  CPPUNIT_TEST(testSaveErrorDownload);
   CPPUNIT_TEST_SUITE_END();
 public:
   void testSave();
+  void testSaveErrorDownload();
 };
 
 
@@ -123,4 +125,26 @@ void SessionSerializerTest::testSave()
 #endif // defined(ENABLE_BITTORRENT) && defined(ENABLE_METALINK)
 }
 
+void SessionSerializerTest::testSaveErrorDownload()
+{
+  SharedHandle<DownloadResult> dr = createDownloadResult(error_code::TIME_OUT,
+                                                         "http://error");
+  dr->fileEntries[0]->getSpentUris().swap
+    (dr->fileEntries[0]->getRemainingUris());
+  SharedHandle<Option> option(new Option());
+  option->put(PREF_MAX_DOWNLOAD_RESULT, "10");
+  SharedHandle<RequestGroupMan> rgman
+    (new RequestGroupMan(std::vector<SharedHandle<RequestGroup> >(), 1,
+                         option.get()));
+  rgman->addDownloadResult(dr);
+  SessionSerializer s(rgman);
+  std::string filename =
+    A2_TEST_OUT_DIR"/aria2_SessionSerializerTest_testSaveErrorDownload";
+  CPPUNIT_ASSERT(s.save(filename));
+  std::ifstream ss(filename.c_str(), std::ios::binary);
+  std::string line;
+  std::getline(ss, line);
+  CPPUNIT_ASSERT_EQUAL(std::string("http://error\t"), line);
+}
+
 } // namespace aria2