Use tempfiles

well done dsimic, you managed to get me to do something that I don't
complete get the point of :)
but ig if it's better it's better

Signed-off-by: Vosjedev <vosje@vosjedev.net>

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):