Exemple #1
0
 def apply(callback) -> 'promise.Promise':
     f = Future()
     try:
         f.set_result(callback())
     except BaseException as e:
         f.set_exception(e)
     return Promise(f)
Exemple #2
0
    def test_future_list(self):
        d = Doc('a')
        f = Future()
        f.set_result([etree.Comment('ccc'), etree.Element('bbb')])
        d.put(f)

        self.assertXmlEqual(d.to_etree_element(), """<?xml version='1.0'?>\n<a><!--ccc--><bbb/></a>""")
Exemple #3
0
    def _start_processing_requests(self):
        while True:
            data = yield gen.Task(self._stream.read_until, '\r\n')
            log.debug('New request: %r', data)
            try:
                msg = json.loads(data)
                key = msg['key']
                method = msg['method']
                args = msg['args']
                kwargs = msg['kwargs']
            except (KeyError, ValueError):
                log.error('Malformed request data: %s', data)
                continue
            try:
                res = self._handler(method, *args, **kwargs)
                if isinstance(res, Future):
                    future = res
                else:
                    future = Future()
                    future.set_result(res)
            except Exception as e:
                log.exception('Failed to handle request: %s', key)
                future = concurrent.TracebackFuture()
                future.set_exception(e)

            future.add_done_callback(partial(self._on_future_finished, key))
  def test_get(self):
    # Mock cassandra response
    async_response = Future()
    async_response.set_result([
      ('a', 'c1', '1'), ('a', 'c2', '2'), ('a', 'c3', '3'),
      ('b', 'c1', '4'), ('b', 'c2', '5'), ('b', 'c3', '6'),
      ('c', 'c1', '7'), ('c', 'c2', '8'), ('c', 'c3', '9'),
    ])
    self.execute_mock.return_value = async_response

    # Call function under test
    keys = ['a', 'b', 'c']
    columns = ['c1', 'c2', 'c3']
    result = yield self.db.batch_get_entity('table', keys, columns)

    # Make sure cassandra interface prepared good query
    query = self.execute_mock.call_args[0][0]
    parameters = self.execute_mock.call_args[1]["parameters"]
    self.assertEqual(
      query.query_string,
      'SELECT * FROM "table" WHERE key IN %s and column1 IN %s')
    self.assertEqual(parameters, ([b'a', b'b', b'c'], ['c1', 'c2', 'c3']) )
    # And result matches expectation
    self.assertEqual(result, {
      'a': {'c1': '1', 'c2': '2', 'c3': '3'},
      'b': {'c1': '4', 'c2': '5', 'c3': '6'},
      'c': {'c1': '7', 'c2': '8', 'c3': '9'}
    })
Exemple #5
0
    def test_both(self):
        marker = object()
        second_marker = object()
        result_marker = object()

        def _mapper(_):
            return marker

        def _exception_mapper(_):
            return second_marker

        first_future = Future()
        folded_future = future_fold(first_future, result_mapper=_mapper, exception_mapper=_exception_mapper)
        folded_future_probe = FutureProbe(folded_future)

        second_future = Future()
        second_folded_future = future_fold(second_future, result_mapper=_mapper, exception_mapper=_exception_mapper)
        second_folded_future_probe = FutureProbe(second_folded_future, stop_cb=self.stop)

        first_future.set_result(result_marker)
        second_future.set_exception(MyException())
        self.wait()

        folded_future_probe.assert_single_result_call(self, marker)
        second_folded_future_probe.assert_single_result_call(self, second_marker)
    def setUp(self):
        AsyncTestCase.setUp(self)

        return_value = {
            'label': 'A Label',
            'comment': 'A Comment',
            'depiction': 'A URL'
        }
        future = Future()
        future.set_result(return_value)
        self._wrapped_fn = Mock(return_value=future)
        self._decorated_fn = cache_facts(self._wrapped_fn)
        
        #setup a function with a value to be cached
        #and decorate it with the decorated under test
        @cache_facts
        def intense_fact_processing(uri):
            return {
                'label': 'A Label',
                'comment': 'A Comment',
                'depiction': 'A URL',
                'uri': uri
            }

        self.intense_fact_processing =  intense_fact_processing
  def test_range_query(self):
    # Mock cassandra response
    async_response = Future()
    async_response.set_result([
      ('keyA', 'c1', '1'), ('keyA', 'c2', '2'),
      ('keyB', 'c1', '4'), ('keyB', 'c2', '5'),
      ('keyC', 'c1', '7'), ('keyC', 'c2', '8')
    ])
    self.execute_mock.return_value = async_response

    # Call function under test
    columns = ['c1', 'c2']
    result = yield self.db.range_query("tableZ", columns, "keyA", "keyC", 5)

    # Make sure cassandra interface prepared good query
    query = self.execute_mock.call_args[0][0]
    parameters = self.execute_mock.call_args[1]["parameters"]
    self.assertEqual(
      query.query_string,
      'SELECT * FROM "tableZ" WHERE '
      'token(key) >= %s AND '
      'token(key) <= %s AND '
      'column1 IN %s '
      'LIMIT 10 '    # 5 * number of columns
      'ALLOW FILTERING')
    self.assertEqual(parameters, (b'keyA', b'keyC', ['c1', 'c2']) )
    # And result matches expectation
    self.assertEqual(result, [
      {'keyA': {'c1': '1', 'c2': '2'}},
      {'keyB': {'c1': '4', 'c2': '5'}},
      {'keyC': {'c1': '7', 'c2': '8'}}
    ])
def queue_get():
    future = Future()
    try:
        future.set_result(queue.get_nowait())
    except QueueEmpty:
        pass
    return future
Exemple #9
0
    def _get_conn(self):  # -> Future[connection]
        now = self.io_loop.time()

        # Try to reuse in free pool
        while self._free_conn:
            conn = self._free_conn.popleft()
            if now - conn.connected_time > self.max_recycle_sec:
                self._close_async(conn)
                continue
            log.debug("Reusing connection from pool: %s", self.stat())
            fut = Future()
            fut.set_result(conn)
            return fut

        # Open new connection
        if self.max_open == 0 or self._opened_conns < self.max_open:
            self._opened_conns += 1
            log.debug("Creating new connection: %s", self.stat())
            fut = connect(**self.connect_kwargs)
            fut.add_done_callback(self._on_connect)  # self._opened_conns -=1 on exception
            return fut

        # Wait to other connection is released.
        fut = Future()
        self._waitings.append(fut)
        return fut
Exemple #10
0
    def put(self, value):
        """Puts an item into the queue.

        Returns a Future that resolves to None once the value has been
        accepted by the queue.
        """
        io_loop = IOLoop.current()
        new_hole = Future()

        new_put = Future()
        new_put.set_result(new_hole)

        with self._lock:
            self._put, put = new_put, self._put

        answer = Future()

        def _on_put(future):
            if future.exception():  # pragma: no cover (never happens)
                return answer.set_exc_info(future.exc_info())

            old_hole = put.result()
            old_hole.set_result(Node(value, new_hole))
            answer.set_result(None)

        io_loop.add_future(put, _on_put)
        return answer
    def fetch(self, request, callback=None, raise_error=True, **kwargs):
        if not isinstance(request, HTTPRequest):
            request = HTTPRequest(url=request, **kwargs)

        key = self.cache.create_key(request)

        # Check and return future if there is a pending request
        pending = self.pending_requests.get(key)
        if pending:
            return pending

        response = self.cache.get_response_and_time(key)
        if response:
            response.cached = True
            if callback:
                self.io_loop.add_callback(callback, response)
            future = Future()
            future.set_result(response)
            return future

        future = orig_fetch(self, request, callback, raise_error, **kwargs)

        self.pending_requests[key] = future

        def cache_response(future):
            exc = future.exception()
            if exc is None:
                self.cache.save_response(key, future.result())

        future.add_done_callback(cache_response)
        return future
Exemple #12
0
    def test_future_etree_element(self):
        d = Doc('a')
        f = Future()
        f.set_result(etree.Element('b'))
        d.put(f)

        self.assertXmlEqual(d.to_etree_element(), """<?xml version='1.0' encoding='utf-8'?>\n<a><b/></a>""")
 def sendmessage(self,topic,msg,key=None):
     if key is None:
         raise Exception("dispatch key is none")
     future = Future()
     res = self.mqtt.sendmessage(topic,msg,key)
     future.set_result('ok')
     return future
    def test_rank_calls_sort_and_returns_output(self):
        #setup the response from the summarum endpoint
        expected_result = [
            ('a', 'b', 10.0),
            ('c', 'd', 1.0)
        ]
        future = Future()
        future.set_result(expected_result)
        self.endpoint.fetch_and_parse = Mock(return_value=future)
        
        #setup the response return value from the sort call
        expected_ranked_facts = {
            'predicate': {},
            'objects': []
        }
        self.ranking_service.sort = Mock(return_value = expected_ranked_facts)
        
        #call the function under test
        facts = {}
        ranked_facts = yield self.ranking_service.rank(facts)

        #check that sort was called
        self.ranking_service.sort.assert_called_once_with(facts)
        #check that rank returns the output from sort
        self.assertEquals(ranked_facts, expected_ranked_facts)
Exemple #15
0
def maybe_future(x):
    if is_future(x):
        return x
    else:
        fut = Future()
        fut.set_result(x)
        return fut
        def wrapper(*args, **kargs):
            future = Future()
            future.set_result(self._response)

            with patch.object(AsyncHTTPClient, "fetch", return_value=future):
                with patch.object(Client, "fetch", return_value=future):
                    yield coroutine(*args, **kargs)
 def update_directories(self,update_dir_list):
     res_future = Future()
     res = {}
     for dirpath in update_dir_list:
         dir_list = []
         file_list = []
         try:
             for i in common.get_dir_contents(self.current_user, dirpath):
                 if i[0].startswith('.'):
                     continue
                 if i[2]:
                     dir_list.append(tuple(list(i)+[i[0].lower()]))
                     # dir_list.append(i)
                 else:
                     file_list.append(tuple(list(i)+[i[0].lower()]))
                     # file_list.append(i)
             # dir_list.sort()
             # file_list.sort()
             dir_list = sorted(dir_list,key=operator.itemgetter(3))
             file_list = sorted(file_list,key=operator.itemgetter(3))
             res.update({dirpath:dir_list+file_list})
         except common.MissingFileError:
             continue
     res_future.set_result(res)
     return res_future
Exemple #18
0
    def write(self, data):
        assert isinstance(data, bytes)
        if self._closed:
            raise StreamClosedError(real_error=self.error)

        if not data:
            if self._write_future:
                return self._write_future
            future = Future()
            future.set_result(None)
            return future

        if self._write_buffer_size:
            self._write_buffer += data
        else:
            self._write_buffer = bytearray(data)
        self._write_buffer_size += len(data)
        future = self._write_future = Future()

        if not self._connecting:
            self._handle_write()
            if self._write_buffer_size:
                if not self._state & self.io_loop.WRITE:
                    self._state = self._state | self.io_loop.WRITE
                    self.io_loop.update_handler(self.fileno(), self._state)

        return future
Exemple #19
0
 def acquire(self, pool_need_log=False):
     """Occupy free connection"""
     future = Future()
     while True:
         if self.free:
             conn = self.free.pop()
             if conn.valid:
                 self.busy.add(conn)
             else:
                 self.dead.add(conn)
                 continue
             future.set_result(conn)
             conn.connection_need_log = pool_need_log
             log.debug("Acquired free connection %s", conn.fileno)
             return future
         elif self.busy:
             log.debug("No free connections, and some are busy - put in waiting queue")
             self.waiting_queue.appendleft(future)
             return future
         elif self.pending:
             log.debug("No free connections, but some are pending - put in waiting queue")
             self.waiting_queue.appendleft(future)
             return future
         else:
             log.debug("All connections are dead")
             return None
Exemple #20
0
    def test_moment(self):
        calls = []

        @gen.coroutine
        def f(name, yieldable):
            for i in range(5):
                calls.append(name)
                yield yieldable

        # First, confirm the behavior without moment: each coroutine
        # monopolizes the event loop until it finishes.
        immediate = Future()  # type: Future[None]
        immediate.set_result(None)
        yield [f("a", immediate), f("b", immediate)]
        self.assertEqual("".join(calls), "aaaaabbbbb")

        # With moment, they take turns.
        calls = []
        yield [f("a", gen.moment), f("b", gen.moment)]
        self.assertEqual("".join(calls), "ababababab")
        self.finished = True

        calls = []
        yield [f("a", gen.moment), f("b", immediate)]
        self.assertEqual("".join(calls), "abbbbbaaaa")
Exemple #21
0
class ManualCapClient(BaseCapClient):
    def capitalize(self, request_data, callback=None):
        logging.debug("capitalize")
        self.request_data = request_data
        self.stream = IOStream(socket.socket())
        self.stream.connect(('127.0.0.1', self.port),
                            callback=self.handle_connect)
        self.future = Future()
        if callback is not None:
            self.future.add_done_callback(
                stack_context.wrap(lambda future: callback(future.result())))
        return self.future

    def handle_connect(self):
        logging.debug("handle_connect")
        self.stream.write(utf8(self.request_data + "\n"))
        self.stream.read_until(b'\n', callback=self.handle_read)

    def handle_read(self, data):
        logging.debug("handle_read")
        self.stream.close()
        try:
            self.future.set_result(self.process_response(data))
        except CapError as e:
            self.future.set_exception(e)
