Cancellable md5sum calculation on thread, Ref. gui#2239

gns3server/compute/base_manager.py

@@ -20,12 +20,17 @@ import os
 import struct
 import stat
 import asyncio
+from asyncio.futures import CancelledError
+
 import aiohttp
 import socket
 import shutil
 import re
 
 import logging
+
+from gns3server.utils.asyncio import cancellable_wait_run_in_executor
+
 log = logging.getLogger(__name__)
 
 from uuid import UUID, uuid4
@@ -557,7 +562,7 @@ class BaseManager:
                     f.write(packet)
             os.chmod(tmp_path, stat.S_IWRITE | stat.S_IREAD | stat.S_IEXEC)
             shutil.move(tmp_path, path)
-            yield from wait_run_in_executor(md5sum, path)
+            yield from cancellable_wait_run_in_executor(md5sum, path)
         except OSError as e:
             raise aiohttp.web.HTTPConflict(text="Could not write image: {} because {}".format(filename, e))
 

gns3server/utils/asyncio/__init__.py

@@ -20,6 +20,9 @@ import functools
 import asyncio
 import sys
 import os
+import threading
+
+from asyncio.futures import CancelledError
 
 
 @asyncio.coroutine
@@ -40,6 +43,25 @@ def wait_run_in_executor(func, *args, **kwargs):
     return future.result()
 
 
+@asyncio.coroutine
+def cancellable_wait_run_in_executor(func, *args, **kwargs):
+    """
+    Run blocking code in a different thread and wait
+    for the result. Support cancellation.
+
+    :param func: Run this function in a different thread
+    :param args: Parameters of the function
+    :param kwargs: Keyword parameters of the function
+    :returns: Return the result of the function
+    """
+    stopped_event = threading.Event()
+    kwargs['stopped_event'] = stopped_event
+    try:
+        yield from wait_run_in_executor(func, *args, **kwargs)
+    except CancelledError:
+        stopped_event.set()
+
+
 @asyncio.coroutine
 def subprocess_check_output(*args, cwd=None, env=None, stderr=False):
     """

gns3server/utils/images.py

@@ -143,11 +143,13 @@ def images_directories(type):
     return [force_unix_path(p) for p in paths if os.path.exists(p)]
 
 
-def md5sum(path):
+def md5sum(path, stopped_event=None):
     """
     Return the md5sum of an image and cache it on disk
 
     :param path: Path to the image
+    :param stopped_event: In case you execute this function on thread and would like to have possibility
+                          to cancel operation pass the `threading.Event`
     :returns: Digest of the image
     """
 
@@ -167,6 +169,9 @@ def md5sum(path):
         m = hashlib.md5()
         with open(path, 'rb') as f:
             while True:
+                if stopped_event is not None and stopped_event.is_set():
+                    log.error("MD5 sum calculation of `{}` has stopped due to cancellation".format(path))
+                    return
                 buf = f.read(128)
                 if not buf:
                     break

tests/utils/test_images.py

@@ -17,6 +17,7 @@
 
 import os
 import sys
+import threading
 from unittest.mock import patch
 
 
@@ -57,6 +58,18 @@ def test_md5sum(tmpdir):
         assert f.read() == '5d41402abc4b2a76b9719d911017c592'
 
 
+def test_md5sum_stopped_event(tmpdir):
+    fake_img = str(tmpdir / 'hello_stopped')
+    with open(fake_img, 'w+') as f:
+        f.write('hello')
+
+    event = threading.Event()
+    event.set()
+
+    assert md5sum(fake_img, stopped_event=event) is None
+    assert not os.path.exists(str(tmpdir / 'hello_stopped.md5sum'))
+
+
 def test_md5sum_existing_digest(tmpdir):
     fake_img = str(tmpdir / 'hello')