diff --git a/synapse/module_api/__init__.py b/synapse/module_api/__init__.py
index 7a25a47ac5..cb09423cb3 100644
--- a/synapse/module_api/__init__.py
+++ b/synapse/module_api/__init__.py
@@ -74,6 +74,11 @@ from synapse.module_api.callbacks.account_validity_callbacks import (
ON_LEGACY_SEND_MAIL_CALLBACK,
ON_USER_REGISTRATION_CALLBACK,
)
+from synapse.module_api.callbacks.background_updater_callbacks import (
+ DEFAULT_BATCH_SIZE_CALLBACK,
+ MIN_BATCH_SIZE_CALLBACK,
+ ON_UPDATE_CALLBACK,
+)
from synapse.module_api.callbacks.presence_router_callbacks import (
GET_INTERESTED_USERS_CALLBACK,
GET_USERS_FOR_STATES_CALLBACK,
@@ -107,11 +112,6 @@ from synapse.module_api.callbacks.third_party_event_rules_callbacks import (
)
from synapse.rest.client.login import LoginResponse
from synapse.storage import DataStore
-from synapse.storage.background_updates import (
- DEFAULT_BATCH_SIZE_CALLBACK,
- MIN_BATCH_SIZE_CALLBACK,
- ON_UPDATE_CALLBACK,
-)
from synapse.storage.database import DatabasePool, LoggingTransaction
from synapse.storage.databases.main.roommember import ProfileInfo
from synapse.types import (
@@ -438,12 +438,11 @@ class ModuleApi:
Added in Synapse v1.49.0.
"""
- for db in self._hs.get_datastores().databases:
- db.updates.register_update_controller_callbacks(
- on_update=on_update,
- default_batch_size=default_batch_size,
- min_batch_size=min_batch_size,
- )
+ self._callbacks.background_updater.register_callbacks(
+ on_update=on_update,
+ default_batch_size=default_batch_size,
+ min_batch_size=min_batch_size,
+ )
def register_account_data_callbacks(
self,
diff --git a/synapse/module_api/callbacks/__init__.py b/synapse/module_api/callbacks/__init__.py
index 50b5f2f4d7..a35e7bf511 100644
--- a/synapse/module_api/callbacks/__init__.py
+++ b/synapse/module_api/callbacks/__init__.py
@@ -13,6 +13,7 @@
# limitations under the License.
from .account_validity_callbacks import AccountValidityModuleApiCallbacks
+from .background_updater_callbacks import BackgroundUpdaterModuleApiCallbacks
from .presence_router_callbacks import PresenceRouterModuleApiCallbacks
from .spam_checker_callbacks import SpamCheckerModuleApiCallbacks
from .third_party_event_rules_callbacks import ThirdPartyEventRulesModuleApiCallbacks
@@ -25,6 +26,7 @@ __all__ = [
class ModuleApiCallbacks:
def __init__(self) -> None:
self.account_validity = AccountValidityModuleApiCallbacks()
+ self.background_updater = BackgroundUpdaterModuleApiCallbacks()
+ self.presence_router = PresenceRouterModuleApiCallbacks()
self.spam_checker = SpamCheckerModuleApiCallbacks()
self.third_party_event_rules = ThirdPartyEventRulesModuleApiCallbacks()
- self.presence_router = PresenceRouterModuleApiCallbacks()
diff --git a/synapse/module_api/callbacks/background_updater_callbacks.py b/synapse/module_api/callbacks/background_updater_callbacks.py
new file mode 100644
index 0000000000..e645cc9106
--- /dev/null
+++ b/synapse/module_api/callbacks/background_updater_callbacks.py
@@ -0,0 +1,54 @@
+# Copyright 2014-2016 OpenMarket Ltd
+# Copyright 2023 The Matrix.org Foundation C.I.C.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+from typing import AsyncContextManager, Awaitable, Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+ON_UPDATE_CALLBACK = Callable[[str, str, bool], AsyncContextManager[int]]
+DEFAULT_BATCH_SIZE_CALLBACK = Callable[[str, str], Awaitable[int]]
+MIN_BATCH_SIZE_CALLBACK = Callable[[str, str], Awaitable[int]]
+
+
+class BackgroundUpdaterModuleApiCallbacks:
+ def __init__(self) -> None:
+ self.on_update_callback: Optional[ON_UPDATE_CALLBACK] = None
+ self.default_batch_size_callback: Optional[DEFAULT_BATCH_SIZE_CALLBACK] = None
+ self.min_batch_size_callback: Optional[MIN_BATCH_SIZE_CALLBACK] = None
+
+ def register_callbacks(
+ self,
+ on_update: ON_UPDATE_CALLBACK,
+ default_batch_size: Optional[DEFAULT_BATCH_SIZE_CALLBACK] = None,
+ min_batch_size: Optional[DEFAULT_BATCH_SIZE_CALLBACK] = None,
+ ) -> None:
+ """Register callbacks from a module for each hook."""
+ if self.on_update_callback is not None:
+ logger.warning(
+ "More than one module tried to register callbacks for controlling"
+ " background updates. Only the callbacks registered by the first module"
+ " (in order of appearance in Synapse's configuration file) that tried to"
+ " do so will be called."
+ )
+
+ return
+
+ self.on_update_callback = on_update
+
+ if default_batch_size is not None:
+ self.default_batch_size_callback = default_batch_size
+
+ if min_batch_size is not None:
+ self.min_batch_size_callback = min_batch_size
diff --git a/synapse/storage/background_updates.py b/synapse/storage/background_updates.py
index a99aea8926..bd673e2fb5 100644
--- a/synapse/storage/background_updates.py
+++ b/synapse/storage/background_updates.py
@@ -42,11 +42,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
-ON_UPDATE_CALLBACK = Callable[[str, str, bool], AsyncContextManager[int]]
-DEFAULT_BATCH_SIZE_CALLBACK = Callable[[str, str], Awaitable[int]]
-MIN_BATCH_SIZE_CALLBACK = Callable[[str, str], Awaitable[int]]
-
-
@attr.s(slots=True, frozen=True, auto_attribs=True)
class _BackgroundUpdateHandler:
"""A handler for a given background update.
@@ -149,13 +144,11 @@ class BackgroundUpdater:
self._database_name = database.name()
+ self._module_api_callbacks = hs.get_module_api_callbacks().background_updater
+
# if a background update is currently running, its name.
self._current_background_update: Optional[str] = None
- self._on_update_callback: Optional[ON_UPDATE_CALLBACK] = None
- self._default_batch_size_callback: Optional[DEFAULT_BATCH_SIZE_CALLBACK] = None
- self._min_batch_size_callback: Optional[MIN_BATCH_SIZE_CALLBACK] = None
-
self._background_update_performance: Dict[str, BackgroundUpdatePerformance] = {}
self._background_update_handlers: Dict[str, _BackgroundUpdateHandler] = {}
self._all_done = False
@@ -175,31 +168,6 @@ class BackgroundUpdater:
self.sleep_duration_ms = hs.config.background_updates.sleep_duration_ms
self.sleep_enabled = hs.config.background_updates.sleep_enabled
- def register_update_controller_callbacks(
- self,
- on_update: ON_UPDATE_CALLBACK,
- default_batch_size: Optional[DEFAULT_BATCH_SIZE_CALLBACK] = None,
- min_batch_size: Optional[DEFAULT_BATCH_SIZE_CALLBACK] = None,
- ) -> None:
- """Register callbacks from a module for each hook."""
- if self._on_update_callback is not None:
- logger.warning(
- "More than one module tried to register callbacks for controlling"
- " background updates. Only the callbacks registered by the first module"
- " (in order of appearance in Synapse's configuration file) that tried to"
- " do so will be called."
- )
-
- return
-
- self._on_update_callback = on_update
-
- if default_batch_size is not None:
- self._default_batch_size_callback = default_batch_size
-
- if min_batch_size is not None:
- self._min_batch_size_callback = min_batch_size
-
def _get_context_manager_for_update(
self,
sleep: bool,
@@ -228,8 +196,10 @@ class BackgroundUpdater:
Note: this is a *target*, and an iteration may take substantially longer or
shorter.
"""
- if self._on_update_callback is not None:
- return self._on_update_callback(update_name, database_name, oneshot)
+ if self._module_api_callbacks.on_update_callback is not None:
+ return self._module_api_callbacks.on_update_callback(
+ update_name, database_name, oneshot
+ )
return _BackgroundUpdateContextManager(
sleep, self._clock, self.sleep_duration_ms, self.update_duration_ms
@@ -239,8 +209,10 @@ class BackgroundUpdater:
"""The batch size to use for the first iteration of a new background
update.
"""
- if self._default_batch_size_callback is not None:
- return await self._default_batch_size_callback(update_name, database_name)
+ if self._module_api_callbacks.default_batch_size_callback is not None:
+ return await self._module_api_callbacks.default_batch_size_callback(
+ update_name, database_name
+ )
return self.default_background_batch_size
@@ -249,8 +221,10 @@ class BackgroundUpdater:
Used to ensure that progress is always made. Must be greater than 0.
"""
- if self._min_batch_size_callback is not None:
- return await self._min_batch_size_callback(update_name, database_name)
+ if self._module_api_callbacks.min_batch_size_callback is not None:
+ return await self._module_api_callbacks.min_batch_size_callback(
+ update_name, database_name
+ )
return self.minimum_background_batch_size
|
