[tor-commits] [onionperf/master] Adds oneshot mode

[irl at torproject.org]

commit 479ec5dd72b9061fa3730d2c8b73d213a1b76807
Author: Ana Custura <ana at nestat.org.uk>
Date:   Mon Jan 28 13:14:09 2019 +0000

    Adds oneshot mode
    
    This adds new logic using the oneshot model to perform a single download
    and then exit gracefully. It can be used via a command line option.
---
 onionperf/measurement.py | 14 ++++++++++++++
 onionperf/onionperf      | 11 ++++++++---
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/onionperf/measurement.py b/onionperf/measurement.py
index cc76ce8..508bd2c 100644
--- a/onionperf/measurement.py
+++ b/onionperf/measurement.py
@@ -166,6 +166,7 @@ class Measurement(object):
         self.tor_bin_path = tor_bin_path
         self.tgen_bin_path = tgen_bin_path
         self.datadir_path = datadir_path
+        self.oneshot = oneshot
         self.nickname = nickname
         self.threads = None
         self.done_event = None
@@ -236,9 +237,20 @@ class Measurement(object):
 
                 logging.info("Bootstrapping finished, entering heartbeat loop")
                 time.sleep(1)
+                logging.info("Onionperf is running in Oneshot mode. It will download a 5M file and shut down gracefully...")
                 while True:
                     # TODO add status update of some kind? maybe the number of files in the www directory?
                     # logging.info("Heartbeat: {0} downloads have completed successfully".format(self.__get_download_count(tgen_client_writable.filename)))
+                    if self.oneshot:
+                        downloads = 0
+                        while True:
+                            downloads = self.__get_download_count(tgen_client_writable.filename)
+                            if downloads == 1:
+                               logging.info("Onionperf has downloaded a 5M file in oneshot mode, and will now shut down.")
+                               break
+                        else:
+                            continue
+                        break
 
                     if self.__is_alive():
                         logging.info("All helper processes seem to be alive :)")
@@ -310,6 +322,8 @@ class Measurement(object):
         if socks_port is None:
             model.ListenModel(tgen_port="{0}".format(tgen_port)).dump_to_file(tgen_confpath)
             logging.info("TGen server running at 0.0.0.0:{0}".format(tgen_port))
+        elif self.oneshot:
+            model.OneshotModel(tgen_port="{0}".format(tgen_port), tgen_servers=server_urls, socksproxy="127.0.0.1:{0}".format(socks_port)).dump_to_file(tgen_confpath)
         else:
             model.TorperfModel(tgen_port="{0}".format(tgen_port), tgen_servers=server_urls, socksproxy="127.0.0.1:{0}".format(socks_port)).dump_to_file(tgen_confpath)
 
diff --git a/onionperf/onionperf b/onionperf/onionperf
index cba93ca..f349942 100755
--- a/onionperf/onionperf
+++ b/onionperf/onionperf
@@ -145,6 +145,11 @@ def main():
         action="store", dest="tgenpath",
         default=util.which("tgen"))
 
+    measure_parser.add_argument('--oneshot',
+        help="""Enables oneshot mode, onionperf closes on successfully downloading a file""",
+        action="store_true", dest="oneshot",
+        default=False)
+
     measure_parser.add_argument('--tgen-connect-ip',
         help="""the TGen client connect IP address ADDR, or 0.0.0.0 to do an external IP lookup; must be Internet-accessible for non-onion downloads to work""",
         metavar="ADDR", type=type_str_ip_in,
@@ -163,12 +168,12 @@ def main():
         action="store", dest="tgenconnectport",
         default=8080)
 
-    measure_parser.add_argument('-i', '--no-inet',
+    measure_parser.add_argument('-o', '--onion-only',
         help="""disable measuring download times over Tor back to {0}""".format(hostname),
         action="store_false", dest="do_inet",
         default=True)
 
-    measure_parser.add_argument('-o', '--no-onion',
+    measure_parser.add_argument('-i', '--inet-only',
         help="""disable measuring download times over Tor to an ephemeral onion service""",
         action="store_false", dest="do_onion",
         default=True)
@@ -363,7 +368,7 @@ def measure(args):
         server_tor_ctl_port = util.get_random_free_port()
         server_tor_socks_port = util.get_random_free_port()
 
-        meas = Measurement(args.torpath, args.tgenpath, args.prefix, args.nickname)
+        meas = Measurement(args.torpath, args.tgenpath, args.prefix, args.nickname, args.oneshot)
         meas.run(do_onion=args.do_onion, do_inet=args.do_inet,
              client_tgen_listen_port=client_tgen_port, client_tgen_connect_ip=client_connect_ip, client_tgen_connect_port=client_connect_port, client_tor_ctl_port=client_tor_ctl_port, client_tor_socks_port=client_tor_socks_port,
              server_tgen_listen_port=server_tgen_port, server_tor_ctl_port=server_tor_ctl_port, server_tor_socks_port=server_tor_socks_port)