Skip to content

Commit

Permalink
less aggressively close sessions
Browse files Browse the repository at this point in the history
  • Loading branch information
agoscinski committed Feb 2, 2025
1 parent 54f8fe3 commit f3f41c1
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 5 deletions.
12 changes: 7 additions & 5 deletions src/aiida/storage/psql_dos/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,16 +186,18 @@ def close(self) -> None:
return # the instance is already closed, and so this is a no-op
# close the connection

# PR QUESTION: can we have multiple sessons? Then we should close all of them
#from sqlalchemy.orm import close_all_sessions
#close_all_sessions()

engine = self._session_factory.bind
if engine is not None:
engine.dispose() # type: ignore[union-attr]
self._session_factory.expunge_all()
self._session_factory.close()
self._session_factory = None

# Without this, sqlalchemy keeps a weakref to a session
# in sqlalchemy.orm.session._sessions
gc.collect()
# IMPORTANT: Dispose engine only after sessions have been closed
if engine is not None:
engine.dispose() # type: ignore[union-attr]

def _clear(self) -> None:
from aiida.storage.psql_dos.models.settings import DbSetting
Expand Down
21 changes: 21 additions & 0 deletions tests/storage/sqlite/test_container.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,26 @@
import psutil, os

def test_file_descriptor_closed(aiida_profile):
"""Checks if the number of open file descriptors change during a reset."""
def list_open_file_descriptors():
process = psutil.Process(os.getpid())
return process.open_files()
# We have some connections active due to aiida profile during the first
# reset these are destroyed. We check the second time changes the number of
# open file descriptors.
from aiida.manage import get_manager
storage_backend = get_manager().get_profile_storage()
migrator_context = storage_backend.migrator_context
open_file_descriptors_before = list_open_file_descriptors()
with migrator_context(aiida_profile) as migrator:
migrator.initialise_repository()
migrator.reset_repository()
open_file_descriptors_after = list_open_file_descriptors()
assert len(open_file_descriptors_before) == len(open_file_descriptors_after), f"Before these file descriptor were open:\n{open_file_descriptors_before}\n Now these are open:\n{open_file_descriptors_after}"


# PR COMMENT this is just a sanity check for me, I don' think that the test should be included in final PR
def test_reset_storage_file_descriptor_closed(aiida_profile):
"""Checks if the number of open file descriptors change during a reset."""
def list_open_file_descriptors():
process = psutil.Process(os.getpid())
Expand All @@ -11,6 +31,7 @@ def list_open_file_descriptors():
# open file descriptors.
# TODO The fix should keep the existing connections alive and just reuse them
# then one does not need to do two resets.
from aiida.manage import get_manager
aiida_profile.reset_storage()
open_file_descriptors_before = list_open_file_descriptors()
aiida_profile.reset_storage()
Expand Down

0 comments on commit f3f41c1

Please sign in to comment.