From d58d7947576a409d8053ea571c86eecb302da70c Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@allmydata.com>
Date: Thu, 20 Sep 2007 18:52:44 -0700
Subject: [PATCH] check_speed: do both upload and download tests

---
 src/allmydata/control.py          | 91 +++++++++++++++++++++++--------
 src/allmydata/interfaces.py       |  7 ++-
 src/allmydata/test/check_speed.py | 45 ++++++++-------
 3 files changed, 100 insertions(+), 43 deletions(-)

diff --git a/src/allmydata/control.py b/src/allmydata/control.py
index 39289638..5d0c75e8 100644
--- a/src/allmydata/control.py
+++ b/src/allmydata/control.py
@@ -47,14 +47,36 @@ class ControlServer(Referenceable, service.Service, testutil.PollMixin):
         d.addCallback(lambda res: filename)
         return d
 
-    def remote_upload_speed_test(self, count, size):
+    def remote_speed_test(self, count, size):
         assert size > 8
-        basedir = os.path.join(self.parent.basedir, "_speed_test_data")
         log.msg("speed_test: count=%d, size=%d" % (count, size))
-        fileutil.make_dirs(basedir)
-        for i in range(count):
-            s = size
-            fn = os.path.join(basedir, str(i))
+        st = SpeedTest(self.parent, count, size)
+        return st.run()
+
+    def remote_get_memory_usage(self):
+        return get_memory_usage()
+
+class SpeedTest:
+    def __init__(self, parent, count, size):
+        self.parent = parent
+        self.count = count
+        self.size = size
+        self.uris = {}
+        self.basedir = os.path.join(self.parent.basedir, "_speed_test_data")
+
+    def run(self):
+        self.create_data()
+        d = self.do_upload()
+        d.addCallback(lambda res: self.do_download())
+        d.addBoth(self.do_cleanup)
+        d.addCallback(lambda res: (self.upload_time, self.download_time))
+        return d
+
+    def create_data(self):
+        fileutil.make_dirs(self.basedir)
+        for i in range(self.count):
+            s = self.size
+            fn = os.path.join(self.basedir, str(i))
             if os.path.exists(fn):
                 os.unlink(fn)
             f = open(fn, "w")
@@ -65,28 +87,53 @@ class ControlServer(Referenceable, service.Service, testutil.PollMixin):
                 f.write("\x00" * chunk)
                 s -= chunk
             f.close()
+
+    def do_upload(self):
         uploader = self.parent.getServiceNamed("uploader")
         start = time.time()
         d = defer.succeed(None)
-        def _do_one_file(uri, i):
-            if i >= count:
+        def _record_uri(uri, i):
+            self.uris[i] = uri
+        def _upload_one_file(ignored, i):
+            if i >= self.count:
                 return
-            fn = os.path.join(basedir, str(i))
+            fn = os.path.join(self.basedir, str(i))
             d1 = uploader.upload_filename(fn)
-            d1.addCallback(_do_one_file, i+1)
+            d1.addCallback(_record_uri, i)
+            d1.addCallback(_upload_one_file, i+1)
             return d1
-        d.addCallback(_do_one_file, 0)
-        def _done(ignored):
+        d.addCallback(_upload_one_file, 0)
+        def _upload_done(ignored):
             stop = time.time()
-            return stop - start
-        d.addCallback(_done)
-        def _cleanup(res):
-            for i in range(count):
-                fn = os.path.join(basedir, str(i))
-                os.unlink(fn)
-            return res
-        d.addBoth(_cleanup)
+            self.upload_time = stop - start
+        d.addCallback(_upload_done)
         return d
 
-    def remote_get_memory_usage(self):
-        return get_memory_usage()
+    def do_download(self):
+        downloader = self.parent.getServiceNamed("downloader")
+        start = time.time()
+        d = defer.succeed(None)
+        def _download_one_file(ignored, i):
+            if i >= self.count:
+                return
+            d1 = downloader.download_to_filehandle(self.uris[i], Discard())
+            d1.addCallback(_download_one_file, i+1)
+            return d1
+        d.addCallback(_download_one_file, 0)
+        def _download_done(ignored):
+            stop = time.time()
+            self.download_time = stop - start
+        d.addCallback(_download_done)
+        return d
+
+    def do_cleanup(self, res):
+        for i in range(self.count):
+            fn = os.path.join(self.basedir, str(i))
+            os.unlink(fn)
+        return res
+
+class Discard:
+    def write(self, data):
+        pass
+    def close(self):
+        pass
diff --git a/src/allmydata/interfaces.py b/src/allmydata/interfaces.py
index 47de6857..efd6962f 100644
--- a/src/allmydata/interfaces.py
+++ b/src/allmydata/interfaces.py
@@ -1054,8 +1054,11 @@ class RIControlClient(RemoteInterface):
         measuring memory consupmtion in bytes."""
         return DictOf(str, int)
 
-    def upload_speed_test(count=int, size=int):
+    def speed_test(count=int, size=int):
         """Write 'count' tempfiles to disk, all of the given size. Measure
         how long (in seconds) it takes to upload them all to the servers.
