<commit_before>"""
xbee.py

By Paul Malmsten, 2010
Inspired by code written by Amit Synderman and Marco Sangalli
pmalmsten@gmail.com

XBee superclass module

This class defines data and methods common to both the Xbee Series 1 and
Series 2 modules. This class should be subclassed in order to provide
series-specific functionality.
"""
import struct, threading, time
from xbee.frame import APIFrame

class ThreadQuitException(Exception):
    pass

class XBeeBase(threading.Thread):
    """
    Abstract base class providing basic API frame generation, validation,
    and data extraction methods for XBee modules
    """

    def __init__(self, ser, shorthand=True, callback=None):
        self.serial = ser
        self.shorthand = shorthand
        self._callback = None
        self._thread_continue = False

        if callback:
            super(XBeeBase, self).__init__()
            self._callback = callback
            self._thread_continue = True
            self._thread_quit = threading.Event()
            self.start()

    def __del__(self):
        """
        In the event that this instance is garbage collected before
        its associated serial port is closed, ensure safe thread
        shutdown.
        """
        self.halt()

    def halt(self):
        """
        halt: None -> None

        If this instance has a separate thread running, it will be
        halted. This method will wait until the thread has cleaned
        up before returning.
        """
        if self._callback:
            self._thread_continue = False
            self._thread_quit.wait()

    def write(self, data):
        """
        write: binary data -> None

        Packages the given binary data in an API frame and writes the
        result to the serial port
        """
        self.serial.write(APIFrame(data).output())

    def run(self):
        """
        run: None -> None

        This method overrides threading.Thread.run() and is automatically
        called when an instance is created with threading enabled.
        """
        while True:
            try:
                self._callback(self.wait_read_frame())
            except ThreadQuitException:
                break
        self._thread_quit.set()

    def wait_for_frame(self):
        """
        wait_for_frame: None -> binary data

        wait_for_frame will read from the serial port until a valid
        API frame arrives. It will then return the binary data
        contained within the frame.

        If this method is called as a separate thread
        and self.thread_continue is set to False, the thread will
        exit by raising a ThreadQuitException.
        """
        WAITING = 0
        PARSING = 1

        data = ''
        state = WAITING

        while True:
            if state == WAITING:
                if self._callback and not self._thread_continue:
                    raise ThreadQuitException

                if self.serial.inWaiting() == 0:
                    time.sleep(.01)
                    continue

                byte = self.serial.read()

                # If a start byte is found, swich states
                if byte == APIFrame.START_BYTE:
                    data += byte
                    state = PARSING
            else:
                # Save all following bytes
                data += self.serial.read()

                if len(data) == 3:
                    # We have the length bytes of the data
                    # Now, wait for the rest to appear
                    data_len = struct.unpack("> h", data[1:3])[0]

                    # Wait for the expected number of bytes to appear
                    # Grab the checksum too
                    data += self.serial.read(data_len + 1)

                    try:
                        # Try to parse and return result
                        return APIFrame.parse(data)
                    except ValueError:
                        # Bad frame, so restart
                        data = ''
                        state = WAITING

    def build_command(self, cmd, **kwargs):
        """
        build_command: string (binary data) ... -> binary data

        build_command will construct a command packet according to the
        specified command's specification in api_commands. It will expect
        named arguments for all fields other than those with a default
        value or a length of 'None'.

        Each field will be written out in the order they are defined
        in the command definition.
        """
        try:
            cmd_spec = self.api_commands[cmd]
        except AttributeError:
            raise NotImplementedError("API command specifications could not be found; use a derived class which defines 'api_commands'.")

        packet = ''

        for field in cmd_spec:
            try:
                # Read this field's name from the function arguments dict
                data = kwargs[field['name']]
            except KeyError:
                # Data wasn't given
                # Only a problem if the field has a specific length
                if field['len'] is not None:
                    # Was a default value specified?
                    default_value = field['default']
                    if default_value:
                        # If so, use it
                        data = default_value
                    else:
                        # Otherwise, fail
                        raise KeyError(
                            "The expected field %s of length %d was not provided"
                            % (field['name'], field['len']))
                else:
                    # No specific length, ignore it
                    data = None

            # Ensure that the proper number of elements will be written
            if field['len'] and len(data) != field['len']:
                raise ValueError(
                    "The data provided for '%s' was not %d bytes long"\
                    % (field['name'], field['len']))

            # Add the data to the packet, if it has been specified
            # Otherwise, the parameter was of variable length, and not
            #  given
            if data:
                packet += data

        return packet

    def split_response(self, data):
        """
        split_response: binary data -> {'id':str,
                                        'param':binary data,
                                        ...}

        split_response takes a data packet received from an XBee device
        and converts it into a dictionary. This dictionary provides
        names for each segment of binary data as specified in the
        api_responses spec.
        """
        # Fetch the first byte, identify the packet
        # If the spec doesn't exist, raise exception
        packet_id = data[0]
        try:
            packet = self.api_responses[packet_id]
        except AttributeError:
            raise NotImplementedError("API response specifications could not be found; use a derived class which defines 'api_responses'.")
        except KeyError:
            raise KeyError(
                "Unrecognized response packet with id byte %s"
                % data[0])

        # Current byte index in the data stream
        index = 1

        # Result info
        info = {'id':packet['name']}
        packet_spec = packet['structure']

        # Parse the packet in the order specified
        for field in packet_spec:
            # Store the number of bytes specified
            # If the data field has no length specified, store any
            #  leftover bytes and quit
            if field['len'] is not None:
                # Are we trying to read beyond the last data element?
                if index + field['len'] > len(data):
                    raise ValueError(
                        "Response packet was shorter than expected")

                field_data = data[index:index + field['len']]
                info[field['name']] = field_data
            else:
                field_data = data[index:]

                # Were there any remaining bytes?
                if field_data:
                    # If so, store them
                    info[field['name']] = field_data
                    index += len(field_data)
                break

            # Move the index
            index += field['len']

        # If there are more bytes than expected, raise an exception
        if index < len(data):
            raise ValueError(
                "Response packet was longer than expected")

        # Check if this packet was an IO sample
        # If so, process the sample data
        if 'parse_as_io_samples' in packet:
            field_to_process = packet['parse_as_io_samples']
            info[field_to_process] = XBeeBase.parse_samples(
                                        info[field_to_process])

        return info

    @staticmethod
    def parse_samples_header(data):
        """
        parse_samples_header: binary data in XBee IO data format ->
                        (int, [int ...], [int ...])

        parse_samples_header will read the first three bytes of the
        binary data given and will return the number of samples which
        follow, a list of enabled digital inputs and a list of enabled
        analog inputs
        """

        ## Parse the header, bytes 0-2
        dio_enabled = []
        adc_enabled = []

        # First byte: number of samples
        len_raw = data[0]
        len_samples = ord(len_raw)

        # Second-third bytes: enabled pin flags
        sources_raw = data[1:3]

        # In order to put the io line names in list positions which
        # match their number (for ease of traversal), the second byte
        # will be read first
        byte_2_data = ord(sources_raw[1])

        # Check each flag
        #  DIO lines 0-7
        i = 1
        for dio in range(0, 8):
            if byte_2_data & i:
                dio_enabled.append(dio)
            i *= 2

        # Byte 1
        byte_1_data = ord(sources_raw[0])

        # Grab DIO8 first
        if byte_1_data & 1:
            dio_enabled.append(8)

        # Check each flag (after the first)
        #  ADC lines 0-5
        i = 2
        for adc in range(0, 6):
            if byte_1_data & i:
                adc_enabled.append(adc)
            i *= 2

        return (len_samples, dio_enabled, adc_enabled)

    @staticmethod
    def parse_samples(data):
        """
        parse_samples: binary data in XBee IO data format ->
                        [ {"dio-0":True,
                           "dio-1":False,
                           "adc-0":100"}, ...]

        parse_samples reads binary data from an XBee device in the IO
        data format specified by the API. It will then return a
        dictionary indicating the status of each enabled IO port.
        """

        ## Parse and store header information
        header_data = XBeeBase.parse_samples_header(data)
        len_samples, dio_enabled, adc_enabled = header_data

        samples = []

        ## Parse the samples
        # Start at byte 3
        byte_pos = 3

        for i in range(0, len_samples):
            sample = {}

            # If one or more DIO lines are set, the first two bytes
            # contain their values
            if dio_enabled:
                # Get two bytes
                values = data[byte_pos:byte_pos + 2]

                # Read and store values for all enabled DIO lines
                for dio in dio_enabled:
                    # Second byte contains values for 0-7
                    # If we want number 8, switch to the first byte, and
                    #  move back to the beginning
                    if dio > 7:
                        sample["dio-%d" % dio] = True if ord(values[0]) & 2 ** (dio - 8) else False
                    else:
                        sample["dio-%d" % dio] = True if ord(values[1]) & 2 ** dio else False

                # Move the starting position for the next new byte
                byte_pos += 2

            # If one or more ADC lines are set, the remaining bytes, in
            # pairs, represent their values, MSB first.
            if adc_enabled:
                for adc in adc_enabled:
                    # Analog reading stored in two bytes
                    value_raw = data[byte_pos:byte_pos + 2]

                    # Unpack the bits
                    value = struct.unpack("> h", value_raw)[0]

                    # Only 10 bits are meaningful
                    value &= 0x3FF

                    # Save the result
                    sample["adc-%d" % adc] = value

                    # Move the starting position for the next new byte
                    byte_pos += 2

            samples.append(sample)

        return samples

    def send(self, cmd, **kwargs):
        """
        send: string param=binary data ... -> None

        When send is called with the proper arguments, an API command
        will be written to the serial port for this XBee device
        containing the proper instructions and data.

        This method must be called with named arguments in accordance
        with the api_command specification. Arguments matching all
        field names other than those in reserved_names (like 'id' and
        'order') should be given, unless they are of variable length
        (of 'None' in the specification. Those are optional).
        """
        # Pass through the keyword arguments
        self.write(self.build_command(cmd, **kwargs))

    def wait_read_frame(self):
        """
        wait_read_frame: None -> frame info dictionary

        wait_read_frame calls XBee.wait_for_frame() and waits until a
        valid frame appears on the serial port. Once it receives a frame,
        wait_read_frame attempts to parse the data contained within it
        and returns the resulting dictionary
        """

        frame = self.wait_for_frame()
        return self.split_response(frame.data)

    def __getattr__(self, name):
        """
        If a method by the name of a valid api command is called,
        the arguments will be automatically sent to an appropriate
        send() call
        """
        # If api_commands is not defined, raise NotImplementedError\
        #  If its not defined, _getattr__ will be called with its name
        if name == 'api_commands':
            raise NotImplementedError("API command specifications could not be found; use a derived class which defines 'api_commands'.")

        # Is shorthand enabled, and is the called name a command?
        if self.shorthand and name in self.api_commands:
            # If so, simply return a function which passes its arguments
            # to an appropriate send() call
            return lambda **kwargs: self.send(name, **kwargs)
        else:
            raise AttributeError("XBee has no attribute '%s'" % name)