Exemple #22
0
    def test_looking_for_driver_no_drivers(self):
        user = {
            'chat_id': 0,
            'current_location': [0., 0.]
        }

        drivers = []
        future_get_drivers = Future()
        future_get_drivers.set_result(drivers)
        self.users.get_drivers_within_distance = mock.MagicMock(
            return_value=future_get_drivers
        )

        yield self.stage.run(user, {})

        self.stage.sender.assert_has_calls([
            mock.call({
                'chat_id': 0,
                'text': 'looking for a driver'
            }),
            mock.call({
                'chat_id': 0,
                'text': 'no available drivers found'
            })
        ])
        self.assertEqual(2,self.stage.sender.call_count)
        self.assertEqual(None, user['proposed_driver'])
Exemple #23
0
    def acquire(
        self, timeout: Union[float, datetime.timedelta] = None
    ) -> "Future[_ReleasingContextManager]":
        """Decrement the counter. Returns a Future.

        Block if the counter is zero and wait for a `.release`. The Future
        raises `.TimeoutError` after the deadline.
        """
        waiter = Future()  # type: Future[_ReleasingContextManager]
        if self._value > 0:
            self._value -= 1
            waiter.set_result(_ReleasingContextManager(self))
        else:
            self._waiters.append(waiter)
            if timeout:

                def on_timeout() -> None:
                    if not waiter.done():
                        waiter.set_exception(gen.TimeoutError())
                    self._garbage_collect()

                io_loop = ioloop.IOLoop.current()
                timeout_handle = io_loop.add_timeout(timeout, on_timeout)
                waiter.add_done_callback(
                    lambda _: io_loop.remove_timeout(timeout_handle)
                )
        return waiter
Exemple #24
0
    def get(self, timeout: Union[float, datetime.timedelta] = None) -> Awaitable[_T]:
        """Remove and return an item from the queue.

        Returns an awaitable which resolves once an item is available, or raises
        `tornado.util.TimeoutError` after a timeout.

        ``timeout`` may be a number denoting a time (on the same
        scale as `tornado.ioloop.IOLoop.time`, normally `time.time`), or a
        `datetime.timedelta` object for a deadline relative to the
        current time.

        .. note::

           The ``timeout`` argument of this method differs from that
           of the standard library's `queue.Queue.get`. That method
           interprets numeric values as relative timeouts; this one
           interprets them as absolute deadlines and requires
           ``timedelta`` objects for relative timeouts (consistent
           with other timeouts in Tornado).

        """
        future = Future()  # type: Future[_T]
        try:
            future.set_result(self.get_nowait())
        except QueueEmpty:
            self._getters.append(future)
            _set_timeout(future, timeout)
        return future
Exemple #25
0
    def wait(self, timeout: Union[float, datetime.timedelta] = None) -> "Future[None]":
        """Block until the internal flag is true.

        Returns a Future, which raises `tornado.util.TimeoutError` after a
        timeout.
        """
        fut = Future()  # type: Future[None]
        if self._value:
            fut.set_result(None)
            return fut
        self._waiters.add(fut)
        fut.add_done_callback(lambda fut: self._waiters.remove(fut))
        if timeout is None:
            return fut
        else:
            timeout_fut = gen.with_timeout(
                timeout, fut, quiet_exceptions=(CancelledError,)
            )
            # This is a slightly clumsy workaround for the fact that
            # gen.with_timeout doesn't cancel its futures. Cancelling
            # fut will remove it from the waiters list.
            timeout_fut.add_done_callback(
                lambda tf: fut.cancel() if not fut.done() else None
            )
            return timeout_fut
def wechat_api_mock(client, request, *args, **kwargs):

    url = urlparse(request.url)
    path = url.path.replace('/cgi-bin/', '').replace('/', '_')
    if path.startswith('_'):
        path = path[1:]
    res_file = os.path.join(_FIXTURE_PATH, '%s.json' % path)
    content = {
        'errcode': 99999,
        'errmsg': 'can not find fixture %s' % res_file,
    }
    headers = {
        'Content-Type': 'application/json'
    }
    try:
        with open(res_file, 'rb') as f:
            content = f.read().decode('utf-8')
    except (IOError, ValueError) as e:
        content['errmsg'] = 'Loads fixture {0} failed, error: {1}'.format(
            res_file,
            e
        )
        content = json.dumps(content)

    buffer = StringIO(content)
    resp = HTTPResponse(
        request,
        200,
        headers=headers,
        buffer=buffer,
    )
    future = Future()
    future.set_result(resp)
    return future
    def get_tweets(username):
        result_future = Future()
        """helper function to fetch 200 tweets for a user with @username
        """
        TWITTER_URL = 'https://api.twitter.com/1.1/statuses/user_timeline.json'
        '''
        curl --get 'https://api.twitter.com/1.1/statuses/user_timeline.json' --data 'count=200&screen_name=twitterapi' --header 'Authorization: OAuth oauth_consumer_key="BlXj0VRgkpUOrN3b6vTyJu8YB", oauth_nonce="9cb4b1aaa1fb1d79e0fbd9bc8b33f82a", oauth_signature="SrJxsOCzOTnudKQMr4nMQ0gDuRk%3D", oauth_signature_method="HMAC-SHA1", oauth_timestamp="1456006969", oauth_token="701166849883373568-bqVfk8vajGxWIlKDe94CRjMJtBvwdQQ", oauth_version="1.0"' --verbose
        '''
        auth = OAuth1('BlXj0VRgkpUOrN3b6vTyJu8YB', 'qzkhGeWIYVXod9umMuinHF2OFmJxiucQspX5JsA7aH8xs5t4DT',
                      '701166849883373568-bqVfk8vajGxWIlKDe94CRjMJtBvwdQQ', 'y3gx0F5fLyIQQFNDev8JtpPKpEUmyy3mMibxCcTK2kbZZ')

        data = {'count': 200,
                'screen_name': username}

        r = requests.get(url=TWITTER_URL, params=data, auth=auth)
        data = r.json()

        if 'errors' in data:
            raise Exception
        res = []
        for item in data:
            if 'retweeted_status' not in item.keys():
                res.append(item)
        result_future.set_result(res)
        return result_future
Exemple #28
0
    def release_connection(self, connection):
        if self._closed:
            return connection.do_close()

        if not connection.open:
            future = Future()
            future.set_result(None)
            return future

        if self.continue_next_wait(connection):
            while self._wait_connections and self._connections:
                connection = self._connections.pop()
                if connection.open:
                    if self.continue_next_wait(connection):
                        self._used_connections[id(connection)] = connection
                    else:
                        self._connections.append(connection)
                        break
        else:
            try:
                del self._used_connections[id(connection)]
                self._connections.append(connection)
                connection.idle_time = time.time()
            except KeyError:
                if connection not in self._connections:
                    IOLoop.current().add_callback(connection.do_close)
                    raise ConnectionNotFoundError("Connection not found.")
                else:
                    raise ConnectionNotUsedError("Connection is not used, you maybe close wrong connection.")

        future = Future()
        future.set_result(None)
        return future
Exemple #29
0
 def wait(self, position):
     future = Future()
     if position != self.count():
         future.set_result(dict(position=self.count(), last_word=self.last()))
     else:
         self.waiters.add(future)
     return future
def test_composite_reporter():
    reporter = jaeger_client.reporter.CompositeReporter(
        jaeger_client.reporter.NullReporter(),
        jaeger_client.reporter.LoggingReporter())
    with mock.patch('jaeger_client.reporter.NullReporter.set_process') \
            as null_mock:
        with mock.patch('jaeger_client.reporter.LoggingReporter.set_process') \
                as log_mock:
            reporter.set_process('x', {}, 123)
            null_mock.assert_called_with('x', {}, 123)
            log_mock.assert_called_with('x', {}, 123)
    with mock.patch('jaeger_client.reporter.NullReporter.report_span') \
            as null_mock:
        with mock.patch('jaeger_client.reporter.LoggingReporter.report_span') \
                as log_mock:
            reporter.report_span({})
            null_mock.assert_called_with({})
            log_mock.assert_called_with({})
    with mock.patch('jaeger_client.reporter.NullReporter.close') \
            as null_mock:
        with mock.patch('jaeger_client.reporter.LoggingReporter.close') \
                as log_mock:

            f1 = Future()
            f2 = Future()
            null_mock.return_value = f1
            log_mock.return_value = f2
            f = reporter.close()
            null_mock.assert_called_once()
            log_mock.assert_called_once()
            assert not f.done()
            f1.set_result(True)
            f2.set_result(True)
            assert f.done()
Exemple #31
0
    def test_already_done(self):
        f1 = Future()
        f2 = Future()
        f3 = Future()
        f1.set_result(24)
        f2.set_result(42)
        f3.set_result(84)

        g = gen.WaitIterator(f1, f2, f3)
        i = 0
        while not g.done():
            r = yield g.next()
            # Order is not guaranteed, but the current implementation
            # preserves ordering of already-done Futures.
            if i == 0:
                self.assertEqual(g.current_index, 0)
                self.assertIs(g.current_future, f1)
                self.assertEqual(r, 24)
            elif i == 1:
                self.assertEqual(g.current_index, 1)
                self.assertIs(g.current_future, f2)
                self.assertEqual(r, 42)
            elif i == 2:
                self.assertEqual(g.current_index, 2)
                self.assertIs(g.current_future, f3)
                self.assertEqual(r, 84)
            i += 1

        self.assertEqual(g.current_index, None, "bad nil current index")
        self.assertEqual(g.current_future, None, "bad nil current future")

        dg = gen.WaitIterator(f1=f1, f2=f2)

        while not dg.done():
            dr = yield dg.next()
            if dg.current_index == "f1":
                self.assertTrue(dg.current_future == f1 and dr == 24,
                                "WaitIterator dict status incorrect")
            elif dg.current_index == "f2":
                self.assertTrue(dg.current_future == f2 and dr == 42,
                                "WaitIterator dict status incorrect")
            else:
                self.fail("got bad WaitIterator index {}".format(
                    dg.current_index))

            i += 1

        self.assertEqual(dg.current_index, None, "bad nil current index")
        self.assertEqual(dg.current_future, None, "bad nil current future")
    def test_nested_future(self):
        j = JsonBuilder()
        f1 = Future()
        f2 = Future()
        f3 = Future()

        f1.set_result({'nested': f2})
        j.put(f1)

        self.assertEqual(j.to_string(), """{"nested": null}""")

        f2.set_result({'a': f3})
        f3.set_result(['b', 'c'])

        self.assertEqual(j.to_string(), """{"nested": {"a": ["b", "c"]}}""")
Exemple #33
0
    def exchange_declare(self,
                         exchange,
                         exchange_type='direct',
                         passive=False,
                         durable=False,
                         auto_delete=False,
                         internal=False,
                         nowait=False,
                         arguments=None,
                         type=None):
        f = Future()

        self.channel.exchange_declare(lambda *a: f.set_result(a),
                                      exchange=exchange,
                                      exchange_type=exchange_type,
                                      passive=passive,
                                      durable=durable,
                                      auto_delete=auto_delete,
                                      internal=internal,
                                      nowait=nowait,
                                      arguments=arguments,
                                      type=type)
        return f
Exemple #34
0
    async def test_call_without_nodes(self, cfg, scanners):
        nodes = []
        future = Future()
        future.set_result(nodes)
        self.thread._get_nodes_for_scanning = MagicMock(return_value=future)

        tcp_scanner = MagicMock()
        udp_scanner = MagicMock()

        scanners.return_value = {
            TransportProtocol.TCP: tcp_scanner,
            TransportProtocol.UDP: udp_scanner
        }

        future = Future()
        future.set_result(MagicMock())
        self.thread.run_scan = MagicMock(return_value=future)

        future_run_scan = Future()
        future_run_scan.set_result(MagicMock())
        self.thread.run_scan.return_value = future_run_scan

        await self.thread.run()
        self.assertFalse(self.thread.run_scan.called)
Exemple #35
0
def sleep(duration):
    f = Future()
    ioloop.IOLoop.current().call_later(duration, lambda: f.set_result(None))
    return f
Exemple #36
0
 def set(self, key, value, *args, **kwargs):
     f = Future()
     f.set_result(None)
     return f
Exemple #37
0
 def get(self, key):
     f = Future()
     f.set_result(self._cache.get(key))
     return f
Exemple #38
0
def daemon(args,
           restart=1,
           first_line=None,
           stream=True,
           timeout=5,
           buffer_size='line',
           **kwargs):
    '''
    This is the same as :py:class:`Subprocess`, but has a few additional checks.

    1. If we have already called :py:class:`Subprocess` with the same arguments,
       re-use the same instance.
    2. Send the process STDOUT and STDERR to this application's STDERR. This
       makes it easy to see what errors the application reports.
    3. Supports retry attempts.
    4. Checks if the first line of output is a matches a string / re -- ensuring
       that the application started properly.
    '''
    arg_str = args if isinstance(args, six.string_types) else ' '.join(args)
    try:
        key = cache_key(arg_str, kwargs)
    except (TypeError, ValueError):
        app_log.error('daemon args must be JSON serializable')
        raise
    # Send the stdout and stderr to (a) stderr AND to (b) a local queue we read
    queue = Queue(maxsize=10)
    for channel in ('stream_stdout', 'stream_stderr'):
        if channel not in kwargs:
            kwargs[channel] = []
        elif not isinstance(kwargs[channel], list):
            kwargs[channel] = [kwargs[channel]]
        if first_line:
            kwargs[channel].append(queue.put)
        if stream is True:
            kwargs[channel].append(_stderr_write)
        elif callable(stream):
            kwargs[channel].append(stream)
    # Buffer by line by default. This is required for the first_line check, not otherwise.
    kwargs['buffer_size'] = buffer_size
    # started is set if we actually call Subprocess as part of this function
    started = False

    # If process was never started, start it
    if key not in _daemons:
        started = _daemons[key] = Subprocess(args, **kwargs)

    # Ensure that process is running. Restart if required
    proc = _daemons[key]
    restart = int(restart)
    while proc.proc.returncode is not None and restart > 0:
        restart -= 1
        proc = started = _daemons[key] = Subprocess(args, **kwargs)
    if proc.proc.returncode is not None:
        raise RuntimeError('Error %d starting %s' %
                           (proc.proc.returncode, arg_str))
    if started:
        app_log.info('Started: %s', arg_str)

    future = Future()
    # If process was started, wait until it has initialized. Else just return the proc
    if first_line and started:
        if isinstance(first_line, six.string_types):

            def check(proc):
                actual = queue.get(timeout=timeout).decode('utf-8')
                if first_line not in actual:
                    raise AssertionError('%s: wrong first line: %s (no "%s")' %
                                         (arg_str, actual, first_line))
        elif isinstance(first_line, _regex_type):

            def check(proc):
                actual = queue.get(timeout=timeout).decode('utf-8')
                if not first_line.search(actual):
                    raise AssertionError('%s: wrong first line: %s' %
                                         (arg_str, actual))
        elif callable(first_line):
            check = first_line
        loop = _get_current_ioloop()

        def checker(proc):
            try:
                check(proc)
            except Exception as e:
                loop.add_callback(future.set_exception, e)
            else:
                loop.add_callback(future.set_result, proc)

        proc._check_thread = t = Thread(target=checker, args=(proc, ))
        t.daemon = True  # Thread dies with the program
        t.start()
    else:
        future.set_result(proc)
    return future
