fix: executor shutdown race, dropped kwargs, chunk revalidation, unknown-class cache

src/uproot/model.py

@@ -1279,14 +1279,17 @@ def new_class(cls, file, version):
             return versioned_cls
 
         else:
-            unknown_cls = uproot.unknown_classes.get(classname)
+            # key on the encoded name, which carries the version and so keeps
+            # distinct versions (and the versionless UnknownClass) apart
+            encoded_classname = classname_encode(classname, version, unknown=True)
+            unknown_cls = uproot.unknown_classes.get(encoded_classname)
             if unknown_cls is None:
                 unknown_cls = uproot._util.new_class(
-                    classname_encode(classname, version, unknown=True),
+                    encoded_classname,
                     (UnknownClassVersion,),
                     {},
                 )
-                uproot.unknown_classes[classname] = unknown_cls
+                uproot.unknown_classes[encoded_classname] = unknown_cls
             return unknown_cls
 
     @classmethod

src/uproot/reading.py

@@ -1131,14 +1131,19 @@ def class_named(self, classname, version=None):
 
         if cls is None:
             if len(streamers) == 0:
-                unknown_cls = uproot.unknown_classes.get(classname)
+                # key on the encoded name, which distinguishes this versionless
+                # UnknownClass from the per-version UnknownClassVersion classes
+                encoded_classname = uproot.model.classname_encode(
+                    classname, unknown=True
+                )
+                unknown_cls = uproot.unknown_classes.get(encoded_classname)
                 if unknown_cls is None:
                     unknown_cls = uproot._util.new_class(
-                        uproot.model.classname_encode(classname, unknown=True),
+                        encoded_classname,
                         (uproot.model.UnknownClass,),
                         {},
                     )
-                    uproot.unknown_classes[classname] = unknown_cls
+                    uproot.unknown_classes[encoded_classname] = unknown_cls
                 return unknown_cls
 
             else:
@@ -1159,16 +1164,17 @@ def class_named(self, classname, version=None):
                 elif version == "min" and len(cls.known_versions) != 0:
                     version = min(cls.known_versions)
                 else:
-                    unknown_cls = uproot.unknown_classes.get(classname)
+                    encoded_classname = uproot.model.classname_encode(
+                        classname, version, unknown=True
+                    )
+                    unknown_cls = uproot.unknown_classes.get(encoded_classname)
                     if unknown_cls is None:
                         unknown_cls = uproot._util.new_class(
-                            uproot.model.classname_encode(
-                                classname, version, unknown=True
-                            ),
+                            encoded_classname,
                             (uproot.model.UnknownClassVersion,),
                             {},
                         )
-                        uproot.unknown_classes[classname] = unknown_cls
+                        uproot.unknown_classes[encoded_classname] = unknown_cls
                     return unknown_cls
 
             versioned_cls = cls.class_of_version(version)

src/uproot/source/chunk.py

@@ -385,21 +385,25 @@ def wait(self, insist: bool = True):
         :ref:`uproot.source.chunk.Chunk.future` completes).
         """
         if self._raw_data is None:
-            self._raw_data = numpy.frombuffer(self._future.result(), dtype=self._dtype)
+            raw_data = numpy.frombuffer(self._future.result(), dtype=self._dtype)
             if insist is True:
-                requirement = len(self._raw_data) == self._stop - self._start
+                requirement = len(raw_data) == self._stop - self._start
             elif isinstance(insist, numbers.Integral):
-                requirement = len(self._raw_data) >= insist - self._start
+                requirement = len(raw_data) >= insist - self._start
             elif insist is False:
                 requirement = True
             else:
                 raise TypeError(f"""insist must be a bool or an int, not {insist!r}
 for file path {self._source.file_path}""")
 
+            # only publish the data once it has passed validation: assigning it
+            # first would make every later access skip this check and hand back
+            # the short buffer instead of raising again
             if not requirement:
                 raise OSError(f"""expected Chunk of length {self._stop - self._start},
-received {len(self._raw_data)} bytes from {type(self._source).__name__}
+received {len(raw_data)} bytes from {type(self._source).__name__}
 for file path {self._source.file_path}""")
+            self._raw_data = raw_data
             self._future = None
 
     @property

src/uproot/source/futures.py

@@ -109,9 +109,10 @@ class Future:
     The :doc:`uproot.source.futures.ResourceFuture` extends this class.
     """
 
-    def __init__(self, task, args):
+    def __init__(self, task, args, kwargs=None):
         self._task = task
         self._args = args
+        self._kwargs = {} if kwargs is None else kwargs
         self._finished = threading.Event()
         self._result = None
         self._excinfo = None
