Python websockets 模块，ConnectionClosed() 实例源码

def _pinger(self):
        '''
        A Controller can time us out if we are silent for too long. This
        is especially true in JaaS, which has a fairly strict timeout.

        To prevent timing out, we send a ping every ten seconds.

        '''
        async def _do_ping():
            try:
                await pinger_facade.Ping()
                await asyncio.sleep(10, loop=self.loop)
            except CancelledError:
                pass

        pinger_facade = client.PingerFacade.from_connection(self)
        while self.monitor.status == Monitor.CONNECTED:
            try:
                await utils.run_with_interrupt(
                    _do_ping(),
                    self.monitor.close_called,
                    loop=self.loop)
            except websockets.ConnectionClosed:
                pass

def block_until(self, *conditions, timeout=None, wait_period=0.5):
        """Return only after all conditions are true.

        Raises `websockets.ConnectionClosed` if disconnected.
        """
        def _disconnected():
            return not (self.is_connected() and self.connection().is_open)

        def done():
            return _disconnected() or all(c() for c in conditions)

        await utils.block_until(done,
                                timeout=timeout,
                                wait_period=wait_period,
                                loop=self.loop)
        if _disconnected():
            raise websockets.ConnectionClosed(1006, 'no reason')

def ping_handler(self):
        ping_interval = self.shark.config['WS_PING']['interval']
        if not ping_interval:
            return
        latency = 0
        while True:
            await asyncio.sleep(ping_interval - latency)
            self.session.log.debug('ping')
            start_time = time.time()
            try:
                ping = await self.websocket.ping()
            except websockets.ConnectionClosed:
                return
            timeout_handler = asyncio.ensure_future(
                    self.ping_timeout_handler(ping))
            await ping
            latency = time.time() - start_time
            self.session.log.debug('pong', latency=round(latency, 3))
            # Return immediately if a ping timeout occurred.
            if not timeout_handler.cancel() and timeout_handler.result():
                return

def consumer_handler(self):
        try:
            ping_handler = asyncio.ensure_future(self.ping_handler())
            try:
                while True:
                    event = await self.websocket.recv()
                    try:
                        data = json.loads(event)
                    except json.decoder.JSONDecodeError:
                        self.session.log.warn('received invalid json')
                        await self.send({
                            "status": "error",
                            "error": c.ERR_INVALID_EVENT,
                        })
                    else:
                        await self.session.on_client_event(data)
            except websockets.ConnectionClosed:
                await self.session.on_close()

            ping_handler.cancel()
        except Exception:
            self.session.log.exception('unhandled error in consumer handler')

def api(ws, path):
    while True:
        try:
#            msg = await ws.recv()
            # get a websockets string
            msg = yield from ws.recv()
            print('msg', msg)
            try:
                msgJ = json.loads(msg)
            except json.decoder.JSONDecodeError:
                print("error decoding msg >{}<".format(msg))
                continue

            print("got json msgJ >{}<".format(msgJ))
            # and handle it...
            retJ = handleMsg(msgJ)

            #print(retJ)
            # and return the response to the client
            yield from ws.send(retJ)
            #            await ws.send(retJ)

        except websockets.ConnectionClosed:
            print('connection closed')
            return

def __aenter__(self):
        self._soc = await websockets.connect(self._ws_url)
        async def loop():
            try:
                while 1:
                    resp = json.loads(await self._soc.recv())
                    if 'id' in resp:
                        self._method_responses[resp['id']] = resp
                        self._recv_data_lock[resp['id']].release()
                    elif 'method' in resp:
                        asyncio.ensure_future(self._run_later(self._handle_event(resp['method'], resp['params'])))
                    else:
                        raise RuntimeError('Unknown data came: {0}'.format(resp))
            except (websockets.ConnectionClosed, concurrent.futures.CancelledError):
                pass
            except Exception as e:
                traceback.print_exc()
        asyncio.ensure_future(self._run_later(loop()))
        return self

def api(ws, path):
    while True:
        try:
#            msg = await ws.recv()
            # get a websockets string
            msg = yield from ws.recv()
            print('msg', msg)
            try:
                msgJ = json.loads(msg)
            except ValueError:
                print("error decoding msg >{}<".format(msg))
                continue

            #print("got json msgJ >{}<".format(msgJ))
            # and handle it...
            retJ = handleMsg(msgJ)

            #print(retJ)
            # and return the response to the client
            yield from ws.send(retJ)
            #            await ws.send(retJ)

        except websockets.ConnectionClosed:
            print('connection closed')
            return

def _read(self):
        """
        Endless read loop that runs until the socket is closed.
        """
        while True:
            try:
                data = await self._read_single()
            except (asyncio.CancelledError, websockets.ConnectionClosed):
                break  # will already be handled
            except Exception as e:
                logger.error("error in interactive read loop", extra=e)
                break

            if isinstance(data, list):
                for item in data:
                    self._handle_recv(item)
            else:
                self._handle_recv(data)

def _recv_loop(self):
        self._done_event.clear()

        while not self._ws_close_event.is_set():
            try:
                data = json.loads(await self._ws.recv())
            except websockets.ConnectionClosed:
                await self._close()

            else:
                message_id = data.get('message-id')
                if message_id is not None:
                    self._message_map.pop(message_id).set_result(data)
                    continue

                type_name = data.get('update-type')
                if type_name is not None:
                    asyncio.ensure_future(
                        self._handle_event(type_name, data), loop=self._loop)

                    continue

                # TODO: Not a response nor an event - log an error maybe?

        self._done_event.set()

def send(self, message):
        """
        Disconnected
        MessageError
        MessageFlowError
        """
        # Pack
        self.log.debug('Packing message: {}', message.type)
        data = message.pack(self)
        self.log.trace('server >> {}', message)

        # Send data
        self.log.debug('Sending message')
        try:
            yield from self._connection.send(data)
        except websockets.ConnectionClosed as exc:
            self.log.debug('Connection closed while sending')
            raise Disconnected(exc.code) from exc

def receive(self):
        """
        Disconnected
        """
        # Receive data
        try:
            data = yield from self._connection.recv()
        except websockets.ConnectionClosed as exc:
            self.log.debug('Connection closed while receiving')
            raise Disconnected(exc.code) from exc
        self.log.debug('Received message')

        # Unpack data and return
        message = unpack(self, data)
        self.log.debug('Unpacked message: {}', message.type)
        self.log.trace('server << {}', message)
        return message

def test_path_full_lite(self, initiator_key, server, client_factory):
        """
        Add 253 fake responders to a path. Then, add a 254th responder
        and check that the correct error code (Path Full) is being
        returned.
        """
        assert len(server.protocols) == 0

        # Get path instance of server
        path = server.paths.get(initiator_key.pk)

        # Add fake clients to path
        clients = [_FakePathClient() for _ in range(0x02, 0x100)]
        for client in clients:
            path.add_responder(client)

        # Now the path is full
        with pytest.raises(websockets.ConnectionClosed) as exc_info:
            yield from client_factory(responder_handshake=True)
        assert exc_info.value.code == CloseCode.path_full_error

        # Remove fake clients from path
        for client in clients:
            path.remove_client(client)
        yield from server.wait_connections_closed()

def test_path_full(self, event_loop, server, client_factory):
        """
        Add 253 responders to a path. Then, add a 254th responder
        and check that the correct error code (Path Full) is being
        returned.
        """
        tasks = [client_factory(responder_handshake=True, timeout=20.0)
                 for _ in range(0x02, 0x100)]
        clients = yield from asyncio.gather(*tasks, loop=event_loop)

        # All clients must be open
        assert all((client.ws_client.open for client, _ in clients))

        # Now the path is full
        with pytest.raises(websockets.ConnectionClosed) as exc_info:
            yield from client_factory(responder_handshake=True)
        assert exc_info.value.code == CloseCode.path_full_error

        # Close all clients
        tasks = [client.close() for client, _ in clients]
        yield from asyncio.wait(tasks, loop=event_loop)
        yield from server.wait_connections_closed()