Exemple #39
0
 def get(self, key):
     f = Future()
     f.set_result(None)
     return f
Exemple #40
0
def future(f):
    r = Future()
    r.set_result(f)
    return r
Exemple #41
0
 def _create_future(self, result):
     future = Future()
     future.set_result(result)
     return future
Exemple #42
0
                except psycopg2.Error as error:
                    log.debug("Method failed Asynchronously")
                    return self._retry(retry, when_available, conn, keep,
                                       future)

                future.set_result(result)
                if not keep:
                    self.putconn(conn)

            self.ioloop.add_future(future_or_result, when_done)

        if not connection:
            self.ioloop.add_future(self.getconn(ping=False), when_available)
        else:
            f = Future()
            f.set_result(connection)
            when_available(f)
        return future

    def _retry(self, retry, what, conn, keep, future):
        if conn.closed:
            if not retry:
                retry.append(conn)
                self.ioloop.add_future(conn.connect(), what)
                return
            else:
                future.set_exception(self._no_conn_available_error)
        else:
            future.set_exc_info(sys.exc_info())
        if not keep:
            self.putconn(conn)
Exemple #43
0
def async_fetch_future(url):
    http_client = AsyncHTTPClient()
    my_future = Future()
    fetch_future = http_client.fetch(url)
    fetch_future.add_done_callback(lambda f: my_future.set_result(f.result()))
    return my_future.result()
Exemple #44
0
class Event(object):
    """An event blocks coroutines until its internal flag is set to True.

    Similar to `threading.Event`.

    A coroutine can wait for an event to be set. Once it is set, calls to
    ``yield event.wait()`` will not block unless the event has been cleared:

    .. testcode::

        from tornado import gen
        from tornado.ioloop import IOLoop
        from tornado.locks import Event

        event = Event()

        @gen.coroutine
        def waiter():
            print("Waiting for event")
            yield event.wait()
            print("Not waiting this time")
            yield event.wait()
            print("Done")

        @gen.coroutine
        def setter():
            print("About to set the event")
            event.set()

        @gen.coroutine
        def runner():
            yield [waiter(), setter()]

        IOLoop.current().run_sync(runner)

    .. testoutput::

        Waiting for event
        About to set the event
        Not waiting this time
        Done
    """
    def __init__(self):
        self._future = Future()

    def __repr__(self):
        return '<%s %s>' % (
            self.__class__.__name__, 'set' if self.is_set() else 'clear')

    def is_set(self):
        """Return ``True`` if the internal flag is true."""
        return self._future.done()

    def set(self):
        """Set the internal flag to ``True``. All waiters are awakened.

        Calling `.wait` once the flag is set will not block.
        """
        if not self._future.done():
            self._future.set_result(None)

    def clear(self):
        """Reset the internal flag to ``False``.

        Calls to `.wait` will block until `.set` is called.
        """
        if self._future.done():
            self._future = Future()

    def wait(self, timeout=None):
        """Block until the internal flag is true.

        Returns a Future, which raises `tornado.gen.TimeoutError` after a
        timeout.
        """
        if timeout is None:
            return self._future
        else:
            return gen.with_timeout(timeout, self._future)
Exemple #45
0
    def _iter_stream_lines(self,
                           chunk_size=ITER_CHUNK_SIZE,
                           decode_unicode=None,
                           delimiter=None):
        """ stream version of iter_lines.

        Basic Usage::

          >>> import trip
          >>> @trip.coroutine
          >>> def main():
          >>>     url = 'http://httpbin.org/get'
          >>>     r = yield trip.get(url, stream=True)
          >>>     for line in r.iter_lines(1):
          >>>         line = yield line
          >>>         if line is not None:
          >>>             print(line)
          >>> trip.IOLoop.current().run_sync(main)
          {
            "args": {},
            "headers": {}
            "origin": "0.0.0.0",
            "url": "http://httpbin.org/get"
          }
        """

        content = {'': []}
        pending = {'': None}

        def handle_content(f):

            chunk = f.result()

            if pending[''] is not None:
                chunk = pending[''] + chunk

            if delimiter:
                lines = chunk.split(delimiter)
            else:
                lines = chunk.splitlines()

            if lines and lines[-1] and chunk and lines[-1][-1] == chunk[-1]:
                pending[''] = lines.pop()
            else:
                pending[''] = None

            content[''] = lines[1:]
            f._result = lines[0] if lines else None

        for future in self.iter_content(chunk_size, decode_unicode):

            future.add_done_callback(handle_content)
            yield future

            for line in content['']:
                future = Future()
                future.set_result(line)
                yield future

        if pending[''] is not None:
            future = Future()
            future.set_result(pending[''])
            yield future
Exemple #46
0
class TestIOStreamStartTLS(AsyncTestCase):
    def setUp(self):
        try:
            super(TestIOStreamStartTLS, self).setUp()
            self.listener, self.port = bind_unused_port()
            self.server_stream = None
            self.server_accepted = Future()  # type: Future[None]
            netutil.add_accept_handler(self.listener, self.accept)
            self.client_stream = IOStream(socket.socket())
            self.io_loop.add_future(
                self.client_stream.connect(("10.0.0.7", self.port)), self.stop)
            self.wait()
            self.io_loop.add_future(self.server_accepted, self.stop)
            self.wait()
        except Exception as e:
            print(e)
            raise

    def tearDown(self):
        if self.server_stream is not None:
            self.server_stream.close()
        if self.client_stream is not None:
            self.client_stream.close()
        self.listener.close()
        super(TestIOStreamStartTLS, self).tearDown()

    def accept(self, connection, address):
        if self.server_stream is not None:
            self.fail("should only get one connection")
        self.server_stream = IOStream(connection)
        self.server_accepted.set_result(None)

    @gen.coroutine
    def client_send_line(self, line):
        self.client_stream.write(line)
        recv_line = yield self.server_stream.read_until(b"\r\n")
        self.assertEqual(line, recv_line)

    @gen.coroutine
    def server_send_line(self, line):
        self.server_stream.write(line)
        recv_line = yield self.client_stream.read_until(b"\r\n")
        self.assertEqual(line, recv_line)

    def client_start_tls(self, ssl_options=None, server_hostname=None):
        client_stream = self.client_stream
        self.client_stream = None
        return client_stream.start_tls(False, ssl_options, server_hostname)

    def server_start_tls(self, ssl_options=None):
        server_stream = self.server_stream
        self.server_stream = None
        return server_stream.start_tls(True, ssl_options)

    @gen_test
    def test_start_tls_smtp(self):
        # This flow is simplified from RFC 3207 section 5.
        # We don't really need all of this, but it helps to make sure
        # that after realistic back-and-forth traffic the buffers end up
        # in a sane state.
        yield self.server_send_line(b"220 mail.example.com ready\r\n")
        yield self.client_send_line(b"EHLO mail.example.com\r\n")
        yield self.server_send_line(b"250-mail.example.com welcome\r\n")
        yield self.server_send_line(b"250 STARTTLS\r\n")
        yield self.client_send_line(b"STARTTLS\r\n")
        yield self.server_send_line(b"220 Go ahead\r\n")
        client_future = self.client_start_tls(dict(cert_reqs=ssl.CERT_NONE))
        server_future = self.server_start_tls(_server_ssl_options())
        self.client_stream = yield client_future
        self.server_stream = yield server_future
        self.assertTrue(isinstance(self.client_stream, SSLIOStream))
        self.assertTrue(isinstance(self.server_stream, SSLIOStream))
        yield self.client_send_line(b"EHLO mail.example.com\r\n")
        yield self.server_send_line(b"250 mail.example.com welcome\r\n")

    @gen_test
    def test_handshake_fail(self):
        server_future = self.server_start_tls(_server_ssl_options())
        # Certificates are verified with the default configuration.
        with ExpectLog(gen_log, "SSL Error"):
            client_future = self.client_start_tls(server_hostname="10.0.0.7")
            with self.assertRaises(ssl.SSLError):
                yield client_future
            with self.assertRaises((ssl.SSLError, socket.error)):
                yield server_future

    @gen_test
    def test_check_hostname(self):
        # Test that server_hostname parameter to start_tls is being used.
        # The check_hostname functionality is only available in python 2.7 and
        # up and in python 3.4 and up.
        server_future = self.server_start_tls(_server_ssl_options())
        with ExpectLog(gen_log, "SSL Error"):
            client_future = self.client_start_tls(ssl.create_default_context(),
                                                  server_hostname="10.0.0.7")
            with self.assertRaises(ssl.SSLError):
                # The client fails to connect with an SSL error.
                yield client_future
            with self.assertRaises(Exception):
                # The server fails to connect, but the exact error is unspecified.
                yield server_future
Exemple #47
0
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
    """WebSocket client connection."""
    def __init__(self, io_loop, request):
        self.connect_future = Future()
        self.read_future = None
        self.read_queue = collections.deque()
        self.key = base64.b64encode(os.urandom(16))

        scheme, sep, rest = request.url.partition(':')
        scheme = {'ws': 'http', 'wss': 'https'}[scheme]
        request.url = scheme + sep + rest
        request.headers.update({
            'Upgrade': 'websocket',
            'Connection': 'Upgrade',
            'Sec-WebSocket-Key': self.key,
            'Sec-WebSocket-Version': '13',
        })

        super(WebSocketClientConnection,
              self).__init__(io_loop, None, request, lambda: None,
                             lambda response: None, 104857600,
                             Resolver(io_loop=io_loop))

    def _on_close(self):
        self.on_message(None)

    def _handle_1xx(self, code):
        assert code == 101
        assert self.headers['Upgrade'].lower() == 'websocket'
        assert self.headers['Connection'].lower() == 'upgrade'
        accept = WebSocketProtocol13.compute_accept_value(self.key)
        assert self.headers['Sec-Websocket-Accept'] == accept

        self.protocol = WebSocketProtocol13(self, mask_outgoing=True)
        self.protocol._receive_frame()

        if self._timeout is not None:
            self.io_loop.remove_timeout(self._timeout)
            self._timeout = None

        self.connect_future.set_result(self)

    def write_message(self, message, binary=False):
        """Sends a message to the WebSocket server."""
        self.protocol.write_message(message, binary)

    def read_message(self, callback=None):
        """Reads a message from the WebSocket server.

        Returns a future whose result is the message, or None
        if the connection is closed.  If a callback argument
        is given it will be called with the future when it is
        ready.
        """
        assert self.read_future is None
        future = Future()
        if self.read_queue:
            future.set_result(self.read_queue.popleft())
        else:
            self.read_future = future
        if callback is not None:
            self.io_loop.add_future(future, callback)
        return future

    def on_message(self, message):
        if self.read_future is not None:
            self.read_future.set_result(message)
            self.read_future = None
        else:
            self.read_queue.append(message)

    def on_pong(self, data):
        pass
