Speed up uploads by avoiding calls to the database

This removes a select across the whole table by ensuring we generate a media ID which is highly unlikely to conflict with the database.

CHANGELOG.md

@@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Fixed federation with some homeserver setups (delegation with ports). Thanks @MatMaul!
 * Fixed the Synapse import script to not skip duplicated media. Thanks @jaywink!
 * Fixed requests to IPv6 hosts. Thanks @MatMaul!
+* Removed excessive calls to the database during upload.
 
 ## [1.1.2] - April 21st, 2020
 

controllers/upload_controller/upload_controller.go

 package upload_controller
 
 import (
-	"database/sql"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"strconv"
+	"time"
 
+	"github.com/patrickmn/go-cache"
 	"github.com/pkg/errors"
 	"github.com/ryanuber/go-glob"
 	"github.com/sirupsen/logrus"
@@ -21,6 +22,8 @@ import (
 
 const NoApplicableUploadUser = ""
 
+var recentMediaIds = cache.New(30 * time.Second, 60 * time.Second)
+
 type AlreadyUploadedFile struct {
 	DS         *datastore.DatastoreRef
 	ObjectInfo *types.ObjectInfo
@@ -94,7 +97,6 @@ func UploadMedia(contents io.ReadCloser, contentLength int64, contentType string
 	}
 
 	metadataDb := storage.GetDatabase().GetMetadataStore(ctx)
-	mediaDb := storage.GetDatabase().GetMediaStore(ctx)
 
 	mediaTaken := true
 	var mediaId string
@@ -110,27 +112,26 @@ func UploadMedia(contents io.ReadCloser, contentLength int64, contentType string
 		if err != nil {
 			return nil, err
 		}
-
-		mediaTaken, err = metadataDb.IsReserved(origin, mediaId)
+		mediaId, err = util.GetSha1OfString(mediaId + strconv.FormatInt(util.NowMillis(), 10))
 		if err != nil {
 			return nil, err
 		}
 
-		if !mediaTaken {
-			// Double check it isn't already in use
-			var media *types.Media
-			media, err = mediaDb.Get(origin, mediaId)
-			if err == sql.ErrNoRows {
-				mediaTaken = false
-				continue
-			}
-			if err != nil {
-				return nil, err
-			}
-			mediaTaken = media != nil
+		// Because we use the current time in the media ID, we don't need to worry about
+		// collisions from the database.
+		if _, present := recentMediaIds.Get(mediaId); present {
+			mediaTaken = true
+			continue
+		}
+
+		mediaTaken, err = metadataDb.IsReserved(origin, mediaId)
+		if err != nil {
+			return nil, err
 		}
 	}
 
+	_ = recentMediaIds.Add(mediaId, true, cache.DefaultExpiration)
+
 	var existingFile *AlreadyUploadedFile = nil
 	ds, err := datastore.PickDatastore(common.KindLocalMedia, ctx)
 	if err != nil {

util/random.go

@@ -25,3 +25,9 @@ func GenerateRandomString(nBytes int) (string, error) {
 	hasher.Write(b)
 	return hex.EncodeToString(hasher.Sum(nil)), nil
 }
+
+func GetSha1OfString(str string) (string, error) {
+	hasher := sha1.New()
+	hasher.Write([]byte(str))
+	return hex.EncodeToString(hasher.Sum(nil)), nil
+}
\ No newline at end of file