[performance] cache account db lookups by public key URI (#795)

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

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

6 files changed, 66 insertions, 19 deletions

diff --git a/internal/cache/account.go b/internal/cache/account.go
index 1f958ebb8..5ba97c7d8 100644
--- a/internal/cache/account.go
+++ b/internal/cache/account.go
@@ -37,6 +37,7 @@ func NewAccountCache() *AccountCache {
 		RegisterLookups: func(lm *cache.LookupMap[string, string]) {
 			lm.RegisterLookup("uri")
 			lm.RegisterLookup("url")
+			lm.RegisterLookup("pubkeyid")
 			lm.RegisterLookup("usernamedomain")
 		},
 
@@ -47,6 +48,7 @@ func NewAccountCache() *AccountCache {
 			if url := acc.URL; url != "" {
 				lm.Set("url", url, acc.ID)
 			}
+			lm.Set("pubkeyid", acc.PublicKeyURI, acc.ID)
 			lm.Set("usernamedomain", usernameDomainKey(acc.Username, acc.Domain), acc.ID)
 		},
 
@@ -57,6 +59,7 @@ func NewAccountCache() *AccountCache {
 			if url := acc.URL; url != "" {
 				lm.Delete("url", url)
 			}
+			lm.Delete("pubkeyid", acc.PublicKeyURI)
 			lm.Delete("usernamedomain", usernameDomainKey(acc.Username, acc.Domain))
 		},
 	})
@@ -80,10 +83,16 @@ func (c *AccountCache) GetByURI(uri string) (*gtsmodel.Account, bool) {
 	return c.cache.GetBy("uri", uri)
 }
 
+// GettByUsernameDomain attempts to fetch an account from the cache by its username@domain combo (or just username), you will receive a copy for thread-safety.
 func (c *AccountCache) GetByUsernameDomain(username string, domain string) (*gtsmodel.Account, bool) {
 	return c.cache.GetBy("usernamedomain", usernameDomainKey(username, domain))
 }
 
+// GetByPubkeyID attempts to fetch an account from the cache by its public key URI (ID), you will receive a copy for thread-safety.
+func (c *AccountCache) GetByPubkeyID(id string) (*gtsmodel.Account, bool) {
+	return c.cache.GetBy("pubkeyid", id)
+}
+
 // Put places a account in the cache, ensuring that the object place is a copy for thread-safety
 func (c *AccountCache) Put(account *gtsmodel.Account) {
 	if account == nil || account.ID == "" {
diff --git a/internal/cache/account_test.go b/internal/cache/account_test.go
index a6d3c6b7d..d373e5f1d 100644
--- a/internal/cache/account_test.go
+++ b/internal/cache/account_test.go
@@ -19,6 +19,7 @@
 package cache_test
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/suite"
@@ -59,19 +60,23 @@ func (suite *AccountCacheTestSuite) TestAccountCache() {
 		// Check we can retrieve
 		check, ok = suite.cache.GetByID(account.ID)
 		if !ok && !accountIs(account, check) {
-			suite.Fail("Failed to fetch expected account with ID: %s", account.ID)
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with ID: %s", account.ID))
 		}
 		check, ok = suite.cache.GetByURI(account.URI)
 		if account.URI != "" && !ok && !accountIs(account, check) {
-			suite.Fail("Failed to fetch expected account with URI: %s", account.URI)
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with URI: %s", account.URI))
 		}
 		check, ok = suite.cache.GetByURL(account.URL)
 		if account.URL != "" && !ok && !accountIs(account, check) {
-			suite.Fail("Failed to fetch expected account with URL: %s", account.URL)
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with URL: %s", account.URL))
+		}
+		check, ok = suite.cache.GetByPubkeyID(account.PublicKeyURI)
+		if account.PublicKeyURI != "" && !ok && !accountIs(account, check) {
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with public key URI: %s", account.PublicKeyURI))
 		}
 		check, ok = suite.cache.GetByUsernameDomain(account.Username, account.Domain)
 		if !ok && !accountIs(account, check) {
-			suite.Fail("Failed to fetch expected account with username/domain: %s/%s", account.Username, account.Domain)
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with username/domain: %s/%s", account.Username, account.Domain))
 		}
 	}
 }
@@ -81,5 +86,11 @@ func TestAccountCache(t *testing.T) {
 }
 
 func accountIs(account1, account2 *gtsmodel.Account) bool {
-	return account1.ID == account2.ID && account1.URI == account2.URI && account1.URL == account2.URL
+	if account1 == nil || account2 == nil {
+		return account1 == account2
+	}
+	return account1.ID == account2.ID &&
+		account1.URI == account2.URI &&
+		account1.URL == account2.URL &&
+		account1.PublicKeyURI == account2.PublicKeyURI
 }
diff --git a/internal/cache/status_test.go b/internal/cache/status_test.go
index 8b0621182..c1c4173fb 100644
--- a/internal/cache/status_test.go
+++ b/internal/cache/status_test.go
@@ -19,6 +19,7 @@
 package cache_test
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/suite"
@@ -59,15 +60,15 @@ func (suite *StatusCacheTestSuite) TestStatusCache() {
 		// Check we can retrieve
 		check, ok = suite.cache.GetByID(status.ID)
 		if !ok && !statusIs(status, check) {
-			suite.Fail("Failed to fetch expected account with ID: %s", status.ID)
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with ID: %s", status.ID))
 		}
 		check, ok = suite.cache.GetByURI(status.URI)
 		if status.URI != "" && !ok && !statusIs(status, check) {
-			suite.Fail("Failed to fetch expected account with URI: %s", status.URI)
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with URI: %s", status.URI))
 		}
 		check, ok = suite.cache.GetByURL(status.URL)
 		if status.URL != "" && !ok && !statusIs(status, check) {
-			suite.Fail("Failed to fetch expected account with URL: %s", status.URL)
+			suite.Fail(fmt.Sprintf("Failed to fetch expected account with URL: %s", status.URL))
 		}
 	}
 }
