add cache wrapper

only cache token if not to rotate

Co-authored-by: Kalle Persson <kalle.persson@grafana.com>

anticipate next rotation

Co-authored-by: ievaVasiljeva <ieva.vasiljeva@grafana.com>
Co-authored-by: Kalle Persson <kalle.persson@grafana.com>
Co-authored-by: Gabriel MABILLE <gamab@users.noreply.github.com>

pkg/services/auth/auth_token.go

@@ -4,10 +4,12 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
+	"errors"
 	"net"
 	"strings"
 	"time"
 
+	"github.com/grafana/grafana/pkg/infra/remotecache"
 	"github.com/grafana/grafana/pkg/infra/serverlock"
 
 	"github.com/grafana/grafana/pkg/infra/log"
@@ -22,16 +24,25 @@ const ServiceName = "UserAuthTokenService"
 
 var getTime = time.Now
 
-const urgentRotateTime = 1 * time.Minute
+const (
+	ttl              = 15 * time.Second
+	urgentRotateTime = 1 * time.Minute
+)
 
-func ProvideUserAuthTokenService(sqlStore *sqlstore.SQLStore, serverLockService *serverlock.ServerLockService,
+func ProvideUserAuthTokenService(sqlStore *sqlstore.SQLStore,
+	serverLockService *serverlock.ServerLockService,
+	remoteCache *remotecache.RemoteCache,
 	cfg *setting.Cfg) *UserAuthTokenService {
 	s := &UserAuthTokenService{
 		SQLStore:          sqlStore,
 		ServerLockService: serverLockService,
 		Cfg:               cfg,
 		log:               log.New("auth"),
+		remoteCache:       remoteCache,
 	}
+
+	remotecache.Register(models.UserToken{})
+
 	return s
 }
 
@@ -40,6 +51,7 @@ type UserAuthTokenService struct {
 	ServerLockService *serverlock.ServerLockService
 	Cfg               *setting.Cfg
 	log               log.Logger
+	remoteCache       *remotecache.RemoteCache
 }
 
 type ActiveAuthTokenService struct {
@@ -118,7 +130,59 @@ func (s *UserAuthTokenService) CreateToken(ctx context.Context, user *user.User,
 	return &userToken, err
 }
 
+func (s *UserAuthTokenService) lookupTokenWithCache(ctx context.Context, unhashedToken string) (*models.UserToken, error) {
+	hashedToken := hashToken(unhashedToken)
+	cacheKey := "auth_token:" + hashedToken
+
+	session, errCache := s.remoteCache.Get(ctx, cacheKey)
+	if errCache == nil {
+		token := session.(models.UserToken)
+		return &token, nil
+	} else {
+		if errors.Is(errCache, remotecache.ErrCacheItemNotFound) {
+			s.log.Debug("user auth token not found in cache",
+				"cacheKey", cacheKey)
+		} else {
+			s.log.Warn("failed to get user auth token from cache",
+				"cacheKey", cacheKey, "error", errCache)
+		}
+	}
+
+	token, err := s.lookupToken(ctx, unhashedToken)
+	if err != nil {
+		return nil, err
+	}
+
+	// only cache tokens if they are not going to be rotated soon
+	nextRotation := time.Unix(token.RotatedAt, 0).Add(time.Duration(s.Cfg.TokenRotationIntervalMinutes) * time.Minute)
+	now := getTime()
+	if now.Before(nextRotation) {
+		ttlSet := min(ttl, nextRotation.Sub(now))
+		if err := s.remoteCache.Set(ctx, cacheKey, *token, ttlSet); err != nil {
+			s.log.Warn("could not cache token", "error", err, "cacheKey", cacheKey, "userId", token.UserId)
+		}
+	}
+
+	return token, nil
+
+}
+
+func min(x, y time.Duration) time.Duration {
+	if x < y {
+		return x
+	}
+	return y
+}
+
 func (s *UserAuthTokenService) LookupToken(ctx context.Context, unhashedToken string) (*models.UserToken, error) {
+	// replace with feature flag
+	if true {
+		return s.lookupTokenWithCache(ctx, unhashedToken)
+	}
+	return s.lookupToken(ctx, unhashedToken)
+}
+
+func (s *UserAuthTokenService) lookupToken(ctx context.Context, unhashedToken string) (*models.UserToken, error) {
 	hashedToken := hashToken(unhashedToken)
 	var model userAuthToken
 	var exists bool