@@ -135,12 +136,13 @@ def _run(self):
         try:
             if self._task is None:
                 raise RuntimeError("cannot run Future twice")
-            self._result = self._task(*self._args)
+            self._result = self._task(*self._args, **self._kwargs)
         except Exception as err:
             self._excinfo = err
         self._finished.set()
         self._task = None
         self._args = ()
+        self._kwargs = {}
 
 
 class Worker(threading.Thread):
@@ -201,6 +203,7 @@ class ThreadPoolExecutor(Executor):
 
     def __init__(self, max_workers: int | None = None):
         self._max_workers = max_workers or os.cpu_count()
+        self._closed = False
 
         self._work_queue = queue.Queue()
         self._workers = []
@@ -235,21 +238,35 @@ def workers(self) -> list[Worker]:
 
     def submit(self, task, /, *args, **kwargs):
         """
-        Pass the ``task`` and ``args`` onto the workers'
+        Pass the ``task``, ``args``, and ``kwargs`` onto the workers'
         :ref:`uproot.source.futures.Worker.work_queue` as a
         :doc:`uproot.source.futures.Future` so that it will be executed when
         one is available.
         """
-        future = Future(task, args)
+        if self.closed:
+            raise OSError("executor is closed")
+        future = Future(task, args, kwargs)
         self._work_queue.put(future)
         return future
 
+    @property
+    def closed(self) -> bool:
+        """
+        True if :ref:`uproot.source.futures.ThreadPoolExecutor.shutdown` has
+        been started; False otherwise.
+        """
+        return self._closed
+
     def shutdown(self, wait: bool = True):
         """
         Stop every :doc:`uproot.source.futures.Worker` by putting one None per
         worker on the :ref:`uproot.source.futures.Worker.work_queue` and
         joining each worker thread.
         """
+        # mark this executor closed *before* queuing the sentinels: a submit
+        # accepted after a sentinel is queued would sit behind it in the queue,
+        # no worker would ever reach it, and its Future would block forever
+        self._closed = True
         for _ in self._workers:
             self._work_queue.put(None)
         for worker in self._workers:
@@ -390,9 +407,10 @@ def close(self):
     @property
     def closed(self) -> bool:
         """
-        True if the :doc:`uproot.source.futures.ResourceWorker` threads have
-        been stopped and their
-        :ref:`uproot.source.futures.ResourceWorker.resource` freed.
+        True once the :doc:`uproot.source.futures.ResourceWorker` threads have
+        started stopping; their
+        :ref:`uproot.source.futures.ResourceWorker.resource` is freed by the
+        time teardown returns.
         """
         return self._closed
 
@@ -401,10 +419,11 @@ def __enter__(self):
             worker.resource.__enter__()
 
     def __exit__(self, exception_type, exception_value, traceback):
+        # shutdown sets self._closed before it queues the sentinels, so a
+        # concurrent submit is rejected instead of being orphaned behind one
         self.shutdown()
         for worker in self._workers:
             worker.resource.__exit__(exception_type, exception_value, traceback)
-        self._closed = True
 
 
 ##################### use-case 4: resources for I/O with trivial executor

tests/test_1688_source_hygiene.py

