From: Michael Tremer <michael.tremer@ipfire.org>
Date: Thu, 14 Jul 2022 15:49:38 +0000 (+0000)
Subject: uploads: Avoid copying the entire uploaded file again
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=efb81883ada0d52c3ecc0ccce4c8765205634ca4;p=pbs.git

uploads: Avoid copying the entire uploaded file again

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---

diff --git a/src/buildservice/uploads.py b/src/buildservice/uploads.py
index b8cb7d2c..d79c0f71 100644
--- a/src/buildservice/uploads.py
+++ b/src/buildservice/uploads.py
@@ -34,19 +34,20 @@ class Uploads(base.Object):
 		return self._get_upload("SELECT * FROM uploads \
 			WHERE uuid = %s AND expires_at > CURRENT_TIMESTAMP", uuid)
 
-	def create(self, filename, f, builder=None, user=None):
+	def allocate_file(self):
+		"""
+			Returns a file handle which can be used to write temporary data to.
+		"""
+		return tempfile.NamedTemporaryFile(dir=UPLOADS_DIR, delete=False)
+
+	def create(self, filename, path, size=None, builder=None, user=None):
 		# Check if either builder or user are set
 		if not builder and not user:
 			raise ValueError("builder or user must be set")
 
-		# Reset f
-		f.seek(0)
-
-		# Create a new temporary file
-		t = tempfile.NamedTemporaryFile(dir=UPLOADS_DIR, delete=False)
-
-		# Copy all content from f
-		shutil.copyfileobj(f, t)
+		# Fetch size if none given
+		if size is None:
+			size = os.path.getsize(path)
 
 		upload = self._get_upload("""
 			INSERT INTO
@@ -68,15 +69,12 @@ class Uploads(base.Object):
 			)
 			RETURNING *""",
 			filename,
-			t.name,
-			t.tell(),
-			builder.id if builder else None,
-			user.id if user else None,
+			path,
+			size,
+			builder,
+			user,
 		)
 
-		# Close the temporary file
-		t.close()
-
 		# Return the newly created upload object
 		return upload
 
diff --git a/src/hub/handlers.py b/src/hub/handlers.py
index b7d30164..f8009fd7 100644
--- a/src/hub/handlers.py
+++ b/src/hub/handlers.py
@@ -161,7 +161,7 @@ class UploadHandler(BaseHandler):
 		self.bytes_read = 0
 
 		# Allocate a temporary file
-		self.f = tempfile.SpooledTemporaryFile(max_size=10485760)
+		self.f = self.backend.uploads.allocate_file()
 
 		self.h, self.hexdigest = self._setup_digest()
 
@@ -182,7 +182,7 @@ class UploadHandler(BaseHandler):
 		# Write payload
 		self.f.write(data)
 
-	def put(self):
+	async def put(self):
 		"""
 			Called after the entire file has been received
 		"""
@@ -201,11 +201,14 @@ class UploadHandler(BaseHandler):
 		with self.db.transaction():
 			upload = self.backend.uploads.create(
 				self.filename,
-				self.f,
+				self.f.name,
 				builder=self.builder,
 				user=self.user,
 			)
 
+			# Free the temporary file (to prevent cleanup)
+			self.f = None
+
 		# Send the ID of the upload back to the client
 		self.finish({
 			"id"         : upload.uuid,
@@ -226,6 +229,15 @@ class UploadHandler(BaseHandler):
 
 		return h, hexdigest
 
+	def on_connection_close(self):
+		"""
+			Called when a connection was unexpectedly closed
+		"""
+		# Try deleting the file
+		if self.f:
+			log.debug("Deleting temporary file %s" % self.f.name)
+			os.unlink(self.f.name)
+
 
 # Builds