#

# Module providing the `SyncManager` class for dealing

# with shared objects

#

# multiprocessing/managers.py

#

# Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt

#

__all__ = [ 'BaseManager', 'SyncManager', 'BaseProxy', 'Token' ]

#

# Imports

#

import os

import sys

import weakref

import threading

import array

import Queue

from traceback import format_exc

from multiprocessing import Process, current_process, active_children, Pool, util, connection

from multiprocessing.process import AuthenticationString

from multiprocessing.forking import exit, Popen, assert_spawning, ForkingPickler

from multiprocessing.util import Finalize, info

try:

    from cPickle import PicklingError

except ImportError:

    from pickle import PicklingError

#

# Register some things for pickling

#

def reduce_array(a):

    return array.array, (a.typecode, a.tostring())

ForkingPickler.register(array.array, reduce_array)

view_types = [type(getattr({}, name)()) for name in ('items','keys','values')]

#

# Type for identifying shared objects

#

class Token(object):

    '''

    Type to uniquely indentify a shared object

    '''

    __slots__ = ('typeid', 'address', 'id')

    def __init__(self, typeid, address, id):

        (self.typeid, self.address, self.id) = (typeid, address, id)

    def __getstate__(self):

        return (self.typeid, self.address, self.id)

    def __setstate__(self, state):

        (self.typeid, self.address, self.id) = state

    def __repr__(self):

        return 'Token(typeid=%r, address=%r, id=%r)' % \

               (self.typeid, self.address, self.id)

#

# Function for communication with a manager's server process

#

def dispatch(c, id, methodname, args=(), kwds={}):

    '''

    Send a message to manager using connection `c` and return response

    '''

    c.send((id, methodname, args, kwds))

    kind, result = c.recv()

    if kind == '#RETURN':

        return result

    raise convert_to_error(kind, result)

def convert_to_error(kind, result):

    if kind == '#ERROR':

        return result

    elif kind == '#TRACEBACK':

        assert type(result) is str

        return  RemoteError(result)

    elif kind == '#UNSERIALIZABLE':

        assert type(result) is str

        return RemoteError('Unserializable message: %s\n' % result)

    else:

        return ValueError('Unrecognized message type')

class RemoteError(Exception):

    def __str__(self):

        return ('\n' + '-'*75 + '\n' + str(self.args[0]) + '-'*75)

#

# Functions for finding the method names of an object

#

def all_methods(obj):

    '''

    Return a list of names of methods of `obj`

    '''

    temp = []

    for name in dir(obj):

        func = getattr(obj, name)

        if hasattr(func, '__call__'):

            temp.append(name)

    return temp

def public_methods(obj):

    '''

    Return a list of names of methods of `obj` which do not start with '_'

    '''

    return [name for name in all_methods(obj) if name[0] != '_']

#

# Server which is run in a process controlled by a manager

#