Exemple #48
0
class Runner(object):
    """Internal implementation of `tornado.gen.engine`.

    Maintains information about pending callbacks and their results.

    The results of the generator are stored in ``result_future`` (a
    `.Future`)
    """
    def __init__(self, gen, result_future, first_yielded):
        self.gen = gen
        self.result_future = result_future
        self.future = _null_future
        self.yield_point = None
        self.pending_callbacks = None
        self.results = None
        self.running = False
        self.finished = False
        self.had_exception = False
        self.io_loop = IOLoop.current()
        # For efficiency, we do not create a stack context until we
        # reach a YieldPoint (stack contexts are required for the historical
        # semantics of YieldPoints, but not for Futures).  When we have
        # done so, this field will be set and must be called at the end
        # of the coroutine.
        self.stack_context_deactivate = None
        if self.handle_yield(first_yielded):
            gen = result_future = first_yielded = None
            self.run()

    def register_callback(self, key):
        """Adds ``key`` to the list of callbacks."""
        if self.pending_callbacks is None:
            # Lazily initialize the old-style YieldPoint data structures.
            self.pending_callbacks = set()
            self.results = {}
        if key in self.pending_callbacks:
            raise KeyReuseError("key %r is already pending" % (key, ))
        self.pending_callbacks.add(key)

    def is_ready(self, key):
        """Returns true if a result is available for ``key``."""
        if self.pending_callbacks is None or key not in self.pending_callbacks:
            raise UnknownKeyError("key %r is not pending" % (key, ))
        return key in self.results

    def set_result(self, key, result):
        """Sets the result for ``key`` and attempts to resume the generator."""
        self.results[key] = result
        if self.yield_point is not None and self.yield_point.is_ready():
            try:
                self.future.set_result(self.yield_point.get_result())
            except:
                future_set_exc_info(self.future, sys.exc_info())
            self.yield_point = None
            self.run()

    def pop_result(self, key):
        """Returns the result for ``key`` and unregisters it."""
        self.pending_callbacks.remove(key)
        return self.results.pop(key)

    def run(self):
        """Starts or resumes the generator, running until it reaches a
        yield point that is not ready.
        """
        if self.running or self.finished:
            return
        try:
            self.running = True
            while True:
                future = self.future
                if not future.done():
                    return
                self.future = None
                try:
                    orig_stack_contexts = stack_context._state.contexts
                    exc_info = None

                    try:
                        value = future.result()
                    except Exception:
                        self.had_exception = True
                        exc_info = sys.exc_info()
                    future = None

                    if exc_info is not None:
                        try:
                            yielded = self.gen.throw(*exc_info)
                        finally:
                            # Break up a reference to itself
                            # for faster GC on CPython.
                            exc_info = None
                    else:
                        yielded = self.gen.send(value)

                    if stack_context._state.contexts is not orig_stack_contexts:
                        self.gen.throw(
                            stack_context.StackContextInconsistentError(
                                'stack_context inconsistency (probably caused '
                                'by yield within a "with StackContext" block)')
                        )
                except (StopIteration, Return) as e:
                    self.finished = True
                    self.future = _null_future
                    if self.pending_callbacks and not self.had_exception:
                        # If we ran cleanly without waiting on all callbacks
                        # raise an error (really more of a warning).  If we
                        # had an exception then some callbacks may have been
                        # orphaned, so skip the check in that case.
                        raise LeakedCallbackError(
                            "finished without waiting for callbacks %r" %
                            self.pending_callbacks)
                    self.result_future.set_result(_value_from_stopiteration(e))
                    self.result_future = None
                    self._deactivate_stack_context()
                    return
                except Exception:
                    self.finished = True
                    self.future = _null_future
                    future_set_exc_info(self.result_future, sys.exc_info())
                    self.result_future = None
                    self._deactivate_stack_context()
                    return
                if not self.handle_yield(yielded):
                    return
                yielded = None
        finally:
            self.running = False

    def handle_yield(self, yielded):
        # Lists containing YieldPoints require stack contexts;
        # other lists are handled in convert_yielded.
        if _contains_yieldpoint(yielded):
            yielded = multi(yielded)

        if isinstance(yielded, YieldPoint):
            # YieldPoints are too closely coupled to the Runner to go
            # through the generic convert_yielded mechanism.
            self.future = Future()

            def start_yield_point():
                try:
                    yielded.start(self)
                    if yielded.is_ready():
                        self.future.set_result(yielded.get_result())
                    else:
                        self.yield_point = yielded
                except Exception:
                    self.future = Future()
                    future_set_exc_info(self.future, sys.exc_info())

            if self.stack_context_deactivate is None:
                # Start a stack context if this is the first
                # YieldPoint we've seen.
                with stack_context.ExceptionStackContext(
                        self.handle_exception) as deactivate:
                    self.stack_context_deactivate = deactivate

                    def cb():
                        start_yield_point()
                        self.run()

                    self.io_loop.add_callback(cb)
                    return False
            else:
                start_yield_point()
        else:
            try:
                self.future = convert_yielded(yielded)
            except BadYieldError:
                self.future = Future()
                future_set_exc_info(self.future, sys.exc_info())

        if self.future is moment:
            self.io_loop.add_callback(self.run)
            return False
        elif not self.future.done():

            def inner(f):
                # Break a reference cycle to speed GC.
                f = None  # noqa
                self.run()

            self.io_loop.add_future(self.future, inner)
            return False
        return True

    def result_callback(self, key):
        return stack_context.wrap(
            _argument_adapter(functools.partial(self.set_result, key)))

    def handle_exception(self, typ, value, tb):
        if not self.running and not self.finished:
            self.future = Future()
            future_set_exc_info(self.future, (typ, value, tb))
            self.run()
            return True
        else:
            return False

    def _deactivate_stack_context(self):
        if self.stack_context_deactivate is not None:
            self.stack_context_deactivate()
            self.stack_context_deactivate = None
class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
    '''There is one ZMQChannelsHandler per running kernel and it oversees all
    the sessions.
    '''
    
    # class-level registry of open sessions
    # allows checking for conflict on session-id,
    # which is used as a zmq identity and must be unique.
    _open_sessions = {}

    @property
    def kernel_info_timeout(self):
        km_default = self.kernel_manager.kernel_info_timeout
        return self.settings.get('kernel_info_timeout', km_default)

    @property
    def iopub_msg_rate_limit(self):
        return self.settings.get('iopub_msg_rate_limit', 0)

    @property
    def iopub_data_rate_limit(self):
        return self.settings.get('iopub_data_rate_limit', 0)

    @property
    def rate_limit_window(self):
        return self.settings.get('rate_limit_window', 1.0)

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, getattr(self, 'kernel_id', 'uninitialized'))

    def create_stream(self):
        km = self.kernel_manager
        identity = self.session.bsession
        for channel in ('shell', 'iopub', 'stdin'):
            meth = getattr(km, 'connect_' + channel)
            self.channels[channel] = stream = meth(self.kernel_id, identity=identity)
            stream.channel = channel
    
    def request_kernel_info(self):
        """send a request for kernel_info"""
        km = self.kernel_manager
        kernel = km.get_kernel(self.kernel_id)
        try:
            # check for previous request
            future = kernel._kernel_info_future
        except AttributeError:
            self.log.debug("Requesting kernel info from %s", self.kernel_id)
            # Create a kernel_info channel to query the kernel protocol version.
            # This channel will be closed after the kernel_info reply is received.
            if self.kernel_info_channel is None:
                self.kernel_info_channel = km.connect_shell(self.kernel_id)
            self.kernel_info_channel.on_recv(self._handle_kernel_info_reply)
            self.session.send(self.kernel_info_channel, "kernel_info_request")
            # store the future on the kernel, so only one request is sent
            kernel._kernel_info_future = self._kernel_info_future
        else:
            if not future.done():
                self.log.debug("Waiting for pending kernel_info request")
            future.add_done_callback(lambda f: self._finish_kernel_info(f.result()))
        return self._kernel_info_future
    
    def _handle_kernel_info_reply(self, msg):
        """process the kernel_info_reply
        
        enabling msg spec adaptation, if necessary
        """
        idents,msg = self.session.feed_identities(msg)
        try:
            msg = self.session.deserialize(msg)
        except:
            self.log.error("Bad kernel_info reply", exc_info=True)
            self._kernel_info_future.set_result({})
            return
        else:
            info = msg['content']
            self.log.debug("Received kernel info: %s", info)
            if msg['msg_type'] != 'kernel_info_reply' or 'protocol_version' not in info:
                self.log.error("Kernel info request failed, assuming current %s", info)
                info = {}
            self._finish_kernel_info(info)
        
        # close the kernel_info channel, we don't need it anymore
        if self.kernel_info_channel:
            self.kernel_info_channel.close()
        self.kernel_info_channel = None
    
    def _finish_kernel_info(self, info):
        """Finish handling kernel_info reply
        
        Set up protocol adaptation, if needed,
        and signal that connection can continue.
        """
        protocol_version = info.get('protocol_version', client_protocol_version)
        if protocol_version != client_protocol_version:
            self.session.adapt_version = int(protocol_version.split('.')[0])
            self.log.info("Adapting to protocol v%s for kernel %s", protocol_version, self.kernel_id)
        if not self._kernel_info_future.done():
            self._kernel_info_future.set_result(info)
    
    def initialize(self):
        super(ZMQChannelsHandler, self).initialize()
        self.zmq_stream = None
        self.channels = {}
        self.kernel_id = None
        self.kernel_info_channel = None
        self._kernel_info_future = Future()
        self._close_future = Future()
        self.session_key = ''

        # Rate limiting code
        self._iopub_window_msg_count = 0
        self._iopub_window_byte_count = 0
        self._iopub_msgs_exceeded = False
        self._iopub_data_exceeded = False
        # Queue of (time stamp, byte count)
        # Allows you to specify that the byte count should be lowered
        # by a delta amount at some point in the future.
        self._iopub_window_byte_queue = []

    @gen.coroutine
    def pre_get(self):
        # authenticate first
        super(ZMQChannelsHandler, self).pre_get()
        # check session collision:
        yield self._register_session()
        # then request kernel info, waiting up to a certain time before giving up.
        # We don't want to wait forever, because browsers don't take it well when
        # servers never respond to websocket connection requests.
        kernel = self.kernel_manager.get_kernel(self.kernel_id)
        self.session.key = kernel.session.key
        future = self.request_kernel_info()
        
        def give_up():
            """Don't wait forever for the kernel to reply"""
            if future.done():
                return
            self.log.warning("Timeout waiting for kernel_info reply from %s", self.kernel_id)
            future.set_result({})
        loop = IOLoop.current()
        loop.add_timeout(loop.time() + self.kernel_info_timeout, give_up)
        # actually wait for it
        yield future
    
    @gen.coroutine
    def get(self, kernel_id):
        self.kernel_id = cast_unicode(kernel_id, 'ascii')
        yield super(ZMQChannelsHandler, self).get(kernel_id=kernel_id)
    
    @gen.coroutine
    def _register_session(self):
        """Ensure we aren't creating a duplicate session.
        
        If a previous identical session is still open, close it to avoid collisions.
        This is likely due to a client reconnecting from a lost network connection,
        where the socket on our side has not been cleaned up yet.
        """
        self.session_key = '%s:%s' % (self.kernel_id, self.session.session)
        stale_handler = self._open_sessions.get(self.session_key)
        if stale_handler:
            self.log.warning("Replacing stale connection: %s", self.session_key)
            yield stale_handler.close()
        self._open_sessions[self.session_key] = self

    def open(self, kernel_id):
        super(ZMQChannelsHandler, self).open()
        km = self.kernel_manager
        km.notify_connect(kernel_id)

        # on new connections, flush the message buffer
        buffer_info = km.get_buffer(kernel_id, self.session_key)
        if buffer_info and buffer_info['session_key'] == self.session_key:
            self.log.info("Restoring connection for %s", self.session_key)
            self.channels = buffer_info['channels']
            replay_buffer = buffer_info['buffer']
            if replay_buffer:
                self.log.info("Replaying %s buffered messages", len(replay_buffer))
                for channel, msg_list in replay_buffer:
                    stream = self.channels[channel]
                    self._on_zmq_reply(stream, msg_list)
        else:
            try:
                self.create_stream()
            except web.HTTPError as e:
                self.log.error("Error opening stream: %s", e)
                # WebSockets don't response to traditional error codes so we
                # close the connection.
                for channel, stream in self.channels.items():
                    if not stream.closed():
                        stream.close()
                self.close()
                return

        km.add_restart_callback(self.kernel_id, self.on_kernel_restarted)
        km.add_restart_callback(self.kernel_id, self.on_restart_failed, 'dead')

        for channel, stream in self.channels.items():
            stream.on_recv_stream(self._on_zmq_reply)

    def on_message(self, msg):
        if not self.channels:
            # already closed, ignore the message
            self.log.debug("Received message on closed websocket %r", msg)
            return
        if isinstance(msg, bytes):
            msg = deserialize_binary_message(msg)
        else:
            msg = json.loads(msg)
        channel = msg.pop('channel', None)
        if channel is None:
            self.log.warning("No channel specified, assuming shell: %s", msg)
            channel = 'shell'
        if channel not in self.channels:
            self.log.warning("No such channel: %r", channel)
            return
        stream = self.channels[channel]
        self.session.send(stream, msg)

    def _on_zmq_reply(self, stream, msg_list):
        idents, fed_msg_list = self.session.feed_identities(msg_list)
        msg = self.session.deserialize(fed_msg_list)
        parent = msg['parent_header']
        def write_stderr(error_message):
            self.log.warning(error_message)
            msg = self.session.msg("stream",
                content={"text": error_message + '\n', "name": "stderr"},
                parent=parent
            )
            msg['channel'] = 'iopub'
            self.write_message(json.dumps(msg, default=date_default))
        channel = getattr(stream, 'channel', None)
        msg_type = msg['header']['msg_type']

        if channel == 'iopub' and msg_type == 'status' and msg['content'].get('execution_state') == 'idle':
            # reset rate limit counter on status=idle,
            # to avoid 'Run All' hitting limits prematurely.
            self._iopub_window_byte_queue = []
            self._iopub_window_msg_count = 0
            self._iopub_window_byte_count = 0
            self._iopub_msgs_exceeded = False
            self._iopub_data_exceeded = False

        if channel == 'iopub' and msg_type not in {'status', 'comm_open', 'execute_input'}:
            
            # Remove the counts queued for removal.
            now = IOLoop.current().time()
            while len(self._iopub_window_byte_queue) > 0:
                queued = self._iopub_window_byte_queue[0]
                if (now >= queued[0]):
                    self._iopub_window_byte_count -= queued[1]
                    self._iopub_window_msg_count -= 1
                    del self._iopub_window_byte_queue[0]
                else:
                    # This part of the queue hasn't be reached yet, so we can
                    # abort the loop.
                    break

            # Increment the bytes and message count
            self._iopub_window_msg_count += 1
            if msg_type == 'stream':
                byte_count = sum([len(x) for x in msg_list])
            else:
                byte_count = 0
            self._iopub_window_byte_count += byte_count
            
            # Queue a removal of the byte and message count for a time in the 
            # future, when we are no longer interested in it.
            self._iopub_window_byte_queue.append((now + self.rate_limit_window, byte_count))
            
            # Check the limits, set the limit flags, and reset the
            # message and data counts.
            msg_rate = float(self._iopub_window_msg_count) / self.rate_limit_window
            data_rate = float(self._iopub_window_byte_count) / self.rate_limit_window
            
            # Check the msg rate
            if self.iopub_msg_rate_limit > 0 and msg_rate > self.iopub_msg_rate_limit:
                if not self._iopub_msgs_exceeded:
                    self._iopub_msgs_exceeded = True
                    write_stderr(dedent("""\
                    IOPub message rate exceeded.
                    The notebook server will temporarily stop sending output
                    to the client in order to avoid crashing it.
                    To change this limit, set the config variable
                    `--NotebookApp.iopub_msg_rate_limit`.
                    
                    Current values:
                    NotebookApp.iopub_msg_rate_limit={} (msgs/sec)
                    NotebookApp.rate_limit_window={} (secs)
                    """.format(self.iopub_msg_rate_limit, self.rate_limit_window)))
            else:
                # resume once we've got some headroom below the limit
                if self._iopub_msgs_exceeded and msg_rate < (0.8 * self.iopub_msg_rate_limit):
                    self._iopub_msgs_exceeded = False
                    if not self._iopub_data_exceeded:
                        self.log.warning("iopub messages resumed")

            # Check the data rate
            if self.iopub_data_rate_limit > 0 and data_rate > self.iopub_data_rate_limit:
                if not self._iopub_data_exceeded:
                    self._iopub_data_exceeded = True
                    write_stderr(dedent("""\
                    IOPub data rate exceeded.
                    The notebook server will temporarily stop sending output
                    to the client in order to avoid crashing it.
                    To change this limit, set the config variable
                    `--NotebookApp.iopub_data_rate_limit`.
                    
                    Current values:
                    NotebookApp.iopub_data_rate_limit={} (bytes/sec)
                    NotebookApp.rate_limit_window={} (secs)
                    """.format(self.iopub_data_rate_limit, self.rate_limit_window)))
            else:
                # resume once we've got some headroom below the limit
                if self._iopub_data_exceeded and data_rate < (0.8 * self.iopub_data_rate_limit):
                    self._iopub_data_exceeded = False
                    if not self._iopub_msgs_exceeded:
                        self.log.warning("iopub messages resumed")
        
            # If either of the limit flags are set, do not send the message.
            if self._iopub_msgs_exceeded or self._iopub_data_exceeded:
                # we didn't send it, remove the current message from the calculus
                self._iopub_window_msg_count -= 1
                self._iopub_window_byte_count -= byte_count
                self._iopub_window_byte_queue.pop(-1)
                return
        super(ZMQChannelsHandler, self)._on_zmq_reply(stream, msg)

    def close(self):
        super(ZMQChannelsHandler, self).close()
        return self._close_future

    def on_close(self):
        self.log.debug("Websocket closed %s", self.session_key)
        # unregister myself as an open session (only if it's really me)
        if self._open_sessions.get(self.session_key) is self:
            self._open_sessions.pop(self.session_key)

        km = self.kernel_manager
        if self.kernel_id in km:
            km.notify_disconnect(self.kernel_id)
            km.remove_restart_callback(
                self.kernel_id, self.on_kernel_restarted,
            )
            km.remove_restart_callback(
                self.kernel_id, self.on_restart_failed, 'dead',
            )

            # start buffering instead of closing if this was the last connection
            if km._kernel_connections[self.kernel_id] == 0:
                km.start_buffering(self.kernel_id, self.session_key, self.channels)
                self._close_future.set_result(None)
                return

        # This method can be called twice, once by self.kernel_died and once
        # from the WebSocket close event. If the WebSocket connection is
        # closed before the ZMQ streams are setup, they could be None.
        for channel, stream in self.channels.items():
            if stream is not None and not stream.closed():
                stream.on_recv(None)
                stream.close()

        self.channels = {}
        self._close_future.set_result(None)

    def _send_status_message(self, status):
        iopub = self.channels.get('iopub', None)
        if iopub and not iopub.closed():
            # flush IOPub before sending a restarting/dead status message
            # ensures proper ordering on the IOPub channel
            # that all messages from the stopped kernel have been delivered
            iopub.flush()
        msg = self.session.msg("status",
            {'execution_state': status}
        )
        msg['channel'] = 'iopub'
        self.write_message(json.dumps(msg, default=date_default))

    def on_kernel_restarted(self):
        logging.warn("kernel %s restarted", self.kernel_id)
        self._send_status_message('restarting')

    def on_restart_failed(self):
        logging.error("kernel %s restarted failed!", self.kernel_id)
        self._send_status_message('dead')
