fleet/server/datastore/cached_mysql/cached_mysql.go

package cached_mysql

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"time"

	"github.com/fleetdm/fleet/v4/server/contexts/ctxdb"
	"github.com/fleetdm/fleet/v4/server/fleet"
	"github.com/jinzhu/copier"
	"github.com/patrickmn/go-cache"
)

const (
	appConfigKey                       = "AppConfig:%s"
	defaultAppConfigExpiration         = 1 * time.Second
	packsHostKey                       = "Packs:host:%d"
	defaultPacksExpiration             = 1 * time.Minute
	scheduledQueriesKey                = "ScheduledQueries:pack:%d"
	defaultScheduledQueriesExpiration  = 1 * time.Minute
	teamAgentOptionsKey                = "TeamAgentOptions:team:%d"
	defaultTeamAgentOptionsExpiration  = 1 * time.Minute
	teamFeaturesKey                    = "TeamFeatures:team:%d"
	defaultTeamFeaturesExpiration      = 1 * time.Minute
	teamMDMConfigKey                   = "TeamMDMConfig:team:%d"
	defaultTeamMDMConfigExpiration     = 1 * time.Minute
	queryByNameKey                     = "QueryByName:team:%d:%s"
	defaultQueryByNameExpiration       = 1 * time.Second
	queryResultsCountKey               = "QueryResultsCount:%d"
	defaultQueryResultsCountExpiration = 1 * time.Second
)

// cloner represents any type that can clone itself. Used by types to provide a more efficient clone method.
type cloner interface {
	Clone() (interface{}, error)
}

func clone(v interface{}) (interface{}, error) {
	if cloner, ok := v.(cloner); ok {
		return cloner.Clone()
	}

	// TODO(mna): consider making implementation of the cloner interface
	// mandatory, and panic/fail loudly if not implemented. Reflection-based deep
	// cloning has significant performance issues at scale (better yet - make the
	// cache accept/return cloner types instead of interface{}).

	if v == nil {
		return nil, nil
	}

	// Use reflection to initialize a clone of v of the same type.
	vv := reflect.ValueOf(v)

	// If the value is a pointer, then calling reflect.New on it will result in a double pointer.
	// Instead, dereference the pointer first.
	isPtr := false
	if vv.Kind() == reflect.Ptr {
		isPtr = true
		if vv.IsNil() {
			return nil, nil
		}
		vv = vv.Elem()
	}

	clone := reflect.New(vv.Type())

	err := copier.CopyWithOption(clone.Interface(), v, copier.Option{DeepCopy: true, IgnoreEmpty: true})
	if err != nil {
		return nil, err
	}

	if isPtr {
		return clone.Interface(), nil
	}

	// The value was not a pointer. Need to dereference it before returning.
	return clone.Elem().Interface(), nil
}

// cloneCache wraps the in memory cache with one that clones items before returning them.
type cloneCache struct {
	*cache.Cache
}

func (c *cloneCache) Get(ctx context.Context, k string) (interface{}, bool) {
	if ctxdb.IsCachedMysqlBypassed(ctx) {
		// cache miss if the caller explicitly asked to bypass the cache
		return nil, false
	}

	x, found := c.Cache.Get(k)
	if !found {
		return nil, false
	}

	clone, err := clone(x)
	if err != nil {
		// Unfortunely, we can't return an error here. Return a cache miss instead of panic'ing.
		return nil, false
	}
	return clone, true
}

func (c *cloneCache) Set(ctx context.Context, k string, x interface{}, d time.Duration) {
	clone, err := clone(x)
	if err != nil {
		// Unfortunately, we can't return an error here. Skip caching it if clone
		// fails, but ensure that we clear any existing cached item for this key,
		// as the call to Set indicates the cache is now stale.
		c.Cache.Delete(k)
		return
	}

	c.Cache.Set(k, clone, d)
}

type cachedMysql struct {
	fleet.Datastore

	c *cloneCache

	appConfigExp         time.Duration
	packsExp             time.Duration
	scheduledQueriesExp  time.Duration
	teamAgentOptionsExp  time.Duration
	teamFeaturesExp      time.Duration
	teamMDMConfigExp     time.Duration
	queryByNameExp       time.Duration
	queryResultsCountExp time.Duration
}

