[bugfix] Fix up `error getting account avatar/header` errors, other small fixes (#1496)

* start fiddling with media + account queries a little

* initialize state when pruning

* allow for unsetting remote media
make sure to wait til media loaded
fix silly tiny bug

* move comment a bit for readability

* slight reformat of fetchRemoteAccount{Avatar,Header}

* fix issue after rebase

* slightly neaten up logic of avatar/header media handling

* remove log prefix (callername log field handles this)

---------

Signed-off-by: kim <grufwub@gmail.com>
Co-authored-by: kim <grufwub@gmail.com>

cmd/gotosocial/action/admin/media/prune/common.go

@@ -33,6 +33,7 @@ type prune struct {
 	dbService db.DB
 	storage   *gtsstorage.Driver
 	manager   media.Manager
+	state     *state.State
 }
 
 func setupPrune(ctx context.Context) (*prune, error) {
@@ -44,6 +45,7 @@ func setupPrune(ctx context.Context) (*prune, error) {
 	if err != nil {
 		return nil, fmt.Errorf("error creating dbservice: %w", err)
 	}
+	state.DB = dbService
 
 	//nolint:contextcheck
 	storage, err := gtsstorage.AutoConfig()
@@ -61,6 +63,7 @@ func setupPrune(ctx context.Context) (*prune, error) {
 		dbService: dbService,
 		storage:   storage,
 		manager:   manager,
+		state:     &state,
 	}, nil
 }
 
@@ -73,5 +76,8 @@ func (p *prune) shutdown(ctx context.Context) error {
 		return fmt.Errorf("error closing dbservice: %w", err)
 	}
 
+	p.state.Caches.Stop()
+	p.state.Workers.Stop()
+
 	return nil
 }

internal/db/bundb/account.go

@@ -41,9 +41,7 @@ type accountDB struct {
 func (a *accountDB) newAccountQ(account *gtsmodel.Account) *bun.SelectQuery {
 	return a.conn.
 		NewSelect().
-		Model(account).
+		Model(account)
-		Relation("AvatarMediaAttachment").
-		Relation("HeaderMediaAttachment")
 }
 
 func (a *accountDB) GetAccountByID(ctx context.Context, id string) (*gtsmodel.Account, db.Error) {

internal/db/bundb/bundb.go

@@ -178,7 +178,8 @@ func NewBunDBService(ctx context.Context, state *state.State) (db.DB, error) {
 			conn: conn,
 		},
 		Media: &mediaDB{
-			conn: conn,
+			conn:  conn,
+			state: state,
 		},
 		Mention: &mentionDB{
 			conn:  conn,

internal/db/bundb/media.go

@@ -24,38 +24,58 @@ import (
 
 	"github.com/superseriousbusiness/gotosocial/internal/db"
 	"github.com/superseriousbusiness/gotosocial/internal/gtsmodel"
+	"github.com/superseriousbusiness/gotosocial/internal/log"
+	"github.com/superseriousbusiness/gotosocial/internal/state"
 	"github.com/uptrace/bun"
 )
 
 type mediaDB struct {
-	conn *DBConn
+	conn  *DBConn
+	state *state.State
 }
 
-func (m *mediaDB) newMediaQ(i interface{}) *bun.SelectQuery {
+func (m *mediaDB) newMediaQ(i *gtsmodel.MediaAttachment) *bun.SelectQuery {
 	return m.conn.
 		NewSelect().
-		Model(i).
+		Model(i)
-		Relation("Account")
 }
 
 func (m *mediaDB) GetAttachmentByID(ctx context.Context, id string) (*gtsmodel.MediaAttachment, db.Error) {
-	attachment := &gtsmodel.MediaAttachment{}
+	return m.getAttachment(
+		ctx,
+		"ID",
+		func(attachment *gtsmodel.MediaAttachment) error {
+			return m.newMediaQ(attachment).Where("? = ?", bun.Ident("media_attachment.id"), id).Scan(ctx)
+		},
+		id,
+	)
+}
 
-	q := m.newMediaQ(attachment).
+func (m *mediaDB) getAttachments(ctx context.Context, ids []string) ([]*gtsmodel.MediaAttachment, db.Error) {
-		Where("? = ?", bun.Ident("media_attachment.id"), id)
+	attachments := make([]*gtsmodel.MediaAttachment, 0, len(ids))
 
-	if err := q.Scan(ctx); err != nil {
+	for _, id := range ids {
-		return nil, m.conn.ProcessError(err)
+		// Attempt fetch from DB
+		attachment, err := m.GetAttachmentByID(ctx, id)
+		if err != nil {
+			log.Errorf("error getting attachment %q: %v", id, err)
+			continue
+		}
+
+		// Append attachment
+		attachments = append(attachments, attachment)
 	}
-	return attachment, nil
+
+	return attachments, nil
 }
 
 func (m *mediaDB) GetRemoteOlderThan(ctx context.Context, olderThan time.Time, limit int) ([]*gtsmodel.MediaAttachment, db.Error) {
-	attachments := []*gtsmodel.MediaAttachment{}
+	attachmentIDs := []string{}
 
 	q := m.conn.
 		NewSelect().
-		Model(&attachments).
+		TableExpr("? AS ?", bun.Ident("media_attachments"), bun.Ident("media_attachment")).
+		Column("media_attachment.id").
 		Where("? = ?", bun.Ident("media_attachment.cached"), true).
 		Where("? < ?", bun.Ident("media_attachment.created_at"), olderThan).
 		WhereGroup(" AND ", whereNotEmptyAndNotNull("media_attachment.remote_url")).
@@ -65,11 +85,11 @@ func (m *mediaDB) GetRemoteOlderThan(ctx context.Context, olderThan time.Time, l
 		q = q.Limit(limit)
 	}
 
-	if err := q.Scan(ctx); err != nil {
+	if err := q.Scan(ctx, &attachmentIDs); err != nil {
 		return nil, m.conn.ProcessError(err)
 	}
 
-	return attachments, nil
+	return m.getAttachments(ctx, attachmentIDs)
 }
 
 func (m *mediaDB) CountRemoteOlderThan(ctx context.Context, olderThan time.Time) (int, db.Error) {
@@ -90,9 +110,11 @@ func (m *mediaDB) CountRemoteOlderThan(ctx context.Context, olderThan time.Time)
 }
 
 func (m *mediaDB) GetAvatarsAndHeaders(ctx context.Context, maxID string, limit int) ([]*gtsmodel.MediaAttachment, db.Error) {
-	attachments := []*gtsmodel.MediaAttachment{}
+	attachmentIDs := []string{}
 
-	q := m.newMediaQ(&attachments).
+	q := m.conn.NewSelect().
+		TableExpr("? AS ?", bun.Ident("media_attachments"), bun.Ident("media_attachment")).
+		Column("media_attachment.id").
 		WhereGroup(" AND ", func(innerQ *bun.SelectQuery) *bun.SelectQuery {
 			return innerQ.
 				WhereOr("? = ?", bun.Ident("media_attachment.avatar"), true).
@@ -108,17 +130,20 @@ func (m *mediaDB) GetAvatarsAndHeaders(ctx context.Context, maxID string, limit
 		q = q.Limit(limit)
 	}
 
-	if err := q.Scan(ctx); err != nil {
+	if err := q.Scan(ctx, &attachmentIDs); err != nil {
 		return nil, m.conn.ProcessError(err)
 	}
 
-	return attachments, nil
+	return m.getAttachments(ctx, attachmentIDs)
 }
 
 func (m *mediaDB) GetLocalUnattachedOlderThan(ctx context.Context, olderThan time.Time, limit int) ([]*gtsmodel.MediaAttachment, db.Error) {
-	attachments := []*gtsmodel.MediaAttachment{}
+	attachmentIDs := []string{}
 
-	q := m.newMediaQ(&attachments).
+	q := m.conn.
+		NewSelect().
+		TableExpr("? AS ?", bun.Ident("media_attachments"), bun.Ident("media_attachment")).
+		Column("media_attachment.id").
 		Where("? = ?", bun.Ident("media_attachment.cached"), true).
 		Where("? = ?", bun.Ident("media_attachment.avatar"), false).
 		Where("? = ?", bun.Ident("media_attachment.header"), false).
@@ -131,11 +156,11 @@ func (m *mediaDB) GetLocalUnattachedOlderThan(ctx context.Context, olderThan tim
 		q = q.Limit(limit)
 	}
 
-	if err := q.Scan(ctx); err != nil {
+	if err := q.Scan(ctx, &attachmentIDs); err != nil {
 		return nil, m.conn.ProcessError(err)
 	}
 
-	return attachments, nil
+	return m.getAttachments(ctx, attachmentIDs)
 }
 
 func (m *mediaDB) CountLocalUnattachedOlderThan(ctx context.Context, olderThan time.Time) (int, db.Error) {
@@ -157,3 +182,15 @@ func (m *mediaDB) CountLocalUnattachedOlderThan(ctx context.Context, olderThan t
 
 	return count, nil
 }
+
+func (m *mediaDB) getAttachment(ctx context.Context, lookup string, dbQuery func(*gtsmodel.MediaAttachment) error, keyParts ...any) (*gtsmodel.MediaAttachment, db.Error) {
+	// Fetch attachment from database
+	// todo: cache this lookup
+	attachment := new(gtsmodel.MediaAttachment)
+
+	if err := dbQuery(attachment); err != nil {
+		return nil, m.conn.ProcessError(err)
+	}
+
+	return attachment, nil
+}

internal/federation/dereferencing/account.go

@@ -238,27 +238,45 @@ func (d *deref) enrichAccount(ctx context.Context, requestUser string, uri *url.
 	latestAcc.AvatarMediaAttachmentID = account.AvatarMediaAttachmentID
 	latestAcc.HeaderMediaAttachmentID = account.HeaderMediaAttachmentID
 
-	if latestAcc.AvatarRemoteURL != account.AvatarRemoteURL && latestAcc.AvatarRemoteURL != "" {
+	if latestAcc.AvatarRemoteURL != account.AvatarRemoteURL {
-		// Account avatar URL has changed; fetch up-to-date copy and use new media ID.
+		// Reset the avatar media ID (handles removed).
-		latestAcc.AvatarMediaAttachmentID, err = d.fetchRemoteAccountAvatar(ctx,
+		latestAcc.AvatarMediaAttachmentID = ""
-			transport,
+
-			latestAcc.AvatarRemoteURL,
+		if latestAcc.AvatarRemoteURL != "" {
-			latestAcc.ID,
+			// Avatar has changed to a new one, fetch up-to-date copy and use new ID.
-		)
+			latestAcc.AvatarMediaAttachmentID, err = d.fetchRemoteAccountAvatar(ctx,
-		if err != nil {
+				transport,
-			log.Errorf("error fetching remote avatar for account %s: %v", uri, err)
+				latestAcc.AvatarRemoteURL,
+				latestAcc.ID,
+			)
+			if err != nil {
+				log.Errorf("error fetching remote avatar for account %s: %v", uri, err)
+
+				// Keep old avatar for now, we'll try again in $interval.
+				latestAcc.AvatarMediaAttachmentID = account.AvatarMediaAttachmentID
+				latestAcc.AvatarRemoteURL = account.AvatarRemoteURL
+			}
 		}
 	}
 
-	if latestAcc.HeaderRemoteURL != account.HeaderRemoteURL && latestAcc.HeaderRemoteURL != "" {
+	if latestAcc.HeaderRemoteURL != account.HeaderRemoteURL {
-		// Account header URL has changed; fetch up-to-date copy and use new media ID.
+		// Reset the header media ID (handles removed).
-		latestAcc.HeaderMediaAttachmentID, err = d.fetchRemoteAccountHeader(ctx,
+		latestAcc.HeaderMediaAttachmentID = ""
-			transport,
+
-			latestAcc.HeaderRemoteURL,
+		if latestAcc.HeaderRemoteURL != "" {
-			latestAcc.ID,
+			// Header has changed to a new one, fetch up-to-date copy and use new ID.
-		)
+			latestAcc.HeaderMediaAttachmentID, err = d.fetchRemoteAccountHeader(ctx,
-		if err != nil {
+				transport,
-			log.Errorf("error fetching remote header for account %s: %v", uri, err)
+				latestAcc.HeaderRemoteURL,
+				latestAcc.ID,
+			)
+			if err != nil {
+				log.Errorf("error fetching remote header for account %s: %v", uri, err)
+
+				// Keep old header for now, we'll try again in $interval.
+				latestAcc.HeaderMediaAttachmentID = account.HeaderMediaAttachmentID
+				latestAcc.HeaderRemoteURL = account.HeaderRemoteURL
+			}
 		}
 	}
 
@@ -345,38 +363,40 @@ func (d *deref) fetchRemoteAccountAvatar(ctx context.Context, tsport transport.T
 	unlock := d.derefAvatarsMu.Lock()
 	defer unlock()
 
-	if processing, ok := d.derefAvatars[avatarURL]; ok {
+	// Look for an existing dereference in progress.
-		// we're already dereferencing it, nothing to do.
+	processing, ok := d.derefAvatars[avatarURL]
-		return processing.AttachmentID(), nil
-	}
 
-	// Set the media data function to dereference avatar from URI.
+	if !ok {
-	data := func(ctx context.Context) (io.ReadCloser, int64, error) {
+		var err error
-		return tsport.DereferenceMedia(ctx, avatarURI)
-	}
 
-	// Create new media processing request from the media manager instance.
+		// Set the media data function to dereference avatar from URI.
-	processing, err := d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
+		data := func(ctx context.Context) (io.ReadCloser, int64, error) {
-		Avatar:    func() *bool { v := false; return &v }(),
+			return tsport.DereferenceMedia(ctx, avatarURI)
-		RemoteURL: &avatarURL,
+		}
-	})
-	if err != nil {
-		return "", err
-	}
 
-	// Store media in map to mark as processing.
+		// Create new media processing request from the media manager instance.
-	d.derefAvatars[avatarURL] = processing
+		processing, err = d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
+			Avatar:    func() *bool { v := true; return &v }(),
+			RemoteURL: &avatarURL,
+		})
+		if err != nil {
+			return "", err
+		}
+
+		// Store media in map to mark as processing.
+		d.derefAvatars[avatarURL] = processing
+
+		defer func() {
+			// On exit safely remove media from map.
+			unlock := d.derefAvatarsMu.Lock()
+			delete(d.derefAvatars, avatarURL)
+			unlock()
+		}()
+	}
 
 	// Unlock map.
 	unlock()
 
-	defer func() {
-		// On exit safely remove media from map.
-		unlock := d.derefAvatarsMu.Lock()
-		delete(d.derefAvatars, avatarURL)
-		unlock()
-	}()
-
 	// Start media attachment loading (blocking call).
 	if _, err := processing.LoadAttachment(ctx); err != nil {
 		return "", err
@@ -396,38 +416,40 @@ func (d *deref) fetchRemoteAccountHeader(ctx context.Context, tsport transport.T
 	unlock := d.derefHeadersMu.Lock()
 	defer unlock()
 
-	if processing, ok := d.derefHeaders[headerURL]; ok {
+	// Look for an existing dereference in progress.
-		// we're already dereferencing it, nothing to do.
+	processing, ok := d.derefHeaders[headerURL]
-		return processing.AttachmentID(), nil
-	}
 
-	// Set the media data function to dereference header from URI.
+	if !ok {
-	data := func(ctx context.Context) (io.ReadCloser, int64, error) {
+		var err error
-		return tsport.DereferenceMedia(ctx, headerURI)
-	}
 
-	// Create new media processing request from the media manager instance.
+		// Set the media data function to dereference header from URI.
-	processing, err := d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
+		data := func(ctx context.Context) (io.ReadCloser, int64, error) {
-		Header:    func() *bool { v := true; return &v }(),
+			return tsport.DereferenceMedia(ctx, headerURI)
-		RemoteURL: &headerURL,
+		}
-	})
-	if err != nil {
-		return "", err
-	}
 
-	// Store media in map to mark as processing.
+		// Create new media processing request from the media manager instance.
-	d.derefHeaders[headerURL] = processing
+		processing, err = d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
+			Header:    func() *bool { v := true; return &v }(),
+			RemoteURL: &headerURL,
+		})
+		if err != nil {
+			return "", err
+		}
+
+		// Store media in map to mark as processing.
+		d.derefHeaders[headerURL] = processing
+
+		defer func() {
+			// On exit safely remove media from map.
+			unlock := d.derefHeadersMu.Lock()
+			delete(d.derefHeaders, headerURL)
+			unlock()
+		}()
+	}
 
 	// Unlock map.
 	unlock()
 
-	defer func() {
-		// On exit safely remove media from map.
-		unlock := d.derefHeadersMu.Lock()
-		delete(d.derefHeaders, headerURL)
-		unlock()
-	}()
-
 	// Start media attachment loading (blocking call).
 	if _, err := processing.LoadAttachment(ctx); err != nil {
 		return "", err

internal/gtsmodel/mediaattachment.go

@@ -34,7 +34,6 @@ type MediaAttachment struct {
 	Type              FileType         `validate:"oneof=Image Gifv Audio Video Unknown" bun:",nullzero,notnull"`                       // Type of file (image/gifv/audio/video)
 	FileMeta          FileMeta         `validate:"required" bun:",embed:,nullzero,notnull"`                                            // Metadata about the file
 	AccountID         string           `validate:"required,ulid" bun:"type:CHAR(26),nullzero,notnull"`                                 // To which account does this attachment belong
-	Account           *Account         `validate:"-" bun:"rel:belongs-to,join:account_id=id"`                                          // Account corresponding to accountID
 	Description       string           `validate:"-" bun:""`                                                                           // Description of the attachment (for screenreaders)
 	ScheduledStatusID string           `validate:"omitempty,ulid" bun:"type:CHAR(26),nullzero"`                                        // To which scheduled status does this attachment belong
 	Blurhash          string           `validate:"required_if=Type Image,required_if=Type Gif,required_if=Type Video" bun:",nullzero"` // What is the generated blurhash of this attachment

internal/media/prune.go

@@ -119,14 +119,24 @@ func (m *manager) PruneUnusedRemote(ctx context.Context, dry bool) (int, error)
 	for attachments, err = m.state.DB.GetAvatarsAndHeaders(ctx, maxID, selectPruneLimit); err == nil && len(attachments) != 0; attachments, err = m.state.DB.GetAvatarsAndHeaders(ctx, maxID, selectPruneLimit) {
 		maxID = attachments[len(attachments)-1].ID // use the id of the last attachment in the slice as the next 'maxID' value
 
-		// Prune each attachment that meets one of the following criteria:
-		// - Has no owning account in the database.
-		// - Is a header but isn't the owning account's current header.
-		// - Is an avatar but isn't the owning account's current avatar.
 		for _, attachment := range attachments {
-			if attachment.Account == nil ||
+			// Retrieve owning account if possible.
-				(*attachment.Header && attachment.ID != attachment.Account.HeaderMediaAttachmentID) ||
+			var account *gtsmodel.Account
-				(*attachment.Avatar && attachment.ID != attachment.Account.AvatarMediaAttachmentID) {
+			if accountID := attachment.AccountID; accountID != "" {
+				account, err = m.state.DB.GetAccountByID(ctx, attachment.AccountID)
+				if err != nil && !errors.Is(err, db.ErrNoEntries) {
+					// Only return on a real error.
+					return 0, fmt.Errorf("PruneUnusedRemote: error fetching account with id %s: %w", accountID, err)
+				}
+			}
+
+			// Prune each attachment that meets one of the following criteria:
+			// - Has no owning account in the database.
+			// - Is a header but isn't the owning account's current header.
+			// - Is an avatar but isn't the owning account's current avatar.
+			if account == nil ||
+				(*attachment.Header && attachment.ID != account.HeaderMediaAttachmentID) ||
+				(*attachment.Avatar && attachment.ID != account.AvatarMediaAttachmentID) {
 				if err := f(ctx, attachment); err != nil {
 					return totalPruned, err
 				}