Fix task leak (#1632)

nucliadb/nucliadb/common/maindb/driver.py

@@ -60,12 +60,19 @@ async def count(self, match: str) -> int:
 
 class Driver:
     initialized = False
+    _abort_tasks: List[asyncio.Task] = []
 
     async def initialize(self):
         raise NotImplementedError()
 
     async def finalize(self):
-        raise NotImplementedError()
+        while len(self._abort_tasks) > 0:
+            task = self._abort_tasks.pop()
+            if not task.done():
+                try:
+                    await task
+                except Exception:
+                    pass
 
     async def begin(self) -> Transaction:
         raise NotImplementedError()
@@ -94,4 +101,9 @@ async def transaction(
                 if error or wait_for_abort:
                     await txn.abort()
                 else:
-                    asyncio.create_task(txn.abort())
+                    self._async_abort(txn)
+
+    def _async_abort(self, txn: Transaction):
+        task = asyncio.create_task(txn.abort())
+        task.add_done_callback(lambda task: self._abort_tasks.remove(task))
+        self._abort_tasks.append(task)

nucliadb/nucliadb/tests/unit/common/maindb/test_driver.py

@@ -17,6 +17,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
+import asyncio
 from unittest import mock
 
 import pytest
@@ -34,14 +35,43 @@ async def commit(self, **kw):
 
 
 @pytest.fixture(scope="function")
-def driver() -> Driver:  # type: ignore
+def txn():
+    return TransactionTest()
+
+
+@pytest.fixture(scope="function")
+def driver(txn) -> Driver:  # type: ignore
     driver = Driver()
-    with mock.patch.object(
-        driver, "begin", new=mock.AsyncMock(return_value=TransactionTest())
-    ):
+    with mock.patch.object(driver, "begin", new=mock.AsyncMock(return_value=txn)):
         yield driver
 
 
+@pytest.mark.asyncio
+async def test_driver_async_abort(driver, txn):
+    async with driver.transaction(wait_for_abort=False):
+        pass
+
+    assert len(driver._abort_tasks) == 1
+    await asyncio.sleep(0.1)
+
+    txn.abort.assert_called_once()
+    assert len(driver._abort_tasks) == 0
+
+
+@pytest.mark.asyncio
+async def test_driver_finalize_aborts_transactions(driver, txn):
+    async with driver.transaction(wait_for_abort=False):
+        pass
+
+    assert len(driver._abort_tasks) == 1
+
+    await driver.finalize()
+
+    txn.abort.assert_called_once()
+
+    assert len(driver._abort_tasks) == 0
+
+
 @pytest.mark.asyncio
 async def test_transaction_handles_txn_begin_errors(driver):
     driver.begin.side_effect = ValueError()