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