class Server(object):

    '''

    Server class which runs in a process controlled by a manager object

    '''

    public = ['shutdown', 'create', 'accept_connection', 'get_methods',

              'debug_info', 'number_of_objects', 'dummy', 'incref', 'decref']

    def __init__(self, registry, address, authkey, serializer):

        assert isinstance(authkey, bytes)

        self.registry = registry

        self.authkey = AuthenticationString(authkey)

        Listener, Client = listener_client[serializer]

        # do authentication later

        self.listener = Listener(address=address, backlog=5)

        self.address = self.listener.address

        self.id_to_obj = {0: (None, ())}

        self.id_to_refcount = {}

        self.mutex = threading.RLock()

        self.stop = 0

    def serve_forever(self):

        '''

        Run the server forever

        '''

        current_process()._manager_server = self

        try:

            try:

                while 1:

                    try:

                        c = self.listener.accept()

                    except (OSError, IOError):

                        continue

                    t = threading.Thread(target=self.handle_request, args=(c,))

                    t.daemon = True

                    t.start()

            except (KeyboardInterrupt, SystemExit):

                pass

        finally:

            self.stop = 999

            self.listener.close()

    def handle_request(self, c):

        '''

        Handle a new connection

        '''

        funcname = result = request = None

        try:

            connection.deliver_challenge(c, self.authkey)

            connection.answer_challenge(c, self.authkey)

            request = c.recv()

            ignore, funcname, args, kwds = request

            assert funcname in self.public, '%r unrecognized' % funcname

            func = getattr(self, funcname)

        except Exception:

            msg = ('#TRACEBACK', format_exc())

        else:

            try:

                result = func(c, *args, **kwds)

            except Exception:

                msg = ('#TRACEBACK', format_exc())

            else:

                msg = ('#RETURN', result)

        try:

            c.send(msg)

        except Exception, e:

            try:

                c.send(('#TRACEBACK', format_exc()))

            except Exception:

                pass

            util.info('Failure to send message: %r', msg)

            util.info(' ... request was %r', request)

            util.info(' ... exception was %r', e)

        c.close()

    def serve_client(self, conn):

        '''

        Handle requests from the proxies in a particular process/thread

        '''

        util.debug('starting server thread to service %r',

                   threading.current_thread().name)

        recv = conn.recv

        send = conn.send

        id_to_obj = self.id_to_obj

        while not self.stop:

            try:

                methodname = obj = None

                request = recv()

                ident, methodname, args, kwds = request

                obj, exposed, gettypeid = id_to_obj[ident]

                if methodname not in exposed:

                    raise AttributeError(

                        'method %r of %r object is not in exposed=%r' %

                        (methodname, type(obj), exposed)

                        )

                function = getattr(obj, methodname)

                try:

                    res = function(*args, **kwds)

                except Exception, e:

                    msg = ('#ERROR', e)

                else:

                    typeid = gettypeid and gettypeid.get(methodname, None)

                    if typeid:

                        rident, rexposed = self.create(conn, typeid, res)

                        token = Token(typeid, self.address, rident)

                        msg = ('#PROXY', (rexposed, token))

                    else:

                        msg = ('#RETURN', res)

            except AttributeError:

                if methodname is None:

                    msg = ('#TRACEBACK', format_exc())

                else:

                    try:

                        fallback_func = self.fallback_mapping[methodname]

                        result = fallback_func(

                            self, conn, ident, obj, *args, **kwds

                            )

                        msg = ('#RETURN', result)

                    except Exception:

                        msg = ('#TRACEBACK', format_exc())

            except EOFError:

                util.debug('got EOF -- exiting thread serving %r',

                           threading.current_thread().name)

                sys.exit(0)

            except Exception:

                msg = ('#TRACEBACK', format_exc())

            try:

                try:

                    send(msg)

                except Exception, e:

                    send(('#UNSERIALIZABLE', repr(msg)))

            except Exception, e:

                util.info('exception in thread serving %r',

                        threading.current_thread().name)

                util.info(' ... message was %r', msg)

                util.info(' ... exception was %r', e)

                conn.close()

                sys.exit(1)

    def fallback_getvalue(self, conn, ident, obj):

        return obj

    def fallback_str(self, conn, ident, obj):

        return str(obj)

    def fallback_repr(self, conn, ident, obj):

        return repr(obj)

    fallback_mapping = {

        '__str__':fallback_str,

        '__repr__':fallback_repr,

        '#GETVALUE':fallback_getvalue

        }

    def dummy(self, c):

        pass

    def debug_info(self, c):

        '''

        Return some info --- useful to spot problems with refcounting

        '''

        self.mutex.acquire()

        try:

            result = []

            keys = self.id_to_obj.keys()

            keys.sort()

            for ident in keys:

                if ident != 0:

                    result.append('  %s:       refcount=%s\n    %s' %

                                  (ident, self.id_to_refcount[ident],

                                   str(self.id_to_obj[ident][0])[:75]))

            return '\n'.join(result)

        finally:

            self.mutex.release()

    def number_of_objects(self, c):

        '''

        Number of shared objects

        '''

        return len(self.id_to_obj) - 1      # don't count ident=0

    def shutdown(self, c):

        '''

        Shutdown this process

        '''

        try:

            try:

                util.debug('manager received shutdown message')

                c.send(('#RETURN', None))

                if sys.stdout != sys.__stdout__:

                    util.debug('resetting stdout, stderr')

                    sys.stdout = sys.__stdout__

                    sys.stderr = sys.__stderr__

                util._run_finalizers(0)

                for p in active_children():

                    util.debug('terminating a child process of manager')

                    p.terminate()

                for p in active_children():

                    util.debug('terminating a child process of manager')

                    p.join()

                util._run_finalizers()

                util.info('manager exiting with exitcode 0')

            except:

                import traceback

                traceback.print_exc()

        finally:

            exit(0)

    def create(self, c, typeid, *args, **kwds):

        '''

        Create a new shared object and return its id

        '''

        self.mutex.acquire()

        try:

            callable, exposed, method_to_typeid, proxytype = \

                      self.registry[typeid]

            if callable is None:

                assert len(args) == 1 and not kwds

                obj = args[0]

            else:

                obj = callable(*args, **kwds)

            if exposed is None:

                exposed = public_methods(obj)

            if method_to_typeid is not None:

                assert type(method_to_typeid) is dict

                exposed = list(exposed) + list(method_to_typeid)

            ident = '%x' % id(obj)  # convert to string because xmlrpclib

                                    # only has 32 bit signed integers

            util.debug('%r callable returned object with id %r', typeid, ident)

            self.id_to_obj[ident] = (obj, set(exposed), method_to_typeid)

            if ident not in self.id_to_refcount:

                self.id_to_refcount[ident] = 0

            # increment the reference count immediately, to avoid

            # this object being garbage collected before a Proxy

            # object for it can be created.  The caller of create()

            # is responsible for doing a decref once the Proxy object

            # has been created.

            self.incref(c, ident)

            return ident, tuple(exposed)

        finally:

            self.mutex.release()

    def get_methods(self, c, token):

        '''

        Return the methods of the shared object indicated by token

        '''

        return tuple(self.id_to_obj[token.id][1])

    def accept_connection(self, c, name):

        '''

        Spawn a new thread to serve this connection

        '''

        threading.current_thread().name = name

        c.send(('#RETURN', None))

        self.serve_client(c)

    def incref(self, c, ident):

        self.mutex.acquire()

        try:

            self.id_to_refcount[ident] += 1

        finally:

            self.mutex.release()

    def decref(self, c, ident):

        self.mutex.acquire()

        try:

            assert self.id_to_refcount[ident] >= 1

            self.id_to_refcount[ident] -= 1

            if self.id_to_refcount[ident] == 0:

                del self.id_to_obj[ident], self.id_to_refcount[ident]

                util.debug('disposing of obj with id %d', ident)

        finally:

            self.mutex.release()

