ENH: spatial: ensure thread-safety for KDTree

[andfoy]

Reference issue

What does this implement/fix?

This PR adds a quick check to ensure that concurrent calls to KDTree do not cause any unexpected errors. This is specially handy in the context of No-GIL CPython.

Additional information

In principle, concurrent access to KDTree is allowed out-of-the-box, since there are no methods that can mutate the state of the tree besides the constructor.

[sturlamolden]

I am a bit uncertain if property .tree is thread-safe or not. It might be sufficiently serialized by the GIL though. It can mutate the inner state.

[sturlamolden]

In context of No-GIL CPython, there are several things that are not thread-safe. For example:

• Calling query methods or pickling before __init__ has completed.
• Property .tree
• Serialization (depickling), I think
• Rouge thread making call to __init__ (even extra calls cause problems, even now)

[andfoy]

Thanks for the pointers @sturlamolden. I have the following questions and comments:

1. Just to clarify, should we be locking the access to .tree before it gets initialized? If I understood correctly, the potential race condition is that T1 creates a KDTree instance with many points so the construction is deferred, then other thread T2 is created with a reference to the same tree as T1 and tries to call any method while the construction takes place in T1? I don't know if this is possible under No-GIL Python, maybe @ngoldbaum has more ideas about this kind of scenario?
2. Is there a mechanism by which the tree state gets modified during a query call? I took a peek at some of the function definitions, and while concurrent calls would be sharing access to the tree values, they are not being modified on any capacity, unless there is some detail that we are overlooking here?
3. Serialization is a good point regarding concurrency, however, since there's not a method that is able to add new points or modify the tree (assuming that 2 holds), there shouldn't be any issue regarding state storage. The only problem would be multiple threads trying to write to the same file, which should be (in my opinion) responsibility of the end user

[sturlamolden]

The issue with .tree is that it is an expensive structure to create, and was added for the purpose of introspection and testing, i.e. to verify that a kd-tree is created as expected or for education (show students how it partitions space and data). Since it is normally not needed, its creation is deferred to when someone actually asks for it.

[sturlamolden]

The more serious issue is that __init__ is serialized by the GIL. Also it needs to complete before query functions are called. If we do not have the GIL, we need to serialize it in such a way that is has completed before queries are allowed. However, we should not serialize query functions, only have them wait for the __init__ to finish, so I guess we could set a threading.Event upon finishing __init__ and have query functions wait for it. Normally it will be set before they are called. We must also use a threading.Lock within __init__. It can also be used to serialize access to .tree.

[andfoy]

Now I understand everything, so, basically we should lock on tree, and also add an atomic flag that signals if the tree structure is ready to use.

[sturlamolden]

__init__ is a callable Python function. So it must be serialized by a lock and it must set an atomic flag upon completion.

.tree must wait for the flag and also be internally serialized by an rlock.

Query functions must wait for the atomic flag set by __init__.

[andfoy]

Thanks for the helpful discussion @sturlamolden!

[sturlamolden]

No problem. When I wrote the code I depended on the GIL for serialization. For example it ensures that __init__ has completed before queries can be called.

[sturlamolden]

Actually it is already broken as __init__ now calls Python functions instead of the NumPy C API directly, so context switch on the GIL is possible while running __init__. It is hard to create though, but it could result in a race conditon or segfault. Which means it needs the serialization described for No-GIL CPython.

[andfoy]

Great! Then this PR will fix the concurrency issues in KDTree, let me add the changes and update the title

[sturlamolden]

Note that my comments refer to cKDTree. KDTree is a Python wrapper on top of it so it will be necessary to fix both.

[ngoldbaum]

I don't know if this is possible under No-GIL Python, maybe @ngoldbaum has more ideas about this kind of scenario?

Yes, in free-threaded python when the GIL is disabled then something needs to explicitly lock access to prevent data races. If the data race happens in python-level code, then a threading.Lock or other python synchronization primitive is needed.

[sturlamolden]

We already make use of the threading module in cKDTree´s Cython code so it is not just for Python code. The problems I mentioned happens in Cython level code, which should have been protected by the GIL.

[ngoldbaum]

Right, C-level code will need to use C-level locking primitives.

[sturlamolden]

Use the primitives in threading like the rest of cKDTree, not cpython.pythread.

[sturlamolden]

Also you cannot use a bool as an atomic flag. It must be an atomic variable, which a bool is not. Either you use platform-dependet C++ code for atomic read and write on a std::atomic<bool>, or you use threading.Event which takes care of these details. Basically to the latter, because we do not need that kind of microoptimizations here. Also if you use atomic read and write, you must provide your own spinlocks in the query functions, which I am certain we do not want to mess with. Just use threading.Event.

[sturlamolden]

I am roughly -inf on the changes to _ckdtree.pyx.

Just use the threading module. The important optmisations are in the C++ code.

[andfoy]

@sturlamolden After an internal discussion with the Python free-threaded team, it has been decided that all the free-threaded work assumes that __init__ is effectively serialized and no reference leaks cannot take place (unless the implementation is changed manually) before __init__ has stopped execution, thus the proposed race condition related to initialization doesn't hold up.

Given this, the only potential race condition would be accessing the tree property concurrently.

cc @ngoldbaum @colesbury