def _recv(self, request_id):
        if not self.is_open:
            raise websockets.exceptions.ConnectionClosed(0, 'websocket closed')
        return await self.messages.get(request_id)

def _receiver(self):
        try:
            while self.is_open:
                result = await utils.run_with_interrupt(
                    self.ws.recv(),
                    self.monitor.close_called,
                    loop=self.loop)
                if self.monitor.close_called.is_set():
                    break
                if result is not None:
                    result = json.loads(result)
                    await self.messages.put(result['request-id'], result)
        except CancelledError:
            pass
        except websockets.ConnectionClosed as e:
            log.warning('Receiver: Connection closed, reconnecting')
            await self.messages.put_all(e)
            # the reconnect has to be done as a task because the receiver will
            # be cancelled by the reconnect and we don't want the reconnect
            # to be aborted half-way through
            self.loop.create_task(self.reconnect())
            return
        except Exception as e:
            log.exception("Error in receiver")
            # make pending listeners aware of the error
            await self.messages.put_all(e)
            raise

def _recv_loop(self) -> None:
        async with self._ws as connection:
            self._connected = True
            self.connection = connection
            while self._connected:
                try:
                    resp = await self.connection.recv()
                    if resp:
                        self._on_message(resp)
                except websockets.ConnectionClosed:
                    logger.info('connection closed')
                    break

def send(self, event):
        try:
            await self.websocket.send(json.dumps(event))
        except websockets.ConnectionClosed:
            self.session.log.warn('attempted to send to closed socket')

def mocked_websocket(*args, **kwargs):
        class MockedWebsocket(mock.Mock):
            def __init__(self, **kwargs):
                super().__init__(**kwargs)
                self.count = 0

            @asyncio.coroutine
            def recv(self, *args, **kwargs):
                self.count += 1
                if self.count > 3:
                    raise ConnectionClosed(1000, "Peace out homie!")
                return "Hello World!"

        return MockedWebsocket()

def send(self, body):
        if isinstance(body, str):
            body = body.encode()
        try:
            self.writer.write(body)
            self.writer.write(b'\n')
        except BrokenPipeError:
            raise ConnectionClosed()

def connection_handler(self, websocket, path):
        """
        Internal asyncio.coroutine function for handling one websocket request.

        :param websocket: Socket with request
        :param path: Requested path of socket (not used)
        :return: Returns when socket is closed or poison pill is found in message queue
            from ClientConnections.
        """
        wanted_id = None
        try:
            wanted_id = yield from websocket.recv()
            queue = self._connections.add_client(wanted_id)
            self._logger.debug("websocket server: got client for channel '{}'".format(wanted_id))
            while True:
                # wait for message
                result = yield from queue.get()
                if not result:
                    break
                self._logger.debug("websocket server: message '{}' for channel '{}'".format(result, wanted_id))
                # send message to client
                yield from websocket.send(result)
                self._logger.debug("websocket server: message sent to channel '{}'".format(wanted_id))
        except websockets.ConnectionClosed:
            self._logger.info("websocket server: connection closed for channel '{}'". format(wanted_id))
        finally:
            self._connections.remove_client(wanted_id, queue)

def run(self):
        while True:
            try:
                data = yield from self.websocket.recv()
            except websockets.ConnectionClosed:
                break
            self._on_data(data)

        self.close()

def _accept(self, websocket, path):
        print("Connection opened.")
        try:
            while True:
                await self._process(websocket, path)
        except websockets.ConnectionClosed:
            print("Closing connection.")