type Option func(*cachedMysql)

func WithAppConfigExpiration(d time.Duration) Option {
	return func(o *cachedMysql) {
		o.appConfigExp = d
	}
}

func WithPacksExpiration(d time.Duration) Option {
	return func(o *cachedMysql) {
		o.packsExp = d
	}
}

func WithScheduledQueriesExpiration(d time.Duration) Option {
	return func(o *cachedMysql) {
		o.scheduledQueriesExp = d
	}
}

func WithTeamAgentOptionsExpiration(d time.Duration) Option {
	return func(o *cachedMysql) {
		o.teamAgentOptionsExp = d
	}
}

func WithTeamFeaturesExpiration(d time.Duration) Option {
	return func(o *cachedMysql) {
		o.teamFeaturesExp = d
	}
}

func WithTeamMDMConfigExpiration(d time.Duration) Option {
	return func(o *cachedMysql) {
		o.teamMDMConfigExp = d
	}
}

func New(ds fleet.Datastore, opts ...Option) fleet.Datastore {
	c := &cachedMysql{
		Datastore:            ds,
		c:                    &cloneCache{cache.New(5*time.Minute, 10*time.Minute)},
		appConfigExp:         defaultAppConfigExpiration,
		packsExp:             defaultPacksExpiration,
		scheduledQueriesExp:  defaultScheduledQueriesExpiration,
		teamAgentOptionsExp:  defaultTeamAgentOptionsExpiration,
		teamFeaturesExp:      defaultTeamFeaturesExpiration,
		teamMDMConfigExp:     defaultTeamMDMConfigExpiration,
		queryByNameExp:       defaultQueryByNameExpiration,
		queryResultsCountExp: defaultQueryResultsCountExpiration,
	}
	for _, fn := range opts {
		fn(c)
	}
	return c
}

func (ds *cachedMysql) NewAppConfig(ctx context.Context, info *fleet.AppConfig) (*fleet.AppConfig, error) {
	ac, err := ds.Datastore.NewAppConfig(ctx, info)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, appConfigKey, ac, ds.appConfigExp)

	return ac, nil
}

func (ds *cachedMysql) AppConfig(ctx context.Context) (*fleet.AppConfig, error) {
	if x, found := ds.c.Get(ctx, appConfigKey); found {
		ac, ok := x.(*fleet.AppConfig)
		if ok {
			return ac, nil
		}
	}

	ac, err := ds.Datastore.AppConfig(ctx)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, appConfigKey, ac, ds.appConfigExp)

	return ac, nil
}

func (ds *cachedMysql) SaveAppConfig(ctx context.Context, info *fleet.AppConfig) error {
	err := ds.Datastore.SaveAppConfig(ctx, info)
	if err != nil {
		return err
	}

	ds.c.Set(ctx, appConfigKey, info, ds.appConfigExp)

	return nil
}

func (ds *cachedMysql) ListPacksForHost(ctx context.Context, hid uint) ([]*fleet.Pack, error) {
	key := fmt.Sprintf(packsHostKey, hid)
	if x, found := ds.c.Get(ctx, key); found {
		cachedPacks, ok := x.([]*fleet.Pack)
		if ok {
			return cachedPacks, nil
		}
	}

	packs, err := ds.Datastore.ListPacksForHost(ctx, hid)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, key, packs, ds.packsExp)

	return packs, nil
}

func (ds *cachedMysql) ListScheduledQueriesInPack(ctx context.Context, packID uint) (fleet.ScheduledQueryList, error) {
	key := fmt.Sprintf(scheduledQueriesKey, packID)
	if x, found := ds.c.Get(ctx, key); found {
		scheduledQueries, ok := x.(fleet.ScheduledQueryList)
		if ok {
			return scheduledQueries, nil
		}
	}

	scheduledQueries, err := ds.Datastore.ListScheduledQueriesInPack(ctx, packID)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, key, scheduledQueries, ds.scheduledQueriesExp)

	return scheduledQueries, nil
}

