Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

perf: avoid returning future if requester wants #2955

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 19 additions & 12 deletions tornado/websocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,8 @@ def max_message_size(self) -> int:
)

def write_message(
self, message: Union[bytes, str, Dict[str, Any]], binary: bool = False
self, message: Union[bytes, str, Dict[str, Any]], binary: bool = False,
return_future: bool = True
) -> "Future[None]":
"""Sends the given message to the client of this Web Socket.

Expand All @@ -337,7 +338,8 @@ def write_message(
raise WebSocketClosedError()
if isinstance(message, dict):
message = tornado.escape.json_encode(message)
return self.ws_connection.write_message(message, binary=binary)
return self.ws_connection.write_message(message, binary=binary,
return_future=return_future)

def select_subprotocol(self, subprotocols: List[str]) -> Optional[str]:
"""Override to implement subprotocol negotiation.
Expand Down Expand Up @@ -680,7 +682,8 @@ async def accept_connection(self, handler: WebSocketHandler) -> None:

@abc.abstractmethod
def write_message(
self, message: Union[str, bytes], binary: bool = False
self, message: Union[str, bytes], binary: bool = False,
return_future: bool = True
) -> "Future[None]":
raise NotImplementedError()

Expand Down Expand Up @@ -1072,7 +1075,8 @@ def _write_frame(
return self.stream.write(frame)

def write_message(
self, message: Union[str, bytes], binary: bool = False
self, message: Union[str, bytes], binary: bool = False,
return_future: bool = True
) -> "Future[None]":
"""Sends the given message to the client of this Web Socket."""
if binary:
Expand All @@ -1096,13 +1100,14 @@ def write_message(
except StreamClosedError:
raise WebSocketClosedError()

async def wrapper() -> None:
try:
await fut
except StreamClosedError:
raise WebSocketClosedError()
if return_future:
async def wrapper() -> None:
try:
await fut
except StreamClosedError:
raise WebSocketClosedError()

return asyncio.ensure_future(wrapper())
return asyncio.ensure_future(wrapper())

def write_ping(self, data: bytes) -> None:
"""Send ping frame."""
Expand Down Expand Up @@ -1493,7 +1498,8 @@ async def headers_received(
future_set_result_unless_cancelled(self.connect_future, self)

def write_message(
self, message: Union[str, bytes], binary: bool = False
self, message: Union[str, bytes], binary: bool = False,
return_future: bool = True,
) -> "Future[None]":
"""Sends a message to the WebSocket server.

Expand All @@ -1504,7 +1510,8 @@ def write_message(
Exception raised on a closed stream changed from `.StreamClosedError`
to `WebSocketClosedError`.
"""
return self.protocol.write_message(message, binary=binary)
return self.protocol.write_message(message, binary=binary,
return_future=return_future)

def read_message(
self,
Expand Down