You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
mitogen/econtext.py

758 lines
21 KiB
Python

#!/usr/bin/env python2.5
'''
Python External Execution Contexts.
'''
import atexit
import cPickle
import cStringIO
import commands
import getpass
import hmac
import imp
import inspect
import os
import select
import sha
import signal
import socket
import struct
import sys
import syslog
import textwrap
import threading
import traceback
import types
import zlib
#
# Module-level data.
#
GET_MODULE = 0L
CALL_FUNCTION = 1L
DEBUG = True
#
# Exceptions.
#
class ContextError(Exception):
'Raised when a problem occurs with a context.'
def __init__(self, fmt, *args):
Exception.__init__(self, fmt % args)
class ChannelError(ContextError):
'Raised when a channel dies or has been closed.'
class StreamError(ContextError):
'Raised when a stream cannot be established.'
class CorruptMessageError(StreamError):
'Raised when a corrupt message is received on a stream.'
class TimeoutError(StreamError):
'Raised when a timeout occurs on a stream.'
#
# Helpers.
#
def Log(fmt, *args):
if DEBUG:
sys.stderr.write('%d (%d): %s\n' % (os.getpid(), os.getppid(),
(fmt%args).replace('econtext.', '')))
def CreateChild(*args):
'''
Create a child process whos stdin/stdout is connected to a socket.
Returns:
pid, socket
'''
sock1, sock2 = socket.socketpair()
pid = os.fork()
if not pid:
for pair in ((0, sock1), (1, sock2)):
os.dup2(sock2.fileno(), pair[0])
os.close(pair[1].fileno())
os.execvp(args[0], args)
raise SystemExit
return pid, sock1
class PartialFunction(object):
def __init__(self, fn, *partial_args):
self.fn = fn
self.partial_args = partial_args
def __call__(self, *args, **kwargs):
return self.fn(*(self.partial_args+args), **kwargs)
def __repr__(self):
return 'PartialFunction(%r, *%r)' % (self.fn, self.partial_args)
class Channel(object):
def __init__(self, stream, handle):
self._stream = stream
self._handle = handle
self._wake_event = threading.Event()
self._queue_lock = threading.Lock()
self._queue = []
self._stream.AddHandleCB(self._InternalReceive, handle)
def _InternalReceive(self, killed, data):
'''
Callback from the stream object; appends a tuple of
(killed-or-closed, data) to the internal queue and wakes the internal
event.
Args:
# Has the Stream object lost its connection?
killed: bool
data: (
# Has the remote Channel had Close() called?
bool,
# The object passed to the remote Send()
object
)
'''
Log('%r._InternalReceive(%r, %r)', self, killed, data)
self._queue_lock.acquire()
try:
self._queue.append((killed or data[0], killed or data[1]))
self._wake_event.set()
finally:
self._queue_lock.release()
def Close(self):
'''
Indicate this channel is closed to the remote side.
'''
Log('%r.Close()', self)
self._stream.Enqueue(handle, (True, None))
def Send(self, data):
'''
Send the given object to the remote side.
'''
Log('%r.Send(%r)', self, data)
self._stream.Enqueue(handle, (False, data))
def Receive(self, timeout=None):
'''
Receive the next object to arrive on this channel, or return if the
optional timeout is reached.
Args:
timeout: float
Returns:
object
'''
Log('%r.Receive(%r)', self, timeout)
if not self._queue:
self._wake_event.wait(timeout)
if not self._wake_event.isSet():
return
self._queue_lock.acquire()
try:
self._wake_event.clear()
Log('%r.Receive() queue is %r', self, self._queue)
closed, data = self._queue.pop(0)
Log('%r.Receive() got closed=%r, data=%r', self, closed, data)
if closed:
raise ChannelError('Channel is closed.')
return data
finally:
self._queue_lock.release()
def __iter__(self):
'''
Return an iterator that yields objects arriving on this channel, until the
channel dies or is closed.
'''
while True:
try:
yield self.Receive()
except ChannelError:
return
def __repr__(self):
return 'econtext.Channel(%r, %r)' % (self._stream, self._handle)
class SlaveModuleImporter(object):
'''
This implements the import protocol described in PEP 302. It works like so:
- Python asks it if it can import a module.
- It asks Python (via imp module) if it can import the module.
- If Python says yes, it says no.
- If Python says no, it asks the parent context for the module.
- If the module isn't returned by the parent, asplode, otherwise ask Python
to load the returned module.
This roundabout crap is necessary because the built-in importer is tried only
after custom hooks are. A class method is provided for the parent context to
satisfy the module request; it will only return modules that have been loaded
in the parent context.
'''
def __init__(self, context):
self._context = context
def find_module(self, fullname, path=None):
if not imp.find_module(fullname):
return self
def load_module(self, fullname):
kind, data = self._context.EnqueueAwaitReply(GET_MODULE, fullname)
def GetModule(cls, killed, fullname):
Log('%r.GetModule(%r, %r)', cls, killed, fullname)
if killed:
return
if fullname in sys.modules:
pass
GetModule = classmethod(GetModule)
#
# Stream implementations.
#
class Stream(object):
def __init__(self, context):
self._context = context
self._input_buf = self._output_buf = ''
self._input_buf_lock = threading.Lock()
self._output_buf_lock = threading.Lock()
self._rhmac = hmac.new(context.key, digestmod=sha.new)
self._whmac = self._rhmac.copy()
self._last_handle = 0
self._handle_map = {}
self._handle_lock = threading.Lock()
self._func_refs = {}
self._func_ref_lock = threading.Lock()
self._pickler_file = cStringIO.StringIO()
self._pickler = cPickle.Pickler(self._pickler_file)
self._pickler.persistent_id = self._CheckFunctionPerID
self._unpickler_file = cStringIO.StringIO()
self._unpickler = cPickle.Unpickler(self._unpickler_file)
self._unpickler.persistent_load = self._LoadFunctionFromPerID
def Pickle(self, obj):
self._pickler.dump(obj)
data = self._pickler_file.getvalue()
self._pickler_file.seek(0)
self._pickler_file.truncate(0)
return data
def Unpickle(self, data):
Log('%r.Unpickle(%r)', self, data)
self._unpickler_file.write(data)
self._unpickler_file.seek(0)
data = self._unpickler.load()
self._unpickler_file.seek(0)
self._unpickler_file.truncate(0)
return data
def _CheckFunctionPerID(self, obj):
'''
Please see the cPickle documentation. Given an object, return None
indicating normal pickle processing or a string 'persistent ID'.
Args:
obj: object
Returns:
str
'''
if isinstance(obj, (types.FunctionType, types.MethodType)):
pid = 'FUNC:' + repr(obj)
self._func_refs[per_id] = obj
return pid
def _LoadFunctionFromPerID(self, pid):
'''
Please see the cPickle documentation. Given a string created by
_CheckFunctionPerID, turn it into an object again.
Args:
pid: str
Returns:
object
'''
if not pid.startswith('FUNC:'):
raise CorruptMessageError('unrecognized persistent ID received: %r', pid)
return PartialFunction(self._CallPersistentWhatsit, pid)
def AllocHandle(self):
'''
Allocate a unique handle for this stream.
Returns:
long
'''
self._handle_lock.acquire()
try:
self._last_handle += 1L
finally:
self._handle_lock.release()
return self._last_handle
def AddHandleCB(self, fn, handle, persist=True):
'''
Arrange to invoke the given function for all messages tagged with the given
handle. By default, process one message and discard this arrangement.
Args:
fn: callable
handle: long
persist: bool
'''
Log('%r.AddHandleCB(%r, %r, persist=%r)', self, fn, handle, persist)
self._handle_lock.acquire()
try:
self._handle_map[handle] = persist, fn
finally:
self._handle_lock.release()
def Receive(self):
'''
Handle the next complete message on the stream. Raise CorruptMessageError
or IOError on failure.
'''
Log('%r.Receive()', self)
self._input_buf += os.read(self._rfd, 4096)
if len(self._input_buf) < 24:
return
msg_mac = self._input_buf[:20]
msg_len = struct.unpack('>L', self._input_buf[20:24])[0]
if len(self._input_buf) < msg_len-24:
return
self._rhmac.update(self._input_buf[20:msg_len+24])
expected_mac = self._rhmac.digest()
if msg_mac != expected_mac:
raise CorruptMessageError('%r got invalid MAC: expected %r, got %r',
self, msg_mac.encode('hex'),
expected_mac.encode('hex'))
try:
handle, data = self.Unpickle(self._input_buf[24:msg_len+24])
self._input_buf = self._input_buf[msg_len+4:]
handle = long(handle)
Log('%r.Receive(): decoded handle=%r; data=%r', self, handle, data)
persist, fn = self._handle_map[handle]
if not persist:
del self._handle_map[handle]
except KeyError, ex:
raise CorruptMessageError('%r got invalid handle: %r', self, handle)
except (TypeError, ValueError), ex:
raise CorruptMessageError('%r got invalid message: %s', self, ex)
fn(False, data)
def Transmit(self):
'''
Transmit pending messages. Raises IOError on failure. Return value
indicates whether there is still data buffered.
Returns:
bool
'''
Log('%r.Transmit()', self)
written = os.write(self._fd, self._output_buf[:4096])
self._output_buf = self._output_buf[written:]
return bool(self._output_buf)
def Enqueue(self, handle, data):
Log('%r.Enqueue(%r, %r)', self, handle, data)
self._output_buf_lock.acquire()
try:
encoded = self.Pickle((handle, data))
msg = struct.pack('>L', len(encoded)) + encoded
self._whmac.update(msg)
self._output_buf += self._whmac.digest() + msg
finally:
self._output_buf_lock.release()
self._context.broker.Register(self._context)
def Disconnect(self):
'''
Called to handle disconnects.
'''
Log('%r.Disconnect()', self)
try:
os.close(self._fd)
except OSError, e:
Log('WARNING: %s', e)
for handle, (persist, fn) in self._handle_map.iteritems():
Log('%r.Disconnect(): killing stale callback handle=%r; fn=%r',
self, handle, fn)
fn(True, None)
@classmethod
def Accept(cls, broker, sock):
context = Context(broker)
stream = cls(context)
context.SetStream(stream)
broker.Register(context)
def Connect(self):
'''
Connect to a Broker at the address given in the Context instance.
'''
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._fd = sock.fileno()
sock.connect(self._context.parent_addr)
self.Enqueue(0, self._context.name)
def fileno(self):
return self._fd
def __repr__(self):
return 'econtext.%s(<context=%r>)' %\
(self.__class__.__name__, self._context)
class LocalStream(Stream):
'''
Base for streams capable of starting new slaves.
'''
python_path = property(
lambda self: getattr(self, '_python_path', sys.executable),
lambda self, path: setattr(self, '_python_path', path),
doc='The path to the remote Python interpreter.')
def __init__(self, context):
super(LocalStream, self).__init__(context)
self._permitted_modules = {}
self._unpickler.find_global = self._FindGlobal
self.AddHandleCB(SlaveModuleImporter.GetModule, handle=GET_MODULE)
def _FindGlobal(self, module_name, class_name):
'''
See the cPickle documentation: given a module and class name, determine
whether class referred to is safe for unpickling.
Args:
module_name: str
class_name: str
Returns:
classobj or type
'''
if module_name not in self._permitted_modules:
raise StreamError('context %r attempted to unpickle %r in module %r',
self._context, class_name, module_name)
return getattr(sys.modules[module_name], class_name)
def AllowModule(self, module_name):
'''
Add the given module to the list of permitted modules.
Args:
module_name: str
'''
self._permitted_modules.add(module_name)
# Hexed and passed to 'python -c'. It forks, dups 0->100, creates a pipe,
# then execs a new interpreter with a custom argv. CONTEXT_NAME is replaced
# with the context name. Optimized for source size.
def _FirstStage():
import os,sys,zlib
R,W=os.pipe()
if os.fork():
os.dup2(0,100)
os.dup2(R,0)
os.close(R)
os.close(W)
os.execv(sys.executable,(CONTEXT_NAME,))
else:
os.fdopen(W,'wb',0).write(zlib.decompress(sys.stdin.read(input())))
print 'OK'
sys.exit(0)
def GetBootCommand(self):
source = inspect.getsource(self._FirstStage)
source = textwrap.dedent('\n'.join(source.strip().split('\n')[1:]))
source = source.replace(' ', '\t')
source = source.replace('CONTEXT_NAME', repr(self._context.name))
return [ self.python_path, '-c',
'exec "%s".decode("hex")' % (source.encode('hex'),) ]
def __repr__(self):
return '%s(%s)' % (self.__class__.__name__, self._context)
def Connect(self):
Log('%r.Connect()', self)
pid, sock = CreateChild(*self.GetBootCommand())
self._fd = sock.fileno()
Log('%r.Connect(): chlid process stdin/stdout=%r', self, self._fd)
source = inspect.getsource(sys.modules[__name__])
source += '\nExternalContextMain(%r, %r, %r)\n' %\
(self._context.name, self._context.broker._listen_addr,
self._context.key)
compressed = zlib.compress(source)
preamble = str(len(compressed)) + '\n' + compressed
sock.sendall(preamble)
assert os.read(self._fd, 3) == 'OK\n'
class SSHStream(LocalStream):
ssh_path = property(
lambda self: getattr(self, '_ssh_path', 'ssh'),
lambda self, path: setattr(self, '_ssh_path', path),
doc='The path to the SSH binary.')
def GetBootCommand(self):
bits = [self.ssh_path]
if self._context.username:
bits += ['-l', self._context.username]
bits.append(self._context.hostname)
return bits + map(commands.mkarg, super(SSHStream, self).GetBootCommand())
class Context(object):
'''
Represents a remote context regardless of connection method.
'''
def __init__(self, broker, name=None, hostname=None, username=None, key=None,
parent_addr=None):
self.broker = broker
self.name = name
self.hostname = hostname
self.username = username
self.parent_addr = parent_addr
if key:
self.key = key
else:
self.key = file('/dev/urandom', 'rb').read(16).encode('hex')
def GetStream(self):
return self._stream
def SetStream(self, stream):
self._stream = stream
return stream
def EnqueueAwaitReply(self, handle, deadline, data):
'''
Send a message to the given handle and wait for a response with an optional
timeout. The message contains (reply_handle, data), where reply_handle is
the handle on which this function expects its reply.
'''
Log('%r.EnqueueAwaitReply(%r, %r, %r)', self, handle, deadline, data)
reply_handle = self._stream.AllocHandle()
reply_event = threading.Event()
container = []
def _Receive(killed, data):
Log('%r._Receive(%r, %r)', self, killed, data)
container.extend([killed, data])
reply_event.set()
self._stream.AddHandleCB(_Receive, reply_handle, persist=False)
self._stream.Enqueue(CALL_FUNCTION, (False, (reply_handle,) + data))
reply_event.wait(deadline)
if not reply_event.isSet():
self.Disconnect()
raise TimeoutError('deadline exceeded.')
killed, data = container
if killed:
raise StreamError('lost connection during call.')
Log('%r._EnqueueAwaitReply(): got reply: %r', self, data)
return data
def CallWithDeadline(self, fn, deadline, *args, **kwargs):
Log('%r.CallWithDeadline(%r, %r, *%r, **%r)', self, fn, deadline, args,
kwargs)
call = (fn.__module__, fn.__name__, args, kwargs)
success, result = self.EnqueueAwaitReply(CALL_FUNCTION, deadline, call)
if success:
return result
else:
exc_obj, traceback = result
exc_obj.real_traceback = traceback
raise exc_obj
def Call(self, fn, *args, **kwargs):
return self.CallWithDeadline(fn, None, *args, **kwargs)
def __repr__(self):
bits = map(repr, filter(None, [self.name, self.hostname, self.username]))
return 'Context(%s)' % ', '.join(bits)
class Broker(object):
'''
Context broker: this is responsible for keeping track of contexts, any
stream that is associated with them, and for I/O multiplexing.
'''
def __init__(self):
self._dead = False
self._poller = select.poll()
self._poller_fd_map = {}
self._poller_lock = threading.Lock()
self._contexts = {}
self._wake_rfd, self._wake_wfd = os.pipe()
self._poller.register(self._wake_rfd)
self._listen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._listen_sock.bind(('0.0.0.0', 0)) # plz 2 allocate 4 me kthx.
self._listen_sock.listen(5)
self._listen_addr = self._listen_sock.getsockname()
self._poller.register(self._listen_sock)
self._thread = threading.Thread(target=self.Loop, name='Broker')
self._thread.setDaemon(True)
self._thread.start()
def Register(self, context):
'''
Put a context under control of this broker.
'''
Log('%r.Register(%r)', self, context)
self._poller_lock.acquire()
os.write(self._wake_wfd, ' ')
try:
self._contexts[context.name] = context
self._poller.register(context.GetStream())
self._poller_fd_map[context.GetStream().fileno()] = context.GetStream()
finally:
self._poller_lock.release()
return context
def GetLocal(self, name):
'''
Return the named local context, or create it if it doesn't exist.
Args:
name: 'my-local-context'
Returns:
econtext.Context
'''
context = Context(self, name)
context.SetStream(LocalStream(context)).Connect()
return self.Register(context)
def GetRemote(self, hostname, username, name=None):
'''
Return the named remote context, or create it if it doesn't exist.
'''
if name is None:
name = 'econtext[%s@%s:%d]' %\
(username, os.getenv('HOSTNAME'), os.getpid())
context = Context(self, name, hostname, username)
context.SetStream(SSHStream(context)).Connect()
return self.Register(context)
def Loop(self):
'''
Handle stream events until Finalize() is called.
'''
while not self._dead:
Log('%r.Loop()', self)
self._poller_lock.acquire()
self._poller_lock.release()
for fd, event in self._poller.poll():
if fd == self._wake_rfd:
Log('%r: got event on wake_rfd=%d.', self, self._wake_rfd)
os.read(self._wake_rfd, 1)
break
elif fd == self._listen_sock.fileno():
Stream.Accept(self, self._listen_sock.accept())
continue
obj = self._poller_fd_map[fd]
if event & select.POLLHUP:
Log('%r: POLLHUP on %r', self, obj)
obj.Disconnect()
elif event & select.POLLIN:
Log('%r: POLLIN on %r', self, obj)
obj.Receive()
elif event & select.POLLOUT:
Log('%r: POLLOUT on %r', self, obj)
if not obj.Transmit(): # If no output buffered, unset POLLOUT.
self._poller.unregister(obj)
self._poller.register(obj, select.POLLIN)
elif event & select.POLLNVAL:
Log('%r: POLLNVAL for %r', self, obj)
obj.Disconnect()
self._poller.unregister(obj)
def Finalize(self):
'''
Tell all active streams to disconnect.
'''
self._dead = True
self._poller_lock.acquire()
try:
for name, context in self._contexts.iteritems():
context.GetStream().Disconnect()
finally:
self._poller_lock.release()
def __repr__(self):
return 'econtext.Broker(<contexts=%s>)' % (self._contexts.keys(),)
def ExternalContextMain(context_name, parent_addr, key):
syslog.openlog('%s:%s' % (getpass.getuser(), context_name), syslog.LOG_PID)
syslog.syslog('initializing (parent=%s)' % (os.getenv('SSH_CLIENT'),))
Log('ExternalContextMain(%r, %r, %r)', context_name, parent_addr, key)
os.wait() # Reap the first stage.
os.dup2(100, 0)
os.close(100)
broker = Broker()
context = Context(broker, 'parent', parent_addr=parent_addr, key=key)
stream = context.SetStream(Stream(context))
stream.Connect()
broker.Register(context)
for call_info in Channel(stream, CALL_FUNCTION):
Log('ExternalContextMain(): CALL_FUNCTION %r', call_info)
reply_handle, mod_name, func_name, args, kwargs = call_info
fn = getattr(__import__(mod_name), func_name)
try:
stream.Enqueue(reply_handle, (True, fn(*args, **kwargs)))
except Exception, e:
stram.Enqueue(reply_handle, (False, (e, traceback.extract_stack())))