#

# Class to represent state of a manager

#

class State(object):

    __slots__ = ['value']

    INITIAL = 0

    STARTED = 1

    SHUTDOWN = 2

#

# Mapping from serializer name to Listener and Client types

#

listener_client = {

    'pickle' : (connection.Listener, connection.Client),

    'xmlrpclib' : (connection.XmlListener, connection.XmlClient)

    }

#

# Definition of BaseManager

#

class BaseManager(object):

    '''

    Base class for managers

    '''

    _registry = {}

    _Server = Server

    def __init__(self, address=None, authkey=None, serializer='pickle'):

        if authkey is None:

            authkey = current_process().authkey

        self._address = address     # XXX not final address if eg ('', 0)

        self._authkey = AuthenticationString(authkey)

        self._state = State()

        self._state.value = State.INITIAL

        self._serializer = serializer

        self._Listener, self._Client = listener_client[serializer]

    def __reduce__(self):

        return type(self).from_address, \

               (self._address, self._authkey, self._serializer)

    def get_server(self):

        '''

        Return server object with serve_forever() method and address attribute

        '''

        assert self._state.value == State.INITIAL

        return Server(self._registry, self._address,

                      self._authkey, self._serializer)

    def connect(self):

        '''

        Connect manager object to the server process

        '''

        Listener, Client = listener_client[self._serializer]

        conn = Client(self._address, authkey=self._authkey)

        dispatch(conn, None, 'dummy')

        self._state.value = State.STARTED

    def start(self):

        '''

        Spawn a server process for this manager object

        '''

        assert self._state.value == State.INITIAL

        # pipe over which we will retrieve address of server

        reader, writer = connection.Pipe(duplex=False)

        # spawn process which runs a server

        self._process = Process(

            target=type(self)._run_server,

            args=(self._registry, self._address, self._authkey,

                  self._serializer, writer),

            )

        ident = ':'.join(str(i) for i in self._process._identity)

        self._process.name = type(self).__name__  + '-' + ident

        self._process.start()

        # get address of server

        writer.close()

        self._address = reader.recv()

        reader.close()

        # register a finalizer

        self._state.value = State.STARTED

        self.shutdown = util.Finalize(

            self, type(self)._finalize_manager,

            args=(self._process, self._address, self._authkey,

                  self._state, self._Client),

            exitpriority=0

            )

    @classmethod

    def _run_server(cls, registry, address, authkey, serializer, writer):

        '''

        Create a server, report its address and run it

        '''

        # create server

        server = cls._Server(registry, address, authkey, serializer)

        # inform parent process of the server's address

        writer.send(server.address)

        writer.close()

        # run the manager

        util.info('manager serving at %r', server.address)

        server.serve_forever()

    def _create(self, typeid, *args, **kwds):

        '''

        Create a new shared object; return the token and exposed tuple