Properly handle multiple requests to threaded REST API

[CasperWA]

Fixes #3743

This PR implements the same fix as was used for aiida-optimade, wherein it closes the SQLAlchemy session used by QueryBuilder after having run a method/function to handle a GET request to the REST API.

If this is not done, the session should(?) still eventually close, releasing connections made in the SQLA connection pool. However, this may take such a long time that when a new GET request comes in (or more precisely, whenever the QueryBuilder is used), the newly created thread that handles the request cannot use the singleton session that AiiDA has, since the connections have not been released, and so it fails with the following message:

sqlalchemy.exc.TimeoutError: QueuePool limit of size 5 overflow 10 reached, connection timed out, timeout 30 (Background on this error at: http://sqlalche.me/e/3o7r)

At least that is how I remember the explanation @sphuber and I came up with some time ago. Hence @sphuber also cleaned up the whole session handling code in AiiDA and provided a public interface for closing the given session.

The major issue here is threading; the AiiDA SQLA session (that the QueryBuilder uses no matter the actual profile's database backend) is a singleton, and only meant to be used within the same (main) thread. Since the REST API runs threaded, as it should to accommodate multiple incoming requests effectively, the singleton AiiDA SQLA sessions design breaks down. However, by closing the session at the end of every request, the connections are immediately released and only in extreme cases will the QueuePool ever reach is limit. In fact, one may argue that in those cases, it should complain, since something is probably wrong/you're flooding the database with too many requests.

This issue was never actually tested with the current tests, since a threaded server is never started.
I have added a new test that starts the REST API server in a new Thread, issues 100 requests in serial, and finally shuts the server down.
This resulted in a reproduction of the QueuePool error, which was then fixed after implementing my fix.

A side note: Since the REST API (as far as I know, please back me up here @waychal) only serves data from the database, and never alters or creates data in the database, there is no risk of corrupting the database by bombarding the REST API with requests.

Last thing: @sphuber, should the close_session decorator function be placed somewhere else in the code? I think I would like to perhaps use this directly in aiida-optimade as well, instead of re-implementing it.
Also, would I then have to move the from aiida.manage.manager import get_manager import line into the finally: section of the close_session() function again, where it was originally placed, in order to not unintentionally load the AiiDA profile before it's truly time?

[codecov]

Codecov Report

Merging #3974 into develop will increase coverage by 0.07%.
The diff coverage is 100.00%.

@@             Coverage Diff             @@
##           develop    #3974      +/-   ##
===========================================
+ Coverage    78.50%   78.57%   +0.07%     
===========================================
  Files          461      461              
  Lines        34094    34095       +1     
===========================================
+ Hits         26763    26787      +24     
+ Misses        7331     7308      -23     

Flag	Coverage Δ
#django	70.59% <85.19%> (+0.05%)	⬆️
#sqlalchemy	71.44% <81.49%> (+0.04%)	⬆️

Impacted Files	Coverage Δ
aiida/manage/tests/unittest_classes.py	0.00% <ø> (ø)
aiida/backends/djsite/__init__.py	100.00% <100.00%> (+33.34%)	⬆️
aiida/backends/djsite/manager.py	82.03% <100.00%> (+2.72%)	⬆️
aiida/backends/manager.py	72.42% <100.00%> (ø)
aiida/backends/sqlalchemy/__init__.py	100.00% <100.00%> (+33.34%)	⬆️
aiida/backends/sqlalchemy/manager.py	78.00% <100.00%> (+2.00%)	⬆️
aiida/restapi/common/identifiers.py	71.52% <100.00%> (ø)
aiida/restapi/common/utils.py	71.70% <100.00%> (+0.46%)	⬆️
aiida/restapi/resources.py	94.56% <100.00%> (-0.15%)	⬇️
aiida/restapi/run_api.py	90.70% <100.00%> (-0.21%)	⬇️
... and 6 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 68fabf0...11130a4. Read the comment docs.

[sphuber]

Thanks a lot @CasperWA , looking good in general with a few minor things and a design question/suggestion.

Last thing: @sphuber, should the close_session decorator function be placed somewhere else in the code?

Any code path that will cause the QueryBuilder to be constructed by a thread should have this. I have the feeling that the get method would be a good place to catch all of these, but I am not sure. We should check with @waychal and @asle85

Also, would I then have to move the from aiida.manage.manager import get_manager import line into the finally: section of the close_session() function again, where it was originally placed, in order to not unintentionally load the AiiDA profile before it's truly time?

You can import the get_manager function safely without having to have loaded a profile.

[CasperWA]

Last thing: @sphuber, should the close_session decorator function be placed somewhere else in the code?

Any code path that will cause the QueryBuilder to be constructed by a thread should have this. I have the feeling that the get method would be a good place to catch all of these, but I am not sure. We should check with @waychal and @asle85

I didn't mean where to use the decorater, I believe this is already the best place where it is now.
I meant more the actual decorator code, i.e., the close_session() function. Should it be moved to aiida.backends.sqlalchemy instead, together with the rest of the session-related code?

Also, would I then have to move the from aiida.manage.manager import get_manager import line into the finally: section of the close_session() function again, where it was originally placed, in order to not unintentionally load the AiiDA profile before it's truly time?

You can import the get_manager function safely without having to have loaded a profile.

Great! Thanks 👍

[sphuber]

I didn't mean where to use the decorater, I believe this is already the best place where it is now.
I meant more the actual decorator code, i.e., the close_session() function. Should it be moved to aiida.backends.sqlalchemy instead, together with the rest of the session-related code?

Ah I see, no it shouldn't go to aiida.backends.sqlalchemy because it is also need when you run with a django backend. This is why I created the get_session on the Backend class, such that it can be gotten in a backend independent way. Since this is the offical API to dealing with the API, we don't need to add additional code.

[CasperWA]

I didn't mean where to use the decorater, I believe this is already the best place where it is now.
I meant more the actual decorator code, i.e., the close_session() function. Should it be moved to aiida.backends.sqlalchemy instead, together with the rest of the session-related code?

Ah I see, no it shouldn't go to aiida.backends.sqlalchemy because it is also need when you run with a django backend. This is why I created the get_session on the Backend class, such that it can be gotten in a backend independent way. Since this is the offical API to dealing with the API, we don't need to add additional code.

I'll try and ask again then: Do you think the close_session() should be moved somewhere else in the code than where it is now? (And if yes, where?) Because I would not consider it an intrinsic part of the REST API, but rather belonging to AiiDA core and its session handling.

[sphuber]

I'll try and ask again then: Do you think the close_session() should be moved somewhere else in the code than where it is now? (And if yes, where?) Because I would not consider it an intrinsic part of the REST API, but rather belonging to AiiDA core and its session handling.

I have answered this in the PR comments themselves. I suggested defining the decorator on the BaseResource class itself, as it is there where it is used. When I refactored the session handling code, I decided to provide get_session and not just close_session because maybe the caller would need to do have access to the session. Since closing the session is as simple as get_session().close() I didn't feel it was necessary to have an explicit close_session method. Since the context manager you wrote is just one of the many potential ways to deal with the problem at hand, I don't really see why this has to be some standard utility in AiiDA.

[CasperWA]

I have answered this in the PR comments themselves. (...)

LazyMan (aka. me ...) strikes again!

[CasperWA]

With the latest commits I have moved to the pytest way of testing, as well as addressed all other review comments by @sphuber.
Furthermore, an extra commit adds another test - I thought this was relevant to actually validate the result of the "first" test, in the sense that it tests a REST API endpoint/resource without the close_session decorator.
The issue with this test is the default time-out for QueuePool connections (30 s), since it means the test will at minimum always take this long ... I haven't found a way to set the time-out without also getting a setup that doesn't properly reflect how the REST API will actually run.

[CasperWA]

@sphuber, it seems the fixture aiida_localhost doesn't work as expected?

[ltalirz]

This is likely related to the discussion here.

In essence, I think the issue is rooted in interference between the pytest fixtures and the remaining AiiDATestCase implementation.
In the PR I linked we decided to remove the change again since we wanted to release. I wrote

Note, however, that we will likely run into this issue again soon, namely when people start to use pytest fixtures in the test suite more actively.

[sphuber]

Almost there @CasperWA thanks a lot. Just a few more minor things

[CasperWA]

This is likely related to the discussion here.

In essence, I think the issue is rooted in interference between the pytest fixtures and the remaining AiiDATestCase implementation.
In the PR I linked we decided to remove the change again since we wanted to release. I wrote

Note, however, that we will likely run into this issue again soon, namely when people start to use pytest fixtures in the test suite more actively.

Well that's crap.

So what now? Should I undo my conversion to pytest or should I make and store a Computer "manually" or what is the solution, do you think?

[greschd]

So what now?

You could try using the clear_database_before_test fixture. My understanding is that some of the "old-style" tests are trashing the database (specifically here, the user). The "good" solution here would be finding and fixing those tests, but the workaround is to reset the database before running "new-style" tests.

Should I undo my conversion to pytest

Please don't 😅

[sphuber]

But isn't this a bit too hacky?

Not at all. There is a reason that SqlAlchemy provides a way to tweak these parameters. We just haven't exposed the same interface in AiiDA because so far it had not been necessary yet. It is fine to add this for the purposes of the tests. It would of course also open it up to users changing, but the get_scoped_session is not part of the public API so this should serve as enough of a warning.

[CasperWA]

All right!
It seems that this is now (finally) working "reliably" for Py3.8 (edit: run a couple of times now).
However, for Py3.5, shutting down the server, doesn't seem to close the Thread.

Also, I am unsure whether it was the switch to check the capfd instead of capsys or have a small sleep period before shutting down the server that made sure the correct output could be found in the traceback thrown for the test that doesn't apply the fix.
I want to test this by switching back to using capsys, run the tests again, and then I guess I'll have to setup a local Py3.5 environment to adapt the tests to work with older py3 version. NOTE: The fix will work for all supported versions, this is merely an issue with the test for reproducing the original error, i.e., I suggest (based on an offline comment by @ltalirz) to eventually skip the test (perhaps just for Py3.5, or all-together) and merge the PR, since the fix is somewhat critical.

Finally, @sphuber, please review specifically whether my trailing of **kwargs through the code base to end up being passed to sqlalchemy.create_engine fits with your intended pathway. Note, I wanted to make it part of a public function, since this seems to be the best practice, right?

[CasperWA]

Update:
I have set up a local Py3.5 environment and could reproduce the never-ending tests also seen here in the GH Actions CI.
Doing a bit of research I found out that the socket was stuck in an infinite loop trying to read.
I then found this and simply setting a timeout together with the get requests, seemed to work and get the socket out of the loop, and catching the subsequent exception (and passing on it) made the tests run locally.

Now to see if this is also the case with GH Actions ...

[CasperWA]

https://github.com/aiidateam/aiida-core/runs/636184358

🤩 ❤️ 😍 🎉 🥳

[CasperWA]

https://github.com/aiidateam/aiida-core/runs/636184358

🤩 ❤️ 😍 🎉 🥳

However, it does seem that the whole traceback sometimes still doesn't get registered and captured in capfd: https://github.com/CasperWA/aiida_core/runs/636183960

[sphuber]

Thanks a lot @CasperWA . Last thing I would ask is to adapt the docstrings as indicated. Once you have done that, please create a first separate commit that introduces the new kwargs piping and then the second commit with the actual fix and tests on top, then I will merge it.

[CasperWA]

@sphuber Wait a second with merging - I just realized I should use our "new" get_decorators argument when initiating the mock REST API class in one of my tests, instead of also creating a mock resource.
This way the get_decorators argument will be tested also.

[CasperWA]

@sphuber All right, it seems to be all fixed with docs building etc. now. Also, I have updated the test and it seems to work.
It was a good thing as well, I found out the actual way to pass this kwarg to the Resource.
Aaaand out-of-date ....

[sphuber]

Thanks a lot @CasperWA !