From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Mon, 17 Oct 2022 02:12:24 +0000 (+0100)
Subject: feat: add wait_select() function
X-Git-Tag: 3.1.5~7^2~13
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=136c097ab1a5983646ab27060b4ace26671b1c48;p=thirdparty%2Fpsycopg.git

feat: add wait_select() function

A few tests show that using select() from Python has a lower overhead.
See #414.

Use this function wherever defined and epoll() is the default select
strategy instead. This is mostly a draft to figure out where this
strategy can be used with success.
---

diff --git a/psycopg/psycopg/waiting.py b/psycopg/psycopg/waiting.py
index 88e8c6fd3..a784aaab4 100644
--- a/psycopg/psycopg/waiting.py
+++ b/psycopg/psycopg/waiting.py
@@ -12,7 +12,7 @@ These functions are designed to consume the generators returned by the
 import select
 import selectors
 from enum import IntEnum
-from typing import Optional
+from typing import Dict, Optional
 from asyncio import get_event_loop, wait_for, Event, TimeoutError
 from selectors import DefaultSelector, EVENT_READ, EVENT_WRITE
 
@@ -214,6 +214,49 @@ async def wait_conn_async(gen: PQGenConn[RV], timeout: Optional[float] = None) -
         return rv
 
 
+def wait_select(gen: PQGen[RV], fileno: int, timeout: Optional[float] = None) -> RV:
+    """
+    Wait for a generator using select where supported.
+    """
+    try:
+        s = next(gen)
+
+        empty = ()
+        fnlist = (fileno,)
+        while True:
+            rl, wl, xl = select.select(
+                fnlist if s & WAIT_R else empty,
+                fnlist if s & WAIT_W else empty,
+                fnlist,
+                timeout,
+            )
+            ready = 0
+            if rl:
+                ready = READY_R
+            if wl:
+                ready |= READY_W
+            if not ready:
+                continue
+            assert s & ready
+            s = gen.send(ready)  # type: ignore
+
+    except StopIteration as ex:
+        rv: RV = ex.args[0] if ex.args else None
+        return rv
+
+
+poll_evmasks: Dict[Wait, int]
+
+if hasattr(selectors, "EpollSelector"):
+    poll_evmasks = {
+        WAIT_R: select.EPOLLONESHOT | select.EPOLLIN,
+        WAIT_W: select.EPOLLONESHOT | select.EPOLLOUT,
+        WAIT_RW: select.EPOLLONESHOT | select.EPOLLIN | select.EPOLLOUT,
+    }
+else:
+    poll_evmasks = {}
+
+
 def wait_epoll(gen: PQGen[RV], fileno: int, timeout: Optional[float] = None) -> RV:
     """
     Wait for a generator using epoll where supported.
@@ -255,13 +298,13 @@ def wait_epoll(gen: PQGen[RV], fileno: int, timeout: Optional[float] = None) ->
 
 
 if selectors.DefaultSelector is getattr(selectors, "EpollSelector", None):
-    wait = wait_epoll
-
-    poll_evmasks = {
-        Wait.R: select.EPOLLONESHOT | select.EPOLLIN,
-        Wait.W: select.EPOLLONESHOT | select.EPOLLOUT,
-        Wait.RW: select.EPOLLONESHOT | select.EPOLLIN | select.EPOLLOUT,
-    }
+    # NOTE: select seems more performing than epoll. It is admittedly unlikely
+    # that a platform has epoll but not select, so maybe we could kill
+    # wait_epoll altogether(). More testing to do.
+    if hasattr(selectors, "SelectSelector"):
+        wait = wait_select
+    else:
+        wait = wait_epoll
 
 else:
     wait = wait_selector
diff --git a/tests/conftest.py b/tests/conftest.py
index b1882a246..4d83f7ca5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,5 +1,6 @@
 import sys
 import asyncio
+import selectors
 from typing import List
 
 import pytest
@@ -58,11 +59,14 @@ def pytest_addoption(parser):
 
 
 def pytest_report_header(config):
+    rv = []
+
+    rv.append(f"default selector: {selectors.DefaultSelector.__name__}")
     loop = config.getoption("--loop")
-    if loop == "default":
-        return []
+    if loop != "default":
+        rv.append(f"asyncio loop: {loop}")
 
-    return [f"asyncio loop: {loop}"]
+    return rv
 
 
 def pytest_sessionstart(session):