Use Futures to deduplicate incoming requests (#37)

Fixes #23 

As of #19, we cache (most) scan results in memory to avoid having to
download and scan the same file twice (with the caveat that we don't
cache the contents of big files, but it still saves scanning them).
However, it does not prevent duplicate work if multiple requests for the
same file happen simultaneously, since we only update this cache when a
scan completes. This is a legitimate concern, since we can easily
imagine a situation where a user posts a large file into a room and
multiple users try to download it at the same time.

As of this change,

* when a request for a given file with given encryption and thumbnailing
parameters comes in, and we're not already running a download/scan for
it (with those parameters), we create a `Future` that we store into a
new `_current_scans` cache, and run the scan. Once the scan completes,
we resolve the `Future` with the result, and remove it from the cache.
* when a request for the same file comes in _before_ the previous
request completed, we don't download and/or scan the file - instead we
look up the matching `Future` from the `_current_scans` cache and
`await` it. We then return its result.

Author: babolivier

matrix_content_scanner/scanner/scanner.py

@@ -11,10 +11,12 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+import asyncio
 import hashlib
 import logging
 import os
 import subprocess
+from asyncio import Future
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
@@ -85,6 +87,10 @@ def __init__(self, mcs: "MatrixContentScanner"):
         # for unencrypted files).
         self._allowed_mimetypes = mcs.config.scan.allowed_mimetypes
 
+        # Cache of futures for files that are currently scanning and downloading, so that
+        # concurrent requests don't cause a file to be downloaded and scanned twice.
+        self._current_scans: Dict[str, Future[MediaDescription]] = {}
+
     async def scan_file(
         self,
         media_path: str,
@@ -100,6 +106,9 @@ async def scan_file(
         downloading the file again (unless we purposefully did not cache the file's
         content to save up on memory).
 
+        If a file is currently already being downloaded or scanned as a result of another
+        request, don't download it again and use the result from the first request.
+
         Args:
             media_path: The `server_name/media_id` path for the media.
             metadata: The metadata attached to the file (e.g. decryption key), or None if
@@ -115,9 +124,72 @@ async def scan_file(
             FileDirtyError if the result of the scan said that the file is dirty, or if
                 the media path is malformed.
         """
-        # Compute the cache key for the media.
+        # Compute the key to use when caching, both in the current scans cache and in the
+        # results cache.
         cache_key = self._get_cache_key_for_file(media_path, metadata, thumbnail_params)
+        if cache_key not in self._current_scans:
+            # Create a future in the context of the current event loop.
+            loop = asyncio.get_event_loop()
+            f = loop.create_future()
+            # Register the future in the current scans cache so that subsequent queries
+            # can use it.
+            self._current_scans[cache_key] = f
+            # Try to download and scan the file.
+            try:
+                res = await self._scan_file(
+                    cache_key, media_path, metadata, thumbnail_params
+                )
+                # Set the future's result, and mark it as done.
+                f.set_result(res)
+                # Return the result.
+                return res
+            except Exception as e:
+                # If there's an exception, catch it, pass it on to the future, and raise
+                # it.
+                f.set_exception(e)
+                # We retrieve the exception from the future, because if we don't and no
+                # other request is awaiting on the future, asyncio complains about "Future
+                # exception was never retrieved".
+                f.exception()
+                raise
+            finally:
+                # Remove the future from the cache.
+                del self._current_scans[cache_key]
+
+        return await self._current_scans[cache_key]
+
+    async def _scan_file(
+        self,
+        cache_key: str,
+        media_path: str,
+        metadata: Optional[JsonDict] = None,
+        thumbnail_params: Optional[MultiDictProxy[str]] = None,
+    ) -> MediaDescription:
+        """Download and scan the given media.
+
+        Unless the scan fails with one of the codes listed in `do_not_cache_exit_codes`,
+        also cache the result.
+
+        If the file already has an entry in the result cache, return this value without
+        downloading the file again (unless we purposefully did not cache the file's
+        content to save up on memory).
 
+        Args:
+            cache_key: The key to use to cache the result of the scan in the result cache.
+            media_path: The `server_name/media_id` path for the media.
+            metadata: The metadata attached to the file (e.g. decryption key), or None if
+                the file isn't encrypted.
+            thumbnail_params: If present, then we want to request and scan a thumbnail
+                generated with the provided parameters instead of the full media.
+
+        Returns:
+            A description of the media.
+
+        Raises:
+            ContentScannerRestError if the file could not be downloaded.
+            FileDirtyError if the result of the scan said that the file is dirty, or if
+                the media path is malformed.
+        """
         # The media to scan.
         media: Optional[MediaDescription] = None
 

tests/scanner/test_scanner.py

@@ -11,8 +11,9 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+import asyncio
 import copy
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import Mock
 
@@ -303,6 +304,37 @@ async def test_invalid_media_path(self) -> None:
         with self.assertRaises(FileDirtyError):
             await self.scanner.scan_file(MEDIA_PATH + "/baz")
 
+    async def test_deduplicate_scans(self) -> None:
+        """Tests that if two scan requests come in for the same file and with the same
+        parameter, only one download/scan happens.
+        """
+
+        # Change the Mock's side effect to introduce some delay, to simulate a long
+        # download time. We sleep asynchronously to allow additional scans requests to be
+        # processed.
+        async def _scan_file(*args: Any) -> MediaDescription:
+            await asyncio.sleep(0.2)
+
+            return self.downloader_res
+
+        scan_mock = Mock(side_effect=_scan_file)
+        self.scanner._scan_file = scan_mock  # type: ignore[assignment]
+
+        # Request two scans of the same file at the same time.
+        results = await asyncio.gather(
+            asyncio.create_task(self.scanner.scan_file(MEDIA_PATH)),
+            asyncio.create_task(self.scanner.scan_file(MEDIA_PATH)),
+        )
+
+        # Check that the scanner has only been called once, meaning that the second
+        # call did not trigger a scan.
+        scan_mock.assert_called_once()
+
+        # Check that we got two results, and that we actually got the correct media
+        # description in the second scan.
+        self.assertEqual(len(results), 2, results)
+        self.assertEqual(results[0].content, results[1].content, results)
+
     def _setup_encrypted(self) -> None:
         """Sets up class properties to make the downloader return an encrypted file
         instead of a plain text one.