Skip to content

LCORE-1062: reconnection logic in quota runner #921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tisnik merged 4 commits into from
Dec 16, 2025
Merged

LCORE-1062: reconnection logic in quota runner #921

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5719637
Ability to reconnect to DB
tisnik Dec 16, 2025
1f1bd6f
Reconnection on beginning
tisnik Dec 16, 2025
b1f4700
Minor fixes
tisnik Dec 16, 2025
7e4744f
Avoid resource leaks
tisnik Dec 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 47 additions & 2 deletions src/runners/quota_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,11 @@

logger = get_logger(__name__)

DATABASE_RECONNECTION_COUNT: int = 10
RECONNECTION_DELAY: int = 1


# pylint: disable=R0912
def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
"""
Run the quota scheduler loop that applies configured quota limiters periodically.
Expand Down Expand Up @@ -55,8 +59,15 @@ def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
logger.warning("No limiters are setup, skipping")
return False

connection = connect(config)
if connection is None:
for _ in range(DATABASE_RECONNECTION_COUNT):
try:
connection = connect(config)
if connection is not None:
break
except Exception as e: # pylint: disable=broad-exception-caught
logger.warning("Can not connect to database, will try later: %s", e)
sleep(RECONNECTION_DELAY)
else:
logger.warning("Can not connect to database, skipping")
return False
Comment on lines +62 to 72
Copy link
Contributor

@coderabbitai coderabbitai bot Dec 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Optimize retry loop to avoid unnecessary final delay.

The current loop sleeps after each attempt, including the final failed one. This adds an unnecessary 1-second delay before returning.

Consider moving the sleep to the beginning of the loop after the first iteration:

+    retry_count = 0
-    for _ in range(DATABASE_RECONNECTION_COUNT):
+    while retry_count < DATABASE_RECONNECTION_COUNT:
+        if retry_count > 0:
+            sleep(RECONNECTION_DELAY)
+        retry_count += 1
         try:
             connection = connect(config)
             if connection is not None:
                 break
         except Exception as e:  # pylint: disable=broad-exception-caught
             logger.warning("Can not connect to database, will try later: %s", e)
-        sleep(RECONNECTION_DELAY)
     else:
         logger.warning("Can not connect to database, skipping")
         return False
🤖 Prompt for AI Agents 
In src/runners/quota_scheduler.py around lines 62 to 72, the retry loop
currently calls sleep(RECONNECTION_DELAY) after every attempt, causing an
unnecessary delay after the final failed attempt; change the logic so the delay
happens only between attempts (either move the sleep to the start of the loop
after the first iteration or conditionally call sleep only when the current
attempt is not the last), keeping the same number of retries and preserving
exception handling and logging.


Expand All @@ -83,6 +94,16 @@ def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
logger.info("Quota scheduler sync started")
for limiter in config.limiters:
try:
if not connected(connection):
# the old connection might be closed to avoid resource leaks
try:
connection.close()
except Exception: # pylint: disable=broad-exception-caught
pass # Connection already dead
connection = connect(config)
if connection is None:
logger.warning("Can not connect to database, skipping")
continue
Comment on lines +97 to +106
Copy link
Contributor

@coderabbitai coderabbitai bot Dec 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Reinitialize tables after reconnecting to database.

When the connection is lost and re-established, the quota table may not exist (e.g., if the database was recreated or wiped). Without calling init_tables() after reconnection, subsequent quota operations will fail.

Apply this diff to ensure tables are initialized after reconnection:

                 if not connected(connection):
                     # the old connection might be closed to avoid resource leaks
                     try:
                         connection.close()
                     except Exception:  # pylint: disable=broad-exception-caught
                         pass  # Connection already dead
                     connection = connect(config)
                     if connection is None:
                         logger.warning("Can not connect to database, skipping")
                         continue
+                    # Reinitialize tables after reconnection
+                    if create_quota_table is not None:
+                        init_tables(connection, create_quota_table)
                 quota_revocation(
                     connection, limiter, increase_quota_statement, reset_quota_statement
                 )
🤖 Prompt for AI Agents 
In src/runners/quota_scheduler.py around lines 97 to 106, when reconnecting to
the DB the code does not reinitialize schema/tables; after assigning connection
= connect(config) call the existing init_tables(...) helper using the same
arguments/signature used elsewhere (e.g., init_tables(connection) or
init_tables(config, connection)), catch and log any exceptions from init_tables,
and if initialization fails log a warning and continue the loop so subsequent
quota operations don't run against a missing schema.

quota_revocation(
connection, limiter, increase_quota_statement, reset_quota_statement
)
Expand All @@ -95,6 +116,30 @@ def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
return True


def connected(connection: Any) -> bool:
"""Check if DB is still connected.

Parameters:
connection: Database connection object to verify.

Returns:
bool: True if connection is active, False otherwise.
"""
if connection is None:
logger.warning("Not connected, need to reconnect later")
return False
try:
# for compatibility with SQLite it is not possible to use context manager there
cursor = connection.cursor()
cursor.execute("SELECT 1")
cursor.close()
logger.info("Connection to storage is ok")
return True
except Exception as e: # pylint: disable=broad-exception-caught
logger.error("Disconnected from storage: %s", e)
return False


def get_increase_quota_statement(config: QuotaHandlersConfiguration) -> str:
"""
Select the SQL statement used to increase stored quota according to the database backend.
Expand Down
Loading