CRF: extract add_water_defaults()

* Also fixed bug that set_write_buffer_limits(low=16385) fails

Author: fantix

tests/test_tcp.py

@@ -1149,6 +1149,46 @@ def test_create_server_stream_bittype(self):
             srv.close()
             self.loop.run_until_complete(srv.wait_closed())
 
+    def test_flowcontrol_mixin_set_write_limits(self):
+        async def client(addr):
+            paused = False
+
+            class Protocol(asyncio.Protocol):
+                def pause_writing(self):
+                    nonlocal paused
+                    paused = True
+
+                def resume_writing(self):
+                    nonlocal paused
+                    paused = False
+
+            t, p = await self.loop.create_connection(Protocol, *addr)
+
+            t.write(b'q' * 512)
+            self.assertEqual(t.get_write_buffer_size(), 512)
+
+            t.set_write_buffer_limits(low=16385)
+            self.assertFalse(paused)
+            self.assertEqual(t.get_write_buffer_limits(), (16385, 65540))
+
+            with self.assertRaisesRegex(ValueError, 'high.*must be >= low'):
+                t.set_write_buffer_limits(high=0, low=1)
+
+            t.set_write_buffer_limits(high=1024, low=128)
+            self.assertFalse(paused)
+            self.assertEqual(t.get_write_buffer_limits(), (128, 1024))
+
+            t.set_write_buffer_limits(high=256, low=128)
+            self.assertTrue(paused)
+            self.assertEqual(t.get_write_buffer_limits(), (128, 256))
+
+            t.close()
+
+        with self.tcp_server(lambda sock: sock.recv_all(1),
+                             max_clients=1,
+                             backlog=1) as srv:
+            self.loop.run_until_complete(client(srv.addr))
+
 
 class Test_AIO_TCP(_TestTCP, tb.AIOTestCase):
     pass

uvloop/handles/basetransport.pxd

@@ -21,7 +21,6 @@ cdef class UVBaseTransport(UVSocketHandle):
 
     # All "inline" methods are final
 
-    cdef inline _set_write_buffer_limits(self, int high=*, int low=*)
     cdef inline _maybe_pause_protocol(self)
     cdef inline _maybe_resume_protocol(self)
 

uvloop/handles/basetransport.pyx

@@ -2,8 +2,8 @@ cdef class UVBaseTransport(UVSocketHandle):
 
     def __cinit__(self):
         # Flow control
-        self._high_water = FLOW_CONTROL_HIGH_WATER
-        self._low_water = FLOW_CONTROL_LOW_WATER
+        self._high_water = FLOW_CONTROL_HIGH_WATER * 1024
+        self._low_water = FLOW_CONTROL_HIGH_WATER // 4
 
         self._protocol = None
         self._protocol_connected = 0
@@ -59,25 +59,6 @@ cdef class UVBaseTransport(UVSocketHandle):
                 'protocol': self._protocol,
             })
 
-    cdef inline _set_write_buffer_limits(self, int high=-1, int low=-1):
-        if high == -1:
-            if low == -1:
-                high = FLOW_CONTROL_HIGH_WATER
-            else:
-                high = FLOW_CONTROL_LOW_WATER
-
-        if low == -1:
-            low = high // 4
-
-        if not high >= low >= 0:
-            raise ValueError('high (%r) must be >= low (%r) must be >= 0' %
-                             (high, low))
-
-        self._high_water = high
-        self._low_water = low
-
-        self._maybe_pause_protocol()
-
     cdef inline _maybe_pause_protocol(self):
         cdef:
             size_t size = self._get_write_buffer_size()
@@ -283,12 +264,10 @@ cdef class UVBaseTransport(UVSocketHandle):
     def set_write_buffer_limits(self, high=None, low=None):
         self._ensure_alive()
 
-        if high is None:
-            high = -1
-        if low is None:
-            low = -1
+        self._high_water, self._low_water = add_water_defaults(
+            high, low, FLOW_CONTROL_HIGH_WATER)
 
-        self._set_write_buffer_limits(high, low)
+        self._maybe_pause_protocol()
 
     def get_write_buffer_limits(self):
         return (self._low_water, self._high_water)

uvloop/includes/consts.pxi

@@ -1,7 +1,8 @@
 DEF UV_STREAM_RECV_BUF_SIZE = 256000  # 250kb
 
-DEF FLOW_CONTROL_HIGH_WATER = 65536
-DEF FLOW_CONTROL_LOW_WATER = 16384
+DEF FLOW_CONTROL_HIGH_WATER = 64  # KiB
+DEF FLOW_CONTROL_HIGH_WATER_SSL_READ = 256  # KiB
+DEF FLOW_CONTROL_HIGH_WATER_SSL_WRITE = 512  # KiB
 
 DEF DEFAULT_FREELIST_SIZE = 250
 

uvloop/includes/flowcontrol.pxd

@@ -0,0 +1,20 @@
+cdef inline add_water_defaults(high, low, int kb):
+    cdef int h, l
+    if high is None:
+        if low is None:
+            h = kb * 1024
+        else:
+            l = low
+            h = 4 * l
+    else:
+        h = high
+    if low is None:
+        l = h // 4
+    else:
+        l = low
+
+    if not h >= l >= 0:
+        raise ValueError('high (%r) must be >= low (%r) must be >= 0' %
+                         (h, l))
+
+    return h, l

uvloop/loop.pyx

@@ -21,6 +21,7 @@ from .includes.python cimport PY_VERSION_HEX, \
                               PyContext_Exit, \
                               PyMemoryView_FromMemory, PyBUF_WRITE, \
                               PyMemoryView_FromObject, PyMemoryView_Check
+from .includes.flowcontrol cimport add_water_defaults
 
 from libc.stdint cimport uint64_t
 from libc.string cimport memset, strerror, memcpy

uvloop/sslproto.pyx

@@ -766,7 +766,8 @@ cdef class SSLProtocol:
         return self._outgoing.pending + self._write_buffer_size
 
     cdef _set_write_buffer_limits(self, high=None, low=None):
-        high, low = _add_water_defaults(high, low, 512)
+        high, low = add_water_defaults(high, low,
+                                       FLOW_CONTROL_HIGH_WATER_SSL_WRITE)
         self._outgoing_high_water = high
         self._outgoing_low_water = low
 
@@ -800,7 +801,8 @@ cdef class SSLProtocol:
             self._transport.resume_reading()
 
     cdef _set_read_buffer_limits(self, high=None, low=None):
-        high, low = _add_water_defaults(high, low, 256)
+        high, low = add_water_defaults(high, low,
+                                       FLOW_CONTROL_HIGH_WATER_SSL_READ)
         self._incoming_high_water = high
         self._incoming_low_water = low
 
@@ -840,19 +842,3 @@ cdef class SSLProtocol:
                 'transport': self._transport,
                 'protocol': self,
             })
-
-
-cdef _add_water_defaults(high, low, kb):
-    if high is None:
-        if low is None:
-            high = kb * 1024
-        else:
-            high = 4 * low
-    if low is None:
-        low = high // 4
-
-    if not high >= low >= 0:
-        raise ValueError(
-            f'high ({high!r}) must be >= low ({low!r}) must be >= 0')
-
-    return high, low