<commit_msg>Now __init__()'s XBeeBase as a Thread in all use cases<commit_after>"""
xbee.py

By Paul Malmsten, 2010
Inspired by code written by Amit Synderman and Marco Sangalli
pmalmsten@gmail.com

XBee superclass module

This class defines data and methods common to both the Xbee Series 1 and
Series 2 modules. This class should be subclassed in order to provide
series-specific functionality.
"""
import struct, threading, time
from xbee.frame import APIFrame

class ThreadQuitException(Exception):
    pass

class XBeeBase(threading.Thread):
    """
    Abstract base class providing basic API frame generation, validation,
    and data extraction methods for XBee modules
    """

    def __init__(self, ser, shorthand=True, callback=None):
        super(XBeeBase, self).__init__()
        self.serial = ser
        self.shorthand = shorthand
        self._callback = None
        self._thread_continue = False

        if callback:
            self._callback = callback
            self._thread_continue = True
            self._thread_quit = threading.Event()
            self.start()

    def __del__(self):
        """
        In the event that this instance is garbage collected before
        its associated serial port is closed, ensure safe thread
        shutdown.
        """
        self.halt()

    def halt(self):
        """
        halt: None -> None

        If this instance has a separate thread running, it will be
        halted. This method will wait until the thread has cleaned
        up before returning.
        """
        if self._callback:
            self._thread_continue = False
            self._thread_quit.wait()

    def write(self, data):
        """
        write: binary data -> None

        Packages the given binary data in an API frame and writes the
        result to the serial port
        """
        self.serial.write(APIFrame(data).output())

    def run(self):
        """
        run: None -> None

        This method overrides threading.Thread.run() and is automatically
        called when an instance is created with threading enabled.
        """
        while True:
            try:
                self._callback(self.wait_read_frame())
            except ThreadQuitException:
                break
        self._thread_quit.set()

    def wait_for_frame(self):
        """
        wait_for_frame: None -> binary data

        wait_for_frame will read from the serial port until a valid
        API frame arrives. It will then return the binary data
        contained within the frame.

        If this method is called as a separate thread
        and self.thread_continue is set to False, the thread will
        exit by raising a ThreadQuitException.
        """
        WAITING = 0
        PARSING = 1

        data = ''
        state = WAITING

        while True:
            if state == WAITING:
                if self._callback and not self._thread_continue:
                    raise ThreadQuitException

                if self.serial.inWaiting() == 0:
                    time.sleep(.01)
                    continue

                byte = self.serial.read()

                # If a start byte is found, swich states
                if byte == APIFrame.START_BYTE:
                    data += byte
                    state = PARSING
            else:
                # Save all following bytes
                data += self.serial.read()

                if len(data) == 3:
                    # We have the length bytes of the data
                    # Now, wait for the rest to appear
                    data_len = struct.unpack("> h", data[1:3])[0]

                    # Wait for the expected number of bytes to appear
                    # Grab the checksum too
                    data += self.serial.read(data_len + 1)

                    try:
                        # Try to parse and return result
                        return APIFrame.parse(data)
                    except ValueError:
                        # Bad frame, so restart
                        data = ''
                        state = WAITING

    def build_command(self, cmd, **kwargs):
        """
        build_command: string (binary data) ... -> binary data

        build_command will construct a command packet according to the
        specified command's specification in api_commands. It will expect
        named arguments for all fields other than those with a default
        value or a length of 'None'.

        Each field will be written out in the order they are defined
        in the command definition.
        """
        try:
            cmd_spec = self.api_commands[cmd]
        except AttributeError:
            raise NotImplementedError("API command specifications could not be found; use a derived class which defines 'api_commands'.")

        packet = ''

        for field in cmd_spec:
            try:
                # Read this field's name from the function arguments dict
                data = kwargs[field['name']]
            except KeyError:
                # Data wasn't given
                # Only a problem if the field has a specific length
                if field['len'] is not None:
                    # Was a default value specified?
                    default_value = field['default']
                    if default_value:
                        # If so, use it
                        data = default_value
                    else:
                        # Otherwise, fail
                        raise KeyError(
                            "The expected field %s of length %d was not provided"
                            % (field['name'], field['len']))
                else:
                    # No specific length, ignore it
                    data = None

            # Ensure that the proper number of elements will be written
            if field['len'] and len(data) != field['len']:
                raise ValueError(
                    "The data provided for '%s' was not %d bytes long"\
                    % (field['name'], field['len']))

            # Add the data to the packet, if it has been specified
            # Otherwise, the parameter was of variable length, and not
            #  given
            if data:
                packet += data

        return packet

    def split_response(self, data):
        """
        split_response: binary data -> {'id':str,
                                        'param':binary data,
                                        ...}

        split_response takes a data packet received from an XBee device
        and converts it into a dictionary. This dictionary provides
        names for each segment of binary data as specified in the
        api_responses spec.
        """
        # Fetch the first byte, identify the packet
        # If the spec doesn't exist, raise exception
        packet_id = data[0]
        try:
            packet = self.api_responses[packet_id]
        except AttributeError:
            raise NotImplementedError("API response specifications could not be found; use a derived class which defines 'api_responses'.")
        except KeyError:
            raise KeyError(
                "Unrecognized response packet with id byte %s"
                % data[0])

        # Current byte index in the data stream
        index = 1

        # Result info
        info = {'id':packet['name']}
        packet_spec = packet['structure']

        # Parse the packet in the order specified
        for field in packet_spec:
            # Store the number of bytes specified
            # If the data field has no length specified, store any
            #  leftover bytes and quit
            if field['len'] is not None:
                # Are we trying to read beyond the last data element?
                if index + field['len'] > len(data):
                    raise ValueError(
                        "Response packet was shorter than expected")

                field_data = data[index:index + field['len']]
                info[field['name']] = field_data
            else:
                field_data = data[index:]

                # Were there any remaining bytes?
                if field_data:
                    # If so, store them
                    info[field['name']] = field_data
                    index += len(field_data)
                break

            # Move the index
            index += field['len']

        # If there are more bytes than expected, raise an exception
        if index < len(data):
            raise ValueError(
                "Response packet was longer than expected")

        # Check if this packet was an IO sample
        # If so, process the sample data
        if 'parse_as_io_samples' in packet:
            field_to_process = packet['parse_as_io_samples']
            info[field_to_process] = XBeeBase.parse_samples(
                                        info[field_to_process])

        return info

    @staticmethod
    def parse_samples_header(data):
        """
        parse_samples_header: binary data in XBee IO data format ->
                        (int, [int ...], [int ...])

        parse_samples_header will read the first three bytes of the
        binary data given and will return the number of samples which
        follow, a list of enabled digital inputs and a list of enabled
        analog inputs
        """

        ## Parse the header, bytes 0-2
        dio_enabled = []
        adc_enabled = []

        # First byte: number of samples
        len_raw = data[0]
        len_samples = ord(len_raw)

        # Second-third bytes: enabled pin flags
        sources_raw = data[1:3]

        # In order to put the io line names in list positions which
        # match their number (for ease of traversal), the second byte
        # will be read first
        byte_2_data = ord(sources_raw[1])

        # Check each flag
        #  DIO lines 0-7
        i = 1
        for dio in range(0, 8):
            if byte_2_data & i:
                dio_enabled.append(dio)
            i *= 2

        # Byte 1
        byte_1_data = ord(sources_raw[0])

        # Grab DIO8 first
        if byte_1_data & 1:
            dio_enabled.append(8)

        # Check each flag (after the first)
        #  ADC lines 0-5
        i = 2
        for adc in range(0, 6):
            if byte_1_data & i:
                adc_enabled.append(adc)
            i *= 2

        return (len_samples, dio_enabled, adc_enabled)

    @staticmethod
    def parse_samples(data):
        """
        parse_samples: binary data in XBee IO data format ->
                        [ {"dio-0":True,
                           "dio-1":False,
                           "adc-0":100"}, ...]

        parse_samples reads binary data from an XBee device in the IO
        data format specified by the API. It will then return a
        dictionary indicating the status of each enabled IO port.
        """

        ## Parse and store header information
        header_data = XBeeBase.parse_samples_header(data)
        len_samples, dio_enabled, adc_enabled = header_data

        samples = []

        ## Parse the samples
        # Start at byte 3
        byte_pos = 3

        for i in range(0, len_samples):
            sample = {}

            # If one or more DIO lines are set, the first two bytes
            # contain their values
            if dio_enabled:
                # Get two bytes
                values = data[byte_pos:byte_pos + 2]

                # Read and store values for all enabled DIO lines
                for dio in dio_enabled:
                    # Second byte contains values for 0-7
                    # If we want number 8, switch to the first byte, and
                    #  move back to the beginning
                    if dio > 7:
                        sample["dio-%d" % dio] = True if ord(values[0]) & 2 ** (dio - 8) else False
                    else:
                        sample["dio-%d" % dio] = True if ord(values[1]) & 2 ** dio else False

                # Move the starting position for the next new byte
                byte_pos += 2

            # If one or more ADC lines are set, the remaining bytes, in
            # pairs, represent their values, MSB first.
            if adc_enabled:
                for adc in adc_enabled:
                    # Analog reading stored in two bytes
                    value_raw = data[byte_pos:byte_pos + 2]

                    # Unpack the bits
                    value = struct.unpack("> h", value_raw)[0]

                    # Only 10 bits are meaningful
                    value &= 0x3FF

                    # Save the result
                    sample["adc-%d" % adc] = value

                    # Move the starting position for the next new byte
                    byte_pos += 2

            samples.append(sample)

        return samples

    def send(self, cmd, **kwargs):
        """
        send: string param=binary data ... -> None

        When send is called with the proper arguments, an API command
        will be written to the serial port for this XBee device
        containing the proper instructions and data.

        This method must be called with named arguments in accordance
        with the api_command specification. Arguments matching all
        field names other than those in reserved_names (like 'id' and
        'order') should be given, unless they are of variable length
        (of 'None' in the specification. Those are optional).
        """
        # Pass through the keyword arguments
        self.write(self.build_command(cmd, **kwargs))

    def wait_read_frame(self):
        """
        wait_read_frame: None -> frame info dictionary

        wait_read_frame calls XBee.wait_for_frame() and waits until a
        valid frame appears on the serial port. Once it receives a frame,
        wait_read_frame attempts to parse the data contained within it
        and returns the resulting dictionary
        """

        frame = self.wait_for_frame()
        return self.split_response(frame.data)

    def __getattr__(self, name):
        """
        If a method by the name of a valid api command is called,
        the arguments will be automatically sent to an appropriate
        send() call
        """
        # If api_commands is not defined, raise NotImplementedError\
        #  If its not defined, _getattr__ will be called with its name
        if name == 'api_commands':
            raise NotImplementedError("API command specifications could not be found; use a derived class which defines 'api_commands'.")

        # Is shorthand enabled, and is the called name a command?
        if self.shorthand and name in self.api_commands:
            # If so, simply return a function which passes its arguments
            # to an appropriate send() call
            return lambda **kwargs: self.send(name, **kwargs)
        else:
            raise AttributeError("XBee has no attribute '%s'" % name)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""multi process and multi greenlet spawn"""

from gevent.queue import Queue
from functools import partial
import gevent
import gipc
import random
import logging

import sys; sys.modules.pop("threading", None)
from gevent import monkey; monkey.patch_all()
from .dictqueue import DictQueue

logger = logging.getLogger("pglet")

def id_generator():
    i = 0
    wall = 1 << 31
    while True:
        i += 1
        if i > wall:
            i = 1
        yield i

class Delay(object):
    """future对象: PPool的spawn返回的future对象"""

    def __init__(self, ppool, task_id):
        self.ppool = ppool
        self.task_id = task_id

    def get(self, block=True, timeout=None):
        """在这里,block不能为False"""
        return self.ppool.results.get(self.task_id, block, timeout)

class PPool(object):
    """process pool + gevent spawn
    : 1.必须放在入口文件的最开始处调用PPool实例的init函数, 因为子进程会复制主进程空间(包括已有协程)
    : 2.如果需要子进程也包含某些模块/全局变量,可以在模块/全局变量导入后再调用init函数,原则是不能复制过多主进程空间
    Example:
    from share import ppool; ppool.init()
    import sys; sys.modules.pop("threading", None)
    from gevent import monkey; monkey.patch_all()
    import gevent

    def foo(a):
        print a

    ppool.spawn_sub(foo, "abc")

    gevent.wait()
    """

    def __init__(self, process_size=1):
        """
        @param process_size: the process pool size
        """
        self.id_generator = id_generator()
        self.process_size = process_size
        self.parent_pipe_ends = []         # 存放process_size个双向pipe的父进程端
        self.child_pipe_ends = []          # 存放process_size个双向pipe的子进程端
        self.processes = []                # 存放process_size个子进程
        self.results = DictQueue()         # 在阻塞模式下存放结果, {task_id:DictQueue()}
        self.loop_get_result_glets = []    # 存放 loop_get_result 产生的process_size个greenlet

    def init(self):
        for _ in xrange(self.process_size):
            child_pipe_end, parent_pipe_end =  gipc.pipe(duplex=True)
            self.child_pipe_ends.append(child_pipe_end)
            self.parent_pipe_ends.append(parent_pipe_end)
            p = gipc.start_process(target=self.process_target, args=(child_pipe_end, ), daemon=True)
            self.processes.append(p)
            #p.join()

            self.loop_get_result()

    def loop_get_result(self):
        """循环读取子进程返回结果"""
        n = 0 # 用于测试性能,发布时删除
        def loop(p):
            while 1:
                try:
                    k, v = p.get()
                    if k == -9: # 用于测试性能,发布时删除, =-9时表示该子进程joinall结束
                        n += 1
                        if n == self.process_size:
                            print "join all completed!!!"
                        continue
                    self.results.put(k, v, override=True, timeout=None)
                except Exception as e:
                    logger.warning(str(e))
                    break

        for p in self.parent_pipe_ends:
            g = gevent.spawn(loop, p)
            self.loop_get_result_glets.append(g)

    def process_target(self, child_pipe_end):
        """子进程空间"""
        def callback(g, task_id):
            child_pipe_end.put([task_id, g.value])

        taskqs = [] # 用于测试,发布时删除

        def loop(child_pipe_end):
            while 1:
                try:
                    f, args, kwargs, task_id = child_pipe_end.get()
                    if task_id >= 0:
                        g = gevent.spawn(f, *args, **kwargs)
                        cb = partial(callback, task_id=task_id)
                        g.link(cb)
                    elif task_id == -1: # 用于测试性能,发布时删除, =-1时表示joinall开始
                        taskqs.append([f, args, kwargs])
                    elif task_id == -2: # 用于测试性能,发布时删除, =-2时表示joinall结束
                        gevent.joinall([gevent.spawn(f, *args, **kwargs) for f, args, kwargs in taskqs])
                        child_pipe_end.put([-9, None])
                    else:
                        gevent.spawn(f, *args, **kwargs)
                except Exception as e:
                    logger.warning(str(e))
                    break

        loop(child_pipe_end)

    def _benchmark_join_start(self, f, ts=[]):
        parent_pipe_end = self.select_pipe_writer()
        [parent_pipe_end.put([f, args, None, -1]) for args in ts]

    def _benchmark_join_end(self):
        parent_pipe_end.put([None, None, None, -2])

    def select_pipe_writer(self,  sn=None):
        if sn is not None:
            return self.parent_pipe_ends[sn]
        return random.choice(self.parent_pipe_ends)

    def _spawn(self, f, args=tuple(), kwargs={}, future=False, timeout=None, sn=None):
        """和gevent.spawn的区别是,如果block=False 就不返回任何值(gevent.spawn返回greenlet),如果block=True,返回f的结果值,而不是greenlet
        @param f: func
        @param args, kwargs: f的参数, ***必须能够pickle序列化***
        @param: future: 是否返回future对象(future.get阻塞获取结果)
        """
        parent_pipe_end = self.select_pipe_writer(sn)
        if future:
            task_id = self.id_generator.next()
            parent_pipe_end.put([f, args, kwargs, task_id])
            delay = Delay(self, task_id)
            return delay
        else:
            parent_pipe_end.put([f, args, kwargs, None])

    def spawn(self, f, *args, **kwargs):
        """f的参数和返回值***必须能够pickle序列化***
        @return: 返回future对象,future.get()来阻塞获取结果
        """
        return self._spawn(f, args, kwargs, True, None)

    def spawn_sub(self, f, *args, **kwargs):
        """f的参数***必须能够pickle序列化***
        @return: 无 (即:是过程sub)
        """
        return self._spawn(f, args, kwargs, False, None)

    def close(self):
        """不再使用时,必须手动调用"""
        try:
            [p.terminate() for p in self.processes]          # 结束子进程
            [g.kill() for g in self.loop_get_result_glets]   # 结束取结果的协程
            [p.close() for p in self.parent_pipe_ends]       # 管道一端关闭即可
        except Exception as e:
            logger.warning(str(e))

    def __del__(self):
        self.close()

if __name__ == "__main__":
    def x(a):
        return a

    try:
        ppool = PPool(2)
        ppool.init()
        print ppool.spawn(x, "abc").get()
        print ppool.spawn_sub(x, "def")
        ppool.close()
        gevent.wait()
    except KeyboardInterrupt:
        ppool.close()

        <commit_msg>benchmark<commit_after># -*- coding: utf-8 -*-

"""multi process and multi greenlet spawn"""

from gevent.queue import Queue
from functools import partial
import gevent
import gipc
import random
import logging

import sys; sys.modules.pop("threading", None)
from gevent import monkey; monkey.patch_all()
from .dictqueue import DictQueue

logger = logging.getLogger("pglet")

def id_generator():
    i = 0
    wall = 1 << 31
    while True:
        i += 1
        if i > wall:
            i = 1
        yield i

class Delay(object):
    """future对象: PPool的spawn返回的future对象"""

    def __init__(self, ppool, task_id):
        self.ppool = ppool
        self.task_id = task_id

    def get(self, block=True, timeout=None):
        """在这里,block不能为False"""
        return self.ppool.results.get(self.task_id, block, timeout)

class PPool(object):
    """process pool + gevent spawn
    : 1.必须放在入口文件的最开始处调用PPool实例的init函数, 因为子进程会复制主进程空间(包括已有协程)
    : 2.如果需要子进程也包含某些模块/全局变量,可以在模块/全局变量导入后再调用init函数,原则是不能复制过多主进程空间
    Example:
    from share import ppool; ppool.init()
    import sys; sys.modules.pop("threading", None)
    from gevent import monkey; monkey.patch_all()
    import gevent

    def foo(a):
        print a

    ppool.spawn_sub(foo, "abc")

    gevent.wait()
    """

    def __init__(self, process_size=1):
        """
        @param process_size: the process pool size
        """
        self.id_generator = id_generator()
        self.process_size = process_size
        self.parent_pipe_ends = []         # 存放process_size个双向pipe的父进程端
        self.child_pipe_ends = []          # 存放process_size个双向pipe的子进程端
        self.processes = []                # 存放process_size个子进程
        self.results = DictQueue()         # 在阻塞模式下存放结果, {task_id:DictQueue()}
        self.loop_get_result_glets = []    # 存放 loop_get_result 产生的process_size个greenlet

    def init(self):
        for _ in xrange(self.process_size):
            child_pipe_end, parent_pipe_end =  gipc.pipe(duplex=True)
            self.child_pipe_ends.append(child_pipe_end)
            self.parent_pipe_ends.append(parent_pipe_end)
            p = gipc.start_process(target=self.process_target, args=(child_pipe_end, ), daemon=True)
            self.processes.append(p)
            #p.join()

            self.loop_get_result()

    def loop_get_result(self):
        """循环读取子进程返回结果"""
        n = 0 # 用于测试性能,发布时删除
        def loop(p):
            while 1:
                try:
                    k, v = p.get()
                    if k == -9: # 用于测试性能,发布时删除, =-9时表示该子进程joinall结束
                        n += 1
                        if n == self.process_size:
                            print "join all completed!!!"
                        continue
                    self.results.put(k, v, override=True, timeout=None)
                except Exception as e:
                    logger.warning(str(e))
                    break

        for p in self.parent_pipe_ends:
            g = gevent.spawn(loop, p)
            self.loop_get_result_glets.append(g)

    def process_target(self, child_pipe_end):
        """子进程空间"""
        def callback(g, task_id):
            child_pipe_end.put([task_id, g.value])

        taskqs = [] # 用于测试,发布时删除

        def loop(child_pipe_end):
            while 1:
                try:
                    f, args, kwargs, task_id = child_pipe_end.get()
                    if task_id >= 0:
                        g = gevent.spawn(f, *args, **kwargs)
                        cb = partial(callback, task_id=task_id)
                        g.link(cb)
                    elif task_id == -1: # 用于测试性能,发布时删除, =-1时表示joinall开始
                        taskqs.append([f, args, kwargs])
                    elif task_id == -2: # 用于测试性能,发布时删除, =-2时表示joinall结束
                        gevent.joinall([gevent.spawn(f, *args, **kwargs) for f, args, kwargs in taskqs])
                        child_pipe_end.put([-9, None])
                    else:
                        gevent.spawn(f, *args, **kwargs)
                except Exception as e:
                    logger.warning(str(e))
                    break

        loop(child_pipe_end)

    def _benchmark_join_start(self, f, ts=[]):
        parent_pipe_end = self.select_pipe_writer()
        [parent_pipe_end.put([f, args, None, -1]) for args in ts]

    def _benchmark_join_end(self):
        for p in self.parent_pipe_ends:
            p.put([None, None, None, -2])

    def select_pipe_writer(self,  sn=None):
        if sn is not None:
            return self.parent_pipe_ends[sn]
        return random.choice(self.parent_pipe_ends)

    def _spawn(self, f, args=tuple(), kwargs={}, future=False, timeout=None, sn=None):
        """和gevent.spawn的区别是,如果block=False 就不返回任何值(gevent.spawn返回greenlet),如果block=True,返回f的结果值,而不是greenlet
        @param f: func
        @param args, kwargs: f的参数, ***必须能够pickle序列化***
        @param: future: 是否返回future对象(future.get阻塞获取结果)
        """
        parent_pipe_end = self.select_pipe_writer(sn)
        if future:
            task_id = self.id_generator.next()
            parent_pipe_end.put([f, args, kwargs, task_id])
            delay = Delay(self, task_id)
            return delay
        else:
            parent_pipe_end.put([f, args, kwargs, None])

    def spawn(self, f, *args, **kwargs):
        """f的参数和返回值***必须能够pickle序列化***
        @return: 返回future对象,future.get()来阻塞获取结果
        """
        return self._spawn(f, args, kwargs, True, None)

    def spawn_sub(self, f, *args, **kwargs):
        """f的参数***必须能够pickle序列化***
        @return: 无 (即:是过程sub)
        """
        return self._spawn(f, args, kwargs, False, None)

    def close(self):
        """不再使用时,必须手动调用"""
        try:
            [p.terminate() for p in self.processes]          # 结束子进程
            [g.kill() for g in self.loop_get_result_glets]   # 结束取结果的协程
            [p.close() for p in self.parent_pipe_ends]       # 管道一端关闭即可
        except Exception as e:
            logger.warning(str(e))

    def __del__(self):
        self.close()

if __name__ == "__main__":
    def x(a):
        return a

    try:
        ppool = PPool(2)
        ppool.init()
        print ppool.spawn(x, "abc").get()
        print ppool.spawn_sub(x, "def")
        ppool.close()
        gevent.wait()
    except KeyboardInterrupt:
        ppool.close()

        <|endoftext|>


<commit_before>import json
import os
import shutil
import subprocess
import sys
import time
import tempfile
import unittest

from contextlib import closing
from StringIO import StringIO

import zmq

from volttron.platform import aip
from volttron.platform.control import client, server
from volttron.platform import packaging

RESTRICTED_AVAILABLE = False

try:
    from volttron.restricted import (auth, certs)
    RESTRICTED_AVAILABLE = True

except ImportError:
    RESTRICTED_AVAILABLE = False
    auth = None
    certs = None

#Filenames for the config files which are created during setup and then
#passed on the command line
TMP_PLATFORM_CONFIG_FILENAME = "config"
TMP_SMAP_CONFIG_FILENAME = "test-smap.ini"

#Used to fill in TWISTED_CONFIG template
TEST_CONFIG_FILE = 'base-platform-test.json'

PLATFORM_CONFIG_UNRESTRICTED = """
no-resource-monitor
no-verify
no-mobility

control-socket = {tmpdir}/run/control
"""

PLATFORM_CONFIG_RESTRICTED = """
mobility-address = {mobility-address}
control-socket = {tmpdir}/run/control
resource-monitor = {resource-monitor}
"""

TWISTED_CONFIG = """
[report 0]
ReportDeliveryLocation = {smap-uri}/add/{smap-key}

[/]
type = Collection
Metadata/SourceName = {smap-source}
uuid = {smap-uuid}

[/datalogger]
type = volttron.drivers.data_logger.DataLogger
interval = 1

"""

UNRESTRICTED = 0
VERIFY_ONLY = 1
RESOURCE_CHECK_ONLY = 2
RESTRICTED = 3

MODES = (UNRESTRICTED, VERIFY_ONLY, RESOURCE_CHECK_ONLY, RESTRICTED)

rel_path = './'

VSTART = os.path.join(rel_path, "env/bin/volttron")
VCTRL = os.path.join(rel_path, "env/bin/volttron-ctl")
SEND_AGENT = "send"

RUN_DIR = 'run'
PUBLISH_TO = RUN_DIR+'/publish'
SUBSCRIBE_TO = RUN_DIR+'/subscribe'

class PlatformWrapperError(Exception):
    pass

class PlatformWrapper():

    def __init__(self, volttron_home=None):
        self.tmpdir = tempfile.mkdtemp()
        self.wheelhouse = '/'.join((self.tmpdir, 'wheelhouse'))
        os.makedirs(self.wheelhouse)

        os.makedirs(os.path.join(self.tmpdir, RUN_DIR))

        self.env = os.environ.copy()
        self.env['VOLTTRON_HOME'] = self.tmpdir
        self.env['AGENT_PUB_ADDR'] = "ipc://{}/{}".format(
                                            self.env['VOLTTRON_HOME'],
                                            PUBLISH_TO)
        self.env['AGENT_SUB_ADDR'] = "ipc://{}/{}".format(
                                            self.env['VOLTTRON_HOME'],
                                            SUBSCRIBE_TO)
        print ("Agent Home", self.env['VOLTTRON_HOME'])
        print ("Agent Pub Addr", self.env['AGENT_PUB_ADDR'])
        print ("Agent Sub Addr", self.env['AGENT_SUB_ADDR'])
        self.p_process = None
        self.t_process = None
        self.zmq_context = None
        self.use_twistd = False

        if volttron_home is not None:
            self.initialize_volttron_home(volttron_home)

    def initialize_volttron_home(self, volttron_home):
        '''Copies the configuration of the passed "volttron_home".

        The platform is actually being run in a temporary space that is
        dynamically created.  This function will copy the directory tree
        recursively from volttron_home to the platforms true VOLTTRON_HOME.

        raises ValueError if volttron_home does not exist

        volttron_home is the directory where the configurations will be copied
                      from
        '''
        if not os.path.isdir(volttron_home):
            raise ValueError('Invalid directory specified\n{}'.format(
                                                                volttron_home))
        mergetree(volttron_home, self.tmpdir)

    def startup_platform(self, platform_config, use_twistd = False,
                         mode=UNRESTRICTED):
        try:
            print("PLATFORM CONFIG: ", platform_config)
            config = json.loads(open(platform_config, 'r').read())
        except Exception as e:
            config = None
            sys.stderr.write (str(e))

        assert config != None, 'Invalid configuration file passed {}'.format(
                                                                platform_config)

#         self.tmpdir = tempfile.mkdtemp()
        config['tmpdir'] = self.tmpdir
        pconfig = os.path.join(self.tmpdir, TMP_PLATFORM_CONFIG_FILENAME)
        self.mode = mode

        assert(self.mode in MODES, 'Invalid platform mode set: '+str(mode))
        opts = None

        if self.mode == UNRESTRICTED:
            if RESTRICTED_AVAILABLE:
                config['mobility'] = False
                config['resource-monitor'] = False
                config['verify'] = False
            with closing(open(pconfig, 'w')) as cfg:
                cfg.write(PLATFORM_CONFIG_UNRESTRICTED.format(**config))
            opts = type('Options', (), {'verify_agents': False,
                                        'volttron_home': self.tmpdir})()
        elif self.mode == RESTRICTED:
            if not RESTRICTED_AVAILABLE:
                raise ValueError("restricted is not available.")

            certsdir = os.path.join(os.path.expanduser(self.env['VOLTTRON_HOME']),
                                     'certificates')

            print ("certsdir", certsdir)
            self.certsobj = certs.Certs(certsdir)

            with closing(open(pconfig, 'w')) as cfg:
                cfg.write(PLATFORM_CONFIG_RESTRICTED.format(**config))
            opts = type('Options', (), {'resource-monitor':False,
                                        'verify_agents': True,
                                        'volttron_home': self.tmpdir})()

#                 self.create_certs()
        else:
            raise PlatformWrapperError("Invalid platform mode specified: {}".format(mode))

        self.test_aip = aip.AIPplatform(opts)
        self.test_aip.setup()

        lfile = os.path.join(self.tmpdir, "volttron.log")

        pparams = [VSTART, "-c", pconfig, "-vv", "-l", lfile]
        print pparams

        self.p_process = subprocess.Popen(pparams, env=self.env)

        #Setup connector
        path = '{}/run/control'.format(self.env['VOLTTRON_HOME'])

        time.sleep(5)
        tries = 0
        max_tries = 5
        while(not os.path.exists(path) and tries < max_tries):
            time.sleep(5)
            tries += 1

        self.conn= server.ControlConnector(path)

#         if self.mode == RESTRICTED:
#             self.conn.call.create_cgroups()

        self.use_twistd = use_twistd
        #TODO: Revise this to start twistd with platform.
        if self.use_twistd:
            tconfig = os.path.join(self.tmpdir, TMP_SMAP_CONFIG_FILENAME)

            with closing(open(tconfig, 'w')) as cfg:
                cfg.write(TWISTED_CONFIG.format(**config))

            tparams = ["env/bin/twistd", "-n", "smap", tconfig]
            self.t_process = subprocess.Popen(tparams, env=self.env)
            time.sleep(5)
        #self.t_process = subprocess.Popen(["twistd", "-n", "smap", "test-smap.ini"])

    def publish(self, topic, data):
        '''Publish data to a zmq context.

        The publisher is goint to use the platform that is contained within
        this wrapper to write data to.
        '''
        if not self.zmq_context:
            self.zmq_context = zmq.Context()
        print("binding publisher to: ", self.env['AGENT_PUB_ADDR'])
        pub = zmq.Socket(zmq_context, zmq.PUB)
        pub.bind(self.env['AGENT_PUB_ADDR'])
        pub.send_multipart([topic, data])

    def fillout_file(self, filename, template, config_file):

        try:
            config = json.loads(open(config_file, 'r').read())
        except Exception as e:
            sys.stderr.write (str(e))
            raise PlatformWrapperError("Could not load configuration file for tests")

        config['tmpdir'] = self.tmpdir

        outfile = os.path.join(self.tmpdir, filename)
        with closing(open(outfile, 'w')) as cfg:
            cfg.write(template.format(**config))

        return outfile

    def create_certs(self):
        auth.create_root_ca(self.tmpdir, ca_name)

    def direct_sign_agentpackage_creator(self, package):
        assert (RESTRICTED), "Auth not available"
        print ("wrapper.certsobj", self.certsobj.cert_dir)
        assert(auth.sign_as_creator(package, 'creator', certsobj=self.certsobj)), "Signing as {} failed.".format('creator')

    def direct_sign_agentpackage_soi(self, package):
        assert (RESTRICTED), "Auth not available"
        assert(auth.sign_as_admin(package, 'soi', certsobj=self.certsobj)), "Signing as {} failed.".format('soi')

    def direct_sign_agentpackage_initiator(self, package, config_file, contract):
        assert (RESTRICTED), "Auth not available"
        files = {"config_file":config_file,"contract":contract}
        assert(auth.sign_as_initiator(package, 'initiator', files=files,
                                      certsobj=self.certsobj)), "Signing as {} failed.".format('initiator')

    def direct_build_agentpackage(self, agent_dir):
        wheel_path = packaging.create_package(os.path.join(rel_path, agent_dir), self.wheelhouse)

        return wheel_path

    def direct_send_agent(self, package, target):
        pparams = [VCTRL, SEND_AGENT, target, package]
        print (pparams, "CWD", os.getcwd())
        send_process = subprocess.call(pparams, env=self.env)
        print ("Done sending to", target)

    def direct_configure_agentpackage(self, agent_wheel, config_file):
        packaging.add_files_to_package(agent_wheel, {'config_file':os.path.join(rel_path, config_file)})

    def direct_buid_install_agent(self, agent_dir, config_file):
        agent_wheel = self.direct_build_agentpackage(agent_dir)
        self.direct_configure_agentpackage(agent_wheel, config_file)
        assert(agent_wheel is not None,"Agent wheel was not built")

        uuid = self.test_aip.install_agent(agent_wheel)
        #aip volttron_home, verify_agents
        return uuid
#         conn.call.start_agent()

    def direct_remove_agent(self, uuid):
        uuid = self.test_aip.remove_agent(uuid)

    def direct_build_install_run_agent(self, agent_dir, config_file):
        agent_uuid = self.direct_buid_install_agent(agent_dir, config_file)
        self.direct_start_agent(agent_uuid)
        return agent_uuid

    def direct_build_send_agent(self, agent_dir, config_file, target):
        agent_uuid = self.direct_buid_install_agent(agent_dir, config_file)
        self.direct_start_agent(agent_uuid)
        return agent_uuid

    def direct_start_agent(self, agent_uuid):

        self.conn.call.start_agent(agent_uuid)
        time.sleep(3)

        status = self.conn.call.status_agents()
#         self.test_aip.status_agents()

#         status = self.conn.call.status_agents()
#         self.assertEquals(len(status[0]), 4, 'Unexpected status message')
        status_uuid = status[0][0]
        assert status_uuid == agent_uuid

        assert len(status[0][2]) == 2, 'Unexpected agent status message'
        status_agent_status = status[0][2][1]
        assert not isinstance(status_agent_status, int)
#         self.assertIn("running",status_agent_status, "Agent status shows error")
        print status

    def confirm_agent_running(self, agent_name, max_retries=5, timeout_seconds=2):

#         self.test_aip.status_agents()

#         status = self.conn.call.status_agents()
#         self.assertEquals(len(status[0]), 4, 'Unexpected status message')
        running = False
        retries = 0
        while (not running and retries < max_retries):
            status = self.conn.call.status_agents()
            print ("Status", status)
            if len(status) > 0:
                status_name = status[0][1]
                assert status_name == agent_name

                assert len(status[0][2]) == 2, 'Unexpected agent status message'
                status_agent_status = status[0][2][1]
                running = not isinstance(status_agent_status, int)
            retries += 1
            time.sleep(timeout_seconds)
        return running

    def direct_stop_agent(self, agent_uuid):
        result = self.conn.call.stop_agent(agent_uuid)
        print result

    def shutdown_platform(self, cleanup_temp=True):
        '''Stop platform here'''
        if self.p_process != None:
            if self.conn is not None:
                self.conn.call.shutdown()
            time.sleep(3)
            self.p_process.kill()
        else:
            print "platform process was null"

        if self.use_twistd and self.t_process != None:
            self.t_process.kill()
            self.t_process.wait()
        elif self.use_twistd:
            print "twistd process was null"
        if cleanup_temp:
            if self.tmpdir:
                shutil.rmtree(self.tmpdir, ignore_errors=True)

    def cleanup(self, cleanup_temp=True):
        '''Shuts the platform down and cleans up based upon cleanup_temp
        '''
        try:
            self.shutdown_platform(cleanup_temp=cleanup_temp)
        except Exception as e:
            sys.stderr.write( str(e))
        finally:
            pass

    def do_nothing(self):
        pass

def mergetree(src, dst, symlinks=False, ignore=None):
    if not os.path.exists(dst):
        os.makedirs(dst)
    for item in os.listdir(src):
        s = os.path.join(src, item)
        d = os.path.join(dst, item)
        if os.path.isdir(s):
            mergetree(s, d, symlinks, ignore)
        else:
            if not os.path.exists(d) or os.stat(src).st_mtime - os.stat(dst).st_mtime > 1:
                shutil.copy2(s, d)

#     def build_install_run_agent(self, agent_dir, agent_name):
#         self.build_and_install_agent(agent_dir)
#         results = subprocess.check_output([VCTRL,CTL_START,agent_name])
#         self.assertTrue(len(results) == 0)
#         results = subprocess.check_output([VCTRL,CTL_STATUS])
#         print results
#         self.assertIn('running', results, "Agent was not started")
#

#     def build_and_install_agent(self, agent_dir):
#         agent_wheel = self.build_agentpackage(agent_dir)
#         self.assertIsNotNone(agent_wheel,"Agent wheel was not built")
#
#         sys.stderr = std_err = StringIO()
#
#         results = subprocess.check_output([VCTRL,CTL_INSTALL,agent_wheel])
#         print ("results: "+results)
#         if self.mode == UNRESTRICTED or self.mode == RESOURCE_CHECK_ONLY:
#             self.assertTrue(len(results) == 0)
#             proc_out = std_err.getvalue().split(':')
#             print (proc_out)
#             self.check_default_dir(proc_out[1].strip())
#         elif self.mode == RESTRICTED or self.mode == VERIFY_ONLY:
#             self.assertTrue(results.startswith('Unpacking to: '))
<commit_msg>VNATION-286 Modified to use full path to volttron during running volttron<commit_after>import json
import os
import shutil
import subprocess
import sys
import time
import tempfile
import unittest

from os.path import dirname
from contextlib import closing
from StringIO import StringIO

import zmq

from volttron.platform import aip
from volttron.platform.control import client, server
from volttron.platform import packaging

RESTRICTED_AVAILABLE = False

try:
    from volttron.restricted import (auth, certs)
    RESTRICTED_AVAILABLE = True

except ImportError:
    RESTRICTED_AVAILABLE = False
    auth = None
    certs = None

#Filenames for the config files which are created during setup and then
#passed on the command line
TMP_PLATFORM_CONFIG_FILENAME = "config"
TMP_SMAP_CONFIG_FILENAME = "test-smap.ini"

#Used to fill in TWISTED_CONFIG template
TEST_CONFIG_FILE = 'base-platform-test.json'

PLATFORM_CONFIG_UNRESTRICTED = """
no-resource-monitor
no-verify
no-mobility

control-socket = {tmpdir}/run/control
"""

PLATFORM_CONFIG_RESTRICTED = """
mobility-address = {mobility-address}
control-socket = {tmpdir}/run/control
resource-monitor = {resource-monitor}
"""

TWISTED_CONFIG = """
[report 0]
ReportDeliveryLocation = {smap-uri}/add/{smap-key}

[/]
type = Collection
Metadata/SourceName = {smap-source}
uuid = {smap-uuid}

[/datalogger]
type = volttron.drivers.data_logger.DataLogger
interval = 1

"""

UNRESTRICTED = 0
VERIFY_ONLY = 1
RESOURCE_CHECK_ONLY = 2
RESTRICTED = 3

MODES = (UNRESTRICTED, VERIFY_ONLY, RESOURCE_CHECK_ONLY, RESTRICTED)

_VOLTTRON_ROOT = dirname(dirname(dirname(os.path.realpath(__file__))))

VSTART = os.path.join(_VOLTTRON_ROOT, "env/bin/volttron")
VCTRL = os.path.join(_VOLTTRON_ROOT, "env/bin/volttron-ctl")
SEND_AGENT = "send"

RUN_DIR = 'run'
PUBLISH_TO = RUN_DIR+'/publish'
SUBSCRIBE_TO = RUN_DIR+'/subscribe'

class PlatformWrapperError(Exception):
    pass

class PlatformWrapper():

    def __init__(self, volttron_home=None):
        self.tmpdir = tempfile.mkdtemp()
        self.wheelhouse = '/'.join((self.tmpdir, 'wheelhouse'))
        os.makedirs(self.wheelhouse)

        os.makedirs(os.path.join(self.tmpdir, RUN_DIR))

        self.env = os.environ.copy()
        self.env['VOLTTRON_HOME'] = self.tmpdir
        self.env['AGENT_PUB_ADDR'] = "ipc://{}/{}".format(
                                            self.env['VOLTTRON_HOME'],
                                            PUBLISH_TO)
        self.env['AGENT_SUB_ADDR'] = "ipc://{}/{}".format(
                                            self.env['VOLTTRON_HOME'],
                                            SUBSCRIBE_TO)
        print ("Agent Home", self.env['VOLTTRON_HOME'])
        print ("Agent Pub Addr", self.env['AGENT_PUB_ADDR'])
        print ("Agent Sub Addr", self.env['AGENT_SUB_ADDR'])
        self.p_process = None
        self.t_process = None
        self.zmq_context = None
        self.use_twistd = False

        if volttron_home is not None:
            self.initialize_volttron_home(volttron_home)

    def initialize_volttron_home(self, volttron_home):
        '''Copies the configuration of the passed "volttron_home".

        The platform is actually being run in a temporary space that is
        dynamically created.  This function will copy the directory tree
        recursively from volttron_home to the platforms true VOLTTRON_HOME.

        raises ValueError if volttron_home does not exist

        volttron_home is the directory where the configurations will be copied
                      from
        '''
        if not os.path.isdir(volttron_home):
            raise ValueError('Invalid directory specified\n{}'.format(
                                                                volttron_home))
        mergetree(volttron_home, self.tmpdir)

    def startup_platform(self, platform_config, use_twistd = False,
                         mode=UNRESTRICTED):
        try:
            print("PLATFORM CONFIG: ", platform_config)
            config = json.loads(open(platform_config, 'r').read())
        except Exception as e:
            config = None
            sys.stderr.write (str(e))

        assert config != None, 'Invalid configuration file passed {}'.format(
                                                                platform_config)

#         self.tmpdir = tempfile.mkdtemp()
        config['tmpdir'] = self.tmpdir
        pconfig = os.path.join(self.tmpdir, TMP_PLATFORM_CONFIG_FILENAME)
        self.mode = mode

        assert(self.mode in MODES, 'Invalid platform mode set: '+str(mode))
        opts = None

        if self.mode == UNRESTRICTED:
            if RESTRICTED_AVAILABLE:
                config['mobility'] = False
                config['resource-monitor'] = False
                config['verify'] = False
            with closing(open(pconfig, 'w')) as cfg:
                cfg.write(PLATFORM_CONFIG_UNRESTRICTED.format(**config))
            opts = type('Options', (), {'verify_agents': False,
                                        'volttron_home': self.tmpdir})()
        elif self.mode == RESTRICTED:
            if not RESTRICTED_AVAILABLE:
                raise ValueError("restricted is not available.")

            certsdir = os.path.join(os.path.expanduser(self.env['VOLTTRON_HOME']),
                                     'certificates')

            print ("certsdir", certsdir)
            self.certsobj = certs.Certs(certsdir)

            with closing(open(pconfig, 'w')) as cfg:
                cfg.write(PLATFORM_CONFIG_RESTRICTED.format(**config))
            opts = type('Options', (), {'resource-monitor':False,
                                        'verify_agents': True,
                                        'volttron_home': self.tmpdir})()

#                 self.create_certs()
        else:
            raise PlatformWrapperError("Invalid platform mode specified: {}".format(mode))

        self.test_aip = aip.AIPplatform(opts)
        self.test_aip.setup()

        lfile = os.path.join(self.tmpdir, "volttron.log")
        print("VOLTTRON_ROOT: ", _VOLTTRON_ROOT)

        pparams = [VSTART, "-c", pconfig, "-vv", "-l", lfile]
        print("PARAMS: ", pparams)

        self.p_process = subprocess.Popen(pparams, env=self.env)

        #Setup connector
        path = '{}/run/control'.format(self.env['VOLTTRON_HOME'])

        time.sleep(5)
        tries = 0
        max_tries = 5
        while(not os.path.exists(path) and tries < max_tries):
            time.sleep(5)
            tries += 1

        self.conn= server.ControlConnector(path)

#         if self.mode == RESTRICTED:
#             self.conn.call.create_cgroups()

        self.use_twistd = use_twistd
        #TODO: Revise this to start twistd with platform.
        if self.use_twistd:
            tconfig = os.path.join(self.tmpdir, TMP_SMAP_CONFIG_FILENAME)

            with closing(open(tconfig, 'w')) as cfg:
                cfg.write(TWISTED_CONFIG.format(**config))

            tparams = ["env/bin/twistd", "-n", "smap", tconfig]
            self.t_process = subprocess.Popen(tparams, env=self.env)
            time.sleep(5)
        #self.t_process = subprocess.Popen(["twistd", "-n", "smap", "test-smap.ini"])

    def publish(self, topic, data):
        '''Publish data to a zmq context.

        The publisher is goint to use the platform that is contained within
        this wrapper to write data to.
        '''
        if not self.zmq_context:
            self.zmq_context = zmq.Context()
        print("binding publisher to: ", self.env['AGENT_PUB_ADDR'])
        pub = zmq.Socket(zmq_context, zmq.PUB)
        pub.bind(self.env['AGENT_PUB_ADDR'])
        pub.send_multipart([topic, data])

    def fillout_file(self, filename, template, config_file):

        try:
            config = json.loads(open(config_file, 'r').read())
        except Exception as e:
            sys.stderr.write (str(e))
            raise PlatformWrapperError("Could not load configuration file for tests")

        config['tmpdir'] = self.tmpdir

        outfile = os.path.join(self.tmpdir, filename)
        with closing(open(outfile, 'w')) as cfg:
            cfg.write(template.format(**config))

        return outfile

    def create_certs(self):
        auth.create_root_ca(self.tmpdir, ca_name)

    def direct_sign_agentpackage_creator(self, package):
        assert (RESTRICTED), "Auth not available"
        print ("wrapper.certsobj", self.certsobj.cert_dir)
        assert(auth.sign_as_creator(package, 'creator', certsobj=self.certsobj)), "Signing as {} failed.".format('creator')

    def direct_sign_agentpackage_soi(self, package):
        assert (RESTRICTED), "Auth not available"
        assert(auth.sign_as_admin(package, 'soi', certsobj=self.certsobj)), "Signing as {} failed.".format('soi')

    def direct_sign_agentpackage_initiator(self, package, config_file, contract):
        assert (RESTRICTED), "Auth not available"
        files = {"config_file":config_file,"contract":contract}
        assert(auth.sign_as_initiator(package, 'initiator', files=files,
                                      certsobj=self.certsobj)), "Signing as {} failed.".format('initiator')

    def direct_build_agentpackage(self, agent_dir):
        wheel_path = packaging.create_package(os.path.join(rel_path, agent_dir), self.wheelhouse)

        return wheel_path

    def direct_send_agent(self, package, target):
        pparams = [VCTRL, SEND_AGENT, target, package]
        print (pparams, "CWD", os.getcwd())
        send_process = subprocess.call(pparams, env=self.env)
        print ("Done sending to", target)

    def direct_configure_agentpackage(self, agent_wheel, config_file):
        packaging.add_files_to_package(agent_wheel, {'config_file':os.path.join(rel_path, config_file)})

    def direct_buid_install_agent(self, agent_dir, config_file):
        agent_wheel = self.direct_build_agentpackage(agent_dir)
        self.direct_configure_agentpackage(agent_wheel, config_file)
        assert(agent_wheel is not None,"Agent wheel was not built")

        uuid = self.test_aip.install_agent(agent_wheel)
        #aip volttron_home, verify_agents
        return uuid
#         conn.call.start_agent()

    def direct_remove_agent(self, uuid):
        uuid = self.test_aip.remove_agent(uuid)

    def direct_build_install_run_agent(self, agent_dir, config_file):
        agent_uuid = self.direct_buid_install_agent(agent_dir, config_file)
        self.direct_start_agent(agent_uuid)
        return agent_uuid

    def direct_build_send_agent(self, agent_dir, config_file, target):
        agent_uuid = self.direct_buid_install_agent(agent_dir, config_file)
        self.direct_start_agent(agent_uuid)
        return agent_uuid

    def direct_start_agent(self, agent_uuid):

        self.conn.call.start_agent(agent_uuid)
        time.sleep(3)

        status = self.conn.call.status_agents()
#         self.test_aip.status_agents()

#         status = self.conn.call.status_agents()
#         self.assertEquals(len(status[0]), 4, 'Unexpected status message')
        status_uuid = status[0][0]
        assert status_uuid == agent_uuid

        assert len(status[0][2]) == 2, 'Unexpected agent status message'
        status_agent_status = status[0][2][1]
        assert not isinstance(status_agent_status, int)
#         self.assertIn("running",status_agent_status, "Agent status shows error")
        print status

    def confirm_agent_running(self, agent_name, max_retries=5, timeout_seconds=2):

#         self.test_aip.status_agents()

#         status = self.conn.call.status_agents()
#         self.assertEquals(len(status[0]), 4, 'Unexpected status message')
        running = False
        retries = 0
        while (not running and retries < max_retries):
            status = self.conn.call.status_agents()
            print ("Status", status)
            if len(status) > 0:
                status_name = status[0][1]
                assert status_name == agent_name

                assert len(status[0][2]) == 2, 'Unexpected agent status message'
                status_agent_status = status[0][2][1]
                running = not isinstance(status_agent_status, int)
            retries += 1
            time.sleep(timeout_seconds)
        return running

    def direct_stop_agent(self, agent_uuid):
        result = self.conn.call.stop_agent(agent_uuid)
        print result

    def shutdown_platform(self, cleanup_temp=True):
        '''Stop platform here'''
        if self.p_process != None:
            if self.conn is not None:
                self.conn.call.shutdown()
            time.sleep(3)
            self.p_process.kill()
        else:
            print "platform process was null"

        if self.use_twistd and self.t_process != None:
            self.t_process.kill()
            self.t_process.wait()
        elif self.use_twistd:
            print "twistd process was null"
        if cleanup_temp:
            if self.tmpdir:
                shutil.rmtree(self.tmpdir, ignore_errors=True)

    def cleanup(self, cleanup_temp=True):
        '''Shuts the platform down and cleans up based upon cleanup_temp
        '''
        try:
            self.shutdown_platform(cleanup_temp=cleanup_temp)
        except Exception as e:
            sys.stderr.write( str(e))
        finally:
            pass

    def do_nothing(self):
        pass

def mergetree(src, dst, symlinks=False, ignore=None):
    if not os.path.exists(dst):
        os.makedirs(dst)
    for item in os.listdir(src):
        s = os.path.join(src, item)
        d = os.path.join(dst, item)
        if os.path.isdir(s):
            mergetree(s, d, symlinks, ignore)
        else:
            if not os.path.exists(d) or os.stat(src).st_mtime - os.stat(dst).st_mtime > 1:
                shutil.copy2(s, d)

#     def build_install_run_agent(self, agent_dir, agent_name):
#         self.build_and_install_agent(agent_dir)
#         results = subprocess.check_output([VCTRL,CTL_START,agent_name])
#         self.assertTrue(len(results) == 0)
#         results = subprocess.check_output([VCTRL,CTL_STATUS])
#         print results
#         self.assertIn('running', results, "Agent was not started")
#

#     def build_and_install_agent(self, agent_dir):
#         agent_wheel = self.build_agentpackage(agent_dir)
#         self.assertIsNotNone(agent_wheel,"Agent wheel was not built")
#
#         sys.stderr = std_err = StringIO()
#
#         results = subprocess.check_output([VCTRL,CTL_INSTALL,agent_wheel])
#         print ("results: "+results)
#         if self.mode == UNRESTRICTED or self.mode == RESOURCE_CHECK_ONLY:
#             self.assertTrue(len(results) == 0)
#             proc_out = std_err.getvalue().split(':')
#             print (proc_out)
#             self.check_default_dir(proc_out[1].strip())
#         elif self.mode == RESTRICTED or self.mode == VERIFY_ONLY:
#             self.assertTrue(results.startswith('Unpacking to: '))
<|endoftext|>