diff options
| -rw-r--r-- | src/discord_image_bridge/fsmanager.py | 6 | ||||
| -rw-r--r-- | src/discord_image_bridge/utils.py | 28 |
2 files changed, 23 insertions, 11 deletions
diff --git a/src/discord_image_bridge/fsmanager.py b/src/discord_image_bridge/fsmanager.py index 7313889..4ad8d00 100644 --- a/src/discord_image_bridge/fsmanager.py +++ b/src/discord_image_bridge/fsmanager.py @@ -3,12 +3,12 @@ import json from threading import Event def makedirs(): - for dir in ["cache", "meta", "links"]: + for dir in ["cache", "meta", "links", "temp"]: if not os.path.isdir(dir): os.mkdir(dir) -def hash2fname(hash): - return os.path.join("cache",hash[:2],hash[:10]) +def hash2fname(hash,temp=False): + return os.path.join("temp" if temp else "cache",hash[:2],hash[:10]) def hash2meta_fname(hash): return os.path.join("meta",hash[:10]+".json") diff --git a/src/discord_image_bridge/utils.py b/src/discord_image_bridge/utils.py index f937997..23a4fb0 100644 --- a/src/discord_image_bridge/utils.py +++ b/src/discord_image_bridge/utils.py @@ -1,5 +1,6 @@ import os from time import time +from threading import RLock as Lock import cherrypy import requests from hashlib import sha256 as do_hash @@ -8,6 +9,7 @@ from . import discord from .downloadpool import DownloadPool from . import fsmanager +download_locks={} def download_and_cache(url, filename): cherrypy.log("Downloading attachment %s"%filename) @@ -21,14 +23,16 @@ def download_and_cache(url, filename): return None, None hash=do_hash(resp.content).hexdigest() try: - fname=fsmanager.hash2fname(hash) - dirname=os.path.dirname(fname) + tmpname=fsmanager.hash2fname(hash,temp=True) + dirname=os.path.dirname(tmpname) try: os.mkdir(dirname) except FileExistsError: pass - with fsmanager.DataFile(open(fname,'wb')) as fd: + with fsmanager.DataFile(open(tmpname,'wb')) as fd: fd.write(resp.content) + fname=fsmanager.hash2fname(hash) + os.rename(tmpname, fname) return hash, fname @@ -40,24 +44,32 @@ def download_and_cache(url, filename): return None, None def download_uncached(hash): - with fsmanager.MetaFile(hash) as metafd: + if not os.path.isfile(fsmanager.hash2meta_fname(hash)): + return None, None + with download_locks.setdefault(hash, Lock()), fsmanager.MetaFile(hash) as metafd: + target=fsmanager.hash2fname(hash) + if os.path.isfile(target): + return hash, data=metafd.read() if not "sources" in data: cherrypy.log("No sources available for "+hash) return None, None - for source in data["sources"]: + for source in data["sources"] + data["sources"]: # try everything twice match source["type"]: case "discord": - hash, fname=download_uncached_discord( + newhash, fname=download_uncached_discord( channel=source["channel"], msgid=source["message"], attachmentid=source["attachment"] ) - if hash: return hash, fname + case _: # NOTE: maybe log here? return None, None - cherrypy.log("No sources available for "+hash) + + if newhash.startswith(hash): + return hash, fname + cherrypy.log("No working sources available for "+hash) return None, None def download_uncached_discord(channel, msgid, attachmentid): |