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

Nodelay #680

Merged
merged 11 commits into from
Dec 15, 2015
Prev Previous commit
Next Next commit
Add cork support to StreamWriter
  • Loading branch information
asvetlov committed Dec 14, 2015
commit 69296ec2ee486b5f456937d344897b216ccc6518
24 changes: 24 additions & 0 deletions aiohttp/parsers.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,13 @@ def set_parser(self, parser):

DEFAULT_LIMIT = 2 ** 16

if hasattr(socket, 'TCP_CORK'): # pragma: no cover
CORK = socket.TCP_CORK
elif hasattr(socket, 'TCP_NOPUSH'): # pragma: no cover
CORK = socket.TCP_NOPUSH
else: # pragma: no cover
CORK = None


class StreamParser:
"""StreamParser manages incoming bytes stream and protocol parsers.
Expand Down Expand Up @@ -226,6 +233,7 @@ def __init__(self, transport, protocol, reader, loop):
self._reader = reader
self._loop = loop
self._tcp_nodelay = False
self._tcp_cork = False
self._socket = transport.get_extra_info('socket')

@property
Expand All @@ -243,6 +251,22 @@ def set_tcp_nodelay(self, value):
return
self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, value)

@property
def tcp_cork(self):
return self._tcp_cork

def set_tcp_cork(self, value):
value = bool(value)
if self._tcp_cork == value:
return
self._tcp_cork = value
if self._socket is None:
return
if self._socket.family not in (socket.AF_INET, socket.AF_INET6):
return
if CORK is not None: # pragma: no branch
self._socket.setsockopt(socket.IPPROTO_TCP, CORK, value)


class StreamProtocol(asyncio.streams.FlowControlMixin, asyncio.Protocol):
"""Helper class to adapt between Protocol and StreamReader."""
Expand Down
97 changes: 96 additions & 1 deletion tests/test_stream_writer.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
import pytest
import socket
from aiohttp.parsers import StreamWriter
from aiohttp.parsers import StreamWriter, CORK
from unittest import mock


# nodelay

def test_nodelay_default(loop):
transport = mock.Mock()
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Expand Down Expand Up @@ -86,3 +88,96 @@ def test_set_nodelay_enable_no_socket(loop):
writer.set_tcp_nodelay(True)
assert writer.tcp_nodelay
assert writer._socket is None


# cork

@pytest.mark.skipif(CORK is None, reason="TCP_CORK or TCP_NOPUSH required")
def test_cork_default(loop):
transport = mock.Mock()
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
transport.get_extra_info.return_value = s
proto = mock.Mock()
reader = mock.Mock()
writer = StreamWriter(transport, proto, reader, loop)
assert not writer.tcp_cork
assert not s.getsockopt(socket.IPPROTO_TCP, CORK)


@pytest.mark.skipif(CORK is None, reason="TCP_CORK or TCP_NOPUSH required")
def test_set_cork_no_change(loop):
transport = mock.Mock()
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
transport.get_extra_info.return_value = s
proto = mock.Mock()
reader = mock.Mock()
writer = StreamWriter(transport, proto, reader, loop)
writer.set_tcp_cork(False)
assert not writer.tcp_cork
assert not s.getsockopt(socket.IPPROTO_TCP, CORK)


@pytest.mark.skipif(CORK is None, reason="TCP_CORK or TCP_NOPUSH required")
def test_set_cork_enable(loop):
transport = mock.Mock()
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
transport.get_extra_info.return_value = s
proto = mock.Mock()
reader = mock.Mock()
writer = StreamWriter(transport, proto, reader, loop)
writer.set_tcp_cork(True)
assert writer.tcp_cork
assert s.getsockopt(socket.IPPROTO_TCP, CORK)


@pytest.mark.skipif(CORK is None, reason="TCP_CORK or TCP_NOPUSH required")
def test_set_cork_enable_and_disable(loop):
transport = mock.Mock()
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
transport.get_extra_info.return_value = s
proto = mock.Mock()
reader = mock.Mock()
writer = StreamWriter(transport, proto, reader, loop)
writer.set_tcp_cork(True)
writer.set_tcp_cork(False)
assert not writer.tcp_cork
assert not s.getsockopt(socket.IPPROTO_TCP, CORK)


@pytest.mark.skipif(CORK is None, reason="TCP_CORK or TCP_NOPUSH required")
def test_set_cork_enable_ipv6(loop):
transport = mock.Mock()
s = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
transport.get_extra_info.return_value = s
proto = mock.Mock()
reader = mock.Mock()
writer = StreamWriter(transport, proto, reader, loop)
writer.set_tcp_cork(True)
assert writer.tcp_cork
assert s.getsockopt(socket.IPPROTO_TCP, CORK)


@pytest.mark.skipif(not hasattr(socket, 'AF_UNIX'),
reason="requires unix sockets")
@pytest.mark.skipif(CORK is None, reason="TCP_CORK or TCP_NOPUSH required")
def test_set_cork_enable_unix(loop):
transport = mock.Mock()
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
transport.get_extra_info.return_value = s
proto = mock.Mock()
reader = mock.Mock()
writer = StreamWriter(transport, proto, reader, loop)
writer.set_tcp_cork(True)
assert writer.tcp_cork


@pytest.mark.skipif(CORK is None, reason="TCP_CORK or TCP_NOPUSH required")
def test_set_cork_enable_no_socket(loop):
transport = mock.Mock()
transport.get_extra_info.return_value = None
proto = mock.Mock()
reader = mock.Mock()
writer = StreamWriter(transport, proto, reader, loop)
writer.set_tcp_cork(True)
assert writer.tcp_cork
assert writer._socket is None