From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 29 Dec 2020 03:03:34 +0000 (+0100)
Subject: Added epoll-based wait() implementation
X-Git-Tag: 3.0.dev0~228^2~1
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=a2c4c4bf7ba16878258560a6e3931565b5b35356;p=thirdparty%2Fpsycopg.git

Added epoll-based wait() implementation

Only used if the EpollSelector would (i.e. epoll is available and it's
the best thing on the menu). Saves noticeable time over the
EpollSelector as there's less Python stuff going on and can use ONESHOT
+ MOD_CTL.
---

diff --git a/psycopg3/psycopg3/waiting.py b/psycopg3/psycopg3/waiting.py
index af818fbb8..e227aabc7 100644
--- a/psycopg3/psycopg3/waiting.py
+++ b/psycopg3/psycopg3/waiting.py
@@ -9,6 +9,8 @@ These functions are designed to consume the generators returned by the
 # Copyright (C) 2020 The Psycopg Team
 
 
+import select
+import selectors
 from enum import IntEnum
 from typing import Optional
 from asyncio import get_event_loop, Event
@@ -190,3 +192,55 @@ async def wait_async_conn(gen: PQGenConn[RV]) -> RV:
     except StopIteration as ex:
         rv: RV = ex.args[0] if ex.args else None
         return rv
+
+
+poll_evmasks = {
+    Wait.R: select.EPOLLONESHOT | select.EPOLLIN,
+    Wait.W: select.EPOLLONESHOT | select.EPOLLOUT,
+    Wait.RW: select.EPOLLONESHOT | select.EPOLLIN | select.EPOLLOUT,
+}
+
+
+def wait_epoll(
+    gen: PQGen[RV], fileno: int, timeout: Optional[float] = None
+) -> RV:
+    """
+    Wait for a generator using epoll where supported.
+
+    Parameters are like for `wait()`. If it is detected that the best selector
+    strategy is `epoll` then this function will be used instead of `wait`.
+
+    See also: https://linux.die.net/man/2/epoll_ctl
+    """
+    epoll = select.epoll()
+    if timeout is None or timeout < 0:
+        timeout = 0
+    else:
+        timeout = int(timeout * 1000.0)
+    try:
+        s = next(gen)
+        evmask = poll_evmasks[s]
+        epoll.register(fileno, evmask)
+        while 1:
+            fileevs = None
+            while not fileevs:
+                fileevs = epoll.poll(timeout)
+            ev = fileevs[0][1]
+            if ev & ~select.EPOLLOUT:
+                s = Ready.R
+            else:
+                s = Ready.W
+            s = gen.send(s)
+            evmask = poll_evmasks[s]
+            epoll.modify(fileno, evmask)
+
+    except StopIteration as ex:
+        rv: RV = ex.args[0] if ex.args else None
+        return rv
+
+
+if (
+    selectors.DefaultSelector  # type: ignore[comparison-overlap]
+    is selectors.EpollSelector
+):
+    wait = wait_epoll  # noqa: F811