@@ -0,0 +1,182 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE
+
+"""Regression tests for issue #1688: executor, chunk, and unknown-class fixes.
+
+* a submit racing with shutdown was accepted and then orphaned behind a sentinel
+* ``ThreadPoolExecutor.submit`` accepted keyword arguments and discarded them
+* a Chunk that failed its length check returned the short buffer on re-access
+* unknown model classes were cached by classname alone, ignoring the version
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+
+import pytest
+
+import uproot
+import uproot._util
+import uproot.model
+import uproot.source.chunk
+import uproot.source.futures
+
+
+class DummyResource:
+    file_path = "dummy"
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+def test_submit_during_shutdown_is_rejected_not_orphaned():
+    executor = uproot.source.futures.ResourceThreadPoolExecutor([DummyResource()])
+    executor.__enter__()
+
+    # occupy the only worker so that shutdown's join() blocks for a while
+    executor.submit(
+        uproot.source.futures.ResourceFuture(lambda resource: time.sleep(1))
+    )
+    time.sleep(0.1)
+
+    closer = threading.Thread(target=lambda: executor.__exit__(None, None, None))
+    closer.start()
+    try:
+        time.sleep(0.1)  # the sentinel is queued and join() is blocking
+        assert executor.closed
+        late = uproot.source.futures.ResourceFuture(lambda resource: 42)
+        with pytest.raises(OSError):
+            executor.submit(late)
+    finally:
+        closer.join()
+
+
+def test_thread_pool_executor_rejects_submit_after_shutdown():
+    executor = uproot.source.futures.ThreadPoolExecutor(1)
+    assert not executor.closed
+    executor.shutdown()
+    assert executor.closed
+    with pytest.raises(OSError):
+        executor.submit(lambda: 1)
+
+
+def test_thread_pool_executor_forwards_kwargs():
+    executor = uproot.source.futures.ThreadPoolExecutor(1)
+    try:
+        future = executor.submit(lambda a, x=None: (a, x), 1, x=3)
+        assert future.result(timeout=10) == (1, 3)
+
+        future = executor.submit(lambda *, only_kw: only_kw, only_kw="value")
+        assert future.result(timeout=10) == "value"
+    finally:
+        executor.shutdown()
+
+
+def test_trivial_executor_forwards_kwargs():
+    executor = uproot.source.futures.TrivialExecutor()
+    future = executor.submit(lambda a, x=None: (a, x), 1, x=3)
+    assert future.result() == (1, 3)
+
+
+class _FakeSource:
+    file_path = "fake"
+
+
+def test_short_chunk_raises_every_time():
+    chunk = uproot.source.chunk.Chunk(
+        _FakeSource(), 0, 5, uproot.source.futures.TrivialFuture(b"abc")
+    )
+    for _ in range(3):
+        with pytest.raises(OSError, match="expected Chunk of length 5"):
+            chunk.raw_data
+
+
+def test_chunk_of_expected_length_still_works():
+    chunk = uproot.source.chunk.Chunk(
+        _FakeSource(), 0, 5, uproot.source.futures.TrivialFuture(b"abcde")
+    )
+    assert chunk.raw_data.tobytes() == b"abcde"
+    assert chunk.raw_data.tobytes() == b"abcde"
+
+
+def test_chunk_insist_false_does_not_raise():
+    chunk = uproot.source.chunk.Chunk(
+        _FakeSource(), 0, 5, uproot.source.futures.TrivialFuture(b"abc")
+    )
+    chunk.wait(insist=False)
+    assert chunk.raw_data.tobytes() == b"abc"
+
+
+class _StreamerlessFile:
+    custom_classes = None
+    file_path = "fake"
+    streamers = {}
+
+    def streamer_named(self, classname, version):
+        return None
+
+
+def _make_dispatch():
+    return uproot._util.new_class(
+        uproot.model.classname_encode("MyClass"),
+        (uproot.model.DispatchByVersion,),
+        {"known_versions": {}},
+    )
+
+
+def test_unknown_class_versions_are_distinct():
+    uproot.model.reset_classes()
+    try:
+        dispatch = _make_dispatch()
+
+        v1 = dispatch.new_class(_StreamerlessFile(), 1)
+        v2 = dispatch.new_class(_StreamerlessFile(), 2)
+
+        assert v1.__name__ == "Unknown_MyClass_v1"
+        assert v2.__name__ == "Unknown_MyClass_v2"
+        assert v1 is not v2
+        # asking again returns the cached class rather than building a new one
+        assert dispatch.new_class(_StreamerlessFile(), 1) is v1
+        assert set(uproot.unknown_classes) == {
+            "Unknown_MyClass_v1",
+            "Unknown_MyClass_v2",
+        }
+    finally:
+        uproot.model.reset_classes()
+
+
+class _NoStreamerFile:
+    """Enough of a ReadOnlyFile for class_named to reach the unknown-class paths."""
+
+    def __init__(self, custom_classes=None):
+        self._custom_classes = custom_classes
+
+    def streamers_named(self, classname):
+        return []
+
+    def streamer_named(self, classname, version):
+        return None
+
+
+def test_versioned_and_versionless_unknown_classes_coexist():
+    uproot.model.reset_classes()
+    try:
+        dispatch = _make_dispatch()
+        # a file that knows the class but not the requested version
+        versioned = uproot.reading.ReadOnlyFile.class_named(
+            _NoStreamerFile({"MyClass": dispatch}), "MyClass", "max"
+        )
+        # a file with no streamers at all for the same class
+        versionless = uproot.reading.ReadOnlyFile.class_named(
+            _NoStreamerFile(), "MyClass"
+        )
+
+        assert issubclass(versioned, uproot.model.UnknownClassVersion)
+        assert issubclass(versionless, uproot.model.UnknownClass)
+        assert versioned is not versionless
+        assert versionless.__name__ == "Unknown_MyClass"
+    finally:
+        uproot.model.reset_classes()