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_upload_from_file_to_uri(self, filename, convergence):
  64         uploader = self.parent.getServiceNamed("uploader")
  65         u = upload.FileName(filename, convergence=convergence)
  66         d = uploader.upload(u)
  67         d.addCallback(lambda results: results.uri)
  68         return d
  69
  70     def remote_download_from_uri_to_file(self, uri, filename):
  71         filenode = self.parent.create_node_from_uri(uri, name=filename)
  72         if not IFileNode.providedBy(filenode):
  73             raise AssertionError("The URI does not reference a file.")
  74         c = FileWritingConsumer(filename)
  75         d = filenode.read(c)
  76         d.addCallback(lambda res: filename)
  77         return d
  78
  79     def remote_speed_test(self, count, size, mutable):
  80         assert size > 8
  81         log.msg("speed_test: count=%d, size=%d, mutable=%s" % (count, size,
  82                                                                mutable))
  83         st = SpeedTest(self.parent, count, size, mutable)
  84         return st.run()
  85
  86     def remote_get_memory_usage(self):
  87         return get_memory_usage()
  88
  89     def remote_measure_peer_response_time(self):
  90         # I'd like to average together several pings, but I don't want this
  91         # phase to take more than 10 seconds. Expect worst-case latency to be
  92         # 300ms.
  93         results = {}
  94         sb = self.parent.get_storage_broker()
  95         everyone = sb.get_connected_servers()
  96         num_pings = int(mathutil.div_ceil(10, (len(everyone) * 0.3)))
  97         everyone = list(everyone) * num_pings
  98         d = self._do_one_ping(None, everyone, results)
  99         return d
 100     def _do_one_ping(self, res, everyone_left, results):
 101         if not everyone_left:
 102             return results
 103         server = everyone_left.pop(0)
 104         server_name = server.get_longname()
 105         connection = server.get_rref()
 106         start = time.time()
 107         d = connection.callRemote("get_buckets", "\x00"*16)
 108         def _done(ignored):
 109             stop = time.time()
 110             elapsed = stop - start
 111             if server_name in results:
 112                 results[server_name].append(elapsed)
 113             else:
 114                 results[server_name] = [elapsed]
 115         d.addCallback(_done)
 116         d.addCallback(self._do_one_ping, everyone_left, results)
 117         def _average(res):
 118             averaged = {}
 119             for server_name,times in results.iteritems():
 120                 averaged[server_name] = sum(times) / len(times)
 121             return averaged
 122         d.addCallback(_average)
 123         return d
 124
 125 class SpeedTest:
 126     def __init__(self, parent, count, size, mutable):
 127         self.parent = parent
 128         self.count = count
 129         self.size = size
 130         self.mutable_mode = mutable
 131         self.uris = {}
 132         self.basedir = os.path.join(self.parent.basedir, "_speed_test_data")
 133
 134     def run(self):
 135         self.create_data()
 136         d = self.do_upload()
 137         d.addCallback(lambda res: self.do_download())
 138         d.addBoth(self.do_cleanup)
 139         d.addCallback(lambda res: (self.upload_time, self.download_time))
 140         return d
 141
 142     def create_data(self):
 143         fileutil.make_dirs(self.basedir)
 144         for i in range(self.count):
 145             s = self.size
 146             fn = os.path.join(self.basedir, str(i))
 147             if os.path.exists(fn):
 148                 os.unlink(fn)
 149             f = open(fn, "w")
 150             f.write(os.urandom(8))
 151             s -= 8
 152             while s > 0:
 153                 chunk = min(s, 4096)
 154                 f.write("\x00" * chunk)
 155                 s -= chunk
 156             f.close()
 157
 158     def do_upload(self):
 159         d = defer.succeed(None)
 160         def _create_slot(res):
 161             d1 = self.parent.create_mutable_file("")
 162             def _created(n):
 163                 self._n = n
 164             d1.addCallback(_created)
 165             return d1
 166         if self.mutable_mode == "upload":
 167             d.addCallback(_create_slot)
 168         def _start(res):
 169             self._start = time.time()
 170         d.addCallback(_start)
 171
 172         def _record_uri(uri, i):
 173             self.uris[i] = uri
 174         def _upload_one_file(ignored, i):
 175             if i >= self.count:
 176                 return
 177             fn = os.path.join(self.basedir, str(i))
 178             if self.mutable_mode == "create":
 179                 data = open(fn,"rb").read()
 180                 d1 = self.parent.create_mutable_file(data)
 181                 d1.addCallback(lambda n: n.get_uri())
 182             elif self.mutable_mode == "upload":
 183                 data = open(fn,"rb").read()
 184                 d1 = self._n.overwrite(MutableData(data))
 185                 d1.addCallback(lambda res: self._n.get_uri())
 186             else:
 187                 up = upload.FileName(fn, convergence=None)
 188                 d1 = self.parent.upload(up)
 189                 d1.addCallback(lambda results: results.uri)
 190             d1.addCallback(_record_uri, i)
 191             d1.addCallback(_upload_one_file, i+1)
 192             return d1
 193         d.addCallback(_upload_one_file, 0)
 194         def _upload_done(ignored):
 195             stop = time.time()
 196             self.upload_time = stop - self._start
 197         d.addCallback(_upload_done)
 198         return d
 199
 200     def do_download(self):
 201         start = time.time()
 202         d = defer.succeed(None)
 203         def _download_one_file(ignored, i):
 204             if i >= self.count:
 205                 return
 206             n = self.parent.create_node_from_uri(self.uris[i])
 207             if not IFileNode.providedBy(n):
 208                 raise AssertionError("The URI does not reference a file.")
 209             if n.is_mutable():
 210                 d1 = n.download_best_version()
 211             else:
 212                 d1 = n.read(DiscardingConsumer())
 213             d1.addCallback(_download_one_file, i+1)
 214             return d1
 215         d.addCallback(_download_one_file, 0)
 216         def _download_done(ignored):
 217             stop = time.time()
 218             self.download_time = stop - start
 219         d.addCallback(_download_done)
 220         return d
 221
 222     def do_cleanup(self, res):
 223         for i in range(self.count):
 224             fn = os.path.join(self.basedir, str(i))
 225             os.unlink(fn)
 226         return res
 227
 228 class DiscardingConsumer:
 229     implements(IConsumer)
 230     def __init__(self):
 231         self.done = False
 232     def registerProducer(self, p, streaming):
 233         if streaming:
 234             p.resumeProducing()
 235         else:
 236             while not self.done:
 237                 p.resumeProducing()
 238     def write(self, data):
 239         pass
 240     def unregisterProducer(self):
 241         self.done = True