Improve caching of thumbnail information.

nion/swift/Thumbnails.py

@@ -27,7 +27,6 @@
 from nion.utils import ReferenceCounting
 
 _NDArray = numpy.typing.NDArray[typing.Any]
-_ThumbnailSourceWeakRef = typing.Callable[[], typing.Optional["ThumbnailSource"]]  # Python 3.9+
 
 
 class ThumbnailSource:
@@ -47,8 +46,15 @@ def __init__(self, ui: UserInterface.UserInterface, display_item: DisplayItem.Di
         self.__display_item = display_item
         self.__recompute_lock = threading.RLock()
         self.__recompute_future: typing.Optional[concurrent.futures.Future[typing.Any]] = None
+        # the cache is used to store the thumbnail data persistently. for performance, it is ideal
+        # to minimize calling it and instead use the cached value in this class.
         self.__cache = self.__display_item._display_cache
         self.__cache_property_name = "thumbnail_data"
+        self.__cache_properties_known = False
+        self.__cache_thumbnail_data: typing.Optional[_NDArray] = None
+        self.__cache_is_dirty = False
+
+        self.thumbnail_dirty_event = Event.Event()  # for testing
 
         self.__display_changed_event_listener = display_item.display_changed_event.listen(ReferenceCounting.weak_partial(ThumbnailSource.__thumbnail_changed, self))
         self.__graphics_changed_event_listener = display_item.graphics_changed_event.listen(ReferenceCounting.weak_partial(ThumbnailSource.__graphics_changed, self))
@@ -58,8 +64,18 @@ def __init__(self, ui: UserInterface.UserInterface, display_item: DisplayItem.Di
 
         self.__display_will_close_listener = display_item.display_item_will_close_event.listen(ReferenceCounting.weak_partial(ThumbnailSource.__display_item_will_close, self))
 
+    def __read_cache_properties(self) -> None:
+        if not self.__cache_properties_known:
+            self.__cache_thumbnail_data = typing.cast(typing.Optional[_NDArray], self.__cache.get_cached_value(self.__display_item, self.__cache_property_name)) if self.__display_item else None
+            self.__cache_is_dirty = self.__cache.is_cached_value_dirty(self.__display_item, self.__cache_property_name) if self.__display_item else False
+            self.__cache_properties_known = True
+            self.thumbnail_updated_event.fire()
+
     def __thumbnail_changed(self) -> None:
         self.__cache.set_cached_value_dirty(self.__display_item, self.__cache_property_name)
+        self.thumbnail_dirty_event.fire()
+        self.__cache_is_dirty = True
+        self.__cache_properties_known = True
         self.__recompute_on_thread()
 
     def __recompute_on_thread(self) -> None:
@@ -88,13 +104,10 @@ def __display_item_will_close(self) -> None:
 
     @property
     def thumbnail_data(self) -> typing.Optional[_NDArray]:
-        """Return the cached data for this processor.
-
-        This method is thread safe and always returns quickly, using the cached data.
-        """
-        return typing.cast(typing.Optional[_NDArray], self.__cache.get_cached_value(self.__display_item, self.__cache_property_name)) if self.__display_item else None
+        return self.__cache_thumbnail_data
 
     def __recompute_data_if_needed(self) -> None:
+        self.__read_cache_properties()
         if self._is_thumbnail_dirty:
             self.recompute_data()
 
@@ -126,47 +139,54 @@ def recompute_data(self) -> None:
                 raise
             if calculated_data is None:
                 calculated_data = numpy.zeros((self.height, self.width), dtype=numpy.uint32)
+            self.__cache_thumbnail_data = calculated_data
+            self.__cache_is_dirty = False
+            self.__cache_properties_known = True
             self.__cache.set_cached_value(self.__display_item, self.__cache_property_name, calculated_data)
         self.thumbnail_updated_event.fire()
 
     @property
     def _is_thumbnail_dirty(self) -> bool:
-        assert self.__display_item
-        return self.__cache.is_cached_value_dirty(self.__display_item, self.__cache_property_name)
+        return self.__cache_is_dirty
+
+    @property
+    def _is_valid(self) -> bool:
+        return self.__display_item is not None
 
 
 class ThumbnailManager(metaclass=Utility.Singleton):
     """Manages thumbnail sources for displays."""
 
     def __init__(self) -> None:
-        self.__thumbnail_sources: typing.Dict[uuid.UUID, _ThumbnailSourceWeakRef] = dict()
+        self.__thumbnail_sources: typing.Dict[uuid.UUID, ThumbnailSource] = dict()
         self.__lock = threading.RLock()
 
     def reset(self) -> None:
         with self.__lock:
             self.__thumbnail_sources.clear()
 
+    def clean(self) -> None:
+        with self.__lock:
+            for uuid, thumbnail_source in list(self.__thumbnail_sources.items()):
+                if not thumbnail_source._is_valid:
+                    del self.__thumbnail_sources[uuid]
+
     def thumbnail_source_for_display_item(self, ui: UserInterface.UserInterface, display_item: DisplayItem.DisplayItem) -> ThumbnailSource:
         """Returned ThumbnailSource must be closed."""
         with self.__lock:
-            thumbnail_source_ref = self.__thumbnail_sources.get(display_item.uuid)
-            thumbnail_source = thumbnail_source_ref() if thumbnail_source_ref else None
+            self.clean()
+            thumbnail_source = self.__thumbnail_sources.get(display_item.uuid)
             if not thumbnail_source:
                 thumbnail_source = ThumbnailSource(ui, display_item)
-                self.__thumbnail_sources[display_item.uuid] = weakref.ref(thumbnail_source)
-
-                def pop_thumbnail_source(uuid: uuid.UUID) -> None:
-                    self.__thumbnail_sources.pop(uuid, None)
-
-                weakref.finalize(thumbnail_source, pop_thumbnail_source, display_item.uuid)
+                self.__thumbnail_sources[display_item.uuid] = thumbnail_source
             else:
                 assert thumbnail_source._ui == ui
             return thumbnail_source
 
     def thumbnail_data_for_display_item(self, display_item: typing.Optional[DisplayItem.DisplayItem]) -> typing.Optional[_NDArray]:
         with self.__lock:
-            thumbnail_source_ref = self.__thumbnail_sources.get(display_item.uuid) if display_item else None
-            thumbnail_source = thumbnail_source_ref() if thumbnail_source_ref else None
+            self.clean()
+            thumbnail_source = self.__thumbnail_sources.get(display_item.uuid) if display_item else None
             if thumbnail_source:
                 return thumbnail_source.thumbnail_data
             return None

nion/swift/test/DataItem_test.py

@@ -295,7 +295,7 @@ def test_clear_thumbnail_when_data_item_changed(self):
             self.assertFalse(display_item._display_cache.is_cached_value_dirty(display_item, "thumbnail_data"))
             with display_item.data_item.data_ref() as data_ref:
                 data_ref.data = numpy.zeros((8, 8), numpy.uint32)
-            self.assertTrue(display_item._display_cache.is_cached_value_dirty(display_item, "thumbnail_data"))
+            self.assertTrue(thumbnail_source._is_thumbnail_dirty)
             thumbnail_source = None
 
     def test_thumbnail_2d_handles_small_dimension_without_producing_invalid_thumbnail(self):
@@ -385,15 +385,15 @@ def test_thumbnail_marked_dirty_when_source_data_changed(self):
             thumbnail_source = Thumbnails.ThumbnailManager().thumbnail_source_for_display_item(self.app.ui, inverted_display_item)
             thumbnail_source.recompute_data()
             thumbnail_source.thumbnail_data
-            thumbnail_source = None
             # here the data should be computed and the thumbnail should not be dirty
             self.assertFalse(inverted_display_item._display_cache.is_cached_value_dirty(inverted_display_item, "thumbnail_data"))
             # now the source data changes and the inverted data needs computing.
             # the thumbnail should also be dirty.
             with data_item.data_ref() as data_ref:
                 data_ref.data = data_ref.data + 1.0
             document_model.recompute_all()
-            self.assertTrue(inverted_display_item._display_cache.is_cached_value_dirty(inverted_display_item, "thumbnail_data"))
+            self.assertTrue(thumbnail_source._is_thumbnail_dirty)
+            thumbnail_source = None
 
     def test_thumbnail_widget_when_data_item_has_no_associated_display_item(self):
         with TestContext.create_memory_context() as test_context:

nion/swift/test/Thumbnails_test.py

@@ -61,9 +61,18 @@ def test_thumbnail_marked_dirty_when_display_layers_change(self):
             self.assertFalse(display_item._display_cache.is_cached_value_dirty(display_item, "thumbnail_data"))
             # now the source data changes and the inverted data needs computing.
             # the thumbnail should also be dirty.
+            thumbnail_dirty = False
+
+            def handle_thumbnail_dirty() -> None:
+                nonlocal thumbnail_dirty
+                thumbnail_dirty = True
+
+            listener = thumbnail_source.thumbnail_dirty_event.listen(handle_thumbnail_dirty)
             display_item._set_display_layer_property(0, "fill_color", "teal")
             document_model.recompute_all()
-            self.assertTrue(display_item._display_cache.is_cached_value_dirty(display_item, "thumbnail_data"))
+            # this is a race condition. the thumbnail thread may clear the dirty flag before we check it.
+            # so use the event instead.
+            self.assertTrue(thumbnail_dirty)
 
 
 if __name__ == '__main__':