Exemple #50
0
def create_future(ret_val=None):
    future = Future()
    future.set_result(ret_val)
    return future
Exemple #51
0
class HTTP1Connection(httputil.HTTPConnection):
    """Implements the HTTP/1.x protocol.

    This class can be on its own for clients, or via `HTTP1ServerConnection`
    for servers.
    """
    def __init__(self, stream, is_client, params=None, context=None):
        """
        :arg stream: an `.IOStream`
        :arg bool is_client: client or server
        :arg params: a `.HTTP1ConnectionParameters` instance or ``None``
        :arg context: an opaque application-defined object that can be accessed
            as ``connection.context``.
        """
        self.is_client = is_client
        self.stream = stream
        if params is None:
            params = HTTP1ConnectionParameters()
        self.params = params
        self.context = context
        self.no_keep_alive = params.no_keep_alive
        # The body limits can be altered by the delegate, so save them
        # here instead of just referencing self.params later.
        self._max_body_size = (self.params.max_body_size
                               or self.stream.max_buffer_size)
        self._body_timeout = self.params.body_timeout
        # _write_finished is set to True when finish() has been called,
        # i.e. there will be no more data sent.  Data may still be in the
        # stream's write buffer.
        self._write_finished = False
        # True when we have read the entire incoming body.
        self._read_finished = False
        # _finish_future resolves when all data has been written and flushed
        # to the IOStream.
        self._finish_future = Future()
        # If true, the connection should be closed after this request
        # (after the response has been written in the server side,
        # and after it has been read in the client)
        self._disconnect_on_finish = False
        self._clear_callbacks()
        # Save the start lines after we read or write them; they
        # affect later processing (e.g. 304 responses and HEAD methods
        # have content-length but no bodies)
        self._request_start_line = None
        self._response_start_line = None
        self._request_headers = None
        # True if we are writing output with chunked encoding.
        self._chunking_output = None
        # While reading a body with a content-length, this is the
        # amount left to read.
        self._expected_content_remaining = None
        # A Future for our outgoing writes, returned by IOStream.write.
        self._pending_write = None

    def read_response(self, delegate):
        """Read a single HTTP response.

        Typical client-mode usage is to write a request using `write_headers`,
        `write`, and `finish`, and then call ``read_response``.

        :arg delegate: a `.HTTPMessageDelegate`

        Returns a `.Future` that resolves to None after the full response has
        been read.
        """
        if self.params.decompress:
            delegate = _GzipMessageDelegate(delegate, self.params.chunk_size)
        return self._read_message(delegate)

    @gen.coroutine
    def _read_message(self, delegate):
        need_delegate_close = False
        try:
            header_future = self.stream.read_until_regex(
                b"\r?\n\r?\n", max_bytes=self.params.max_header_size)
            if self.params.header_timeout is None:
                header_data = yield header_future
            else:
                try:
                    header_data = yield gen.with_timeout(
                        self.stream.io_loop.time() +
                        self.params.header_timeout,
                        header_future,
                        io_loop=self.stream.io_loop,
                        quiet_exceptions=iostream.StreamClosedError)
                except gen.TimeoutError:
                    self.close()
                    raise gen.Return(False)
            start_line, headers = self._parse_headers(header_data)
            if self.is_client:
                start_line = httputil.parse_response_start_line(start_line)
                self._response_start_line = start_line
            else:
                start_line = httputil.parse_request_start_line(start_line)
                self._request_start_line = start_line
                self._request_headers = headers

            self._disconnect_on_finish = not self._can_keep_alive(
                start_line, headers)
            need_delegate_close = True
            with _ExceptionLoggingContext(app_log):
                header_future = delegate.headers_received(start_line, headers)
                if header_future is not None:
                    yield header_future
            if self.stream is None:
                # We've been detached.
                need_delegate_close = False
                raise gen.Return(False)
            skip_body = False
            if self.is_client:
                if (self._request_start_line is not None
                        and self._request_start_line.method == 'HEAD'):
                    skip_body = True
                code = start_line.code
                if code == 304:
                    # 304 responses may include the content-length header
                    # but do not actually have a body.
                    # http://tools.ietf.org/html/rfc7230#section-3.3
                    skip_body = True
                if code >= 100 and code < 200:
                    # 1xx responses should never indicate the presence of
                    # a body.
                    if ('Content-Length' in headers
                            or 'Transfer-Encoding' in headers):
                        raise httputil.HTTPInputError(
                            "Response code %d cannot have body" % code)
                    # TODO: client delegates will get headers_received twice
                    # in the case of a 100-continue.  Document or change?
                    yield self._read_message(delegate)
            else:
                if (headers.get("Expect") == "100-continue"
                        and not self._write_finished):
                    self.stream.write(b"HTTP/1.1 100 (Continue)\r\n\r\n")
            if not skip_body:
                body_future = self._read_body(
                    start_line.code if self.is_client else 0, headers,
                    delegate)
                if body_future is not None:
                    if self._body_timeout is None:
                        yield body_future
                    else:
                        try:
                            yield gen.with_timeout(
                                self.stream.io_loop.time() +
                                self._body_timeout,
                                body_future,
                                self.stream.io_loop,
                                quiet_exceptions=iostream.StreamClosedError)
                        except gen.TimeoutError:
                            gen_log.info("Timeout reading body from %s",
                                         self.context)
                            self.stream.close()
                            raise gen.Return(False)
            self._read_finished = True
            if not self._write_finished or self.is_client:
                need_delegate_close = False
                with _ExceptionLoggingContext(app_log):
                    delegate.finish()
            # If we're waiting for the application to produce an asynchronous
            # response, and we're not detached, register a close callback
            # on the stream (we didn't need one while we were reading)
            if (not self._finish_future.done() and self.stream is not None
                    and not self.stream.closed()):
                self.stream.set_close_callback(self._on_connection_close)
                yield self._finish_future
            if self.is_client and self._disconnect_on_finish:
                self.close()
            if self.stream is None:
                raise gen.Return(False)
        except httputil.HTTPInputError as e:
            gen_log.info("Malformed HTTP message from %s: %s", self.context, e)
            self.close()
            raise gen.Return(False)
        finally:
            if need_delegate_close:
                with _ExceptionLoggingContext(app_log):
                    delegate.on_connection_close()
            self._clear_callbacks()
        raise gen.Return(True)

    def _clear_callbacks(self):
        """Clears the callback attributes.

        This allows the request handler to be garbage collected more
        quickly in CPython by breaking up reference cycles.
        """
        self._write_callback = None
        self._write_future = None
        self._close_callback = None
        if self.stream is not None:
            self.stream.set_close_callback(None)

    def set_close_callback(self, callback):
        """Sets a callback that will be run when the connection is closed.

        .. deprecated:: 4.0
            Use `.HTTPMessageDelegate.on_connection_close` instead.
        """
        self._close_callback = stack_context.wrap(callback)

    def _on_connection_close(self):
        # Note that this callback is only registered on the IOStream
        # when we have finished reading the request and are waiting for
        # the application to produce its response.
        if self._close_callback is not None:
            callback = self._close_callback
            self._close_callback = None
            callback()
        if not self._finish_future.done():
            self._finish_future.set_result(None)
        self._clear_callbacks()

    def close(self):
        if self.stream is not None:
            self.stream.close()
        self._clear_callbacks()
        if not self._finish_future.done():
            self._finish_future.set_result(None)

    def detach(self):
        """Take control of the underlying stream.

        Returns the underlying `.IOStream` object and stops all further
        HTTP processing.  May only be called during
        `.HTTPMessageDelegate.headers_received`.  Intended for implementing
        protocols like websockets that tunnel over an HTTP handshake.
        """
        self._clear_callbacks()
        stream = self.stream
        self.stream = None
        if not self._finish_future.done():
            self._finish_future.set_result(None)
        return stream

    def set_body_timeout(self, timeout):
        """Sets the body timeout for a single request.

        Overrides the value from `.HTTP1ConnectionParameters`.
        """
        self._body_timeout = timeout

    def set_max_body_size(self, max_body_size):
        """Sets the body size limit for a single request.

        Overrides the value from `.HTTP1ConnectionParameters`.
        """
        self._max_body_size = max_body_size

    def write_headers(self, start_line, headers, chunk=None, callback=None):
        """Implements `.HTTPConnection.write_headers`."""
        lines = []
        if self.is_client:
            self._request_start_line = start_line
            lines.append(
                utf8('%s %s HTTP/1.1' % (start_line[0], start_line[1])))
            # Client requests with a non-empty body must have either a
            # Content-Length or a Transfer-Encoding.
            self._chunking_output = (start_line.method
                                     in ('POST', 'PUT', 'PATCH')
                                     and 'Content-Length' not in headers
                                     and 'Transfer-Encoding' not in headers)
        else:
            self._response_start_line = start_line
            lines.append(
                utf8('HTTP/1.1 %d %s' % (start_line[1], start_line[2])))
            self._chunking_output = (
                # TODO: should this use
                # self._request_start_line.version or
                # start_line.version?
                self._request_start_line.version == 'HTTP/1.1' and
                # 304 responses have no body (not even a zero-length body), and so
                # should not have either Content-Length or Transfer-Encoding.
                # headers.
                start_line.code != 304 and
                # No need to chunk the output if a Content-Length is specified.
                'Content-Length' not in headers and
                # Applications are discouraged from touching Transfer-Encoding,
                # but if they do, leave it alone.
                'Transfer-Encoding' not in headers)
            # If a 1.0 client asked for keep-alive, add the header.
            if (self._request_start_line.version == 'HTTP/1.0'
                    and (self._request_headers.get('Connection', '').lower()
                         == 'keep-alive')):
                headers['Connection'] = 'Keep-Alive'
        if self._chunking_output:
            headers['Transfer-Encoding'] = 'chunked'
        if (not self.is_client and (self._request_start_line.method == 'HEAD'
                                    or start_line.code == 304)):
            self._expected_content_remaining = 0
        elif 'Content-Length' in headers:
            self._expected_content_remaining = int(headers['Content-Length'])
        else:
            self._expected_content_remaining = None
        lines.extend([utf8(n) + b": " + utf8(v) for n, v in headers.get_all()])
        for line in lines:
            if b'\n' in line:
                raise ValueError('Newline in header: ' + repr(line))
        future = None
        if self.stream.closed():
            future = self._write_future = Future()
            future.set_exception(iostream.StreamClosedError())
            future.exception()
        else:
            if callback is not None:
                self._write_callback = stack_context.wrap(callback)
            else:
                future = self._write_future = Future()
            data = b"\r\n".join(lines) + b"\r\n\r\n"
            if chunk:
                data += self._format_chunk(chunk)
            self._pending_write = self.stream.write(data)
            self._pending_write.add_done_callback(self._on_write_complete)
        return future

    def _format_chunk(self, chunk):
        if self._expected_content_remaining is not None:
            self._expected_content_remaining -= len(chunk)
            if self._expected_content_remaining < 0:
                # Close the stream now to stop further framing errors.
                self.stream.close()
                raise httputil.HTTPOutputError(
                    "Tried to write more data than Content-Length")
        if self._chunking_output and chunk:
            # Don't write out empty chunks because that means END-OF-STREAM
            # with chunked encoding
            return utf8("%x" % len(chunk)) + b"\r\n" + chunk + b"\r\n"
        else:
            return chunk

    def write(self, chunk, callback=None):
        """Implements `.HTTPConnection.write`.

        For backwards compatibility is is allowed but deprecated to
        skip `write_headers` and instead call `write()` with a
        pre-encoded header block.
        """
        future = None
        if self.stream.closed():
            future = self._write_future = Future()
            self._write_future.set_exception(iostream.StreamClosedError())
            self._write_future.exception()
        else:
            if callback is not None:
                self._write_callback = stack_context.wrap(callback)
            else:
                future = self._write_future = Future()
            self._pending_write = self.stream.write(self._format_chunk(chunk))
            self._pending_write.add_done_callback(self._on_write_complete)
        return future

    def finish(self):
        """Implements `.HTTPConnection.finish`."""
        if (self._expected_content_remaining is not None
                and self._expected_content_remaining != 0
                and not self.stream.closed()):
            self.stream.close()
            raise httputil.HTTPOutputError(
                "Tried to write %d bytes less than Content-Length" %
                self._expected_content_remaining)
        if self._chunking_output:
            if not self.stream.closed():
                self._pending_write = self.stream.write(b"0\r\n\r\n")
                self._pending_write.add_done_callback(self._on_write_complete)
        self._write_finished = True
        # If the app finished the request while we're still reading,
        # divert any remaining data away from the delegate and
        # close the connection when we're done sending our response.
        # Closing the connection is the only way to avoid reading the
        # whole input body.
        if not self._read_finished:
            self._disconnect_on_finish = True
        # No more data is coming, so instruct TCP to send any remaining
        # data immediately instead of waiting for a full packet or ack.
        self.stream.set_nodelay(True)
        if self._pending_write is None:
            self._finish_request(None)
        else:
            self._pending_write.add_done_callback(self._finish_request)

    def _on_write_complete(self, future):
        exc = future.exception()
        if exc is not None and not isinstance(exc, iostream.StreamClosedError):
            future.result()
        if self._write_callback is not None:
            callback = self._write_callback
            self._write_callback = None
            self.stream.io_loop.add_callback(callback)
        if self._write_future is not None:
            future = self._write_future
            self._write_future = None
            future.set_result(None)

    def _can_keep_alive(self, start_line, headers):
        if self.params.no_keep_alive:
            return False
        connection_header = headers.get("Connection")
        if connection_header is not None:
            connection_header = connection_header.lower()
        if start_line.version == "HTTP/1.1":
            return connection_header != "close"
        elif ("Content-Length" in headers
              or headers.get("Transfer-Encoding", "").lower() == "chunked"
              or start_line.method in ("HEAD", "GET")):
            return connection_header == "keep-alive"
        return False

    def _finish_request(self, future):
        self._clear_callbacks()
        if not self.is_client and self._disconnect_on_finish:
            self.close()
            return
        # Turn Nagle's algorithm back on, leaving the stream in its
        # default state for the next request.
        self.stream.set_nodelay(False)
        if not self._finish_future.done():
            self._finish_future.set_result(None)

    def _parse_headers(self, data):
        # The lstrip removes newlines that some implementations sometimes
        # insert between messages of a reused connection.  Per RFC 7230,
        # we SHOULD ignore at least one empty line before the request.
        # http://tools.ietf.org/html/rfc7230#section-3.5
        data = native_str(data.decode('latin1')).lstrip("\r\n")
        # RFC 7230 section allows for both CRLF and bare LF.
        eol = data.find("\n")
        start_line = data[:eol].rstrip("\r")
        try:
            headers = httputil.HTTPHeaders.parse(data[eol:])
        except ValueError:
            # probably form split() if there was no ':' in the line
            raise httputil.HTTPInputError("Malformed HTTP headers: %r" %
                                          data[eol:100])
        return start_line, headers

    def _read_body(self, code, headers, delegate):
        if "Content-Length" in headers:
            if "Transfer-Encoding" in headers:
                # Response cannot contain both Content-Length and
                # Transfer-Encoding headers.
                # http://tools.ietf.org/html/rfc7230#section-3.3.3
                raise httputil.HTTPInputError(
                    "Response with both Transfer-Encoding and Content-Length")
            if "," in headers["Content-Length"]:
                # Proxies sometimes cause Content-Length headers to get
                # duplicated.  If all the values are identical then we can
                # use them but if they differ it's an error.
                pieces = re.split(r',\s*', headers["Content-Length"])
                if any(i != pieces[0] for i in pieces):
                    raise httputil.HTTPInputError(
                        "Multiple unequal Content-Lengths: %r" %
                        headers["Content-Length"])
                headers["Content-Length"] = pieces[0]
            content_length = int(headers["Content-Length"])

            if content_length > self._max_body_size:
                raise httputil.HTTPInputError("Content-Length too long")
        else:
            content_length = None

        if code == 204:
            # This response code is not allowed to have a non-empty body,
            # and has an implicit length of zero instead of read-until-close.
            # http://www.w3.org/Protocols/rfc2616/rfc2616-sec4.html#sec4.3
            if ("Transfer-Encoding" in headers
                    or content_length not in (None, 0)):
                raise httputil.HTTPInputError(
                    "Response with code %d should not have body" % code)
            content_length = 0

        if content_length is not None:
            return self._read_fixed_body(content_length, delegate)
        if headers.get("Transfer-Encoding") == "chunked":
            return self._read_chunked_body(delegate)
        if self.is_client:
            return self._read_body_until_close(delegate)
        return None

    @gen.coroutine
    def _read_fixed_body(self, content_length, delegate):
        while content_length > 0:
            body = yield self.stream.read_bytes(min(self.params.chunk_size,
                                                    content_length),
                                                partial=True)
            content_length -= len(body)
            if not self._write_finished or self.is_client:
                with _ExceptionLoggingContext(app_log):
                    ret = delegate.data_received(body)
                    if ret is not None:
                        yield ret

    @gen.coroutine
    def _read_chunked_body(self, delegate):
        # TODO: "chunk extensions" http://tools.ietf.org/html/rfc2616#section-3.6.1
        total_size = 0
        while True:
            chunk_len = yield self.stream.read_until(b"\r\n", max_bytes=64)
            chunk_len = int(chunk_len.strip(), 16)
            if chunk_len == 0:
                return
            total_size += chunk_len
            if total_size > self._max_body_size:
                raise httputil.HTTPInputError("chunked body too large")
            bytes_to_read = chunk_len
            while bytes_to_read:
                chunk = yield self.stream.read_bytes(min(
                    bytes_to_read, self.params.chunk_size),
                                                     partial=True)
                bytes_to_read -= len(chunk)
                if not self._write_finished or self.is_client:
                    with _ExceptionLoggingContext(app_log):
                        ret = delegate.data_received(chunk)
                        if ret is not None:
                            yield ret
            # chunk ends with \r\n
            crlf = yield self.stream.read_bytes(2)
            assert crlf == b"\r\n"

    @gen.coroutine
    def _read_body_until_close(self, delegate):
        body = yield self.stream.read_until_close()
        if not self._write_finished or self.is_client:
            with _ExceptionLoggingContext(app_log):
                delegate.data_received(body)
