src/allmydata/control.py

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