[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

@@ -179,6 +179,7 @@ func NewBunDBService(ctx context.Context, state *state.State) (db.DB, error) {
 		},
 		Media: &mediaDB{
 			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
+	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
 		}
-	return attachment, nil
+
+		// Append attachment
+		attachments = append(attachments, attachment)
+	}
+
+	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,8 +238,12 @@ 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 = ""
+
+		if latestAcc.AvatarRemoteURL != "" {
+			// Avatar has changed to a new one, fetch up-to-date copy and use new ID.
 			latestAcc.AvatarMediaAttachmentID, err = d.fetchRemoteAccountAvatar(ctx,
 				transport,
 				latestAcc.AvatarRemoteURL,
@@ -247,11 +251,20 @@ func (d *deref) enrichAccount(ctx context.Context, requestUser string, uri *url.
 			)
 			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 = ""
+
+		if latestAcc.HeaderRemoteURL != "" {
+			// Header has changed to a new one, fetch up-to-date copy and use new ID.
 			latestAcc.HeaderMediaAttachmentID, err = d.fetchRemoteAccountHeader(ctx,
 				transport,
 				latestAcc.HeaderRemoteURL,
@@ -259,6 +272,11 @@ func (d *deref) enrichAccount(ctx context.Context, requestUser string, uri *url.
 			)
 			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,10 +363,11 @@ 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
+
-	}
+	if !ok {
+		var err error
 
 		// Set the media data function to dereference avatar from URI.
 		data := func(ctx context.Context) (io.ReadCloser, int64, error) {
@@ -356,8 +375,8 @@ func (d *deref) fetchRemoteAccountAvatar(ctx context.Context, tsport transport.T
 		}
 
 		// Create new media processing request from the media manager instance.
-	processing, err := d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
+		processing, err = d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
-		Avatar:    func() *bool { v := false; return &v }(),
+			Avatar:    func() *bool { v := true; return &v }(),
 			RemoteURL: &avatarURL,
 		})
 		if err != nil {
@@ -367,15 +386,16 @@ func (d *deref) fetchRemoteAccountAvatar(ctx context.Context, tsport transport.T
 		// Store media in map to mark as processing.
 		d.derefAvatars[avatarURL] = processing
 
-	// Unlock map.
-	unlock()
-
 		defer func() {
 			// On exit safely remove media from map.
 			unlock := d.derefAvatarsMu.Lock()
 			delete(d.derefAvatars, avatarURL)
 			unlock()
 		}()
+	}
+
+	// Unlock map.
+	unlock()
 
 	// Start media attachment loading (blocking call).
 	if _, err := processing.LoadAttachment(ctx); err != nil {
@@ -396,10 +416,11 @@ 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
+
-	}
+	if !ok {
+		var err error
 
 		// Set the media data function to dereference header from URI.
 		data := func(ctx context.Context) (io.ReadCloser, int64, error) {
@@ -407,7 +428,7 @@ func (d *deref) fetchRemoteAccountHeader(ctx context.Context, tsport transport.T
 		}
 
 		// Create new media processing request from the media manager instance.
-	processing, err := d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
+		processing, err = d.mediaManager.PreProcessMedia(ctx, data, nil, accountID, &media.AdditionalMediaInfo{
 			Header:    func() *bool { v := true; return &v }(),
 			RemoteURL: &headerURL,
 		})
@@ -418,15 +439,16 @@ func (d *deref) fetchRemoteAccountHeader(ctx context.Context, tsport transport.T
 		// Store media in map to mark as processing.
 		d.derefHeaders[headerURL] = processing
 
-	// Unlock map.
-	unlock()
-
 		defer func() {
 			// On exit safely remove media from map.
 			unlock := d.derefHeadersMu.Lock()
 			delete(d.derefHeaders, headerURL)
 			unlock()
 		}()
+	}
+
+	// Unlock map.
+	unlock()
 
 	// Start media attachment loading (blocking call).
 	if _, err := processing.LoadAttachment(ctx); err != nil {

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
 
+		for _, attachment := range attachments {
+			// Retrieve owning account if possible.
+			var account *gtsmodel.Account
+			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.
-		for _, attachment := range attachments {
+			if account == nil ||
-			if attachment.Account == nil ||
+				(*attachment.Header && attachment.ID != account.HeaderMediaAttachmentID) ||
-				(*attachment.Header && attachment.ID != attachment.Account.HeaderMediaAttachmentID) ||
+				(*attachment.Avatar && attachment.ID != account.AvatarMediaAttachmentID) {
-				(*attachment.Avatar && attachment.ID != attachment.Account.AvatarMediaAttachmentID) {
 				if err := f(ctx, attachment); err != nil {
 					return totalPruned, err
 				}