From 027f0a3abc2844c73897887b30469954ce158e63 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 28 Feb 2021 03:55:29 +0100
Subject: [PATCH] Add jitter method to pool

---
 psycopg3/psycopg3/pool/base.py | 17 +++++++++++------
 tests/pool/test_pool.py        |  7 +++++++
 tests/pool/test_pool_async.py  |  9 +++++++++
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/psycopg3/psycopg3/pool/base.py b/psycopg3/psycopg3/pool/base.py
index c999ac113..c56b116c5 100644
--- a/psycopg3/psycopg3/pool/base.py
+++ b/psycopg3/psycopg3/pool/base.py
@@ -4,10 +4,10 @@ psycopg3 connection pool base class and functionalities.
 
 # Copyright (C) 2021 The Psycopg Team
 
-import random
 import logging
 import threading
 from queue import Queue, Empty
+from random import random
 from typing import Any, Callable, Deque, Dict, Generic, List, Optional
 from collections import deque
 
@@ -164,7 +164,7 @@ class BasePool(Generic[ConnectionType]):
         StopWorker is received.
         """
         # Don't make all the workers time out at the same moment
-        timeout = WORKER_TIMEOUT * (0.9 + 0.1 * random.random())
+        timeout = cls._jitter(WORKER_TIMEOUT, -0.1, 0.1)
         while True:
             # Use a timeout to make the wait interruptable
             try:
@@ -187,6 +187,13 @@ class BasePool(Generic[ConnectionType]):
                     "task run %s failed: %s: %s", task, e.__class__.__name__, e
                 )
 
+    @classmethod
+    def _jitter(cls, value: float, min_pc: float, max_pc: float) -> float:
+        """
+        Add a random value to *value* between *min_pc* and *max_pc* percent.
+        """
+        return value * (1.0 + ((max_pc - min_pc) * random()) + min_pc)
+
 
 class ConnectionAttempt:
     """Keep the state of a connection attempt."""
@@ -204,11 +211,9 @@ class ConnectionAttempt:
         """Calculate how long to wait for a new connection attempt"""
         if self.delay == 0.0:
             self.give_up_at = now + self.reconnect_timeout
-            # +/- 10% of the initial delay
-            jitter = self.INITIAL_DELAY * (
-                (2.0 * self.DELAY_JITTER * random.random()) - self.DELAY_JITTER
+            self.delay = BasePool._jitter(
+                self.INITIAL_DELAY, -self.DELAY_JITTER, self.DELAY_JITTER
             )
-            self.delay = self.INITIAL_DELAY + jitter
         else:
             self.delay *= self.DELAY_BACKOFF
 
diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index 1356f8e08..49792724a 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -636,6 +636,13 @@ def test_resize(dsn):
     assert size == [2, 1, 3, 4, 3, 2, 2]
 
 
+def test_jitter():
+    rnds = [pool.ConnectionPool._jitter(30, -0.1, +0.2) for i in range(100)]
+    rnds.sort()
+    assert 27 <= min(rnds) <= 28
+    assert 35 < max(rnds) < 36
+
+
 def delay_connection(monkeypatch, sec):
     """
     Return a _connect_gen function delayed by the amount of seconds
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index 0425947d3..7a19e036b 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -668,6 +668,15 @@ async def test_resize(dsn):
     assert size == [2, 1, 3, 4, 3, 2, 2]
 
 
+def test_jitter():
+    rnds = [
+        pool.AsyncConnectionPool._jitter(30, -0.1, +0.2) for i in range(100)
+    ]
+    rnds.sort()
+    assert 27 <= min(rnds) <= 28
+    assert 35 < max(rnds) < 36
+
+
 def delay_connection(monkeypatch, sec):
     """
     Return a _connect_gen function delayed by the amount of seconds
-- 
2.47.2