Переглянути джерело

Rewritten DHTRoutingTableSerializer using stdio instead of stream.

Tatsuhiro Tsujikawa 14 роки тому
батько
коміт
dea7a7969c

+ 4 - 17
src/DHTAutoSaveCommand.cc

@@ -35,7 +35,6 @@
 #include "DHTAutoSaveCommand.h"
 
 #include <cstring>
-#include <fstream>
 
 #include "DHTRoutingTable.h"
 #include "DHTNode.h"
@@ -86,13 +85,14 @@ void DHTAutoSaveCommand::save()
     get(family_ == AF_INET? PREF_DHT_FILE_PATH : PREF_DHT_FILE_PATH6);
   A2_LOG_INFO(fmt("Saving DHT routing table to %s.", dhtFile.c_str()));
 
-  File tempFile(dhtFile+"__temp");
   // Removing tempFile is unnecessary because the file is truncated on
   // open.  But the bug in 1.10.4 creates directory for this path.
   // Because it is directory, opening directory as file fails.  So we
   // first remove it here.
+  File tempFile(dhtFile+"__temp");
   tempFile.remove();
-  File(tempFile.getDirname()).mkdirs();
+
+  File(File(dhtFile).getDirname()).mkdirs();
   std::vector<SharedHandle<DHTNode> > nodes;
   std::vector<SharedHandle<DHTBucket> > buckets;
   routingTable_->getBuckets(buckets);
@@ -109,20 +109,7 @@ void DHTAutoSaveCommand::save()
   serializer.setNodes(nodes);
 
   try {
-    {
-      std::ofstream o(tempFile.getPath().c_str(),
-                      std::ios::out|std::ios::binary);
-      if(!o) {
-        throw DL_ABORT_EX
-          (fmt("Failed to save DHT routing table to %s.",
-               dhtFile.c_str()));
-      }
-      serializer.serialize(o);
-    }
-    if(!tempFile.renameTo(dhtFile)) {
-      A2_LOG_ERROR(fmt("Cannot move file from %s to %s.",
-                       tempFile.getPath().c_str(), dhtFile.c_str()));
-    }
+    serializer.serialize(dhtFile);
   } catch(RecoverableException& e) {
     A2_LOG_ERROR_EX(fmt("Exception caught while saving DHT routing table to %s",
                         dhtFile.c_str()),

+ 38 - 19
src/DHTRoutingTableSerializer.cc

@@ -35,7 +35,7 @@
 #include "DHTRoutingTableSerializer.h"
 
 #include <cstring>
-#include <ostream>
+#include <cstdio>
 
 #include "DHTNode.h"
 #include "DlAbortEx.h"
@@ -45,6 +45,8 @@
 #include "a2netcompat.h"
 #include "util.h"
 #include "TimeA2.h"
+#include "fmt.h"
+#include "File.h"
 
 namespace aria2 {
 
@@ -65,8 +67,21 @@ void DHTRoutingTableSerializer::setNodes
   nodes_ = nodes;
 }
 
-void DHTRoutingTableSerializer::serialize(std::ostream& o)
+#define FWRITE_CHECK(ptr, count, fp)                                    \
+  if(fwrite((ptr), 1, (count), (fp)) != (count)) {                      \
+    fclose(fp);                                                         \
+    throw DL_ABORT_EX(fmt("Failed to save DHT routing table to %s.",    \
+                          utf8ToNative(filename).c_str()));             \
+  }
+
+void DHTRoutingTableSerializer::serialize(const std::string& filename)
 {
+  std::string filenameTemp = filename+"__temp";
+  FILE* fp = a2fopen(utf8ToWChar(filenameTemp).c_str(), "wb");
+  if(!fp) {
+    throw DL_ABORT_EX(fmt("Failed to save DHT routing table to %s.",
+                          utf8ToNative(filename).c_str()));
+  }
   char header[8];
   memset(header, 0, sizeof(header));
   // magic
@@ -81,24 +96,24 @@ void DHTRoutingTableSerializer::serialize(std::ostream& o)
   char zero[18];
   memset(zero, 0, sizeof(zero));
 
-  o.write(header, 8);
+  FWRITE_CHECK(header, 8, fp);
   // write save date
   uint64_t ntime = hton64(Time().getTime());
-  o.write(reinterpret_cast<const char*>(&ntime), sizeof(ntime));
+  FWRITE_CHECK(&ntime, sizeof(ntime), fp);
 
   // localnode
   // 8bytes reserved
-  o.write(zero, 8);
+  FWRITE_CHECK(zero, 8, fp);
   // 20bytes localnode ID
-  o.write(reinterpret_cast<const char*>(localNode_->getID()), DHT_ID_LENGTH);
+  FWRITE_CHECK(localNode_->getID(), DHT_ID_LENGTH, fp);
   // 4bytes reserved
-  o.write(zero, 4);
+  FWRITE_CHECK(zero, 4, fp);
 
   // number of nodes
   uint32_t numNodes = htonl(nodes_.size());
-  o.write(reinterpret_cast<const char*>(&numNodes), sizeof(uint32_t));
+  FWRITE_CHECK(&numNodes, sizeof(uint32_t), fp);
   // 4bytes reserved
-  o.write(zero, 4);
+  FWRITE_CHECK(zero, 4, fp);
 
   const int clen = bittorrent::getCompactLength(family_);
   // nodes
@@ -112,23 +127,27 @@ void DHTRoutingTableSerializer::serialize(std::ostream& o)
     if(compactlen != clen) {
       memset(compactPeer, 0, clen);
     }
+    uint8_t clen1 = clen;
     // 1byte compact peer format length
-    o << static_cast<uint8_t>(clen);
+    FWRITE_CHECK(&clen1, sizeof(clen1), fp);
     // 7bytes reserved
-    o.write(zero, 7);
+    FWRITE_CHECK(zero, 7, fp);
     // clen bytes compact peer
-    o.write(reinterpret_cast<const char*>(compactPeer), clen);
+    FWRITE_CHECK(compactPeer, static_cast<size_t>(clen), fp);
     // 24-clen bytes reserved
-    o.write(zero, 24-clen);
+    FWRITE_CHECK(zero, 24-clen, fp);
     // 20bytes: node ID
-    o.write(reinterpret_cast<const char*>(node->getID()), DHT_ID_LENGTH);
+    FWRITE_CHECK(node->getID(), DHT_ID_LENGTH, fp);
     // 4bytes reserved
-    o.write(zero, 4);
+    FWRITE_CHECK(zero, 4, fp);
   }
-
-  o.flush();
-  if(!o) {
-    throw DL_ABORT_EX("Failed to save DHT routing table.");
+  if(fclose(fp) == EOF) {
+    throw DL_ABORT_EX(fmt("Failed to save DHT routing table to %s.",
+                          utf8ToNative(filename).c_str()));
+  }
+  if(!File(filenameTemp).renameTo(filename)) {
+    throw DL_ABORT_EX(fmt("Failed to save DHT routing table to %s.",
+                          utf8ToNative(filename).c_str()));
   }
 }
 

+ 2 - 2
src/DHTRoutingTableSerializer.h

@@ -38,7 +38,7 @@
 #include "common.h"
 
 #include <vector>
-#include <iosfwd>
+#include <string>
 
 #include "SharedHandle.h"
 
@@ -62,7 +62,7 @@ public:
 
   void setNodes(const std::vector<SharedHandle<DHTNode> >& nodes);
 
-  void serialize(std::ostream& o);
+  void serialize(const std::string& filename);
 };
 
 } // namespace aria2

