src/allmydata/control.py

   1
   2 import os, time
   3 from zope.interface import implements
   4 from twisted.application import service
   5 from twisted.internet import defer
   6 from twisted.internet.interfaces import IConsumer
   7 from foolscap.api import Referenceable
   8 from allmydata.interfaces import RIControlClient, IFileNode
   9 from allmydata.util import fileutil, mathutil
  10 from allmydata.immutable import upload
  11 from allmydata.mutable.publish import MutableData
  12 from twisted.python import log
  13
  14 def get_memory_usage():
  15     # this is obviously linux-specific
  16     stat_names = ("VmPeak",
  17                   "VmSize",
  18                   #"VmHWM",
  19                   "VmData")
  20     stats = {}
  21     try:
  22         for line in open("/proc/self/status", "r").readlines():
  23             name, right = line.split(":",2)
  24             if name in stat_names:
  25                 assert right.endswith(" kB\n")
  26                 right = right[:-4]
  27                 stats[name] = int(right) * 1024
  28     except:
  29         # Probably not on (a compatible version of) Linux
  30         stats['VmSize'] = 0
  31         stats['VmPeak'] = 0
  32     return stats
  33
  34 def log_memory_usage(where=""):
  35     stats = get_memory_usage()
  36     log.msg("VmSize: %9d  VmPeak: %9d  %s" % (stats["VmSize"],
  37                                               stats["VmPeak"],
  38                                               where))
  39
  40 class FileWritingConsumer:
  41     implements(IConsumer)
  42     def __init__(self, filename):
  43         self.done = False
  44         self.f = open(filename, "wb")
  45     def registerProducer(self, p, streaming):
  46         if streaming:
  47             p.resumeProducing()
  48         else:
  49             while not self.done:
  50                 p.resumeProducing()
  51     def write(self, data):
  52         self.f.write(data)
  53     def unregisterProducer(self):
  54         self.done = True
  55         self.f.close()
  56
  57 class ControlServer(Referenceable, service.Service):
  58     implements(RIControlClient)
  59
  60     def remote_wait_for_client_connections(self, num_clients):
  61         return self.parent.debug_wait_for_client_connections(num_clients)
  62
  63     def remote_speed_test(self, count, size, mutable):
  64         assert size > 8
  65         log.msg("speed_test: count=%d, size=%d, mutable=%s" % (count, size,
  66                                                                mutable))
  67         st = SpeedTest(self.parent, count, size, mutable)
  68         return st.run()
  69
  70     def remote_get_memory_usage(self):
  71         return get_memory_usage()
  72
  73     def remote_measure_peer_response_time(self):
  74         # I'd like to average together several pings, but I don't want this
  75         # phase to take more than 10 seconds. Expect worst-case latency to be
  76         # 300ms.
  77         results = {}
  78         sb = self.parent.get_storage_broker()
  79         everyone = sb.get_connected_servers()
  80         num_pings = int(mathutil.div_ceil(10, (len(everyone) * 0.3)))
  81         everyone = list(everyone) * num_pings
  82         d = self._do_one_ping(None, everyone, results)
  83         return d
  84     def _do_one_ping(self, res, everyone_left, results):
  85         if not everyone_left:
  86             return results
  87         server = everyone_left.pop(0)
  88         server_name = server.get_longname()
  89         connection = server.get_rref()
  90         start = time.time()
  91         d = connection.callRemote("get_buckets", "\x00"*16)
  92         def _done(ignored):
  93             stop = time.time()
  94             elapsed = stop - start
  95             if server_name in results:
  96                 results[server_name].append(elapsed)
  97             else:
  98                 results[server_name] = [elapsed]
  99         d.addCallback(_done)
 100         d.addCallback(self._do_one_ping, everyone_left, results)
 101         def _average(res):
 102             averaged = {}
 103             for server_name,times in results.iteritems():
 104                 averaged[server_name] = sum(times) / len(times)
 105             return averaged
 106         d.addCallback(_average)
 107         return d
 108
 109 class SpeedTest:
 110     def __init__(self, parent, count, size, mutable):
 111         self.parent = parent
 112         self.count = count
 113         self.size = size
 114         self.mutable_mode = mutable
 115         self.uris = {}
 116         self.basedir = os.path.join(self.parent.basedir, "_speed_test_data")
 117
 118     def run(self):
 119         self.create_data()
 120         d = self.do_upload()
 121         d.addCallback(lambda res: self.do_download())
 122         d.addBoth(self.do_cleanup)
 123         d.addCallback(lambda res: (self.upload_time, self.download_time))
 124         return d
 125
 126     def create_data(self):
 127         fileutil.make_dirs(self.basedir)
 128         for i in range(self.count):
 129             s = self.size
 130             fn = os.path.join(self.basedir, str(i))
 131             if os.path.exists(fn):
 132                 os.unlink(fn)
 133             f = open(fn, "w")
 134             f.write(os.urandom(8))
 135             s -= 8
 136             while s > 0:
 137                 chunk = min(s, 4096)
 138                 f.write("\x00" * chunk)
 139                 s -= chunk
 140             f.close()
 141
 142     def do_upload(self):
 143         d = defer.succeed(None)
 144         def _create_slot(res):
 145             d1 = self.parent.create_mutable_file("")
 146             def _created(n):
 147                 self._n = n
 148             d1.addCallback(_created)
 149             return d1
 150         if self.mutable_mode == "upload":
 151             d.addCallback(_create_slot)
 152         def _start(res):
 153             self._start = time.time()
 154         d.addCallback(_start)
 155
 156         def _record_uri(uri, i):
 157             self.uris[i] = uri
 158         def _upload_one_file(ignored, i):
 159             if i >= self.count:
 160                 return
 161             fn = os.path.join(self.basedir, str(i))
 162             if self.mutable_mode == "create":
 163                 data = open(fn,"rb").read()
 164                 d1 = self.parent.create_mutable_file(data)
 165                 d1.addCallback(lambda n: n.get_uri())
 166             elif self.mutable_mode == "upload":
 167                 data = open(fn,"rb").read()
 168                 d1 = self._n.overwrite(MutableData(data))
 169                 d1.addCallback(lambda res: self._n.get_uri())
 170             else:
 171                 up = upload.FileName(fn, convergence=None)
 172                 d1 = self.parent.upload(up)
 173                 d1.addCallback(lambda results: results.get_uri())
 174             d1.addCallback(_record_uri, i)
 175             d1.addCallback(_upload_one_file, i+1)
 176             return d1
 177         d.addCallback(_upload_one_file, 0)
 178         def _upload_done(ignored):
 179             stop = time.time()
 180             self.upload_time = stop - self._start
 181         d.addCallback(_upload_done)
 182         return d
 183
 184     def do_download(self):
 185         start = time.time()
 186         d = defer.succeed(None)
 187         def _download_one_file(ignored, i):
 188             if i >= self.count:
 189                 return
 190             n = self.parent.create_node_from_uri(self.uris[i])
 191             if not IFileNode.providedBy(n):
 192                 raise AssertionError("The URI does not reference a file.")
 193             if n.is_mutable():
 194                 d1 = n.download_best_version()
 195             else:
 196                 d1 = n.read(DiscardingConsumer())
 197             d1.addCallback(_download_one_file, i+1)
 198             return d1
 199         d.addCallback(_download_one_file, 0)
 200         def _download_done(ignored):
 201             stop = time.time()
 202             self.download_time = stop - start
 203         d.addCallback(_download_done)
 204         return d
 205
 206     def do_cleanup(self, res):
 207         for i in range(self.count):
 208             fn = os.path.join(self.basedir, str(i))
 209             os.unlink(fn)
 210         return res
 211
 212 class DiscardingConsumer:
 213     implements(IConsumer)
 214     def __init__(self):
 215         self.done = False
 216     def registerProducer(self, p, streaming):
 217         if streaming:
 218             p.resumeProducing()
 219         else:
 220             while not self.done:
 221                 p.resumeProducing()
 222     def write(self, data):
 223         pass
 224     def unregisterProducer(self):
 225         self.done = True