Improved janitor

fractalis/cleanup.py

@@ -8,10 +8,19 @@ def janitor():
     """Ideally this is maintained by a systemd service to cleanup redis and the
     file system while Fractalis is running.
     """
-    tmp_dir = app.config['FRACTALIS_TMP_DIR']
+    data_dir = os.path.join(app.config['FRACTALIS_TMP_DIR'], 'data')
+    cached_files = [f for f in os.listdir(data_dir)
+                    if os.path.isfile(os.path.join(data_dir, f))]
     tracked_files = [key.split(':')[1] for key in redis.scan_iter('data:*')]
-    cached_files = [f for f in os.listdir(tmp_dir)
-                    if os.path.isfile(os.path.join(tmp_dir, f))]
+
+    # clean cached files
     for cached_file in cached_files:
         if cached_file not in tracked_files:
-            sync.remove_file(os.path.join(tmp_dir, cached_file))
+            sync.remove_file(os.path.join(data_dir, cached_file))
+
+    # clean tracked files
+    for task_id in tracked_files:
+        path = os.path.join(data_dir, task_id)
+        async_result = celery.AsyncResult(task_id)
+        if async_result.state == 'SUCCESS' and not os.path.exists(path):
+            redis.delete('data:{}'.format(task_id))

fractalis/misc/controller.py

@@ -6,6 +6,8 @@ from typing import Tuple
 
 from flask import Blueprint, jsonify, Response
 
+from fractalis.cleanup import janitor
+
 
 misc_blueprint = Blueprint('misc_blueprint', __name__)
 logger = logging.getLogger(__name__)
@@ -14,4 +16,7 @@ logger = logging.getLogger(__name__)
 @misc_blueprint.route('/version', methods=['GET'])
 def get_version() -> Tuple[Response, int]:
     version = '0.6.0'
+    # this is a good place to launch the janitor because /version is one of the
+    # first requests sent by the front-end on initialization
+    janitor.delay()
     return jsonify({'version': version}), 201

manage.py

-from flask_script import Manager
-
-import fractalis.cleanup
-from fractalis import app
-
-
-manager = Manager(app)
-
-
-@manager.command
-def janitor():
-    fractalis.cleanup.janitor.delay()
-
-
-if __name__ == "__main__":
-    manager.run()

setup.py

+flask
 # -*- coding: utf-8 -*-
 
 from distutils.core import setup

tests/unit/test_cleanup.py

 """This module provides tests for the janitor"""
 
 import os
+import json
 from pathlib import Path
+from shutil import rmtree
 
 from fractalis.cleanup import janitor
-from fractalis import app, redis
+from fractalis import app, redis, celery
 
 
 # noinspection PyMissingOrEmptyDocstring,PyMissingTypeHints
-class TestManage:
+class TestCleanup:
+
+    data_dir = os.path.join(app.config['FRACTALIS_TMP_DIR'], 'data')
+
+    def teardown_method(self, method):
+        redis.flushall()
+        rmtree(self.data_dir)
 
     def test_janitor_removes_untracked_files(self):
-        tmp_dir = app.config['FRACTALIS_TMP_DIR']
-        os.makedirs(tmp_dir, exist_ok=True)
-        Path(os.path.join(tmp_dir, 'abc')).touch()
+        os.makedirs(self.data_dir, exist_ok=True)
+        Path(os.path.join(self.data_dir, 'abc')).touch()
         janitor()
-        assert not os.path.exists(os.path.join(tmp_dir, 'abc'))
+        assert not os.path.exists(os.path.join(self.data_dir, 'abc'))
 
     def test_janitor_does_not_remove_tracked_files(self):
-        tmp_dir = app.config['FRACTALIS_TMP_DIR']
-        os.makedirs(tmp_dir, exist_ok=True)
-        Path(os.path.join(tmp_dir, 'abc')).touch()
+        os.makedirs(self.data_dir, exist_ok=True)
+        Path(os.path.join(self.data_dir, 'abc')).touch()
         redis.set('data:abc', '')
         janitor()
-        assert os.path.exists(os.path.join(tmp_dir, 'abc'))
+        assert os.path.exists(os.path.join(self.data_dir, 'abc'))
+
+    def test_janitor_removes_redis_keys_without_file(self, monkeypatch):
+        os.makedirs(self.data_dir, exist_ok=True)
+        path = os.path.join(self.data_dir, '123')
+        data_state = {
+            'file_path': path
+        }
+        redis.set('data:123', json.dumps(data_state))
+
+        class FakeAsyncResult:
+            def __init__(self, *args, **kwargs):
+                self.state = 'SUCCESS'
+
+            def get(self, *args, **kwargs):
+                pass
+        monkeypatch.setattr(celery, 'AsyncResult', FakeAsyncResult)
+        janitor()
+        assert not redis.exists('data:123')
+
+    def test_janitor_doesnt_remove_redis_keys_with_file(self, monkeypatch):
+        os.makedirs(self.data_dir, exist_ok=True)
+        path = os.path.join(self.data_dir, '123')
+        Path(path).touch()
+        data_state = {
+            'file_path': path
+        }
+        redis.set('data:123', json.dumps(data_state))
+
+        class FakeAsyncResult:
+            def __init__(self, *args, **kwargs):
+                self.state = 'SUCCESS'
+
+            def get(self, *args, **kwargs):
+                pass
+        monkeypatch.setattr(celery, 'AsyncResult', FakeAsyncResult)
+        janitor()
+        assert redis.exists('data:123')
+
+    def test_janitor_doesnt_remove_redis_keys_if_not_finished(
+            self, monkeypatch):
+        os.makedirs(self.data_dir, exist_ok=True)
+        path = os.path.join(self.data_dir, '123')
+        data_state = {
+            'file_path': path
+        }
+        redis.set('data:123', json.dumps(data_state))
+
+        class FakeAsyncResult:
+            def __init__(self, *args, **kwargs):
+                self.state = 'SUBMITTED'
+
+            def get(self, *args, **kwargs):
+                pass
+        monkeypatch.setattr(celery, 'AsyncResult', FakeAsyncResult)
+        janitor()
+        assert redis.exists('data:123')