From 9bdf27026ec79d53962503940d2de2aa17fbb15f Mon Sep 17 00:00:00 2001
From: Debdut Chakraborty <debdutdeb@outlook.com>
Date: Wed, 23 Oct 2024 11:59:42 +0530
Subject: [PATCH] appservice: reset recoverer backoff on login

---
 .gitignore                                    |  4 ++++
 synapse/appservice/scheduler.py               | 13 +++++++++++-
 synapse/module_api/callbacks/__init__.py      |  3 +++
 .../callbacks/appservice_login_callbacks.py   | 21 +++++++++++++++++++
 synapse/rest/client/login.py                  | 10 +++++++++
 5 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 synapse/module_api/callbacks/appservice_login_callbacks.py

diff --git a/.gitignore b/.gitignore
index a89f149ec1..61e98352b1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -78,3 +78,7 @@ book/
 
 # Don't include users' poetry configs
 /poetry.toml
+
+data/
+registration.yaml
+registration.yml
diff --git a/synapse/appservice/scheduler.py b/synapse/appservice/scheduler.py
index 7994da0868..387fe13624 100644
--- a/synapse/appservice/scheduler.py
+++ b/synapse/appservice/scheduler.py
@@ -58,6 +58,7 @@ components.
 import logging
 from typing import (
     TYPE_CHECKING,
+    Any,
     Awaitable,
     Callable,
     Collection,
@@ -482,11 +483,21 @@ class _Recoverer:
         self.service = service
         self.callback = callback
         self.backoff_counter = 1
+        self._delayed_call_handler = None
+
+    def reset(self) -> None:
+        if self._delayed_call_handler is not None and self._delayed_call_handler.active():
+            self._delayed_call_handler.cancel()
+            self._delayed_call_handler = None
+
+        self.backoff_counter = 1
+        
+        self.recover()
 
     def recover(self) -> None:
         delay = 2**self.backoff_counter
         logger.info("Scheduling retries on %s in %fs", self.service.id, delay)
-        self.clock.call_later(
+        self._delayed_call_handler = self.clock.call_later(
             delay, run_as_background_process, "as-recoverer", self.retry
         )
 
diff --git a/synapse/module_api/callbacks/__init__.py b/synapse/module_api/callbacks/__init__.py
index c20d9543fb..a8aee341ef 100644
--- a/synapse/module_api/callbacks/__init__.py
+++ b/synapse/module_api/callbacks/__init__.py
@@ -21,6 +21,8 @@
 
 from typing import TYPE_CHECKING
 
+from synapse.module_api.callbacks.appservice_login_callbacks import AppserviceLoginModuleApiCallbacks
+
 if TYPE_CHECKING:
     from synapse.server import HomeServer
 
@@ -40,3 +42,4 @@ class ModuleApiCallbacks:
         self.account_validity = AccountValidityModuleApiCallbacks()
         self.spam_checker = SpamCheckerModuleApiCallbacks(hs)
         self.third_party_event_rules = ThirdPartyEventRulesModuleApiCallbacks(hs)
+        self.appservice_login = AppserviceLoginModuleApiCallbacks(hs)
diff --git a/synapse/module_api/callbacks/appservice_login_callbacks.py b/synapse/module_api/callbacks/appservice_login_callbacks.py
new file mode 100644
index 0000000000..590421d64e
--- /dev/null
+++ b/synapse/module_api/callbacks/appservice_login_callbacks.py
@@ -0,0 +1,21 @@
+from typing import TYPE_CHECKING, Union
+
+
+if TYPE_CHECKING:
+    from synapse.server import HomeServer
+    from synapse.appservice import ApplicationService
+
+class AppserviceLoginModuleApiCallbacks:
+    def __init__(self, hs: "HomeServer") -> None:
+        self.scheduler = hs.get_application_service_scheduler()
+
+    async def reset_recoverer_backoff(self, appservice: "ApplicationService"):
+        print(appservice, "resetting backoff, retrying transactions immediately")
+
+        appservice_id = appservice.id
+
+        if appservice_id not in self.scheduler.txn_ctrl.recoverers:
+            print(appservice, "no recoverer found to reset")
+            return
+
+        self.scheduler.txn_ctrl.recoverers[appservice_id].reset()
diff --git a/synapse/rest/client/login.py b/synapse/rest/client/login.py
index 3271b02d40..20a02f9c43 100644
--- a/synapse/rest/client/login.py
+++ b/synapse/rest/client/login.py
@@ -19,6 +19,7 @@
 #
 #
 
+from copy import copy
 import logging
 import re
 from typing import (
@@ -307,6 +308,8 @@ class LoginRestServlet(RestServlet):
         if not appservice.is_interested_in_user(qualified_user_id):
             raise LoginError(403, "Invalid access_token", errcode=Codes.FORBIDDEN)
 
+        print("logging in appservice")
+
         return await self._complete_login(
             qualified_user_id,
             login_submission,
@@ -316,8 +319,15 @@ class LoginRestServlet(RestServlet):
             # is not actually created in Synapse.
             should_check_deactivated=qualified_user_id != appservice.sender,
             request_info=request_info,
+            callback=self._appservice_login_callback(appservice),
         )
 
+    def _appservice_login_callback(self, appservice: "ApplicationService") -> Optional[Callable[[LoginResponse], Awaitable[None]]]:
+        async def __handle_appservice_login(_: "LoginResponse"):
+            await self.hs.get_module_api_callbacks().appservice_login.reset_recoverer_backoff(appservice)
+
+        return __handle_appservice_login
+
     async def _do_other_login(
         self,
         login_submission: JsonDict,