Exemple #52
0
from cpe import CPE

from os import path
from tornado.concurrent import Future
from tornado.httpclient import HTTPError, HTTPRequest, HTTPResponse
from tornado.testing import gen_test, AsyncTestCase

from fixtures.exploits import Exploit
from structs import Port, Node, Scan, TransportProtocol, Service, PhysicalPort, ScanContext, Vulnerability
from tools.cve_search.exceptions import CVESearchApiException
from tools.cve_search.structs import CVESearchVulnerabilityResults, CVESearchVulnerabilityResult
from tools.cve_search.tasks import CVESearchServiceTask
from utils import Config

future = Future()
future.set_result(True)


@patch('utils.http_client.gen.sleep', MagicMock(return_value=future))
class CVESearchServiceTaskTest(AsyncTestCase):
    @patch('tools.cve_search.tasks.cfg', new_callable=Config)
    def setUp(self, cfg):
        super(CVESearchServiceTaskTest, self).setUp()
        cfg._cfg = {'tools': {'cve-search': {'api': 'localhost:200'}}}

        self.example_output = ''

        with open(
                path.join(path.dirname(path.abspath(__file__)),
                          'example_output.json'), 'rb') as f:
            self.example_output = f.read()
Exemple #53
0
 def close(self):
     fut = Future()
     fut.set_result(True)
     return fut
Exemple #54
0
 def test_already_resolved(self):
     future = Future()
     future.set_result('asdf')
     result = yield gen.with_timeout(datetime.timedelta(seconds=3600),
                                     future)
     self.assertEqual(result, 'asdf')
Exemple #55
0
class DaskClusterManager:
    """
    A class for starting, stopping, and otherwise managing the lifecycle
    of Dask clusters.
    """
    def __init__(self) -> None:
        """ Initialize the cluster manager """
        self._clusters: Dict[str, Cluster] = dict()
        self._adaptives: Dict[str, Adaptive] = dict()
        self._cluster_names: Dict[str, str] = dict()
        self._n_clusters = 0

        self.initialized = Future()

        async def start_clusters():
            for model in dask.config.get('labextension.initial'):
                await self.start_cluster(configuration=model)
            self.initialized.set_result(self)

        IOLoop.current().add_callback(start_clusters)

    async def start_cluster(self,
                            cluster_id: str = "",
                            configuration: dict = {}) -> ClusterModel:
        """
        Start a new Dask cluster.

        Parameters
        ----------
        cluster_id : string
            An optional string id for the cluster. If not given, a random id
            will be chosen.

        Returns
        cluster_model : a dask cluster model.
        """
        if not cluster_id:
            cluster_id = str(uuid4())

        cluster, adaptive = await make_cluster(configuration)
        self._n_clusters += 1

        # Check for a name in the config
        if not configuration.get('name'):
            cluster_type = type(cluster).__name__
            cluster_name = f"{cluster_type} {self._n_clusters}"
        else:
            cluster_name = configuration['name']

        # Check if the cluster was started adaptively
        if adaptive:
            self._adaptives[cluster_id] = adaptive

        self._clusters[cluster_id] = cluster
        self._cluster_names[cluster_id] = cluster_name
        return make_cluster_model(cluster_id,
                                  cluster_name,
                                  cluster,
                                  adaptive=adaptive)

    async def close_cluster(self,
                            cluster_id: str) -> Union[ClusterModel, None]:
        """
        Close a Dask cluster.

        Parameters
        ----------
        cluster_id : string
            A string id for the cluster.

        Returns
        cluster_model : the dask cluster model for the shut down cluster,
            or None if it was not found.
        """
        cluster = self._clusters.get(cluster_id)
        if cluster:
            await cluster.close()
            self._clusters.pop(cluster_id)
            name = self._cluster_names.pop(cluster_id)
            adaptive = self._adaptives.pop(cluster_id, None)
            return make_cluster_model(cluster_id, name, cluster, adaptive)

        else:
            return None

    def get_cluster(self, cluster_id) -> Union[ClusterModel, None]:
        """
        Get a Dask cluster model.

        Parameters
        ----------
        cluster_id : string
            A string id for the cluster.

        Returns
        cluster_model : the dask cluster model for the cluster,
            or None if it was not found.
        """
        cluster = self._clusters.get(cluster_id)
        name = self._cluster_names.get(cluster_id, '')
        adaptive = self._adaptives.get(cluster_id)
        if not cluster:
            return None

        return make_cluster_model(cluster_id, name, cluster, adaptive)

    def list_clusters(self) -> List[ClusterModel]:
        """
        List the Dask cluster models known to the manager.

        Returns
        cluster_models : A list of the dask cluster models known to the manager.
        """
        return [
            make_cluster_model(
                cluster_id,
                self._cluster_names[cluster_id],
                self._clusters[cluster_id],
                self._adaptives.get(cluster_id, None),
            ) for cluster_id in self._clusters
        ]

    def scale_cluster(self, cluster_id: str,
                      n: int) -> Union[ClusterModel, None]:
        cluster = self._clusters.get(cluster_id)
        name = self._cluster_names[cluster_id]
        adaptive = self._adaptives.pop(cluster_id, None)

        # Check if the cluster exists
        if not cluster:
            return None

        # Check if it is actually different.
        model = make_cluster_model(cluster_id, name, cluster, adaptive)
        if model.get("adapt") == None and model["workers"] == n:
            return model

        # Otherwise, rescale the model.
        cluster.scale(n)
        return make_cluster_model(cluster_id, name, cluster, adaptive=None)

    def adapt_cluster(self, cluster_id: str, minimum: int,
                      maximum: int) -> Union[ClusterModel, None]:
        cluster = self._clusters.get(cluster_id)
        name = self._cluster_names[cluster_id]
        adaptive = self._adaptives.pop(cluster_id, None)

        # Check if the cluster exists
        if not cluster:
            return None

        # Check if it is actually different.
        model = make_cluster_model(cluster_id, name, cluster, adaptive)
        if model.get("adapt") != None and model["adapt"][
                "minimum"] == minimum and model["adapt"]["maximum"] == maximum:
            return model

        # Otherwise, rescale the model.
        adaptive = cluster.adapt(minimum=minimum, maximum=maximum)
        self._adaptives[cluster_id] = adaptive
        return make_cluster_model(cluster_id, name, cluster, adaptive)

    async def close(self):
        """ Close all clusters and cleanup """
        for cluster_id in list(self._clusters):
            await self.close_cluster(cluster_id)

    async def __aenter__(self):
        """
        Enter an asynchronous context.
        This waits for any initial clusters specified via configuration to start.
        """
        await self.initialized
        return self

    async def __aexit__(self, exc_type, exc, tb):
        """
        Exit an asynchronous context.
        This closes any extant clusters.
        """
        await self.close()

    def __await__(self):
        """
        Awaiter for the manager to be initialized.
        This waits for any initial clusters specified via configuration to start.
        """
        return self.initialized.__await__()
