Add shutdown method to gracefully quit internal tasks

Author: Syndace

CHANGELOG.md

@@ -6,6 +6,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 
+### Added
+- A new method `shutdown` to gracefully quit internal tasks
+
 ## [1.2.0] - 15th of October, 2024
 
 ### Changed

LICENSE

@@ -1,6 +1,6 @@
 The MIT License
 
-Copyright (c) 2024 Tim Henkes (Syndace)
+Copyright (c) 2025 Tim Henkes (Syndace)
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

omemo/project.py

@@ -4,7 +4,7 @@
     "name"         : "OMEMO",
     "description"  : "A Python implementation of the OMEMO Multi-End Message and Object Encryption protocol.",
     "url"          : "https://github.com/Syndace/python-omemo",
-    "year"         : "2024",
+    "year"         : "2025",
     "author"       : "Tim Henkes (Syndace)",
     "author_email" : "[email protected]",
     "categories"   : [

omemo/session_manager.py

@@ -4,7 +4,8 @@
 import itertools
 import logging
 import secrets
-from typing import Dict, FrozenSet, List, NamedTuple, Optional, Set, Tuple, Type, TypeVar, Union, cast
+from typing import Awaitable, Dict, FrozenSet, List, NamedTuple, Optional, Set, Tuple, Type, TypeVar, Union, \
+    cast
 from typing_extensions import assert_never
 
 try:
@@ -247,6 +248,8 @@ def __init__(self) -> None:
         self.__pre_key_refill_threshold: int
         self.__identity_key_pair: IdentityKeyPair
         self.__synchronizing: bool
+        self.__async_framework: AsyncFramework
+        self.__signed_pre_key_management_task: Awaitable[None]
 
     @classmethod
     async def create(
@@ -349,6 +352,7 @@ async def create(
         self.__pre_key_refill_threshold = pre_key_refill_threshold
         self.__identity_key_pair = await IdentityKeyPair.get(storage)
         self.__synchronizing = True
+        self.__async_framework = async_framework
 
         try:
             self.__own_device_id = (await self.__storage.load_primitive("/own_device_id", int)).from_just()
@@ -473,16 +477,15 @@ async def create(
                 await self.purge_backend(namespace)
 
         # Start signed pre key rotation management "in the background"
+        signed_pre_key_management_coro = self.__manage_signed_pre_key_rotation(
+            signed_pre_key_rotation_period,
+            async_framework
+        )
+
         if async_framework is AsyncFramework.ASYNCIO:
-            asyncio.ensure_future(self.__manage_signed_pre_key_rotation(
-                signed_pre_key_rotation_period,
-                async_framework
-            ))
+            self.__signed_pre_key_management_task = asyncio.ensure_future(signed_pre_key_management_coro)
         elif async_framework is AsyncFramework.TWISTED:
-            defer.ensureDeferred(self.__manage_signed_pre_key_rotation(
-                signed_pre_key_rotation_period,
-                async_framework
-            ))
+            self.__signed_pre_key_management_task = defer.ensureDeferred(signed_pre_key_management_coro)
         else:
             assert_never(async_framework)
 
@@ -492,6 +495,30 @@ async def create(
 
         return self
 
+    async def shutdown(self) -> None:
+        """
+        Gracefully quit internal tasks.
+        """
+
+        if self.__async_framework is AsyncFramework.ASYNCIO:
+            asyncio_task = cast(asyncio.Future[None], self.__signed_pre_key_management_task)
+            try:
+                asyncio_task.cancel()
+                await asyncio_task
+            except asyncio.CancelledError:
+                pass
+
+        elif self.__async_framework is AsyncFramework.TWISTED:
+            twisted_task = cast(defer.Deferred[None], self.__signed_pre_key_management_task)
+            try:
+                twisted_task.cancel()
+                await twisted_task
+            except defer.CancelledError:
+                pass
+
+        else:
+            assert_never(self.__async_framework)
+
     async def purge_backend(self, namespace: str) -> None:
         """
         Purge a backend, removing both the online data (bundle, device list entry) and the offline data that

tests/test_session_manager.py

@@ -169,6 +169,10 @@ async def test_regression0() -> None:
         assert plaintext == b"Hello back, Alice!"
 
 
+    await alice_session_manager.shutdown()
+    await bob_session_manager.shutdown()
+
+
 async def test_oldmemo_migration() -> None:
     """
     Tests the migration of the legacy storage format, which is provided by python-oldmemo.
@@ -217,3 +221,5 @@ async def test_oldmemo_migration() -> None:
     plaintext, _, _ = await session_manager.decrypt(message)
 
     assert plaintext == b"This is a test message"
+
+    await session_manager.shutdown()