+        Then measure how long it takes to download all of them.
+
+        Returns a tuple of (upload_time, download_time).
         """
-        return float
+        return (float, float)
diff --git a/src/allmydata/test/check_speed.py b/src/allmydata/test/check_speed.py
index fe162349..11aceab5 100644
--- a/src/allmydata/test/check_speed.py
+++ b/src/allmydata/test/check_speed.py
@@ -17,7 +17,8 @@ class SpeedTest:
         f.close()
         self.base_service = service.MultiService()
         self.failed = None
-        self.times = {}
+        self.upload_times = {}
+        self.download_times = {}
 
     def run(self):
         print "STARTING"
@@ -57,41 +58,47 @@ class SpeedTest:
         reactor.callLater(delay, d.callback, result)
         return d
 
-    def record_time(self, time, key):
-        print "TIME (%s): %s" % (key, time)
-        self.times[key] = time
+    def record_times(self, times, key):
+        print "TIME (%s): %s up, %s down" % (key, times[0], times[1])
+        self.upload_times[key], self.download_times[key] = times
 
     def one_test(self, res, name, count, size):
-        d = self.client_rref.callRemote("upload_speed_test", count, size)
-        d.addCallback(self.record_time, name)
+        d = self.client_rref.callRemote("speed_test", count, size)
+        d.addCallback(self.record_times, name)
         return d
 
     def do_test(self):
         print "doing test"
         rr = self.client_rref
-        d = rr.callRemote("get_memory_usage")
-        def _got(res):
-            print "MEMORY USAGE:", res
-        d.addCallback(_got)
+        d = defer.succeed(None)
         d.addCallback(self.one_test, "startup", 1, 1000) # ignore this one
         d.addCallback(self.one_test, "1x 200B", 1, 200)
         d.addCallback(self.one_test, "10x 200B", 10, 200)
         #d.addCallback(self.one_test, "100x 200B", 100, 200)
         d.addCallback(self.one_test, "1MB", 1, 1*MB)
         d.addCallback(self.one_test, "10MB", 1, 10*MB)
-        d.addCallback(self.calculate_speed)
+        d.addCallback(self.calculate_speeds)
         return d
 
-    def calculate_speed(self, res):
-        perfile = self.times["1x 200B"]
+    def calculate_speeds(self, res):
         # time = A*size+B
         # we assume that A*200bytes is negligible
-        B = self.times["10x 200B"] / 10
-        print "per-file time: %.3fs" % B
-        A1 = 1*MB / (self.times["1MB"] - B) # in bytes per second
-        print "speed (1MB):", self.number(A1, "Bps")
-        A2 = 10*MB / (self.times["10MB"] - B)
-        print "speed (10MB):", self.number(A2, "Bps")
+
+        # upload
+        B = self.upload_times["10x 200B"] / 10
+        print "upload per-file time: %.3fs" % B
+        A1 = 1*MB / (self.upload_times["1MB"] - B) # in bytes per second
+        print "upload speed (1MB):", self.number(A1, "Bps")
+        A2 = 10*MB / (self.upload_times["10MB"] - B)
+        print "upload speed (10MB):", self.number(A2, "Bps")
+
+        # download
+        B = self.download_times["10x 200B"] / 10
+        print "download per-file time: %.3fs" % B
+        A1 = 1*MB / (self.download_times["1MB"] - B) # in bytes per second
+        print "download speed (1MB):", self.number(A1, "Bps")
+        A2 = 10*MB / (self.download_times["10MB"] - B)
+        print "download speed (10MB):", self.number(A2, "Bps")
 
     def number(self, value, suffix=""):
         scaling = 1
-- 
2.45.2