Exemple #56
0
    if isinstance(future, Future):
        # We know this future will resolve on the IOLoop, so we don't
        # need the extra thread-safety of IOLoop.add_future (and we also
        # don't care about StackContext here.
        future.add_done_callback(
            lambda future: io_loop.remove_timeout(timeout_handle))
    else:
        # concurrent.futures.Futures may resolve on any thread, so we
        # need to route them back to the IOLoop.
        io_loop.add_future(
            future, lambda future: io_loop.remove_timeout(timeout_handle))
    return result


_null_future = Future()
_null_future.set_result(None)

moment = Future()
moment.__doc__ = \
    """A special object which may be yielded to allow the IOLoop to run for
one iteration.

This is not needed in normal use but it can be helpful in long-running
coroutines that are likely to yield Futures that are ready instantly.

Usage: ``yield gen.moment``

.. versionadded:: 4.0
"""
moment.set_result(None)
Exemple #57
0
class WSGIApplication(web.Application):
    """A WSGI equivalent of `tornado.web.Application`.

    .. deprecated:: 4.0

       Use a regular `.Application` and wrap it in `WSGIAdapter` instead.
    """
    def __call__(self, environ, start_response):
        return WSGIAdapter(self)(environ, start_response)


# WSGI has no facilities for flow control, so just return an already-done
# Future when the interface requires it.
_dummy_future = Future()
_dummy_future.set_result(None)


class _WSGIConnection(httputil.HTTPConnection):
    def __init__(self, method, start_response, context):
        self.method = method
        self.start_response = start_response
        self.context = context
        self._write_buffer = []
        self._finished = False
        self._expected_content_remaining = None
        self._error = None

    def set_close_callback(self, callback):
        # WSGI has no facility for detecting a closed connection mid-request,
        # so we can simply ignore the callback.
