OpenFlare/atsf_server/service/node_agent_token_cache.go

package service

import (
	"atsflare/model"
	"errors"
	"time"

	ristretto "github.com/dgraph-io/ristretto/v2"
	"gorm.io/gorm"
)

const (
	agentTokenPositiveCacheTTL = 2 * time.Minute
	agentTokenNegativeCacheTTL = 10 * time.Minute
	agentTokenNegativeCacheCap = 10000
)

type cachedAgentNode struct {
	node      *model.Node
	expiresAt time.Time
}

type cachedMissingAgentToken struct {
	expiresAt time.Time
}

type agentTokenAuthCache struct {
	positive        *ristretto.Cache[string, cachedAgentNode]
	negative        *ristretto.Cache[string, cachedMissingAgentToken]
	now             func() time.Time
	loadNodeByToken func(string) (*model.Node, error)
}

var nodeAgentTokenCache = newAgentTokenAuthCache()

func newAgentTokenAuthCache() *agentTokenAuthCache {
	return &agentTokenAuthCache{
		positive: mustNewAgentTokenPositiveCache(),
		negative: mustNewAgentTokenNegativeCache(),
		now:      time.Now,
		loadNodeByToken: func(token string) (*model.Node, error) {
			return model.GetNodeByAgentToken(token)
		},
	}
}

func mustNewAgentTokenPositiveCache() *ristretto.Cache[string, cachedAgentNode] {
	cache, err := ristretto.NewCache(&ristretto.Config[string, cachedAgentNode]{
		NumCounters: 1e5,
		MaxCost:     2e4,
		BufferItems: 64,
	})
	if err != nil {
		panic(err)
	}
	return cache
}

func mustNewAgentTokenNegativeCache() *ristretto.Cache[string, cachedMissingAgentToken] {
	cache, err := ristretto.NewCache(&ristretto.Config[string, cachedMissingAgentToken]{
		NumCounters: 1e5,
		MaxCost:     agentTokenNegativeCacheCap,
		BufferItems: 64,
	})
	if err != nil {
		panic(err)
	}
	return cache
}

func (c *agentTokenAuthCache) authenticate(token string) (*model.Node, error) {
	now := c.now()
	if node, ok := c.getNode(token, now); ok {
		return node, nil
	}
	if c.isMissing(token, now) {
		return nil, gorm.ErrRecordNotFound
	}

	node, err := c.loadNodeByToken(token)
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			c.storeMissing(token, now.Add(agentTokenNegativeCacheTTL))
		}
		return nil, err
	}

	c.storeNode(token, node, now.Add(agentTokenPositiveCacheTTL))
	return cloneCachedNode(node), nil
}

func (c *agentTokenAuthCache) getNode(token string, now time.Time) (*model.Node, bool) {
	entry, ok := c.positive.Get(token)
	if !ok {
		return nil, false
	}
	if now.After(entry.expiresAt) {
		c.positive.Del(token)
		return nil, false
	}
	return cloneCachedNode(entry.node), true
}

func (c *agentTokenAuthCache) isMissing(token string, now time.Time) bool {
	entry, ok := c.negative.Get(token)
	if !ok {
		return false
	}
	if now.After(entry.expiresAt) {
		c.negative.Del(token)
		return false
	}
	return true
}

func (c *agentTokenAuthCache) storeNode(token string, node *model.Node, expiresAt time.Time) {
	if token == "" || node == nil {
		return
	}
	c.negative.Del(token)
	c.positive.Set(token, cachedAgentNode{
		node:      cloneCachedNode(node),
		expiresAt: expiresAt,
	}, 1)
	c.positive.Wait()
}

func (c *agentTokenAuthCache) storeMissing(token string, expiresAt time.Time) {
	if token == "" {
		return
	}
	c.positive.Del(token)
	c.negative.Set(token, cachedMissingAgentToken{
		expiresAt: expiresAt,
	}, 1)
	c.negative.Wait()
}

func (c *agentTokenAuthCache) invalidate(token string) {
	if token == "" {
		return
	}
	c.positive.Del(token)
	c.negative.Del(token)
}

func (c *agentTokenAuthCache) reset() {
	c.positive.Clear()
	c.negative.Clear()
}

func cloneCachedNode(node *model.Node) *model.Node {
	if node == nil {
		return nil
	}
	cloned := *node
	return &cloned
}

func authenticateAgentTokenWithCache(token string) (*model.Node, error) {
	return nodeAgentTokenCache.authenticate(token)
}

func refreshAgentTokenCache(node *model.Node) {
	if node == nil {
		return
	}
	nodeAgentTokenCache.storeNode(
		node.AgentToken,
		node,
		nodeAgentTokenCache.now().Add(agentTokenPositiveCacheTTL),
	)
}

func invalidateAgentTokenCache(token string) {
	nodeAgentTokenCache.invalidate(token)
}