@@ -103,5 +104,10 @@ func TestStatusCache(t *testing.T) {
 }
 
 func statusIs(status1, status2 *gtsmodel.Status) bool {
-	return status1.ID == status2.ID && status1.URI == status2.URI && status1.URL == status2.URL
+	if status1 == nil || status2 == nil {
+		return status1 == status2
+	}
+	return status1.ID == status2.ID &&
+		status1.URI == status2.URI &&
+		status1.URL == status2.URL
 }
diff --git a/internal/db/account.go b/internal/db/account.go
index 54e24784f..04c76777f 100644
--- a/internal/db/account.go
+++ b/internal/db/account.go
@@ -39,6 +39,9 @@ type Account interface {
 	// GetAccountByUsernameDomain returns one account with the given username and domain, or an error if something goes wrong.
 	GetAccountByUsernameDomain(ctx context.Context, username string, domain string) (*gtsmodel.Account, Error)
 
+	// GetAccountByPubkeyID returns one account with the given public key URI (ID), or an error if something goes wrong.
+	GetAccountByPubkeyID(ctx context.Context, id string) (*gtsmodel.Account, Error)
+
 	// UpdateAccount updates one account by ID.
 	UpdateAccount(ctx context.Context, account *gtsmodel.Account) (*gtsmodel.Account, Error)
 
diff --git a/internal/db/bundb/account.go b/internal/db/bundb/account.go
index 7bee375e3..23030c612 100644
--- a/internal/db/bundb/account.go
+++ b/internal/db/bundb/account.go
@@ -106,6 +106,18 @@ func (a *accountDB) GetAccountByUsernameDomain(ctx context.Context, username str
 	)
 }
 