Exemple #58
0
class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
    """There is one ZMQChannelsHandler per running kernel and it oversees all
    the sessions.
    """

    # class-level registry of open sessions
    # allows checking for conflict on session-id,
    # which is used as a zmq identity and must be unique.
    _open_sessions = {}

    @property
    def kernel_info_timeout(self):
        km_default = self.kernel_manager.kernel_info_timeout
        return self.settings.get("kernel_info_timeout", km_default)

    @property
    def iopub_msg_rate_limit(self):
        return self.settings.get("iopub_msg_rate_limit", 0)

    @property
    def iopub_data_rate_limit(self):
        return self.settings.get("iopub_data_rate_limit", 0)

    @property
    def rate_limit_window(self):
        return self.settings.get("rate_limit_window", 1.0)

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__,
                           getattr(self, "kernel_id", "uninitialized"))

    def create_stream(self):
        km = self.kernel_manager
        identity = self.session.bsession
        for channel in ("iopub", "shell", "control", "stdin"):
            meth = getattr(km, "connect_" + channel)
            self.channels[channel] = stream = meth(self.kernel_id,
                                                   identity=identity)
            stream.channel = channel

    def nudge(self):
        """Nudge the zmq connections with kernel_info_requests
        Returns a Future that will resolve when we have received
        a control reply and at least one iopub message,
        ensuring that zmq subscriptions are established,
        sockets are fully connected, and kernel is responsive.
        Keeps retrying kernel_info_request until these are both received.
        """
        kernel = self.kernel_manager.get_kernel(self.kernel_id)

        # Use a transient control channel to prevent leaking
        # control responses to the front-end.
        control_channel = kernel.connect_control()
        # The IOPub used by the client, whose subscriptions we are verifying.
        iopub_channel = self.channels["iopub"]

        info_future = Future()
        iopub_future = Future()
        both_done = gen.multi([info_future, iopub_future])

        def finish(_=None):
            """Ensure all futures are resolved
            which in turn triggers cleanup
            """
            for f in (info_future, iopub_future):
                if not f.done():
                    f.set_result(None)

        def cleanup(_=None):
            """Common cleanup"""
            loop.remove_timeout(nudge_handle)
            iopub_channel.stop_on_recv()
            if not control_channel.closed():
                control_channel.close()

        # trigger cleanup when both message futures are resolved
        both_done.add_done_callback(cleanup)

        def on_control_reply(msg):
            self.log.debug("Nudge: shell info reply received: %s",
                           self.kernel_id)
            if not info_future.done():
                self.log.debug("Nudge: resolving shell future: %s",
                               self.kernel_id)
                info_future.set_result(None)

        def on_iopub(msg):
            self.log.debug("Nudge: IOPub received: %s", self.kernel_id)
            if not iopub_future.done():
                iopub_channel.stop_on_recv()
                self.log.debug("Nudge: resolving iopub future: %s",
                               self.kernel_id)
                iopub_future.set_result(None)

        iopub_channel.on_recv(on_iopub)
        control_channel.on_recv(on_control_reply)
        loop = IOLoop.current()

        # Nudge the kernel with kernel info requests until we get an IOPub message
        def nudge(count):
            count += 1

            # NOTE: this close check appears to never be True during on_open,
            # even when the peer has closed the connection
            if self.ws_connection is None or self.ws_connection.is_closing():
                self.log.debug("Nudge: cancelling on closed websocket: %s",
                               self.kernel_id)
                finish()
                return

            # check for stopped kernel
            if self.kernel_id not in self.kernel_manager:
                self.log.debug("Nudge: cancelling on stopped kernel: %s",
                               self.kernel_id)
                finish()
                return

            # check for closed zmq socket
            if control_channel.closed():
                self.log.debug("Nudge: cancelling on closed zmq socket: %s",
                               self.kernel_id)
                finish()
                return

            if not both_done.done():
                log = self.log.warning if count % 10 == 0 else self.log.debug
                log("Nudge: attempt %s on kernel %s" % (count, self.kernel_id))
                self.session.send(control_channel, "kernel_info_request")
                nonlocal nudge_handle
                nudge_handle = loop.call_later(0.5, nudge, count)

        nudge_handle = loop.call_later(0, nudge, count=0)

        # resolve with a timeout if we get no response
        future = gen.with_timeout(loop.time() + self.kernel_info_timeout,
                                  both_done)
        # ensure we have no dangling resources or unresolved Futures in case of timeout
        future.add_done_callback(finish)
        return future

    def request_kernel_info(self):
        """send a request for kernel_info"""
        km = self.kernel_manager
        kernel = km.get_kernel(self.kernel_id)
        try:
            # check for previous request
            future = kernel._kernel_info_future
        except AttributeError:
            self.log.debug("Requesting kernel info from %s", self.kernel_id)
            # Create a kernel_info channel to query the kernel protocol version.
            # This channel will be closed after the kernel_info reply is received.
            if self.kernel_info_channel is None:
                self.kernel_info_channel = km.connect_shell(self.kernel_id)
            self.kernel_info_channel.on_recv(self._handle_kernel_info_reply)
            self.session.send(self.kernel_info_channel, "kernel_info_request")
            # store the future on the kernel, so only one request is sent
            kernel._kernel_info_future = self._kernel_info_future
        else:
            if not future.done():
                self.log.debug("Waiting for pending kernel_info request")
            future.add_done_callback(
                lambda f: self._finish_kernel_info(f.result()))
        return self._kernel_info_future

    def _handle_kernel_info_reply(self, msg):
        """process the kernel_info_reply

        enabling msg spec adaptation, if necessary
        """
        idents, msg = self.session.feed_identities(msg)
        try:
            msg = self.session.deserialize(msg)
        except:
            self.log.error("Bad kernel_info reply", exc_info=True)
            self._kernel_info_future.set_result({})
            return
        else:
            info = msg["content"]
            self.log.debug("Received kernel info: %s", info)
            if msg["msg_type"] != "kernel_info_reply" or "protocol_version" not in info:
                self.log.error(
                    "Kernel info request failed, assuming current %s", info)
                info = {}
            self._finish_kernel_info(info)

        # close the kernel_info channel, we don't need it anymore
        if self.kernel_info_channel:
            self.kernel_info_channel.close()
        self.kernel_info_channel = None

    def _finish_kernel_info(self, info):
        """Finish handling kernel_info reply

        Set up protocol adaptation, if needed,
        and signal that connection can continue.
        """
        protocol_version = info.get("protocol_version",
                                    client_protocol_version)
        if protocol_version != client_protocol_version:
            self.session.adapt_version = int(protocol_version.split(".")[0])
            self.log.info(
                "Adapting from protocol version {protocol_version} (kernel {kernel_id}) to {client_protocol_version} (client)."
                .format(
                    protocol_version=protocol_version,
                    kernel_id=self.kernel_id,
                    client_protocol_version=client_protocol_version,
                ))
        if not self._kernel_info_future.done():
            self._kernel_info_future.set_result(info)

    def initialize(self):
        super(ZMQChannelsHandler, self).initialize()
        self.zmq_stream = None
        self.channels = {}
        self.kernel_id = None
        self.kernel_info_channel = None
        self._kernel_info_future = Future()
        self._close_future = Future()
        self.session_key = ""

        # Rate limiting code
        self._iopub_window_msg_count = 0
        self._iopub_window_byte_count = 0
        self._iopub_msgs_exceeded = False
        self._iopub_data_exceeded = False
        # Queue of (time stamp, byte count)
        # Allows you to specify that the byte count should be lowered
        # by a delta amount at some point in the future.
        self._iopub_window_byte_queue = []

    async def pre_get(self):
        # authenticate first
        super(ZMQChannelsHandler, self).pre_get()
        # check session collision:
        await self._register_session()
        # then request kernel info, waiting up to a certain time before giving up.
        # We don't want to wait forever, because browsers don't take it well when
        # servers never respond to websocket connection requests.
        kernel = self.kernel_manager.get_kernel(self.kernel_id)

        if hasattr(kernel, "ready"):
            try:
                await kernel.ready
            except Exception as e:
                kernel.execution_state = "dead"
                kernel.reason = str(e)
                raise web.HTTPError(500, str(e)) from e

        self.session.key = kernel.session.key
        future = self.request_kernel_info()

        def give_up():
            """Don't wait forever for the kernel to reply"""
            if future.done():
                return
            self.log.warning("Timeout waiting for kernel_info reply from %s",
                             self.kernel_id)
            future.set_result({})

        loop = IOLoop.current()
        loop.add_timeout(loop.time() + self.kernel_info_timeout, give_up)
        # actually wait for it
        await future

    async def get(self, kernel_id):
        self.kernel_id = cast_unicode(kernel_id, "ascii")
        await super(ZMQChannelsHandler, self).get(kernel_id=kernel_id)

    async def _register_session(self):
        """Ensure we aren't creating a duplicate session.

        If a previous identical session is still open, close it to avoid collisions.
        This is likely due to a client reconnecting from a lost network connection,
        where the socket on our side has not been cleaned up yet.
        """
        self.session_key = "%s:%s" % (self.kernel_id, self.session.session)
        stale_handler = self._open_sessions.get(self.session_key)
        if stale_handler:
            self.log.warning("Replacing stale connection: %s",
                             self.session_key)
            await stale_handler.close()
        self._open_sessions[self.session_key] = self

    def open(self, kernel_id):
        super(ZMQChannelsHandler, self).open()
        km = self.kernel_manager
        km.notify_connect(kernel_id)

        # on new connections, flush the message buffer
        buffer_info = km.get_buffer(kernel_id, self.session_key)
        if buffer_info and buffer_info["session_key"] == self.session_key:
            self.log.info("Restoring connection for %s", self.session_key)
            if km.ports_changed(kernel_id):
                # If the kernel's ports have changed (some restarts trigger this)
                # then reset the channels so nudge() is using the correct iopub channel
                self.create_stream()
            else:
                # The kernel's ports have not changed; use the channels captured in the buffer
                self.channels = buffer_info["channels"]

            connected = self.nudge()

            def replay(value):
                replay_buffer = buffer_info["buffer"]
                if replay_buffer:
                    self.log.info("Replaying %s buffered messages",
                                  len(replay_buffer))
                    for channel, msg_list in replay_buffer:
                        stream = self.channels[channel]
                        self._on_zmq_reply(stream, msg_list)

            connected.add_done_callback(replay)
        else:
            try:
                self.create_stream()
                connected = self.nudge()
            except web.HTTPError as e:
                # Do not log error if the kernel is already shutdown,
                # as it's normal that it's not responding
                try:
                    self.kernel_manager.get_kernel(kernel_id)

                    self.log.error("Error opening stream: %s", e)
                except KeyError:
                    pass
                # WebSockets don't respond to traditional error codes so we
                # close the connection.
                for channel, stream in self.channels.items():
                    if not stream.closed():
                        stream.close()
                self.close()
                return

        km.add_restart_callback(self.kernel_id, self.on_kernel_restarted)
        km.add_restart_callback(self.kernel_id, self.on_restart_failed, "dead")

        def subscribe(value):
            for channel, stream in self.channels.items():
                stream.on_recv_stream(self._on_zmq_reply)

        connected.add_done_callback(subscribe)

        return connected

    def on_message(self, msg):
        if not self.channels:
            # already closed, ignore the message
            self.log.debug("Received message on closed websocket %r", msg)
            return
        if isinstance(msg, bytes):
            msg = deserialize_binary_message(msg)
        else:
            msg = json.loads(msg)
        channel = msg.pop("channel", None)
        if channel is None:
            self.log.warning("No channel specified, assuming shell: %s", msg)
            channel = "shell"
        if channel not in self.channels:
            self.log.warning("No such channel: %r", channel)
            return
        am = self.kernel_manager.allowed_message_types
        mt = msg["header"]["msg_type"]
        if am and mt not in am:
            self.log.warning(
                'Received message of type "%s", which is not allowed. Ignoring.'
                % mt)
        else:
            stream = self.channels[channel]
            self.session.send(stream, msg)

    def _on_zmq_reply(self, stream, msg_list):
        idents, fed_msg_list = self.session.feed_identities(msg_list)
        msg = self.session.deserialize(fed_msg_list)

        parent = msg["parent_header"]

        def write_stderr(error_message):
            self.log.warning(error_message)
            msg = self.session.msg("stream",
                                   content={
                                       "text": error_message + "\n",
                                       "name": "stderr"
                                   },
                                   parent=parent)
            msg["channel"] = "iopub"
            self.write_message(json.dumps(msg, default=json_default))

        channel = getattr(stream, "channel", None)
        msg_type = msg["header"]["msg_type"]

        if channel == "iopub" and msg_type == "error":
            self._on_error(msg)

        if (channel == "iopub" and msg_type == "status"
                and msg["content"].get("execution_state") == "idle"):
            # reset rate limit counter on status=idle,
            # to avoid 'Run All' hitting limits prematurely.
            self._iopub_window_byte_queue = []
            self._iopub_window_msg_count = 0
            self._iopub_window_byte_count = 0
            self._iopub_msgs_exceeded = False
            self._iopub_data_exceeded = False

        if channel == "iopub" and msg_type not in {
                "status", "comm_open", "execute_input"
        }:

            # Remove the counts queued for removal.
            now = IOLoop.current().time()
            while len(self._iopub_window_byte_queue) > 0:
                queued = self._iopub_window_byte_queue[0]
                if now >= queued[0]:
                    self._iopub_window_byte_count -= queued[1]
                    self._iopub_window_msg_count -= 1
                    del self._iopub_window_byte_queue[0]
                else:
                    # This part of the queue hasn't be reached yet, so we can
                    # abort the loop.
                    break

            # Increment the bytes and message count
            self._iopub_window_msg_count += 1
            if msg_type == "stream":
                byte_count = sum([len(x) for x in msg_list])
            else:
                byte_count = 0
            self._iopub_window_byte_count += byte_count

            # Queue a removal of the byte and message count for a time in the
            # future, when we are no longer interested in it.
            self._iopub_window_byte_queue.append(
                (now + self.rate_limit_window, byte_count))

            # Check the limits, set the limit flags, and reset the
            # message and data counts.
            msg_rate = float(
                self._iopub_window_msg_count) / self.rate_limit_window
            data_rate = float(
                self._iopub_window_byte_count) / self.rate_limit_window

            # Check the msg rate
            if self.iopub_msg_rate_limit > 0 and msg_rate > self.iopub_msg_rate_limit:
                if not self._iopub_msgs_exceeded:
                    self._iopub_msgs_exceeded = True
                    write_stderr(
                        dedent("""\
                    IOPub message rate exceeded.
                    The Jupyter server will temporarily stop sending output
                    to the client in order to avoid crashing it.
                    To change this limit, set the config variable
                    `--ServerApp.iopub_msg_rate_limit`.

                    Current values:
                    ServerApp.iopub_msg_rate_limit={} (msgs/sec)
                    ServerApp.rate_limit_window={} (secs)
                    """.format(self.iopub_msg_rate_limit,
                               self.rate_limit_window)))
            else:
                # resume once we've got some headroom below the limit
                if self._iopub_msgs_exceeded and msg_rate < (
                        0.8 * self.iopub_msg_rate_limit):
                    self._iopub_msgs_exceeded = False
                    if not self._iopub_data_exceeded:
                        self.log.warning("iopub messages resumed")

            # Check the data rate
            if self.iopub_data_rate_limit > 0 and data_rate > self.iopub_data_rate_limit:
                if not self._iopub_data_exceeded:
                    self._iopub_data_exceeded = True
                    write_stderr(
                        dedent("""\
                    IOPub data rate exceeded.
                    The Jupyter server will temporarily stop sending output
                    to the client in order to avoid crashing it.
                    To change this limit, set the config variable
                    `--ServerApp.iopub_data_rate_limit`.

                    Current values:
                    ServerApp.iopub_data_rate_limit={} (bytes/sec)
                    ServerApp.rate_limit_window={} (secs)
                    """.format(self.iopub_data_rate_limit,
                               self.rate_limit_window)))
            else:
                # resume once we've got some headroom below the limit
                if self._iopub_data_exceeded and data_rate < (
                        0.8 * self.iopub_data_rate_limit):
                    self._iopub_data_exceeded = False
                    if not self._iopub_msgs_exceeded:
                        self.log.warning("iopub messages resumed")

            # If either of the limit flags are set, do not send the message.
            if self._iopub_msgs_exceeded or self._iopub_data_exceeded:
                # we didn't send it, remove the current message from the calculus
                self._iopub_window_msg_count -= 1
                self._iopub_window_byte_count -= byte_count
                self._iopub_window_byte_queue.pop(-1)
                return
        super(ZMQChannelsHandler, self)._on_zmq_reply(stream, msg)

    def close(self):
        super(ZMQChannelsHandler, self).close()
        return self._close_future

    def on_close(self):
        self.log.debug("Websocket closed %s", self.session_key)
        # unregister myself as an open session (only if it's really me)
        if self._open_sessions.get(self.session_key) is self:
            self._open_sessions.pop(self.session_key)

        km = self.kernel_manager
        if self.kernel_id in km:
            km.notify_disconnect(self.kernel_id)
            km.remove_restart_callback(
                self.kernel_id,
                self.on_kernel_restarted,
            )
            km.remove_restart_callback(
                self.kernel_id,
                self.on_restart_failed,
                "dead",
            )

            # start buffering instead of closing if this was the last connection
            if km._kernel_connections[self.kernel_id] == 0:
                km.start_buffering(self.kernel_id, self.session_key,
                                   self.channels)
                self._close_future.set_result(None)
                return

        # This method can be called twice, once by self.kernel_died and once
        # from the WebSocket close event. If the WebSocket connection is
        # closed before the ZMQ streams are setup, they could be None.
        for channel, stream in self.channels.items():
            if stream is not None and not stream.closed():
                stream.on_recv(None)
                stream.close()

        self.channels = {}
        self._close_future.set_result(None)

    def _send_status_message(self, status):
        iopub = self.channels.get("iopub", None)
        if iopub and not iopub.closed():
            # flush IOPub before sending a restarting/dead status message
            # ensures proper ordering on the IOPub channel
            # that all messages from the stopped kernel have been delivered
            iopub.flush()
        msg = self.session.msg("status", {"execution_state": status})
        msg["channel"] = "iopub"
        self.write_message(json.dumps(msg, default=json_default))

    def on_kernel_restarted(self):
        self.log.warning("kernel %s restarted", self.kernel_id)
        self._send_status_message("restarting")

    def on_restart_failed(self):
        self.log.error("kernel %s restarted failed!", self.kernel_id)
        self._send_status_message("dead")

    def _on_error(self, msg):
        if self.kernel_manager.allow_tracebacks:
            return
        msg["content"]["ename"] = "ExecutionError"
        msg["content"]["evalue"] = "Execution error"
        msg["content"]["traceback"] = [
            self.kernel_manager.traceback_replacement_message
        ]
Exemple #59
0
class _Connector(object):
    """A stateless implementation of the "Happy Eyeballs" algorithm.

    "Happy Eyeballs" is documented in RFC6555 as the recommended practice
    for when both IPv4 and IPv6 addresses are available.

    In this implementation, we partition the addresses by family, and
    make the first connection attempt to whichever address was
    returned first by ``getaddrinfo``.  If that connection fails or
    times out, we begin a connection in parallel to the first address
    of the other family.  If there are additional failures we retry
    with other addresses, keeping one connection attempt per family
    in flight at a time.

    http://tools.ietf.org/html/rfc6555

    """
    def __init__(self, addrinfo, connect):
        self.io_loop = IOLoop.current()
        self.connect = connect

        self.future = Future()
        self.timeout = None
        self.connect_timeout = None
        self.last_error = None
        self.remaining = len(addrinfo)
        self.primary_addrs, self.secondary_addrs = self.split(addrinfo)
        self.streams = set()

    @staticmethod
    def split(addrinfo):
        """Partition the ``addrinfo`` list by address family.

        Returns two lists.  The first list contains the first entry from
        ``addrinfo`` and all others with the same family, and the
        second list contains all other addresses (normally one list will
        be AF_INET and the other AF_INET6, although non-standard resolvers
        may return additional families).
        """
        primary = []
        secondary = []
        primary_af = addrinfo[0][0]
        for af, addr in addrinfo:
            if af == primary_af:
                primary.append((af, addr))
            else:
                secondary.append((af, addr))
        return primary, secondary

    def start(self, timeout=_INITIAL_CONNECT_TIMEOUT, connect_timeout=None):
        self.try_connect(iter(self.primary_addrs))
        self.set_timeout(timeout)
        if connect_timeout is not None:
            self.set_connect_timeout(connect_timeout)
        return self.future

    def try_connect(self, addrs):
        try:
            af, addr = next(addrs)
        except StopIteration:
            # We've reached the end of our queue, but the other queue
            # might still be working.  Send a final error on the future
            # only when both queues are finished.
            if self.remaining == 0 and not self.future.done():
                self.future.set_exception(self.last_error
                                          or IOError("connection failed"))
            return
        stream, future = self.connect(af, addr)
        self.streams.add(stream)
        future_add_done_callback(
            future, functools.partial(self.on_connect_done, addrs, af, addr))

    def on_connect_done(self, addrs, af, addr, future):
        self.remaining -= 1
        try:
            stream = future.result()
        except Exception as e:
            if self.future.done():
                return
            # Error: try again (but remember what happened so we have an
            # error to raise in the end)
            self.last_error = e
            self.try_connect(addrs)
            if self.timeout is not None:
                # If the first attempt failed, don't wait for the
                # timeout to try an address from the secondary queue.
                self.io_loop.remove_timeout(self.timeout)
                self.on_timeout()
            return
        self.clear_timeouts()
        if self.future.done():
            # This is a late arrival; just drop it.
            stream.close()
        else:
            self.streams.discard(stream)
            self.future.set_result((af, addr, stream))
            self.close_streams()

    def set_timeout(self, timeout):
        self.timeout = self.io_loop.add_timeout(self.io_loop.time() + timeout,
                                                self.on_timeout)

    def on_timeout(self):
        self.timeout = None
        if not self.future.done():
            self.try_connect(iter(self.secondary_addrs))

    def clear_timeout(self):
        if self.timeout is not None:
            self.io_loop.remove_timeout(self.timeout)

    def set_connect_timeout(self, connect_timeout):
        self.connect_timeout = self.io_loop.add_timeout(
            connect_timeout, self.on_connect_timeout)

    def on_connect_timeout(self):
        if not self.future.done():
            self.future.set_exception(TimeoutError())
        self.close_streams()

    def clear_timeouts(self):
        if self.timeout is not None:
            self.io_loop.remove_timeout(self.timeout)
        if self.connect_timeout is not None:
            self.io_loop.remove_timeout(self.connect_timeout)

    def close_streams(self):
        for stream in self.streams:
            stream.close()
def not_blocking():
    future = Future()
    future.set_result('not blocking func result')
    result = yield future
    raise gen.Return(result)