Add context to the mediaapi database (#250)

src/github.com/matrix-org/dendrite/mediaapi/storage/media_repository_table.go

@@ -15,6 +15,7 @@
 package storage
 
 import (
+	"context"
 	"database/sql"
 	"time"
 
@@ -74,9 +75,12 @@ func (s *mediaStatements) prepare(db *sql.DB) (err error) {
 	}.prepare(db)
 }
 
-func (s *mediaStatements) insertMedia(mediaMetadata *types.MediaMetadata) error {
+func (s *mediaStatements) insertMedia(
+	ctx context.Context, mediaMetadata *types.MediaMetadata,
+) error {
 	mediaMetadata.CreationTimestamp = types.UnixMs(time.Now().UnixNano() / 1000000)
-	_, err := s.insertMediaStmt.Exec(
+	_, err := s.insertMediaStmt.ExecContext(
+		ctx,
 		mediaMetadata.MediaID,
 		mediaMetadata.Origin,
 		mediaMetadata.ContentType,
@@ -89,13 +93,15 @@ func (s *mediaStatements) insertMedia(mediaMetadata *types.MediaMetadata) error
 	return err
 }
 
-func (s *mediaStatements) selectMedia(mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName) (*types.MediaMetadata, error) {
+func (s *mediaStatements) selectMedia(
+	ctx context.Context, mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName,
+) (*types.MediaMetadata, error) {
 	mediaMetadata := types.MediaMetadata{
 		MediaID: mediaID,
 		Origin:  mediaOrigin,
 	}
-	err := s.selectMediaStmt.QueryRow(
+	err := s.selectMediaStmt.QueryRowContext(
-		mediaMetadata.MediaID, mediaMetadata.Origin,
+		ctx, mediaMetadata.MediaID, mediaMetadata.Origin,
 	).Scan(
 		&mediaMetadata.ContentType,
 		&mediaMetadata.FileSizeBytes,

src/github.com/matrix-org/dendrite/mediaapi/storage/storage.go

@@ -15,6 +15,7 @@
 package storage
 
 import (
+	"context"
 	"database/sql"
 
 	// Import the postgres database driver.
@@ -44,15 +45,19 @@ func Open(dataSourceName string) (*Database, error) {
 
 // StoreMediaMetadata inserts the metadata about the uploaded media into the database.
 // Returns an error if the combination of MediaID and Origin are not unique in the table.
-func (d *Database) StoreMediaMetadata(mediaMetadata *types.MediaMetadata) error {
+func (d *Database) StoreMediaMetadata(
-	return d.statements.media.insertMedia(mediaMetadata)
+	ctx context.Context, mediaMetadata *types.MediaMetadata,
+) error {
+	return d.statements.media.insertMedia(ctx, mediaMetadata)
 }
 
 // GetMediaMetadata returns metadata about media stored on this server.
 // The media could have been uploaded to this server or fetched from another server and cached here.
 // Returns nil metadata if there is no metadata associated with this media.
-func (d *Database) GetMediaMetadata(mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName) (*types.MediaMetadata, error) {
+func (d *Database) GetMediaMetadata(
-	mediaMetadata, err := d.statements.media.selectMedia(mediaID, mediaOrigin)
+	ctx context.Context, mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName,
+) (*types.MediaMetadata, error) {
+	mediaMetadata, err := d.statements.media.selectMedia(ctx, mediaID, mediaOrigin)
 	if err != nil && err == sql.ErrNoRows {
 		return nil, nil
 	}
@@ -61,15 +66,25 @@ func (d *Database) GetMediaMetadata(mediaID types.MediaID, mediaOrigin gomatrixs
 
 // StoreThumbnail inserts the metadata about the thumbnail into the database.
 // Returns an error if the combination of MediaID and Origin are not unique in the table.
-func (d *Database) StoreThumbnail(thumbnailMetadata *types.ThumbnailMetadata) error {
+func (d *Database) StoreThumbnail(
-	return d.statements.thumbnail.insertThumbnail(thumbnailMetadata)
+	ctx context.Context, thumbnailMetadata *types.ThumbnailMetadata,
+) error {
+	return d.statements.thumbnail.insertThumbnail(ctx, thumbnailMetadata)
 }
 
 // GetThumbnail returns metadata about a specific thumbnail.
 // The media could have been uploaded to this server or fetched from another server and cached here.
 // Returns nil metadata if there is no metadata associated with this thumbnail.
-func (d *Database) GetThumbnail(mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName, width, height int, resizeMethod string) (*types.ThumbnailMetadata, error) {
+func (d *Database) GetThumbnail(
-	thumbnailMetadata, err := d.statements.thumbnail.selectThumbnail(mediaID, mediaOrigin, width, height, resizeMethod)
+	ctx context.Context,
+	mediaID types.MediaID,
+	mediaOrigin gomatrixserverlib.ServerName,
+	width, height int,
+	resizeMethod string,
+) (*types.ThumbnailMetadata, error) {
+	thumbnailMetadata, err := d.statements.thumbnail.selectThumbnail(
+		ctx, mediaID, mediaOrigin, width, height, resizeMethod,
+	)
 	if err != nil && err == sql.ErrNoRows {
 		return nil, nil
 	}
@@ -79,8 +94,10 @@ func (d *Database) GetThumbnail(mediaID types.MediaID, mediaOrigin gomatrixserve
 // GetThumbnails returns metadata about all thumbnails for a specific media stored on this server.
 // The media could have been uploaded to this server or fetched from another server and cached here.
 // Returns nil metadata if there are no thumbnails associated with this media.
-func (d *Database) GetThumbnails(mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName) ([]*types.ThumbnailMetadata, error) {
+func (d *Database) GetThumbnails(
-	thumbnails, err := d.statements.thumbnail.selectThumbnails(mediaID, mediaOrigin)
+	ctx context.Context, mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName,
+) ([]*types.ThumbnailMetadata, error) {
+	thumbnails, err := d.statements.thumbnail.selectThumbnails(ctx, mediaID, mediaOrigin)
 	if err != nil && err == sql.ErrNoRows {
 		return nil, nil
 	}

src/github.com/matrix-org/dendrite/mediaapi/storage/thumbnail_table.go

@@ -15,6 +15,7 @@
 package storage
 
 import (
+	"context"
 	"database/sql"
 	"time"
 
@@ -82,9 +83,12 @@ func (s *thumbnailStatements) prepare(db *sql.DB) (err error) {
 	}.prepare(db)
 }
 
-func (s *thumbnailStatements) insertThumbnail(thumbnailMetadata *types.ThumbnailMetadata) error {
+func (s *thumbnailStatements) insertThumbnail(
+	ctx context.Context, thumbnailMetadata *types.ThumbnailMetadata,
+) error {
 	thumbnailMetadata.MediaMetadata.CreationTimestamp = types.UnixMs(time.Now().UnixNano() / 1000000)
-	_, err := s.insertThumbnailStmt.Exec(
+	_, err := s.insertThumbnailStmt.ExecContext(
+		ctx,
 		thumbnailMetadata.MediaMetadata.MediaID,
 		thumbnailMetadata.MediaMetadata.Origin,
 		thumbnailMetadata.MediaMetadata.ContentType,
@@ -97,7 +101,13 @@ func (s *thumbnailStatements) insertThumbnail(thumbnailMetadata *types.Thumbnail
 	return err
 }
 
-func (s *thumbnailStatements) selectThumbnail(mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName, width, height int, resizeMethod string) (*types.ThumbnailMetadata, error) {
+func (s *thumbnailStatements) selectThumbnail(
+	ctx context.Context,
+	mediaID types.MediaID,
+	mediaOrigin gomatrixserverlib.ServerName,
+	width, height int,
+	resizeMethod string,
+) (*types.ThumbnailMetadata, error) {
 	thumbnailMetadata := types.ThumbnailMetadata{
 		MediaMetadata: &types.MediaMetadata{
 			MediaID: mediaID,
@@ -109,7 +119,8 @@ func (s *thumbnailStatements) selectThumbnail(mediaID types.MediaID, mediaOrigin
 			ResizeMethod: resizeMethod,
 		},
 	}
-	err := s.selectThumbnailStmt.QueryRow(
+	err := s.selectThumbnailStmt.QueryRowContext(
+		ctx,
 		thumbnailMetadata.MediaMetadata.MediaID,
 		thumbnailMetadata.MediaMetadata.Origin,
 		thumbnailMetadata.ThumbnailSize.Width,
@@ -123,9 +134,11 @@ func (s *thumbnailStatements) selectThumbnail(mediaID types.MediaID, mediaOrigin
 	return &thumbnailMetadata, err
 }
 
-func (s *thumbnailStatements) selectThumbnails(mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName) ([]*types.ThumbnailMetadata, error) {
+func (s *thumbnailStatements) selectThumbnails(
-	rows, err := s.selectThumbnailsStmt.Query(
+	ctx context.Context, mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName,
-		mediaID, mediaOrigin,
+) ([]*types.ThumbnailMetadata, error) {
+	rows, err := s.selectThumbnailsStmt.QueryContext(
+		ctx, mediaID, mediaOrigin,
 	)
 	if err != nil {
 		return nil, err

src/github.com/matrix-org/dendrite/mediaapi/thumbnailer/thumbnailer.go

@@ -15,6 +15,7 @@
 package thumbnailer
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"os"
@@ -130,8 +131,18 @@ func broadcastGeneration(dst types.Path, activeThumbnailGeneration *types.Active
 	delete(activeThumbnailGeneration.PathToResult, string(dst))
 }
 
-func isThumbnailExists(dst types.Path, config types.ThumbnailSize, mediaMetadata *types.MediaMetadata, db *storage.Database, logger *log.Entry) (bool, error) {
+func isThumbnailExists(
-	thumbnailMetadata, err := db.GetThumbnail(mediaMetadata.MediaID, mediaMetadata.Origin, config.Width, config.Height, config.ResizeMethod)
+	ctx context.Context,
+	dst types.Path,
+	config types.ThumbnailSize,
+	mediaMetadata *types.MediaMetadata,
+	db *storage.Database,
+	logger *log.Entry,
+) (bool, error) {
+	thumbnailMetadata, err := db.GetThumbnail(
+		ctx, mediaMetadata.MediaID, mediaMetadata.Origin,
+		config.Width, config.Height, config.ResizeMethod,
+	)
 	if err != nil {
 		logger.Error("Failed to query database for thumbnail.")
 		return false, err

src/github.com/matrix-org/dendrite/mediaapi/thumbnailer/thumbnailer_bimg.go

@@ -17,6 +17,7 @@
 package thumbnailer
 
 import (
+	"context"
 	"os"
 	"time"
 
@@ -28,7 +29,16 @@ import (
 )
 
 // GenerateThumbnails generates the configured thumbnail sizes for the source file
-func GenerateThumbnails(src types.Path, configs []config.ThumbnailSize, mediaMetadata *types.MediaMetadata, activeThumbnailGeneration *types.ActiveThumbnailGeneration, maxThumbnailGenerators int, db *storage.Database, logger *log.Entry) (busy bool, errorReturn error) {
+func GenerateThumbnails(
+	ctx context.Context,
+	src types.Path,
+	configs []config.ThumbnailSize,
+	mediaMetadata *types.MediaMetadata,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+	maxThumbnailGenerators int,
+	db *storage.Database,
+	logger *log.Entry,
+) (busy bool, errorReturn error) {
 	buffer, err := bimg.Read(string(src))
 	if err != nil {
 		logger.WithError(err).WithField("src", src).Error("Failed to read src file")
@@ -37,7 +47,10 @@ func GenerateThumbnails(src types.Path, configs []config.ThumbnailSize, mediaMet
 	img := bimg.NewImage(buffer)
 	for _, config := range configs {
 		// Note: createThumbnail does locking based on activeThumbnailGeneration
-		busy, err = createThumbnail(src, img, config, mediaMetadata, activeThumbnailGeneration, maxThumbnailGenerators, db, logger)
+		busy, err = createThumbnail(
+			ctx, src, img, config, mediaMetadata, activeThumbnailGeneration,
+			maxThumbnailGenerators, db, logger,
+		)
 		if err != nil {
 			logger.WithError(err).WithField("src", src).Error("Failed to generate thumbnails")
 			return false, err
@@ -50,7 +63,16 @@ func GenerateThumbnails(src types.Path, configs []config.ThumbnailSize, mediaMet
 }
 
 // GenerateThumbnail generates the configured thumbnail size for the source file
-func GenerateThumbnail(src types.Path, config types.ThumbnailSize, mediaMetadata *types.MediaMetadata, activeThumbnailGeneration *types.ActiveThumbnailGeneration, maxThumbnailGenerators int, db *storage.Database, logger *log.Entry) (busy bool, errorReturn error) {
+func GenerateThumbnail(
+	ctx context.Context,
+	src types.Path,
+	config types.ThumbnailSize,
+	mediaMetadata *types.MediaMetadata,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+	maxThumbnailGenerators int,
+	db *storage.Database,
+	logger *log.Entry,
+) (busy bool, errorReturn error) {
 	buffer, err := bimg.Read(string(src))
 	if err != nil {
 		logger.WithError(err).WithFields(log.Fields{
@@ -60,7 +82,10 @@ func GenerateThumbnail(src types.Path, config types.ThumbnailSize, mediaMetadata
 	}
 	img := bimg.NewImage(buffer)
 	// Note: createThumbnail does locking based on activeThumbnailGeneration
-	busy, err = createThumbnail(src, img, config, mediaMetadata, activeThumbnailGeneration, maxThumbnailGenerators, db, logger)
+	busy, err = createThumbnail(
+		ctx, src, img, config, mediaMetadata, activeThumbnailGeneration,
+		maxThumbnailGenerators, db, logger,
+	)
 	if err != nil {
 		logger.WithError(err).WithFields(log.Fields{
 			"src": src,
@@ -75,7 +100,17 @@ func GenerateThumbnail(src types.Path, config types.ThumbnailSize, mediaMetadata
 
 // createThumbnail checks if the thumbnail exists, and if not, generates it
 // Thumbnail generation is only done once for each non-existing thumbnail.
-func createThumbnail(src types.Path, img *bimg.Image, config types.ThumbnailSize, mediaMetadata *types.MediaMetadata, activeThumbnailGeneration *types.ActiveThumbnailGeneration, maxThumbnailGenerators int, db *storage.Database, logger *log.Entry) (busy bool, errorReturn error) {
+func createThumbnail(
+	ctx context.Context,
+	src types.Path,
+	img *bimg.Image,
+	config types.ThumbnailSize,
+	mediaMetadata *types.MediaMetadata,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+	maxThumbnailGenerators int,
+	db *storage.Database,
+	logger *log.Entry,
+) (busy bool, errorReturn error) {
 	logger = logger.WithFields(log.Fields{
 		"Width":        config.Width,
 		"Height":       config.Height,
@@ -111,7 +146,7 @@ func createThumbnail(src types.Path, img *bimg.Image, config types.ThumbnailSize
 		}()
 	}
 
-	exists, err := isThumbnailExists(dst, config, mediaMetadata, db, logger)
+	exists, err := isThumbnailExists(ctx, dst, config, mediaMetadata, db, logger)
 	if err != nil || exists {
 		return false, err
 	}
@@ -147,7 +182,7 @@ func createThumbnail(src types.Path, img *bimg.Image, config types.ThumbnailSize
 		},
 	}
 
-	err = db.StoreThumbnail(thumbnailMetadata)
+	err = db.StoreThumbnail(ctx, thumbnailMetadata)
 	if err != nil {
 		logger.WithError(err).WithFields(log.Fields{
 			"ActualWidth":  width,

src/github.com/matrix-org/dendrite/mediaapi/thumbnailer/thumbnailer_nfnt.go

@@ -17,6 +17,7 @@
 package thumbnailer
 
 import (
+	"context"
 	"image"
 	"image/draw"
 	// Imported for gif codec
@@ -35,7 +36,16 @@ import (
 )
 
 // GenerateThumbnails generates the configured thumbnail sizes for the source file
-func GenerateThumbnails(src types.Path, configs []config.ThumbnailSize, mediaMetadata *types.MediaMetadata, activeThumbnailGeneration *types.ActiveThumbnailGeneration, maxThumbnailGenerators int, db *storage.Database, logger *log.Entry) (busy bool, errorReturn error) {
+func GenerateThumbnails(
+	ctx context.Context,
+	src types.Path,
+	configs []config.ThumbnailSize,
+	mediaMetadata *types.MediaMetadata,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+	maxThumbnailGenerators int,
+	db *storage.Database,
+	logger *log.Entry,
+) (busy bool, errorReturn error) {
 	img, err := readFile(string(src))
 	if err != nil {
 		logger.WithError(err).WithField("src", src).Error("Failed to read src file")
@@ -43,7 +53,10 @@ func GenerateThumbnails(src types.Path, configs []config.ThumbnailSize, mediaMet
 	}
 	for _, config := range configs {
 		// Note: createThumbnail does locking based on activeThumbnailGeneration
-		busy, err = createThumbnail(src, img, types.ThumbnailSize(config), mediaMetadata, activeThumbnailGeneration, maxThumbnailGenerators, db, logger)
+		busy, err = createThumbnail(
+			ctx, src, img, types.ThumbnailSize(config), mediaMetadata,
+			activeThumbnailGeneration, maxThumbnailGenerators, db, logger,
+		)
 		if err != nil {
 			logger.WithError(err).WithField("src", src).Error("Failed to generate thumbnails")
 			return false, err
@@ -56,7 +69,16 @@ func GenerateThumbnails(src types.Path, configs []config.ThumbnailSize, mediaMet
 }
 
 // GenerateThumbnail generates the configured thumbnail size for the source file
-func GenerateThumbnail(src types.Path, config types.ThumbnailSize, mediaMetadata *types.MediaMetadata, activeThumbnailGeneration *types.ActiveThumbnailGeneration, maxThumbnailGenerators int, db *storage.Database, logger *log.Entry) (busy bool, errorReturn error) {
+func GenerateThumbnail(
+	ctx context.Context,
+	src types.Path,
+	config types.ThumbnailSize,
+	mediaMetadata *types.MediaMetadata,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+	maxThumbnailGenerators int,
+	db *storage.Database,
+	logger *log.Entry,
+) (busy bool, errorReturn error) {
 	img, err := readFile(string(src))
 	if err != nil {
 		logger.WithError(err).WithFields(log.Fields{
@@ -65,7 +87,10 @@ func GenerateThumbnail(src types.Path, config types.ThumbnailSize, mediaMetadata
 		return false, err
 	}
 	// Note: createThumbnail does locking based on activeThumbnailGeneration
-	busy, err = createThumbnail(src, img, config, mediaMetadata, activeThumbnailGeneration, maxThumbnailGenerators, db, logger)
+	busy, err = createThumbnail(
+		ctx, src, img, config, mediaMetadata, activeThumbnailGeneration,
+		maxThumbnailGenerators, db, logger,
+	)
 	if err != nil {
 		logger.WithError(err).WithFields(log.Fields{
 			"src": src,
@@ -107,7 +132,17 @@ func writeFile(img image.Image, dst string) (err error) {
 
 // createThumbnail checks if the thumbnail exists, and if not, generates it
 // Thumbnail generation is only done once for each non-existing thumbnail.
-func createThumbnail(src types.Path, img image.Image, config types.ThumbnailSize, mediaMetadata *types.MediaMetadata, activeThumbnailGeneration *types.ActiveThumbnailGeneration, maxThumbnailGenerators int, db *storage.Database, logger *log.Entry) (busy bool, errorReturn error) {
+func createThumbnail(
+	ctx context.Context,
+	src types.Path,
+	img image.Image,
+	config types.ThumbnailSize,
+	mediaMetadata *types.MediaMetadata,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+	maxThumbnailGenerators int,
+	db *storage.Database,
+	logger *log.Entry,
+) (busy bool, errorReturn error) {
 	logger = logger.WithFields(log.Fields{
 		"Width":        config.Width,
 		"Height":       config.Height,
@@ -143,7 +178,7 @@ func createThumbnail(src types.Path, img image.Image, config types.ThumbnailSize
 		}()
 	}
 
-	exists, err := isThumbnailExists(dst, config, mediaMetadata, db, logger)
+	exists, err := isThumbnailExists(ctx, dst, config, mediaMetadata, db, logger)
 	if err != nil || exists {
 		return false, err
 	}
@@ -179,7 +214,7 @@ func createThumbnail(src types.Path, img image.Image, config types.ThumbnailSize
 		},
 	}
 
-	err = db.StoreThumbnail(thumbnailMetadata)
+	err = db.StoreThumbnail(ctx, thumbnailMetadata)
 	if err != nil {
 		logger.WithError(err).WithFields(log.Fields{
 			"ActualWidth":  width,

src/github.com/matrix-org/dendrite/mediaapi/writers/download.go

@@ -15,6 +15,7 @@
 package writers
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -118,7 +119,9 @@ func Download(
 		return
 	}
 
-	metadata, err := dReq.doDownload(w, cfg, db, activeRemoteRequests, activeThumbnailGeneration)
+	metadata, err := dReq.doDownload(
+		req.Context(), w, cfg, db, activeRemoteRequests, activeThumbnailGeneration,
+	)
 	if err != nil {
 		// TODO: Handle the fact we might have started writing the response
 		dReq.jsonErrorResponse(w, util.ErrorResponse(err))
@@ -192,6 +195,7 @@ func (r *downloadRequest) Validate() *util.JSONResponse {
 }
 
 func (r *downloadRequest) doDownload(
+	ctx context.Context,
 	w http.ResponseWriter,
 	cfg *config.Dendrite,
 	db *storage.Database,
@@ -199,7 +203,9 @@ func (r *downloadRequest) doDownload(
 	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
 ) (*types.MediaMetadata, error) {
 	// check if we have a record of the media in our database
-	mediaMetadata, err := db.GetMediaMetadata(r.MediaMetadata.MediaID, r.MediaMetadata.Origin)
+	mediaMetadata, err := db.GetMediaMetadata(
+		ctx, r.MediaMetadata.MediaID, r.MediaMetadata.Origin,
+	)
 	if err != nil {
 		return nil, errors.Wrap(err, "error querying the database")
 	}
@@ -209,7 +215,9 @@ func (r *downloadRequest) doDownload(
 			return nil, nil
 		}
 		// If we do not have a record and the origin is remote, we need to fetch it and respond with that file
-		resErr := r.getRemoteFile(cfg, db, activeRemoteRequests, activeThumbnailGeneration)
+		resErr := r.getRemoteFile(
+			ctx, cfg, db, activeRemoteRequests, activeThumbnailGeneration,
+		)
 		if resErr != nil {
 			return nil, resErr
 		}
@@ -217,12 +225,17 @@ func (r *downloadRequest) doDownload(
 		// If we have a record, we can respond from the local file
 		r.MediaMetadata = mediaMetadata
 	}
-	return r.respondFromLocalFile(w, cfg.Media.AbsBasePath, activeThumbnailGeneration, cfg.Media.MaxThumbnailGenerators, db, cfg.Media.DynamicThumbnails, cfg.Media.ThumbnailSizes)
+	return r.respondFromLocalFile(
+		ctx, w, cfg.Media.AbsBasePath, activeThumbnailGeneration,
+		cfg.Media.MaxThumbnailGenerators, db,
+		cfg.Media.DynamicThumbnails, cfg.Media.ThumbnailSizes,
+	)
 }
 
 // respondFromLocalFile reads a file from local storage and writes it to the http.ResponseWriter
 // If no file was found then returns nil, nil
 func (r *downloadRequest) respondFromLocalFile(
+	ctx context.Context,
 	w http.ResponseWriter,
 	absBasePath config.Path,
 	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
@@ -256,7 +269,10 @@ func (r *downloadRequest) respondFromLocalFile(
 	var responseFile *os.File
 	var responseMetadata *types.MediaMetadata
 	if r.IsThumbnailRequest {
-		thumbFile, thumbMetadata, resErr := r.getThumbnailFile(types.Path(filePath), activeThumbnailGeneration, maxThumbnailGenerators, db, dynamicThumbnails, thumbnailSizes)
+		thumbFile, thumbMetadata, resErr := r.getThumbnailFile(
+			ctx, types.Path(filePath), activeThumbnailGeneration, maxThumbnailGenerators,
+			db, dynamicThumbnails, thumbnailSizes,
+		)
 		if thumbFile != nil {
 			defer thumbFile.Close() // nolint: errcheck
 		}
@@ -306,6 +322,7 @@ func (r *downloadRequest) respondFromLocalFile(
 // Note: Thumbnail generation may be ongoing asynchronously.
 // If no thumbnail was found then returns nil, nil, nil
 func (r *downloadRequest) getThumbnailFile(
+	ctx context.Context,
 	filePath types.Path,
 	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
 	maxThumbnailGenerators int,
@@ -317,7 +334,10 @@ func (r *downloadRequest) getThumbnailFile(
 	var err error
 
 	if dynamicThumbnails {
-		thumbnail, err = r.generateThumbnail(filePath, r.ThumbnailSize, activeThumbnailGeneration, maxThumbnailGenerators, db)
+		thumbnail, err = r.generateThumbnail(
+			ctx, filePath, r.ThumbnailSize, activeThumbnailGeneration,
+			maxThumbnailGenerators, db,
+		)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -326,7 +346,9 @@ func (r *downloadRequest) getThumbnailFile(
 	// to trying to use a pre-generated thumbnail
 	if thumbnail == nil {
 		var thumbnails []*types.ThumbnailMetadata
-		thumbnails, err = db.GetThumbnails(r.MediaMetadata.MediaID, r.MediaMetadata.Origin)
+		thumbnails, err = db.GetThumbnails(
+			ctx, r.MediaMetadata.MediaID, r.MediaMetadata.Origin,
+		)
 		if err != nil {
 			return nil, nil, errors.Wrap(err, "error looking up thumbnails")
 		}
@@ -343,7 +365,10 @@ func (r *downloadRequest) getThumbnailFile(
 				"Height":       thumbnailSize.Height,
 				"ResizeMethod": thumbnailSize.ResizeMethod,
 			}).Info("Pre-generating thumbnail for immediate response.")
-			thumbnail, err = r.generateThumbnail(filePath, *thumbnailSize, activeThumbnailGeneration, maxThumbnailGenerators, db)
+			thumbnail, err = r.generateThumbnail(
+				ctx, filePath, *thumbnailSize, activeThumbnailGeneration,
+				maxThumbnailGenerators, db,
+			)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -378,6 +403,7 @@ func (r *downloadRequest) getThumbnailFile(
 }
 
 func (r *downloadRequest) generateThumbnail(
+	ctx context.Context,
 	filePath types.Path,
 	thumbnailSize types.ThumbnailSize,
 	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
@@ -389,7 +415,10 @@ func (r *downloadRequest) generateThumbnail(
 		"Height":       thumbnailSize.Height,
 		"ResizeMethod": thumbnailSize.ResizeMethod,
 	})
-	busy, err := thumbnailer.GenerateThumbnail(filePath, thumbnailSize, r.MediaMetadata, activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger)
+	busy, err := thumbnailer.GenerateThumbnail(
+		ctx, filePath, thumbnailSize, r.MediaMetadata,
+		activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger,
+	)
 	if err != nil {
 		return nil, errors.Wrap(err, "error creating thumbnail")
 	}
@@ -397,7 +426,10 @@ func (r *downloadRequest) generateThumbnail(
 		return nil, nil
 	}
 	var thumbnail *types.ThumbnailMetadata
-	thumbnail, err = db.GetThumbnail(r.MediaMetadata.MediaID, r.MediaMetadata.Origin, thumbnailSize.Width, thumbnailSize.Height, thumbnailSize.ResizeMethod)
+	thumbnail, err = db.GetThumbnail(
+		ctx, r.MediaMetadata.MediaID, r.MediaMetadata.Origin,
+		thumbnailSize.Width, thumbnailSize.Height, thumbnailSize.ResizeMethod,
+	)
 	if err != nil {
 		return nil, errors.Wrap(err, "error looking up thumbnail")
 	}
@@ -409,6 +441,7 @@ func (r *downloadRequest) generateThumbnail(
 // regardless of how many download requests are received.
 // Note: The named errorResponse return variable is used in a deferred broadcast of the metadata and error response to waiting goroutines.
 func (r *downloadRequest) getRemoteFile(
+	ctx context.Context,
 	cfg *config.Dendrite,
 	db *storage.Database,
 	activeRemoteRequests *types.ActiveRemoteRequests,
@@ -434,14 +467,20 @@ func (r *downloadRequest) getRemoteFile(
 		}()
 
 		// check if we have a record of the media in our database
-		mediaMetadata, err := db.GetMediaMetadata(r.MediaMetadata.MediaID, r.MediaMetadata.Origin)
+		mediaMetadata, err := db.GetMediaMetadata(
+			ctx, r.MediaMetadata.MediaID, r.MediaMetadata.Origin,
+		)
 		if err != nil {
 			return errors.Wrap(err, "error querying the database.")
 		}
 
 		if mediaMetadata == nil {
 			// If we do not have a record, we need to fetch the remote file first and then respond from the local file
-			err := r.fetchRemoteFileAndStoreMetadata(cfg.Media.AbsBasePath, *cfg.Media.MaxFileSizeBytes, db, cfg.Media.ThumbnailSizes, activeThumbnailGeneration, cfg.Media.MaxThumbnailGenerators)
+			err := r.fetchRemoteFileAndStoreMetadata(
+				ctx, cfg.Media.AbsBasePath, *cfg.Media.MaxFileSizeBytes, db,
+				cfg.Media.ThumbnailSizes, activeThumbnailGeneration,
+				cfg.Media.MaxThumbnailGenerators,
+			)
 			if err != nil {
 				return errors.Wrap(err, "error querying the database.")
 			}
@@ -501,6 +540,7 @@ func (r *downloadRequest) broadcastMediaMetadata(activeRemoteRequests *types.Act
 
 // fetchRemoteFileAndStoreMetadata fetches the file from the remote server and stores its metadata in the database
 func (r *downloadRequest) fetchRemoteFileAndStoreMetadata(
+	ctx context.Context,
 	absBasePath config.Path,
 	maxFileSizeBytes config.FileSizeBytes,
 	db *storage.Database,
@@ -521,7 +561,7 @@ func (r *downloadRequest) fetchRemoteFileAndStoreMetadata(
 	}).Info("Storing file metadata to media repository database")
 
 	// FIXME: timeout db request
-	if err := db.StoreMediaMetadata(r.MediaMetadata); err != nil {
+	if err := db.StoreMediaMetadata(ctx, r.MediaMetadata); err != nil {
 		// If the file is a duplicate (has the same hash as an existing file) then
 		// there is valid metadata in the database for that file. As such we only
 		// remove the file if it is not a duplicate.
@@ -535,7 +575,10 @@ func (r *downloadRequest) fetchRemoteFileAndStoreMetadata(
 	}
 
 	go func() {
-		busy, err := thumbnailer.GenerateThumbnails(finalPath, thumbnailSizes, r.MediaMetadata, activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger)
+		busy, err := thumbnailer.GenerateThumbnails(
+			context.Background(), finalPath, thumbnailSizes, r.MediaMetadata,
+			activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger,
+		)
 		if err != nil {
 			r.Logger.WithError(err).Warn("Error generating thumbnails")
 		}

src/github.com/matrix-org/dendrite/mediaapi/writers/upload.go

@@ -15,6 +15,7 @@
 package writers
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -57,7 +58,7 @@ func Upload(req *http.Request, cfg *config.Dendrite, db *storage.Database, activ
 		return *resErr
 	}
 
-	if resErr = r.doUpload(req.Body, cfg, db, activeThumbnailGeneration); resErr != nil {
+	if resErr = r.doUpload(req.Context(), req.Body, cfg, db, activeThumbnailGeneration); resErr != nil {
 		return *resErr
 	}
 
@@ -97,7 +98,13 @@ func parseAndValidateRequest(req *http.Request, cfg *config.Dendrite) (*uploadRe
 	return r, nil
 }
 
-func (r *uploadRequest) doUpload(reqReader io.Reader, cfg *config.Dendrite, db *storage.Database, activeThumbnailGeneration *types.ActiveThumbnailGeneration) *util.JSONResponse {
+func (r *uploadRequest) doUpload(
+	ctx context.Context,
+	reqReader io.Reader,
+	cfg *config.Dendrite,
+	db *storage.Database,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+) *util.JSONResponse {
 	r.Logger.WithFields(log.Fields{
 		"UploadName":    r.MediaMetadata.UploadName,
 		"FileSizeBytes": r.MediaMetadata.FileSizeBytes,
@@ -134,7 +141,9 @@ func (r *uploadRequest) doUpload(reqReader io.Reader, cfg *config.Dendrite, db *
 	}).Info("File uploaded")
 
 	// check if we already have a record of the media in our database and if so, we can remove the temporary directory
-	mediaMetadata, err := db.GetMediaMetadata(r.MediaMetadata.MediaID, r.MediaMetadata.Origin)
+	mediaMetadata, err := db.GetMediaMetadata(
+		ctx, r.MediaMetadata.MediaID, r.MediaMetadata.Origin,
+	)
 	if err != nil {
 		r.Logger.WithError(err).Error("Error querying the database.")
 		resErr := jsonerror.InternalServerError()
@@ -152,7 +161,10 @@ func (r *uploadRequest) doUpload(reqReader io.Reader, cfg *config.Dendrite, db *
 		}
 	}
 
-	if resErr := r.storeFileAndMetadata(tmpDir, cfg.Media.AbsBasePath, db, cfg.Media.ThumbnailSizes, activeThumbnailGeneration, cfg.Media.MaxThumbnailGenerators); resErr != nil {
+	if resErr := r.storeFileAndMetadata(
+		ctx, tmpDir, cfg.Media.AbsBasePath, db, cfg.Media.ThumbnailSizes,
+		activeThumbnailGeneration, cfg.Media.MaxThumbnailGenerators,
+	); resErr != nil {
 		return resErr
 	}
 
@@ -208,7 +220,15 @@ func (r *uploadRequest) Validate(maxFileSizeBytes config.FileSizeBytes) *util.JS
 // The order of operations is important as it avoids metadata entering the database before the file
 // is ready, and if we fail to move the file, it never gets added to the database.
 // Returns a util.JSONResponse error and cleans up directories in case of error.
-func (r *uploadRequest) storeFileAndMetadata(tmpDir types.Path, absBasePath config.Path, db *storage.Database, thumbnailSizes []config.ThumbnailSize, activeThumbnailGeneration *types.ActiveThumbnailGeneration, maxThumbnailGenerators int) *util.JSONResponse {
+func (r *uploadRequest) storeFileAndMetadata(
+	ctx context.Context,
+	tmpDir types.Path,
+	absBasePath config.Path,
+	db *storage.Database,
+	thumbnailSizes []config.ThumbnailSize,
+	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
+	maxThumbnailGenerators int,
+) *util.JSONResponse {
 	finalPath, duplicate, err := fileutils.MoveFileWithHashCheck(tmpDir, r.MediaMetadata, absBasePath, r.Logger)
 	if err != nil {
 		r.Logger.WithError(err).Error("Failed to move file.")
@@ -221,7 +241,7 @@ func (r *uploadRequest) storeFileAndMetadata(tmpDir types.Path, absBasePath conf
 		r.Logger.WithField("dst", finalPath).Info("File was stored previously - discarding duplicate")
 	}
 
-	if err = db.StoreMediaMetadata(r.MediaMetadata); err != nil {
+	if err = db.StoreMediaMetadata(ctx, r.MediaMetadata); err != nil {
 		r.Logger.WithError(err).Warn("Failed to store metadata")
 		// If the file is a duplicate (has the same hash as an existing file) then
 		// there is valid metadata in the database for that file. As such we only
@@ -236,7 +256,10 @@ func (r *uploadRequest) storeFileAndMetadata(tmpDir types.Path, absBasePath conf
 	}
 
 	go func() {
-		busy, err := thumbnailer.GenerateThumbnails(finalPath, thumbnailSizes, r.MediaMetadata, activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger)
+		busy, err := thumbnailer.GenerateThumbnails(
+			context.Background(), finalPath, thumbnailSizes, r.MediaMetadata,
+			activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger,
+		)
 		if err != nil {
 			r.Logger.WithError(err).Warn("Error generating thumbnails")
 		}