From a9efbfe585e381ea40b21b778cfa99bd026bb23d Mon Sep 17 00:00:00 2001
From: Nejc Saje <nejc@saje.info>
Date: Tue, 16 Jul 2024 11:20:48 +0200
Subject: [PATCH] Add heartbeat jitter (#184)

---
 socketshark/subscription.py | 6 +++++-
 tests/test_basic.py         | 8 +++++---
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/socketshark/subscription.py b/socketshark/subscription.py
index ad07cb1..db8c54a 100644
--- a/socketshark/subscription.py
+++ b/socketshark/subscription.py
@@ -1,4 +1,5 @@
 import asyncio
+import random
 import time
 from typing import Any, Dict, Optional
 
@@ -184,8 +185,10 @@ async def periodic_heartbeat(self):
             subscription=self.name,
             period=period,
         )
+        # Add some randomness to the first heartbeat period to avoid all
+        # subscriptions sending heartbeats at the same time upon redeploy.
+        await asyncio.sleep(period * random.random())
         while True:
-            await asyncio.sleep(period)
             try:
                 self.session.log.debug(
                     'sending heartbeat', subscription=self.name
@@ -202,6 +205,7 @@ async def periodic_heartbeat(self):
                     subscription=self.name,
                     error=e.error,
                 )
+            await asyncio.sleep(period)
 
     async def before_subscribe(self):
         return await self.perform_service_request('before_subscribe')
diff --git a/tests/test_basic.py b/tests/test_basic.py
index d3b8d08..ddc68ca 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -1,6 +1,7 @@
 import asyncio
 import json
 import os
+import random
 import time
 from unittest.mock import patch
 
@@ -1097,6 +1098,7 @@ async def test_subscription_periodic_heartbeat(self):
                 },
             )
 
+            random.seed(1)  # due to a random sleep at the beginning
             await session.on_client_event(
                 {
                     'event': 'subscribe',
@@ -1105,10 +1107,10 @@ async def test_subscription_periodic_heartbeat(self):
             )
 
             mock.assert_not_called()
-
-            await asyncio.sleep(0.4)
-
+            await asyncio.sleep(0.2)
             mock.assert_called_once()
+            await asyncio.sleep(0.2)
+            assert len(list(mock.requests.values())[0]) == 2
 
         await shark.shutdown()