Skip to content

federation/cache: Add Context parameter to <X>Cache #503

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
timedoutuk wants to merge 3 commits into
base: main
Choose a base branch
from
+36 −28
Open

federation/cache: Add Context parameter to <X>Cache #503

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a7b6d56
federation/cache: Add Context parameter to <X>Cache
timedoutuk May 17, 2026
064f629
Fix build errors
timedoutuk May 17, 2026
17b1e8d
Allow returning an error in StoreResolution, StoreKeys, and ShouldReQ…
timedoutuk May 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 25 additions & 20 deletions federation/cache.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
package federation

import (
"context"
"errors"
"fmt"
"math"
Expand All @@ -16,17 +17,17 @@ import (

// ResolutionCache is an interface for caching resolved server names.
type ResolutionCache interface {
StoreResolution(*ResolvedServerName)
StoreResolution(context.Context, *ResolvedServerName) error
// LoadResolution loads a resolved server name from the cache.
// Expired entries MUST NOT be returned.
LoadResolution(serverName string) (*ResolvedServerName, error)
LoadResolution(ctx context.Context, serverName string) (*ResolvedServerName, error)
}

type KeyCache interface {
StoreKeys(*ServerKeyResponse)
StoreFetchError(serverName string, err error)
ShouldReQuery(serverName string) bool
LoadKeys(serverName string) (*ServerKeyResponse, error)
StoreKeys(context.Context, *ServerKeyResponse) error
StoreFetchError(ctx context.Context, serverName string, err error)
ShouldReQuery(ctx context.Context, serverName string) (bool, error)
LoadKeys(ctx context.Context, serverName string) (*ServerKeyResponse, error)
}

type InMemoryCache struct {
Expand Down Expand Up @@ -55,13 +56,14 @@ func NewInMemoryCache() *InMemoryCache {
}
}

func (c *InMemoryCache) StoreResolution(resolution *ResolvedServerName) {
func (c *InMemoryCache) StoreResolution(_ context.Context, resolution *ResolvedServerName) error {
c.resolutionsLock.Lock()
defer c.resolutionsLock.Unlock()
c.resolutions[resolution.ServerName] = resolution
return nil
}

func (c *InMemoryCache) LoadResolution(serverName string) (*ResolvedServerName, error) {
func (c *InMemoryCache) LoadResolution(_ context.Context, serverName string) (*ResolvedServerName, error) {
c.resolutionsLock.RLock()
defer c.resolutionsLock.RUnlock()
resolution, ok := c.resolutions[serverName]
Expand All @@ -71,11 +73,12 @@ func (c *InMemoryCache) LoadResolution(serverName string) (*ResolvedServerName,
return resolution, nil
}

func (c *InMemoryCache) StoreKeys(keys *ServerKeyResponse) {
func (c *InMemoryCache) StoreKeys(_ context.Context, keys *ServerKeyResponse) error {
c.keysLock.Lock()
defer c.keysLock.Unlock()
c.keys[keys.ServerName] = keys
delete(c.lastError, keys.ServerName)
return nil
}

type resolutionErrorCache struct {
Expand All @@ -93,7 +96,7 @@ func (rec *resolutionErrorCache) ShouldRetry() bool {

var ErrRecentKeyQueryFailed = errors.New("last retry was too recent")

func (c *InMemoryCache) LoadKeys(serverName string) (*ServerKeyResponse, error) {
func (c *InMemoryCache) LoadKeys(_ context.Context, serverName string) (*ServerKeyResponse, error) {
c.keysLock.RLock()
defer c.keysLock.RUnlock()
keys, ok := c.keys[serverName]
Expand All @@ -112,7 +115,7 @@ func (c *InMemoryCache) LoadKeys(serverName string) (*ServerKeyResponse, error)
return keys, nil
}

func (c *InMemoryCache) StoreFetchError(serverName string, err error) {
func (c *InMemoryCache) StoreFetchError(_ context.Context, serverName string, err error) {
c.keysLock.Lock()
defer c.keysLock.Unlock()
errorCache, ok := c.lastError[serverName]
Expand All @@ -125,25 +128,27 @@ func (c *InMemoryCache) StoreFetchError(serverName string, err error) {
}
}

func (c *InMemoryCache) ShouldReQuery(serverName string) bool {
func (c *InMemoryCache) ShouldReQuery(_ context.Context, serverName string) (bool, error) {
c.keysLock.Lock()
defer c.keysLock.Unlock()
lastQuery, ok := c.lastReQueryAt[serverName]
if ok && time.Since(lastQuery) < c.MinKeyRefetchDelay {
return false
return false, nil
}
c.lastReQueryAt[serverName] = time.Now()
return true
return true, nil
}

type noopCache struct{}

func (*noopCache) StoreKeys(_ *ServerKeyResponse) {}
func (*noopCache) LoadKeys(_ string) (*ServerKeyResponse, error) { return nil, nil }
func (*noopCache) StoreFetchError(_ string, _ error) {}
func (*noopCache) ShouldReQuery(_ string) bool { return true }
func (*noopCache) StoreResolution(_ *ResolvedServerName) {}
func (*noopCache) LoadResolution(_ string) (*ResolvedServerName, error) { return nil, nil }
func (*noopCache) StoreKeys(_ context.Context, _ *ServerKeyResponse) error { return nil }
func (*noopCache) LoadKeys(_ context.Context, _ string) (*ServerKeyResponse, error) { return nil, nil }
func (*noopCache) StoreFetchError(_ context.Context, _ string, _ error) {}
func (*noopCache) ShouldReQuery(_ context.Context, _ string) (bool, error) { return true, nil }
func (*noopCache) StoreResolution(_ context.Context, _ *ResolvedServerName) error { return nil }
func (*noopCache) LoadResolution(_ context.Context, _ string) (*ResolvedServerName, error) {
return nil, nil
}

var (
_ ResolutionCache = (*noopCache)(nil)
Expand Down
4 changes: 2 additions & 2 deletions federation/httpclient.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,15 +80,15 @@ func (srt *ServerResolvingTransport) resolve(ctx context.Context, serverName str

lock.Lock()
defer lock.Unlock()
res, err := srt.cache.LoadResolution(serverName)
res, err := srt.cache.LoadResolution(ctx, serverName)
if err != nil {
return nil, fmt.Errorf("failed to read cache: %w", err)
} else if res != nil {
return res, nil
} else if res, err = ResolveServerName(ctx, serverName, srt.ResolveOpts); err != nil {
return nil, err
} else {
srt.cache.StoreResolution(res)
srt.cache.StoreResolution(ctx, res)
return res, nil
}
}
Expand Down
2 changes: 1 addition & 1 deletion federation/keyserver.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ func (ks *KeyServer) GetQueryKeys(w http.ResponseWriter, r *http.Request) {
resp.ServerKeys = append(resp.ServerKeys, key.GenerateKeyResponse(serverName, nil))
}
} else if ks.OtherKeys != nil {
otherKey, err := ks.OtherKeys.LoadKeys(serverName)
otherKey, err := ks.OtherKeys.LoadKeys(r.Context(), serverName)
if err != nil {
mautrix.MUnknown.WithMessage("Failed to load keys from cache").Write(w)
return
Expand Down
13 changes: 8 additions & 5 deletions federation/serverauth.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ func ParseXMatrixAuth(auth string) (xma XMatrixAuth) {
}

func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, keyID id.KeyID) (*ServerKeyResponse, error) {
res, err := sa.Keys.LoadKeys(serverName)
res, err := sa.Keys.LoadKeys(ctx, serverName)
if err != nil {
return nil, fmt.Errorf("failed to read cache: %w", err)
} else if res.HasKey(keyID) {
Expand All @@ -120,14 +120,15 @@ func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, k

lock.Lock()
defer lock.Unlock()
res, err = sa.Keys.LoadKeys(serverName)
res, err = sa.Keys.LoadKeys(ctx, serverName)
if err != nil {
return nil, fmt.Errorf("failed to read cache: %w", err)
} else if res != nil {
if res.HasKey(keyID) {
return res, nil
} else if !sa.Keys.ShouldReQuery(serverName) {
} else if shouldQuery, err := sa.Keys.ShouldReQuery(ctx, serverName); !shouldQuery {
zerolog.Ctx(ctx).Trace().
Err(err).
Str("server_name", serverName).
Stringer("key_id", keyID).
Msg("Not sending key request for missing key ID, last query was too recent")
Expand All @@ -136,10 +137,12 @@ func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, k
}
res, err = sa.Client.ServerKeys(ctx, serverName)
if err != nil {
sa.Keys.StoreFetchError(serverName, err)
sa.Keys.StoreFetchError(ctx, serverName, err)
return nil, err
}
if err := sa.Keys.StoreKeys(ctx, res); err != nil {
return nil, err
}
sa.Keys.StoreKeys(res)
return res, nil
}

Expand Down
Loading