From 1439042104bd3fd66c51299913f201ff691460ac Mon Sep 17 00:00:00 2001
From: kim <89579420+NyaaaWhatsUpDoc@users.noreply.github.com>
Date: Sat, 13 Apr 2024 10:36:10 +0100
Subject: [PATCH] [performance] update GetAccountsByIDs() to use the new multi
 cache loader endpoint (#2828)

* update GetAccountsByIDs() to use the new multi cache loader endpoint

* fix broken import
---
 internal/db/bundb/account.go | 53 +++++++++++++++++++++++++++---------
 1 file changed, 40 insertions(+), 13 deletions(-)

diff --git a/internal/db/bundb/account.go b/internal/db/bundb/account.go
index 870c2ff55..1ecf28e42 100644
--- a/internal/db/bundb/account.go
+++ b/internal/db/bundb/account.go
@@ -20,6 +20,7 @@
 import (
 	"context"
 	"errors"
+	"slices"
 	"strings"
 	"time"
 
@@ -56,23 +57,49 @@ func(account *gtsmodel.Account) error {
 }
 
 func (a *accountDB) GetAccountsByIDs(ctx context.Context, ids []string) ([]*gtsmodel.Account, error) {
-	accounts := make([]*gtsmodel.Account, 0, len(ids))
+	// Load all input account IDs via cache loader callback.
+	accounts, err := a.state.Caches.GTS.Account.LoadIDs("ID",
+		ids,
+		func(uncached []string) ([]*gtsmodel.Account, error) {
+			// Preallocate expected length of uncached accounts.
+			accounts := make([]*gtsmodel.Account, 0, len(uncached))
 
-	for _, id := range ids {
-		// Attempt to fetch account from DB.
-		account, err := a.GetAccountByID(
-			gtscontext.SetBarebones(ctx),
-			id,
-		)
-		if err != nil {
-			log.Errorf(ctx, "error getting account %q: %v", id, err)
-			continue
-		}
+			// Perform database query scanning
+			// the remaining (uncached) account IDs.
+			if err := a.db.NewSelect().
+				Model(&accounts).
+				Where("? IN (?)", bun.Ident("id"), bun.In(uncached)).
+				Scan(ctx); err != nil {
+				return nil, err
+			}
 
-		// Append account to return slice.
-		accounts = append(accounts, account)
+			return accounts, nil
+		},
+	)
+	if err != nil {
+		return nil, err
 	}
 
+	// Reorder the statuses by their
+	// IDs to ensure in correct order.
+	getID := func(a *gtsmodel.Account) string { return a.ID }
+	util.OrderBy(accounts, ids, getID)
+
+	if gtscontext.Barebones(ctx) {
+		// no need to fully populate.
+		return accounts, nil
+	}
+
+	// Populate all loaded accounts, removing those we fail to
+	// populate (removes needing so many nil checks everywhere).
+	accounts = slices.DeleteFunc(accounts, func(account *gtsmodel.Account) bool {
+		if err := a.PopulateAccount(ctx, account); err != nil {
+			log.Errorf(ctx, "error populating account %s: %v", account.ID, err)
+			return true
+		}
+		return false
+	})
+
 	return accounts, nil
 }