borg/darc/remote.py

206 lines
6.5 KiB
Python
Raw Normal View History

2011-08-03 15:54:03 +00:00
from __future__ import with_statement
2010-11-15 21:18:47 +00:00
import fcntl
import msgpack
import os
import select
from subprocess import Popen, PIPE
2010-11-15 21:18:47 +00:00
import sys
import getpass
from .store import Store
from .lrucache import LRUCache
2010-11-15 21:18:47 +00:00
BUFSIZE = 10 * 1024 * 1024
2011-07-05 19:29:15 +00:00
2010-11-15 21:18:47 +00:00
class StoreServer(object):
def __init__(self):
self.store = None
def serve(self):
# Make stdin non-blocking
fl = fcntl.fcntl(sys.stdin.fileno(), fcntl.F_GETFL)
fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, fl | os.O_NONBLOCK)
# Make stdout blocking
fl = fcntl.fcntl(sys.stdout.fileno(), fcntl.F_GETFL)
fcntl.fcntl(sys.stdout.fileno(), fcntl.F_SETFL, fl & ~os.O_NONBLOCK)
2010-11-15 21:18:47 +00:00
unpacker = msgpack.Unpacker()
while True:
r, w, es = select.select([sys.stdin], [], [], 10)
if r:
data = os.read(sys.stdin.fileno(), BUFSIZE)
if not data:
return
unpacker.feed(data)
for type, msgid, method, args in unpacker:
try:
2010-11-17 21:40:39 +00:00
try:
f = getattr(self, method)
except AttributeError:
f = getattr(self.store, method)
res = f(*args)
2010-11-15 21:18:47 +00:00
except Exception, e:
sys.stdout.write(msgpack.packb((1, msgid, e.__class__.__name__, None)))
else:
2011-07-17 20:31:37 +00:00
sys.stdout.write(msgpack.packb((1, msgid, None, res)))
2010-11-15 21:18:47 +00:00
sys.stdout.flush()
if es:
return
2011-09-12 19:51:23 +00:00
def negotiate(self, versions):
return 1
2010-11-17 21:40:39 +00:00
def open(self, path, create=False):
if path.startswith('/~'):
path = path[1:]
self.store = Store(os.path.expanduser(path), create)
return self.store.id
2010-11-17 21:40:39 +00:00
2010-11-15 21:18:47 +00:00
class RemoteStore(object):
class DoesNotExist(Exception):
pass
class AlreadyExists(Exception):
pass
class RPCError(Exception):
def __init__(self, name):
self.name = name
def __init__(self, location, create=False):
self.cache = LRUCache(200)
self.to_send = ''
self.extra = {}
self.pending_cache = {}
2010-11-15 21:18:47 +00:00
self.unpacker = msgpack.Unpacker()
self.msgid = 0
self.received_msgid = 0
args = ['ssh', '-p', str(location.port), '%s@%s' % (location.user or getpass.getuser(), location.host), 'darc', 'serve']
self.p = Popen(args, bufsize=0, stdin=PIPE, stdout=PIPE)
self.stdin_fd = self.p.stdin.fileno()
self.stdout_fd = self.p.stdout.fileno()
self.r_fds = [self.stdout_fd]
self.x_fds = [self.stdin_fd, self.stdout_fd]
version = self.call('negotiate', (1,))
2011-09-12 19:51:23 +00:00
if version != 1:
raise Exception('Server insisted on using unsupported protocol version %d' % version)
self.id = self.call('open', (location.path, create))
def __del__(self):
self.p.stdin.close()
self.p.stdout.close()
self.p.wait()
def _read(self):
data = os.read(self.stdout_fd, BUFSIZE)
if not data:
raise Exception('EOF')
self.unpacker.feed(data)
to_yield = []
for type, msgid, error, res in self.unpacker:
self.received_msgid = msgid
if error:
raise self.RPCError(error)
if msgid in self.pending_cache:
args = self.pending_cache.pop(msgid)
self.cache[args] = msgid, res
else:
print 'unknown response'
for args in self.extra.pop(msgid, []):
to_yield.append(self.cache[args][1])
for res in to_yield:
yield res
def call(self, cmd, args, wait=True):
for res in self.call_multi(cmd, [args], wait=wait):
return res
def gen_request(self, cmd, argsv):
data = []
m = self.received_msgid
for args in argsv:
if not args in self.cache:
self.msgid += 1
msgid = self.msgid
self.pending_cache[msgid] = args
self.cache[args] = msgid, None
data.append(msgpack.packb((1, msgid, cmd, args)))
msgid, resp = self.cache[args]
m = max(m, msgid)
self.extra.setdefault(m, []).append(args)
return ''.join(data)
def gen_cache_requests(self, cmd, peek):
data = []
while True:
try:
args = (peek()[0],)
except StopIteration:
break
if args in self.cache:
continue
self.msgid += 1
msgid = self.msgid
self.pending_cache[msgid] = args
self.cache[args] = msgid, None
data.append(msgpack.packb((1, msgid, cmd, args)))
return ''.join(data)
def call_multi(self, cmd, argsv, wait=True, peek=None):
w_fds = [self.stdin_fd]
left = len(argsv)
data = self.gen_request(cmd, argsv)
self.to_send += data
for args in self.extra.pop(self.received_msgid, []):
left -= 1
yield self.cache[args][1]
while left:
r, w, x = select.select(self.r_fds, w_fds, self.x_fds, 1)
if x:
raise Exception('FD exception occured')
if r:
for res in self._read():
left -= 1
yield res
if w:
if not self.to_send and peek:
self.to_send = self.gen_cache_requests(cmd, peek)
if self.to_send:
n = os.write(self.stdin_fd, self.to_send)
assert n > 0
self.to_send = self.to_send[n:]
else:
w_fds = []
2011-09-05 19:20:17 +00:00
def commit(self, *args):
self.call('commit', args)
2010-11-15 21:18:47 +00:00
def rollback(self, *args):
return self.call('rollback', args)
2010-11-15 21:18:47 +00:00
def get(self, id):
2010-11-15 21:18:47 +00:00
try:
return self.call('get', (id, ))
except self.RPCError, e:
if e.name == 'DoesNotExist':
raise self.DoesNotExist
raise
def get_many(self, ids, peek=None):
return self.call_multi('get', [(id, ) for id in ids], peek=peek)
def put(self, id, data, wait=True):
2010-11-15 21:18:47 +00:00
try:
return self.call('put', (id, data), wait=wait)
2010-11-15 21:18:47 +00:00
except self.RPCError, e:
if e.name == 'AlreadyExists':
raise self.AlreadyExists
def delete(self, id, wait=True):
return self.call('delete', (id, ), wait=wait)