func (ds *cachedMysql) TeamAgentOptions(ctx context.Context, teamID uint) (*json.RawMessage, error) {
	key := fmt.Sprintf(teamAgentOptionsKey, teamID)
	if x, found := ds.c.Get(ctx, key); found {
		if agentOptions, ok := x.(*json.RawMessage); ok {
			return agentOptions, nil
		}
	}

	agentOptions, err := ds.Datastore.TeamAgentOptions(ctx, teamID)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, key, agentOptions, ds.teamAgentOptionsExp)

	return agentOptions, nil
}

func (ds *cachedMysql) TeamFeatures(ctx context.Context, teamID uint) (*fleet.Features, error) {
	key := fmt.Sprintf(teamFeaturesKey, teamID)
	if x, found := ds.c.Get(ctx, key); found {
		if features, ok := x.(*fleet.Features); ok {
			return features, nil
		}
	}

	features, err := ds.Datastore.TeamFeatures(ctx, teamID)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, key, features, ds.teamFeaturesExp)

	return features, nil
}

func (ds *cachedMysql) TeamMDMConfig(ctx context.Context, teamID uint) (*fleet.TeamMDM, error) {
	key := fmt.Sprintf(teamMDMConfigKey, teamID)
	if x, found := ds.c.Get(ctx, key); found {
		if cfg, ok := x.(*fleet.TeamMDM); ok {
			return cfg, nil
		}
	}

	cfg, err := ds.Datastore.TeamMDMConfig(ctx, teamID)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, key, cfg, ds.teamMDMConfigExp)

	return cfg, nil
}

func (ds *cachedMysql) SaveTeam(ctx context.Context, team *fleet.Team) (*fleet.Team, error) {
	team, err := ds.Datastore.SaveTeam(ctx, team)
	if err != nil {
		return nil, err
	}

	agentOptionsKey := fmt.Sprintf(teamAgentOptionsKey, team.ID)
	featuresKey := fmt.Sprintf(teamFeaturesKey, team.ID)
	mdmConfigKey := fmt.Sprintf(teamMDMConfigKey, team.ID)

	ds.c.Set(ctx, agentOptionsKey, team.Config.AgentOptions, ds.teamAgentOptionsExp)
	ds.c.Set(ctx, featuresKey, &team.Config.Features, ds.teamFeaturesExp)
	ds.c.Set(ctx, mdmConfigKey, &team.Config.MDM, ds.teamMDMConfigExp)

	return team, nil
}

func (ds *cachedMysql) DeleteTeam(ctx context.Context, teamID uint) error {
	err := ds.Datastore.DeleteTeam(ctx, teamID)
	if err != nil {
		return err
	}

	agentOptionsKey := fmt.Sprintf(teamAgentOptionsKey, teamID)
	featuresKey := fmt.Sprintf(teamFeaturesKey, teamID)
	mdmConfigKey := fmt.Sprintf(teamMDMConfigKey, teamID)

	ds.c.Delete(agentOptionsKey)
	ds.c.Delete(featuresKey)
	ds.c.Delete(mdmConfigKey)

	return nil
}

func (ds *cachedMysql) QueryByName(ctx context.Context, teamID *uint, name string) (*fleet.Query, error) {
	teamID_ := uint(0) // global team is 0
	if teamID != nil {
		teamID_ = *teamID
	}
	key := fmt.Sprintf(queryByNameKey, teamID_, name)

	if x, found := ds.c.Get(ctx, key); found {
		if query, ok := x.(*fleet.Query); ok {
			return query, nil
		}
	}

	query, err := ds.Datastore.QueryByName(ctx, teamID, name)
	if err != nil {
		return nil, err
	}

	ds.c.Set(ctx, key, query, ds.queryByNameExp)

	return query, nil
}

func (ds *cachedMysql) ResultCountForQuery(ctx context.Context, queryID uint) (int, error) {
	key := fmt.Sprintf(queryResultsCountKey, queryID)

	if x, found := ds.c.Get(ctx, key); found {
		if count, ok := x.(int); ok {
			return count, nil
		}
	}

	count, err := ds.Datastore.ResultCountForQuery(ctx, queryID)
	if err != nil {
		return 0, err
	}

	ds.c.Set(ctx, key, count, ds.queryResultsCountExp)

	return count, nil
}