Python SocketManager.bind_to_random_port方法代码示例

# 需要导入模块: from agentzero.core import SocketManager [as 别名]
# 或者: from agentzero.core.SocketManager import bind_to_random_port [as 别名]
def test_socket_bind_to_random_port(get_by_name, register_socket, engage):
    ("SocketManager().bind_to_random_port() should get by name, register then "
     "bind to a random port")

    # Given a zmq mock
    zmq = Mock()

    # And a context
    context = Mock()

    # And a socket manager
    manager = SocketManager(zmq, context)

    # And that the socket is mocked to return a random port
    socket = get_by_name.return_value
    socket.bind_to_random_port.return_value = 42000

    # When I call bind
    result = manager.bind_to_random_port('foobar', 'some-mechanism')

    # Then it should have returned a socket from get_by_name
    result.should.equal((socket, 'tcp://0.0.0.0:42000'))

    # And get_by_name should have been called with the name
    get_by_name.assert_called_once_with('foobar')

    # And register_socket should have been called
    register_socket.assert_called_once_with(
        socket,
        "some-mechanism"
    )

    # And it should have called bind on the given address
    socket.bind_to_random_port.assert_called_once_with('tcp://0.0.0.0')

    # And it should have engaged the manager with zero timeout
    engage.assert_called_once_with(0)

# 需要导入模块: from agentzero.core import SocketManager [as 别名]
# 或者: from agentzero.core.SocketManager import bind_to_random_port [as 别名]
class Step(object):
    def __init__(self, pull_bind_address='tcp://127.0.0.1', subscriber_connect_address='tcp://127.0.0.1:6000', concurrency=100, timeout=1):
        self.context = zmq.Context()
        self.sockets = SocketManager(zmq, self.context)
        self.sockets.create('pull-in', zmq.PULL)
        # self.sockets.set_socket_option('pull-in', zmq.RCVHWM, concurrency)
        self.sockets.create('events', zmq.PUB)
        self.name = self.__class__.__name__
        self.subscriber_connect_address = subscriber_connect_address
        self._allowed_to_run = True
        self.pool = gevent.pool.Pool(concurrency + 1)
        self.timeout = timeout
        self.pull_bind_address = pull_bind_address
        self.id = str(uuid.uuid4())
        self.logger = self.sockets.get_logger('events', 'logs', 'logs')

    def listen(self):
        _, self.address = self.sockets.bind_to_random_port('pull-in', zmq.POLLIN, local_address=self.pull_bind_address)
        gevent.sleep(1)

    def connect(self):
        self.sockets.connect('events', self.subscriber_connect_address, zmq.POLLOUT)
        gevent.sleep(1)
        self.notify_available()

    def notify_available(self):
        self.send_event('available', self.to_dict())

    def should_run(self):
        gevent.sleep(0.1)
        return self._allowed_to_run

    def send_event(self, name, data):
        self.sockets.publish_safe('events', name, data)

    def dispatch(self, job):
        try:
            start = time.time()
            job['job_started_at'] = start
            self.send_event('job:started', job)
            job['instructions'] = self.execute(job['instructions'])
            job['job_finished_at'] = time.time()
            self.logger.info("done processing %s", job)

        except Exception as e:
            job['job_finished_at'] = time.time()
            job['error'] = serialized_exception(e)
            self.send_event('job:error', job)
            job['success'] = False
            self.send_event('job:failed', job)
            self.logger.exception('failed to execute job {id}'.format(**dict(job)))
        finally:
            end = time.time()
            job['end'] = end

        if job.get('instructions'):
            self.send_event('job:success', job)

        else:
            self.send_event('job:failed', job)

    def to_dict(self):
        return dict(
            id=self.id,
            address=self.address,
            job_type=self.job_type,
            step_name=self.name,
        )

    def loop(self):
        self.listen()
        self.connect()
        self.logger.info('listening for jobs on %s', self.address)
        while self.should_run():
            if self.pool.free_count() == 0:
                self.logger.info('waiting for an execution slot')
                self.pool.wait_available()

            job = self.sockets.recv_safe('pull-in')

            if job:
                self.logger.info('received job')
                self.pool.spawn(self.dispatch, job)
            else:
                self.notify_available()
                gevent.sleep(1)