def _run(self):
        """Enter an infinite loop waiting for websocket packets"""
        try:
            while True:
                payload = await self.recv()
                await self.process(payload)
        except (PayloadLengthExceeded, earl.DecodeError, json.JSONDecodeError):
            await self.ws.close(CloseCodes.DECODE_ERROR, 'Decoding Error')
        except asyncio.CancelledError:
            log.info('[ws] Run task was cancelled')
            await self.ws.close(1006, 'Task was cancelled')
        except StopConnection as sc:
            log.info('[ws] StopConncection: %r', sc)

            sc_args = sc.args
            c_code = sc.args[0]
            if len(sc_args) == 1:
                await self.ws.close(c_code, reason=CloseReasons.get(c_code))
            elif len(sc_args) == 2:
                await self.ws.close(c_code, reason=sc.args[1])

        except websockets.ConnectionClosed as err:
            log.info('[ws] Closed with %d, %r', err.code, err.reason)
        except InvalidateSession as err:
            resumable = err.args[0]
            if not resumable:
                await self._clean()
            pass
        except Exception as err:
            log.error('Error while running', exc_info=True)
            await self.ws.close(4000, f'Unexpected error: {err!r}')
            await self._clean()
            return

        await self._clean()
        if self.ws.open:
            await self.ws.close(1000)

def _read_single(self):
        """
        Reads a single event off the websocket.
        """
        try:
            raw_data = await self._socket.recv()
        except (asyncio.CancelledError, websockets.ConnectionClosed) as e:
            if self._recv_await is None:
                self._recv_await = asyncio.Future(loop=self._loop)
            self._recv_await.set_result(False)
            raise e

        return json.loads(self._decode(raw_data))

def test_handles_connection_closed(self):
        yield from self._connection.connect()

        def raise_closed():
            raise websockets.ConnectionClosed(4000, "")

        yield from asyncio.sleep(0)
        self._mock_socket.recv = raise_closed
        self._queue.put_nowait(sample_method)

        has_packet = yield from self._connection.has_packet()
        self.assertTrue(has_packet)  # reads what we pushed to get unblocked
        has_packet = yield from self._connection.has_packet()
        self.assertFalse(has_packet)  # gets a connection closed

def ping(self):
        """
        Disconnected
        """
        self.log.debug('Sending ping')
        try:
            return (yield from self._connection.ping())
        except websockets.ConnectionClosed as exc:
            self.log.debug('Connection closed while pinging')
            raise Disconnected(exc.code) from exc

def test_explicit_invalid_permanent_key(
            self, server, client_factory
    ):
        """
        Check that the server rejects a permanent key it doesn't have.
        """
        key = libnacl.public.SecretKey()

        # Expect invalid key
        with pytest.raises(websockets.ConnectionClosed) as exc_info:
            yield from client_factory(
                permanent_key=key.pk, explicit_permanent_key=True,
                initiator_handshake=True)
        assert exc_info.value.code == CloseCode.invalid_key
        yield from server.wait_connections_closed()

def on_connection_factory(execute_cmd, base_dispatcher):

    async def on_connection(reader, writer):
        context = ClientConnectionContext(reader, writer)
        client_dispatcher = client_dispatcher_factory(context)
        dispatcher = ComposedDispatcher([base_dispatcher, client_dispatcher])
        context.logger.info('Connection started')
        # Wait for two things:
        # - User's command (incomming request)
        # - Event subscribed by user (pushed to client requests)
        # Note user's command should have been replied before sending an event notification
        get_event = asyncio.ensure_future(context.queued_pushed_events.get())
        get_cmd = asyncio.ensure_future(context.recv())
        try:
            while True:
                done, pending = await asyncio.wait((get_event, get_cmd),
                                                   return_when='FIRST_COMPLETED')
                if get_event in done:
                    payload = get_event.result()
                    context.logger.debug('Got event: %s' % payload)
                    await context.send(payload)
                    # Restart watch on incoming notifications
                    get_event = asyncio.ensure_future(context.queued_pushed_events.get())
                else:
                    raw_cmd = get_cmd.result()
                    if not raw_cmd:
                        context.logger.debug('Connection stopped')
                        return
                    context.logger.debug('Received: %r' % raw_cmd)
                    intent = execute_cmd(raw_cmd)
                    raw_resp = await asyncio_perform(dispatcher, intent)
                    context.logger.debug('Replied: %r' % raw_resp)
                    await context.send(raw_resp)
                    # Restart watch on incoming messages
                    get_cmd = asyncio.ensure_future(context.recv())
        except ConnectionClosed:
            context.logger.info('Connection closed')
        finally:
            get_event.cancel()
            get_cmd.cancel()

    return on_connection