Fix transferring same attachment multiple times in parallel

attachments.go

@@ -17,6 +17,7 @@ import (
 	"maunium.net/go/mautrix/crypto/attachment"
 	"maunium.net/go/mautrix/event"
 	"maunium.net/go/mautrix/id"
+	"maunium.net/go/mautrix/util"
 
 	"go.mau.fi/mautrix-discord/database"
 )
@@ -146,32 +147,51 @@ func (br *DiscordBridge) uploadMatrixAttachment(intent *appservice.IntentAPI, da
 }
 
 type AttachmentMeta struct {
-	AttachmentID string
-	MimeType     string
-	EmojiName    string
+	AttachmentID  string
+	MimeType      string
+	EmojiName     string
+	CopyIfMissing bool
 }
 
-func (br *DiscordBridge) copyAttachmentToMatrix(intent *appservice.IntentAPI, url string, encrypt bool, meta *AttachmentMeta) (*database.File, error) {
-	dbFile := br.DB.File.Get(url, encrypt)
-	if dbFile == nil {
-		data, err := downloadDiscordAttachment(url)
-		if err != nil {
-			return nil, err
-		}
+var NoMeta = AttachmentMeta{}
 
-		if meta == nil {
-			meta = &AttachmentMeta{}
-		}
-		dbFile, err = br.uploadMatrixAttachment(intent, data, url, encrypt, *meta)
-		if err != nil {
-			return nil, err
-		}
-		// TODO add option to cache encrypted files too?
-		if !dbFile.Encrypted {
-			dbFile.Insert(nil)
-		}
+type attachmentKey struct {
+	URL     string
+	Encrypt bool
+}
+
+func (br *DiscordBridge) copyAttachmentToMatrix(intent *appservice.IntentAPI, url string, encrypt bool, meta AttachmentMeta) (returnDBFile *database.File, returnErr error) {
+	isCacheable := !encrypt
+	returnDBFile = br.DB.File.Get(url, encrypt)
+	if returnDBFile == nil {
+		transferKey := attachmentKey{url, encrypt}
+		once, _ := br.attachmentTransfers.GetOrSet(transferKey, &util.ReturnableOnce[*database.File]{})
+		returnDBFile, returnErr = once.Do(func() (onceDBFile *database.File, onceErr error) {
+			if isCacheable {
+				onceDBFile = br.DB.File.Get(url, encrypt)
+				if onceDBFile != nil {
+					return
+				}
+			}
+
+			var data []byte
+			data, onceErr = downloadDiscordAttachment(url)
+			if onceErr != nil {
+				return
+			}
+
+			onceDBFile, onceErr = br.uploadMatrixAttachment(intent, data, url, encrypt, meta)
+			if onceErr != nil {
+				return
+			}
+			if isCacheable {
+				onceDBFile.Insert(nil)
+			}
+			br.attachmentTransfers.Delete(transferKey)
+			return
+		})
 	}
-	return dbFile, nil
+	return
 }
 
 func (portal *Portal) getEmojiMXCByDiscordID(emojiID, name string, animated bool) id.ContentURI {
@@ -183,7 +203,7 @@ func (portal *Portal) getEmojiMXCByDiscordID(emojiID, name string, animated bool
 		url = discordgo.EndpointEmoji(emojiID)
 		mimeType = "image/png"
 	}
-	dbFile, err := portal.bridge.copyAttachmentToMatrix(portal.MainIntent(), url, false, &AttachmentMeta{
+	dbFile, err := portal.bridge.copyAttachmentToMatrix(portal.MainIntent(), url, false, AttachmentMeta{
 		AttachmentID: emojiID,
 		MimeType:     mimeType,
 		EmojiName:    name,

go.mod

@@ -14,7 +14,7 @@ require (
 	github.com/stretchr/testify v1.8.1
 	github.com/yuin/goldmark v1.5.3
 	maunium.net/go/maulogger/v2 v2.3.2
-	maunium.net/go/mautrix v0.13.1-0.20230129151130-9eb38c70fff2
+	maunium.net/go/mautrix v0.13.1-0.20230131110946-41fd713d1765
 )
 
 require (

go.sum

@@ -77,5 +77,5 @@ maunium.net/go/mauflag v1.0.0 h1:YiaRc0tEI3toYtJMRIfjP+jklH45uDHtT80nUamyD4M=
 maunium.net/go/mauflag v1.0.0/go.mod h1:nLivPOpTpHnpzEh8jEdSL9UqO9+/KBJFmNRlwKfkPeA=
 maunium.net/go/maulogger/v2 v2.3.2 h1:1XmIYmMd3PoQfp9J+PaHhpt80zpfmMqaShzUTC7FwY0=
 maunium.net/go/maulogger/v2 v2.3.2/go.mod h1:TYWy7wKwz/tIXTpsx8G3mZseIRiC5DoMxSZazOHy68A=
-maunium.net/go/mautrix v0.13.1-0.20230129151130-9eb38c70fff2 h1:/09m+KWf2fjxJMSpnbVudv4hlBaJVU8oou8TMnvKK0I=
-maunium.net/go/mautrix v0.13.1-0.20230129151130-9eb38c70fff2/go.mod h1:gYMQPsZ9lQpyKlVp+DGwOuc9LIcE/c8GZW2CvKHISgM=
+maunium.net/go/mautrix v0.13.1-0.20230131110946-41fd713d1765 h1:A9OYPQ5okmWrU4zMnU21UxZtXATuBLEgQ9FGZrFhNS0=
+maunium.net/go/mautrix v0.13.1-0.20230131110946-41fd713d1765/go.mod h1:gYMQPsZ9lQpyKlVp+DGwOuc9LIcE/c8GZW2CvKHISgM=

main.go

@@ -23,6 +23,7 @@ import (
 	"maunium.net/go/mautrix/bridge"
 	"maunium.net/go/mautrix/bridge/commands"
 	"maunium.net/go/mautrix/id"
+	"maunium.net/go/mautrix/util"
 	"maunium.net/go/mautrix/util/configupgrade"
 
 	"go.mau.fi/mautrix-discord/config"
@@ -71,6 +72,8 @@ type DiscordBridge struct {
 	puppets             map[string]*Puppet
 	puppetsByCustomMXID map[id.UserID]*Puppet
 	puppetsLock         sync.Mutex
+
+	attachmentTransfers *util.SyncMap[attachmentKey, *util.ReturnableOnce[*database.File]]
 }
 
 func (br *DiscordBridge) GetExampleConfig() string {
@@ -163,6 +166,8 @@ func main() {
 
 		puppets:             make(map[string]*Puppet),
 		puppetsByCustomMXID: make(map[id.UserID]*Puppet),
+
+		attachmentTransfers: util.NewSyncMap[attachmentKey, *util.ReturnableOnce[*database.File]](),
 	}
 	br.Bridge = bridge.Bridge{
 		Name:         "mautrix-discord",

portal.go

@@ -555,7 +555,7 @@ func (portal *Portal) sendMediaFailedMessage(intent *appservice.IntentAPI, bridg
 const DiscordStickerSize = 160
 
 func (portal *Portal) handleDiscordFile(typeName string, intent *appservice.IntentAPI, id, url string, content *event.MessageEventContent, ts time.Time, threadRelation *event.RelatesTo) *database.MessagePart {
-	dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, url, portal.Encrypted, &AttachmentMeta{AttachmentID: id, MimeType: content.Info.MimeType})
+	dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, url, portal.Encrypted, AttachmentMeta{AttachmentID: id, MimeType: content.Info.MimeType})
 	if err != nil {
 		errorEventID := portal.sendMediaFailedMessage(intent, err)
 		if errorEventID != "" {
@@ -675,7 +675,7 @@ type ConvertedMessage struct {
 }
 
 func (portal *Portal) convertDiscordVideoEmbed(intent *appservice.IntentAPI, embed *discordgo.MessageEmbed) *ConvertedMessage {
-	dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Video.ProxyURL, portal.Encrypted, nil)
+	dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Video.ProxyURL, portal.Encrypted, NoMeta)
 	if err != nil {
 		return &ConvertedMessage{Content: portal.createMediaFailedMessage(err)}
 	}
@@ -768,7 +768,7 @@ func (portal *Portal) convertDiscordRichEmbed(intent *appservice.IntentAPI, embe
 		}
 		authorHTML = fmt.Sprintf(embedHTMLAuthorPlain, authorNameHTML)
 		if embed.Author.ProxyIconURL != "" {
-			dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Author.ProxyIconURL, false, nil)
+			dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Author.ProxyIconURL, false, NoMeta)
 			if err != nil {
 				portal.log.Warnfln("Failed to reupload author icon in embed #%d of message %s: %v", index+1, msgID, err)
 			} else {
@@ -818,7 +818,7 @@ func (portal *Portal) convertDiscordRichEmbed(intent *appservice.IntentAPI, embe
 		}
 	}
 	if embed.Image != nil {
-		dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Image.ProxyURL, false, nil)
+		dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Image.ProxyURL, false, NoMeta)
 		if err != nil {
 			portal.log.Warnfln("Failed to reupload image in embed #%d of message %s: %v", index+1, msgID, err)
 		} else {
@@ -844,7 +844,7 @@ func (portal *Portal) convertDiscordRichEmbed(intent *appservice.IntentAPI, embe
 		}
 		footerHTML = fmt.Sprintf(embedHTMLFooterPlain, html.EscapeString(embed.Footer.Text), datePart)
 		if embed.Footer.ProxyIconURL != "" {
-			dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Footer.ProxyIconURL, false, nil)
+			dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, embed.Footer.ProxyIconURL, false, NoMeta)
 			if err != nil {
 				portal.log.Warnfln("Failed to reupload footer icon in embed #%d of message %s: %v", index+1, msgID, err)
 			} else {
@@ -876,7 +876,7 @@ type BeeperLinkPreview struct {
 }
 
 func (portal *Portal) convertDiscordLinkEmbedImage(intent *appservice.IntentAPI, url string, width, height int, preview *BeeperLinkPreview) {
-	dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, url, portal.Encrypted, nil)
+	dbFile, err := portal.bridge.copyAttachmentToMatrix(intent, url, portal.Encrypted, NoMeta)
 	if err != nil {
 		portal.log.Warnfln("Failed to copy image in URL preview: %v", err)
 	} else {