Merge pull request #106 from Erazor2/queuePersistence

Queue persistence

2 files changed, 72 insertions, 22 deletions

diff --git a/app/main.py b/app/main.py
index 8b2a57f..36c0e94 100644
--- a/app/main.py
+++ b/app/main.py
@@ -16,6 +16,7 @@ class Config:
     _DEFAULTS = {

         'DOWNLOAD_DIR': '.',

         'AUDIO_DOWNLOAD_DIR': '%%DOWNLOAD_DIR',

+        'STATE_DIR': '.',

         'URL_PREFIX': '',

         'OUTPUT_TEMPLATE': '%(title)s.%(ext)s',

         'YTDL_OPTIONS': '{}',

@@ -68,6 +69,7 @@ class Notifier(DownloadQueueNotifier):
         await sio.emit('cleared', serializer.encode(id))

 

 dqueue = DownloadQueue(config, Notifier())

+app.on_startup.append(lambda app: dqueue.initialize())

 

 @routes.post(config.URL_PREFIX + 'add')

 async def add(request):

diff --git a/app/ytdl.py b/app/ytdl.py
index a5e60b7..deb745c 100644
--- a/app/ytdl.py
+++ b/app/ytdl.py
@@ -1,6 +1,8 @@
 import os

 import yt_dlp

 from collections import OrderedDict

+import shelve

+import time

 import asyncio

 import multiprocessing

 import logging

@@ -31,6 +33,7 @@ class DownloadInfo:
         self.format = format

         self.status = self.msg = self.percent = self.speed = self.eta = None

         self.filename = None

+        self.timestamp = time.time_ns()

 

 class Download:

     manager = None

@@ -131,19 +134,65 @@ class Download:
             self.info.eta = status.get('eta')

             await self.notifier.updated(self.info)

 

+class PersistentQueue:

+    def __init__(self, path):

+        self.path = path

+        with shelve.open(path, 'c'):

+            pass

+        self.dict = OrderedDict()

+    

+    def load(self):

+        for k, v in self.saved_items():

+            self.dict[k] = Download(None, None, None, None, {}, v)

+

+    def exists(self, key):

+        return key in self.dict

+    

+    def get(self, key):

+        return self.dict[key]

+    

+    def items(self):

+        return self.dict.items()

+

+    def saved_items(self):

+        with shelve.open(self.path, 'r') as shelf:

+            return sorted(shelf.items(), key=lambda item: item[1].timestamp)

+

+    def put(self, value):

+        key = value.info.id

+        self.dict[key] = value

+        with shelve.open(self.path, 'w') as shelf:

+            shelf[key] = value.info

+    

+    def delete(self, key):

+        del self.dict[key]

+        with shelve.open(self.path, 'w') as shelf:

+            shelf.pop(key)

+

+    def next(self):

+        k, v = next(iter(self.dict.items()))

+        return k, v

+    

+    def empty(self):

+        return not bool(self.dict)

+

+

 class DownloadQueue:

     def __init__(self, config, notifier):

         self.config = config

         self.notifier = notifier

-        self.queue = OrderedDict()

-        self.done = OrderedDict()

-        self.initialized = False

+        self.queue = PersistentQueue(self.config.STATE_DIR + '/queue')

+        self.done = PersistentQueue(self.config.STATE_DIR + '/completed')

+        self.done.load()

+    

+    async def __import_queue(self):

+        for k, v in self.queue.saved_items():

+            await self.add(v.url, v.quality, v.format)

 

-    def __initialize(self):

-        if not self.initialized:

-            self.initialized = True

-            self.event = asyncio.Event()

-            asyncio.create_task(self.__download())

+    async def initialize(self):

+        self.event = asyncio.Event()

+        asyncio.create_task(self.__download())

+        asyncio.create_task(self.__import_queue())

 

     def __extract_info(self, url):

         return yt_dlp.YoutubeDL(params={

@@ -165,10 +214,10 @@ class DownloadQueue:
                 return {'status': 'error', 'msg': ', '.join(res['msg'] for res in results if res['status'] == 'error' and 'msg' in res)}

             return {'status': 'ok'}

         elif etype == 'video' or etype.startswith('url') and 'id' in entry and 'title' in entry:

-            if entry['id'] not in self.queue:

+            if not self.queue.exists(entry['id']):

                 dl = DownloadInfo(entry['id'], entry['title'], entry.get('webpage_url') or entry['url'], quality, format)

                 dldirectory = self.config.DOWNLOAD_DIR if (quality != 'audio' and format != 'mp3') else self.config.AUDIO_DOWNLOAD_DIR

-                self.queue[entry['id']] = Download(dldirectory, self.config.OUTPUT_TEMPLATE, quality, format, self.config.YTDL_OPTIONS, dl)

+                self.queue.put(Download(dldirectory, self.config.OUTPUT_TEMPLATE, quality, format, self.config.YTDL_OPTIONS, dl))

                 self.event.set()

                 await self.notifier.added(dl)

             return {'status': 'ok'}

@@ -178,7 +227,6 @@ class DownloadQueue:
 

     async def add(self, url, quality, format, already=None):

         log.info(f'adding {url}')

-        self.__initialize()

         already = set() if already is None else already

         if url in already:

             log.info('recursion detected, skipping')

@@ -193,22 +241,22 @@ class DownloadQueue:
 

     async def cancel(self, ids):

         for id in ids:

-            if id not in self.queue:

+            if not self.queue.exists(id):

                 log.warn(f'requested cancel for non-existent download {id}')

                 continue

-            if self.queue[id].started():

-                self.queue[id].cancel()

+            if self.queue.get(id).started():

+                self.queue.get(id).cancel()

             else:

-                del self.queue[id]

+                self.queue.delete(id)

                 await self.notifier.canceled(id)

         return {'status': 'ok'}

 

     async def clear(self, ids):

         for id in ids:

-            if id not in self.done:

+            if not self.done.exists(id):

                 log.warn(f'requested delete for non-existent download {id}')

                 continue

-            del self.done[id]

+            self.done.delete(id)

             await self.notifier.cleared(id)

         return {'status': 'ok'}

 

@@ -218,11 +266,11 @@ class DownloadQueue:
 

     async def __download(self):

         while True:

-            while not self.queue:

+            while self.queue.empty():

                 log.info('waiting for item to download')

                 await self.event.wait()

                 self.event.clear()

-            id, entry = next(iter(self.queue.items()))

+            id, entry = self.queue.next()

             log.info(f'downloading {entry.info.title}')

             await entry.start(self.notifier)

             if entry.info.status != 'finished':

@@ -233,10 +281,10 @@ class DownloadQueue:
                         pass

                 entry.info.status = 'error'

             entry.close()

-            if id in self.queue:

-                del self.queue[id]

+            if self.queue.exists(id):

+                self.queue.delete(id)

                 if entry.canceled:

                     await self.notifier.canceled(id)

                 else:

-                    self.done[id] = entry

+                    self.done.put(entry)

                     await self.notifier.completed(entry.info)