+ 2 - 6
test/DHTRoutingTableDeserializerTest.cc

@@ -55,9 +55,7 @@ void DHTRoutingTableDeserializerTest::testDeserialize()
   s.setNodes(nodes);
 
   std::string filename = A2_TEST_OUT_DIR"/aria2_DHTRoutingTableDeserializerTest_testDeserialize";
-  std::ofstream outfile(filename.c_str(), std::ios::binary);
-  s.serialize(outfile);
-  outfile.close();
+  s.serialize(filename);
 
   DHTRoutingTableDeserializer d(AF_INET);
   d.deserialize(filename);
@@ -96,9 +94,7 @@ void DHTRoutingTableDeserializerTest::testDeserialize6()
   s.setNodes(nodes);
 
   std::string filename = A2_TEST_OUT_DIR"/aria2_DHTRoutingTableDeserializerTest_testDeserialize6";
-  std::ofstream outfile(filename.c_str(), std::ios::binary);
-  s.serialize(outfile);
-  outfile.close();
+  s.serialize(filename);
 
   DHTRoutingTableDeserializer d(AF_INET6);
   d.deserialize(filename);

+ 9 - 5
test/DHTRoutingTableSerializerTest.cc

@@ -1,7 +1,7 @@
 #include "DHTRoutingTableSerializer.h"
 
 #include <cstring>
-#include <sstream>
+#include <fstream>
 #include <iostream>
 
 #include <cppunit/extensions/HelperMacros.h>
@@ -115,8 +115,10 @@ void DHTRoutingTableSerializerTest::testSerialize()
   s.setLocalNode(localNode);
   s.setNodes(nodes);
 
-  std::stringstream ss;
-  s.serialize(ss);
+  std::string filename = A2_TEST_OUT_DIR"/aria2_DHTRoutingTableSerializerTest_testSerialize";
+  s.serialize(filename);
+
+  std::ifstream ss(filename.c_str(), std::ios::binary);
 
   checkToLocalnode(ss, localNode);
   size_t numNodes = 3;
@@ -242,8 +244,10 @@ void DHTRoutingTableSerializerTest::testSerialize6()
   s.setLocalNode(localNode);
   s.setNodes(nodes);
 
-  std::stringstream ss;
-  s.serialize(ss);
+  std::string filename = A2_TEST_OUT_DIR"/aria2_DHTRoutingTableSerializerTest_testSerialize6";
+  s.serialize(filename);
+
+  std::ifstream ss(filename.c_str(), std::ios::binary);
 
   checkToLocalnode(ss, localNode);
   size_t numNodes = 2;