lib/bup/client.py

   1 import re, struct, errno, time, zlib
   2 from bup import git, ssh
   3 from bup.helpers import *
   4
   5 bwlimit = None
   6
   7
   8 class ClientError(Exception):
   9     pass
  10
  11
  12 def _raw_write_bwlimit(f, buf, bwcount, bwtime):
  13     if not bwlimit:
  14         f.write(buf)
  15         return (len(buf), time.time())
  16     else:
  17         # We want to write in reasonably large blocks, but not so large that
  18         # they're likely to overflow a router's queue.  So our bwlimit timing
  19         # has to be pretty granular.  Also, if it takes too long from one
  20         # transmit to the next, we can't just make up for lost time to bring
  21         # the average back up to bwlimit - that will risk overflowing the
  22         # outbound queue, which defeats the purpose.  So if we fall behind
  23         # by more than one block delay, we shouldn't ever try to catch up.
  24         for i in xrange(0,len(buf),4096):
  25             now = time.time()
  26             next = max(now, bwtime + 1.0*bwcount/bwlimit)
  27             time.sleep(next-now)
  28             sub = buf[i:i+4096]
  29             f.write(sub)
  30             bwcount = len(sub)  # might be less than 4096
  31             bwtime = next
  32         return (bwcount, bwtime)
  33
  34
  35 def parse_remote(remote):
  36     protocol = r'([a-z]+)://'
  37     host = r'(?P<sb>\[)?((?(sb)[0-9a-f:]+|[^:/]+))(?(sb)\])'
  38     port = r'(?::(\d+))?'
  39     path = r'(/.*)?'
  40     url_match = re.match(
  41             '%s(?:%s%s)?%s' % (protocol, host, port, path), remote, re.I)
  42     if url_match:
  43         if not url_match.group(1) in ('ssh', 'bup', 'file'):
  44             raise ClientError, 'unexpected protocol: %s' % url_match.group(1)
  45         return url_match.group(1,3,4,5)
  46     else:
  47         rs = remote.split(':', 1)
  48         if len(rs) == 1 or rs[0] in ('', '-'):
  49             return 'file', None, None, rs[-1]
  50         else:
  51             return 'ssh', rs[0], None, rs[1]
  52
  53
  54 class Client:
  55     def __init__(self, remote, create=False, compression_level=1):
  56         self._busy = self.conn = None
  57         self.sock = self.p = self.pout = self.pin = None
  58         self.compression_level = compression_level
  59         is_reverse = os.environ.get('BUP_SERVER_REVERSE')
  60         if is_reverse:
  61             assert(not remote)
  62             remote = '%s:' % is_reverse
  63         (self.protocol, self.host, self.port, self.dir) = parse_remote(remote)
  64         self.cachedir = git.repo('index-cache/%s'
  65                                  % re.sub(r'[^@\w]', '_',
  66                                           "%s:%s" % (self.host, self.dir)))
  67         if is_reverse:
  68             self.pout = os.fdopen(3, 'rb')
  69             self.pin = os.fdopen(4, 'wb')
  70             self.conn = Conn(self.pout, self.pin)
  71         else:
  72             if self.protocol in ('ssh', 'file'):
  73                 try:
  74                     # FIXME: ssh and file shouldn't use the same module
  75                     self.p = ssh.connect(self.host, self.port, 'server')
  76                     self.pout = self.p.stdout
  77                     self.pin = self.p.stdin
  78                     self.conn = Conn(self.pout, self.pin)
  79                 except OSError, e:
  80                     raise ClientError, 'connect: %s' % e, sys.exc_info()[2]
  81             elif self.protocol == 'bup':
  82                 self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  83                 self.sock.connect((self.host, atoi(self.port) or 1982))
  84                 self.sockw = self.sock.makefile('wb')
  85                 self.conn = DemuxConn(self.sock.fileno(), self.sockw)
  86         if self.dir:
  87             self.dir = re.sub(r'[\r\n]', ' ', self.dir)
  88             if create:
  89                 self.conn.write('init-dir %s\n' % self.dir)
  90             else:
  91                 self.conn.write('set-dir %s\n' % self.dir)
  92             self.check_ok()
  93         self.sync_indexes()
  94
  95     def __del__(self):
  96         try:
  97             self.close()
  98         except IOError, e:
  99             if e.errno == errno.EPIPE:
 100                 pass
 101             else:
 102                 raise
 103
 104     def close(self):
 105         if self.conn and not self._busy:
 106             self.conn.write('quit\n')
 107         if self.pin:
 108             self.pin.close()
 109         if self.sock and self.sockw:
 110             self.sockw.close()
 111             self.sock.shutdown(socket.SHUT_WR)
 112         if self.conn:
 113             self.conn.close()
 114         if self.pout:
 115             self.pout.close()
 116         if self.sock:
 117             self.sock.close()
 118         if self.p:
 119             self.p.wait()
 120             rv = self.p.wait()
 121             if rv:
 122                 raise ClientError('server tunnel returned exit code %d' % rv)
 123         self.conn = None
 124         self.sock = self.p = self.pin = self.pout = None
 125
 126     def check_ok(self):
 127         if self.p:
 128             rv = self.p.poll()
 129             if rv != None:
 130                 raise ClientError('server exited unexpectedly with code %r'
 131                                   % rv)
 132         try:
 133             return self.conn.check_ok()
 134         except Exception, e:
 135             raise ClientError, e, sys.exc_info()[2]
 136
 137     def check_busy(self):
 138         if self._busy:
 139             raise ClientError('already busy with command %r' % self._busy)
 140
 141     def ensure_busy(self):
 142         if not self._busy:
 143             raise ClientError('expected to be busy, but not busy?!')
 144
 145     def _not_busy(self):
 146         self._busy = None
 147
 148     def sync_indexes(self):
 149         self.check_busy()
 150         conn = self.conn
 151         mkdirp(self.cachedir)
 152         # All cached idxs are extra until proven otherwise
 153         extra = set()
 154         for f in os.listdir(self.cachedir):
 155             debug1('%s\n' % f)
 156             if f.endswith('.idx'):
 157                 extra.add(f)
 158         needed = set()
 159         conn.write('list-indexes\n')
 160         for line in linereader(conn):
 161             if not line:
 162                 break
 163             assert(line.find('/') < 0)
 164             parts = line.split(' ')
 165             idx = parts[0]
 166             if len(parts) == 2 and parts[1] == 'load' and idx not in extra:
 167                 # If the server requests that we load an idx and we don't
 168                 # already have a copy of it, it is needed
 169                 needed.add(idx)
 170             # Any idx that the server has heard of is proven not extra
 171             extra.discard(idx)
 172
 173         self.check_ok()
 174         debug1('client: removing extra indexes: %s\n' % extra)
 175         for idx in extra:
 176             os.unlink(os.path.join(self.cachedir, idx))
 177         debug1('client: server requested load of: %s\n' % needed)
 178         for idx in needed:
 179             self.sync_index(idx)
 180         git.auto_midx(self.cachedir)
 181
 182     def sync_index(self, name):
 183         #debug1('requesting %r\n' % name)
 184         self.check_busy()
 185         mkdirp(self.cachedir)
 186         fn = os.path.join(self.cachedir, name)
 187         if os.path.exists(fn):
 188             msg = "won't request existing .idx, try `bup bloom --check %s`" % fn
 189             raise ClientError(msg)
 190         self.conn.write('send-index %s\n' % name)
 191         n = struct.unpack('!I', self.conn.read(4))[0]
 192         assert(n)
 193         f = open(fn + '.tmp', 'w')
 194         count = 0
 195         progress('Receiving index from server: %d/%d\r' % (count, n))
 196         for b in chunkyreader(self.conn, n):
 197             f.write(b)
 198             count += len(b)
 199             qprogress('Receiving index from server: %d/%d\r' % (count, n))
 200         progress('Receiving index from server: %d/%d, done.\n' % (count, n))
 201         self.check_ok()
 202         f.close()
 203         os.rename(fn + '.tmp', fn)
 204
 205     def _make_objcache(self):
 206         return git.PackIdxList(self.cachedir)
 207
 208     def _suggest_packs(self):
 209         ob = self._busy
 210         if ob:
 211             assert(ob == 'receive-objects-v2')
 212             self.conn.write('\xff\xff\xff\xff')  # suspend receive-objects-v2
 213         suggested = []
 214         for line in linereader(self.conn):
 215             if not line:
 216                 break
 217             debug2('%s\n' % line)
 218             if line.startswith('index '):
 219                 idx = line[6:]
 220                 debug1('client: received index suggestion: %s\n'
 221                        % git.shorten_hash(idx))
 222                 suggested.append(idx)
 223             else:
 224                 assert(line.endswith('.idx'))
 225                 debug1('client: completed writing pack, idx: %s\n'
 226                        % git.shorten_hash(line))
 227                 suggested.append(line)
 228         self.check_ok()
 229         if ob:
 230             self._busy = None
 231         idx = None
 232         for idx in suggested:
 233             self.sync_index(idx)
 234         git.auto_midx(self.cachedir)
 235         if ob:
 236             self._busy = ob
 237             self.conn.write('%s\n' % ob)
 238         return idx
 239
 240     def new_packwriter(self):
 241         self.check_busy()
 242         def _set_busy():
 243             self._busy = 'receive-objects-v2'
 244             self.conn.write('receive-objects-v2\n')
 245         return PackWriter_Remote(self.conn,
 246                                  objcache_maker = self._make_objcache,
 247                                  suggest_packs = self._suggest_packs,
 248                                  onopen = _set_busy,
 249                                  onclose = self._not_busy,
 250                                  ensure_busy = self.ensure_busy,
 251                                  compression_level = self.compression_level)
 252
 253     def read_ref(self, refname):
 254         self.check_busy()
 255         self.conn.write('read-ref %s\n' % refname)
 256         r = self.conn.readline().strip()
 257         self.check_ok()
 258         if r:
 259             assert(len(r) == 40)   # hexified sha
 260             return r.decode('hex')
 261         else:
 262             return None   # nonexistent ref
 263
 264     def update_ref(self, refname, newval, oldval):
 265         self.check_busy()
 266         self.conn.write('update-ref %s\n%s\n%s\n'
 267                         % (refname, newval.encode('hex'),
 268                            (oldval or '').encode('hex')))
 269         self.check_ok()
 270
 271     def cat(self, id):
 272         self.check_busy()
 273         self._busy = 'cat'
 274         self.conn.write('cat %s\n' % re.sub(r'[\n\r]', '_', id))
 275         while 1:
 276             sz = struct.unpack('!I', self.conn.read(4))[0]
 277             if not sz: break
 278             yield self.conn.read(sz)
 279         e = self.check_ok()
 280         self._not_busy()
 281         if e:
 282             raise KeyError(str(e))
 283
 284
 285 class PackWriter_Remote(git.PackWriter):
 286     def __init__(self, conn, objcache_maker, suggest_packs,
 287                  onopen, onclose,
 288                  ensure_busy,
 289                  compression_level=1):
 290         git.PackWriter.__init__(self, objcache_maker)
 291         self.file = conn
 292         self.filename = 'remote socket'
 293         self.suggest_packs = suggest_packs
 294         self.onopen = onopen
 295         self.onclose = onclose
 296         self.ensure_busy = ensure_busy
 297         self._packopen = False
 298         self._bwcount = 0
 299         self._bwtime = time.time()
 300         self.compression_level = compression_level
 301
 302     def _open(self):
 303         if not self._packopen:
 304             self.onopen()
 305             self._packopen = True
 306
 307     def _end(self):
 308         if self._packopen and self.file:
 309             self.file.write('\0\0\0\0')
 310             self._packopen = False
 311             self.onclose() # Unbusy
 312             self.objcache = None
 313             return self.suggest_packs() # Returns last idx received
 314
 315     def close(self):
 316         id = self._end()
 317         self.file = None
 318         return id
 319
 320     def abort(self):
 321         raise ClientError("don't know how to abort remote pack writing")
 322
 323     def _raw_write(self, datalist, sha):
 324         assert(self.file)
 325         if not self._packopen:
 326             self._open()
 327         self.ensure_busy()
 328         data = ''.join(datalist)
 329         assert(data)
 330         assert(sha)
 331         crc = zlib.crc32(data) & 0xffffffff
 332         outbuf = ''.join((struct.pack('!I', len(data) + 20 + 4),
 333                           sha,
 334                           struct.pack('!I', crc),
 335                           data))
 336         try:
 337             (self._bwcount, self._bwtime) = _raw_write_bwlimit(
 338                     self.file, outbuf, self._bwcount, self._bwtime)
 339         except IOError, e:
 340             raise ClientError, e, sys.exc_info()[2]
 341         self.outbytes += len(data)
 342         self.count += 1
 343
 344         if self.file.has_input():
 345             self.suggest_packs()
 346             self.objcache.refresh()
 347
 348         return sha, crc