+func (a *accountDB) GetAccountByPubkeyID(ctx context.Context, id string) (*gtsmodel.Account, db.Error) {
+	return a.getAccount(
+		ctx,
+		func() (*gtsmodel.Account, bool) {
+			return a.cache.GetByPubkeyID(id)
+		},
+		func(account *gtsmodel.Account) error {
+			return a.newAccountQ(account).Where("account.public_key_uri = ?", id).Scan(ctx)
+		},
+	)
+}
+
 func (a *accountDB) getAccount(ctx context.Context, cacheGet func() (*gtsmodel.Account, bool), dbQuery func(*gtsmodel.Account) error) (*gtsmodel.Account, db.Error) {
 	// Attempt to fetch cached account
 	account, cached := cacheGet()
diff --git a/internal/federation/authenticate.go b/internal/federation/authenticate.go
index cb693892e..ab93fbeaf 100644
--- a/internal/federation/authenticate.go
+++ b/internal/federation/authenticate.go
@@ -34,23 +34,22 @@ import (
 	"github.com/superseriousbusiness/activity/streams/vocab"
 	"github.com/superseriousbusiness/gotosocial/internal/ap"
 	"github.com/superseriousbusiness/gotosocial/internal/config"
-	"github.com/superseriousbusiness/gotosocial/internal/db"
 	"github.com/superseriousbusiness/gotosocial/internal/gtserror"
 	"github.com/superseriousbusiness/gotosocial/internal/gtsmodel"
 	"github.com/superseriousbusiness/gotosocial/internal/log"
 )
 
 /*
-	publicKeyer is BORROWED DIRECTLY FROM https://github.com/go-fed/apcore/blob/master/ap/util.go
-	Thank you @cj@mastodon.technology ! <3
+publicKeyer is BORROWED DIRECTLY FROM https://github.com/go-fed/apcore/blob/master/ap/util.go
+Thank you @cj@mastodon.technology ! <3
 */
 type publicKeyer interface {
 	GetW3IDSecurityV1PublicKey() vocab.W3IDSecurityV1PublicKeyProperty
 }
 
 /*
-	getPublicKeyFromResponse is adapted from https://github.com/go-fed/apcore/blob/master/ap/util.go
-	Thank you @cj@mastodon.technology ! <3
+getPublicKeyFromResponse is adapted from https://github.com/go-fed/apcore/blob/master/ap/util.go
+Thank you @cj@mastodon.technology ! <3
 */
 func getPublicKeyFromResponse(c context.Context, b []byte, keyID *url.URL) (vocab.W3IDSecurityV1PublicKey, error) {
 	m := make(map[string]interface{})
@@ -161,26 +160,33 @@ func (f *federator) AuthenticateFederatedRequest(ctx context.Context, requestedU
 		return nil, errWithCode
 	}
 
-	requestingRemoteAccount := &gtsmodel.Account{}
-	requestingLocalAccount := &gtsmodel.Account{}
-	requestingHost := requestingPublicKeyID.Host
+	var (
+		requestingLocalAccount  *gtsmodel.Account
+		requestingRemoteAccount *gtsmodel.Account
+		requestingHost          = requestingPublicKeyID.Host
+	)
+
 	if host := config.GetHost(); strings.EqualFold(requestingHost, host) {
 		// LOCAL ACCOUNT REQUEST
 		// the request is coming from INSIDE THE HOUSE so skip the remote dereferencing
 		log.Tracef("proceeding without dereference for local public key %s", requestingPublicKeyID)
-		if err := f.db.GetWhere(ctx, []db.Where{{Key: "public_key_uri", Value: requestingPublicKeyID.String()}}, requestingLocalAccount); err != nil {
+
+		requestingLocalAccount, err = f.db.GetAccountByPubkeyID(ctx, requestingPublicKeyID.String())
+		if err != nil {
 			errWithCode := gtserror.NewErrorInternalError(fmt.Errorf("couldn't get account with public key uri %s from the database: %s", requestingPublicKeyID.String(), err))
 			log.Debug(errWithCode)
 			return nil, errWithCode
 		}
+
 		publicKey = requestingLocalAccount.PublicKey
+
 		pkOwnerURI, err = url.Parse(requestingLocalAccount.URI)
 		if err != nil {
 			errWithCode := gtserror.NewErrorBadRequest(err, fmt.Sprintf("couldn't parse public key owner URL %s", requestingLocalAccount.URI))
 			log.Debug(errWithCode)
 			return nil, errWithCode
 		}
-	} else if err := f.db.GetWhere(ctx, []db.Where{{Key: "public_key_uri", Value: requestingPublicKeyID.String()}}, requestingRemoteAccount); err == nil {
+	} else if requestingRemoteAccount, err = f.db.GetAccountByPubkeyID(ctx, requestingPublicKeyID.String()); err == nil {
 		// REMOTE ACCOUNT REQUEST WITH KEY CACHED LOCALLY
 		// this is a remote account and we already have the public key for it so use that
 		log.Tracef("proceeding without dereference for cached public key %s", requestingPublicKeyID)