示例#1
0
文件: server.py 项目: openre/openre
    def run(self):
        def event_done(event):
            if not event.address:
                return
            if event.message.get('no_reply'):
                return
            if event.is_success:
                ret = {
                    'success': event.is_success,
                    'data': event.result
                }
            else:
                ret = {
                    'success': event.is_success,
                    'data': event.result,
                    'error': event.error,
                    'traceback': event.traceback
                }

            self.reply(event.address, ret)
        poll_timeout = 0
        event_pool = EventPool()
        event_pool.context['server'] = self
        # init tasks
        # run proxy and broker
        run_proxy_event = Event('proxy_start', 'server',
                                {'exit_on_error': True})
        event_pool.register(run_proxy_event)

        run_broker_event = Event('broker_start', 'server',
                                 {'exit_on_error': True})
        event_pool.register(run_broker_event)

        while True:
            was_message = False
            socks = dict(self.poller.poll(poll_timeout))
            if socks.get(self.responder) == zmq.POLLIN:
                was_message = True
                message = self.responder.recv_multipart()
                if len(message) < 3:
                    logging.warn('Broken message: %s', message)
                    continue
                address = message[0]
                data = self.from_json(message[2])
                logging.debug('Received message: %s', data)

                if not isinstance(data, dict) or 'action' not in data:
                    logging.warn(
                        'Malformed data in message ' \
                        '(should be dict with \'action\' key): %s', data)
                    ret = {
                        'success': False,
                        'data': None,
                        'error': 'Malformed message: %s' % data,
                        'traceback': 'Malformed message: %s' % data
                    }
                    self.reply(address, ret)
                    continue

                bytes = None
                if data.get('bytes'):
                    bytes = message[-data['bytes']:]
                event = ServerEvent(data['action'], 'server', data, bytes,
                                    address)
                event.set_priority(data.get('priority', 0))
                event.done_callback(event_done)
                event_pool.register(event)
            if socks.get(self.broker_socket) == zmq.POLLIN:
                was_message = True
                message = self.broker_socket.recv_multipart()
                if len(message) < 3:
                    logging.warn('Broken broker response message: %s', message)
                    continue
                if message[0] != '':
                    continue
                data = self.from_json(message[2])
                logging.debug('Received response message from broker: %s', data)
                event_id = (data.get('context', {}) or {}).get('event_id')
                if event_id:
                    for event in event_pool.event_list:
                        if event.context.get('event_id') == event_id:
                            if event.is_done:
                                break
                            event.is_success = data['success']
                            event.result = data['data']
                            if not data['success']:
                                event.error = data.get('error')
                                event.traceback = data.get('traceback')
                            event.done()
                            break
            # get all messages, then proceed
            if was_message:
                poll_timeout = 0
                continue
            event_pool.tick()
            # if no events - then wait for new events without timeout
            poll_timeout = event_pool.poll_timeout()
示例#2
0
文件: domain.py 项目: openre/openre
    def run(self):
        def event_done(event):
            # if we don't need the result of the event
            if not event.data.get('wait') or event.data.get('no_reply'):
                return
            if event.is_success:
                ret = {
                    'success': event.is_success,
                    'data': event.result
                }
            else:
                logging.warn(event.traceback)
                ret = {
                    'success': event.is_success,
                    'data': event.result,
                    'error': event.error,
                    'traceback': event.traceback
                }

            if event.data.get('context'):
                ret['context'] = event.data['context']
            self.reply(event.address, ret)
        poll_timeout = 0
        event_pool = EventPool()
        event_pool.context['agent'] = self

        process_name = self.config['name']
        if not process_name:
            process_name = 'domain'
        if process_name != 'domain':
            process_name = 'domain.%s' % self.config['name']
        self.send_server('process_state', {
            'name':  process_name
        })
        self.send_server('domain_state', {
            'state': 'blank',
            'status': 'done',
            'name': self.config['name']
        })

        register_spike_pack_cache = None
        # main loop
        while True:
            # receive all messages in while loop
            while True:
                was_message = False
                max_priority = 0
                socks = dict(self.poller.poll(poll_timeout))
                if socks.get(self.backend) == zmq.POLLIN:
                    was_message = True
                    message = self.backend.recv_multipart()
                    if len(message) > 4:
                        logging.debug("in binary: %s", message[:4])
                    else:
                        logging.debug("in: %s", message)


                    data = self.from_json(message[3])
                    address = [message[0], '', message[2]]
                    if not isinstance(data, dict) or 'action' not in data:
                        logging.warn(
                            'Malformed data in message ' \
                            '(should be dict with \'action\' key): %s', data)
                        ret = {
                            'success': False,
                            'data': None,
                            'error': 'Malformed message: %s' % message,
                            'traceback': 'Malformed message: %s' % message
                        }
                        self.reply(address, ret)
                    else:
                        bytes = None
                        if data.get('bytes'):
                            bytes = message[-data['bytes']:]
                        event = DomainEvent(data['action'], 'domain', data,
                                            bytes, address)
                        priority = data.get('priority', 0)
                        event.set_priority(priority)
                        event.done_callback(event_done)
                        event_pool.register(event)
                        if priority > max_priority:
                            max_priority = priority
                        # send response immediately
                        if not data.get('wait') and not data.get('no_reply'):
                            if is_registered_action(data['action'], 'domain'):
                                ret = {
                                    'success': True,
                                    'data': None,
                                }
                            else:
                                ret = {
                                    'success': False,
                                    'data': None,
                                    'error': 'Action "%s" in not registered'
                                        % data['action'],
                                    'traceback': 'Action "%s" in not registered'
                                        % data['action'],
                                }
                            if event.data.get('context'):
                                ret['context'] = event.data['context']
                            self.reply(address, ret)

                if socks.get(self.sub) == zmq.POLLIN:
                    was_message = True
                    message = self.sub.recv_multipart()
                    agent_id = message[0]
                    data_type = message[1]
                    data = message[2]
                    if self.id.bytes == agent_id and data_type == 'S':
                        try:
                            register_spike_pack_cache(bytes=data)
                        except TypeError:
                            if 'local_domain' in self.context:
                                register_spike_pack_cache = \
                                    self.context['local_domain'] \
                                        .register_spike_pack
                                register_spike_pack_cache(bytes=data)
                    elif self.id.bytes == agent_id:
                        do_action(data_type, 'domain', self, *message[2:])
                # receive all messages, and only then process them
                if was_message:
                    poll_timeout = 0
                    # we will proccess messages with priority grater than 0
                    if max_priority <= 0:
                        continue
                event_pool.tick()
                # if no events - then wait for new events without timeout
                poll_timeout = event_pool.poll_timeout()