Elgato_dark/fleet
1
0
Fork 0
mirror of https://github.com/fleetdm/fleet synced 2026-05-22 00:18:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 106

Updated cache strategy on queries used in GetClientConfig (#12815)

Browse source 
1. Cached results of `svc.ds.Team`
2. Cached results of `svc.ds.ListQueries` too for scheduled queries
only.
3. Do not load aggregated stats on `svc.ds.ListQueries` insde
`GetClientConfig`
This commit is contained in:
Juan Fernandez 2023-07-20 08:06:43 -04:00 committed by GitHub
parent 8d55966553
commit b0c1dba44c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 312 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
server/datastore/cached_mysql/cached_mysql.go
View file

@ -25,6 +25,8 @@ const (
defaultTeamFeaturesExpiration = 1 * time.Minute
teamMDMConfigKey = "TeamMDMConfig:team:%d"
defaultTeamMDMConfigExpiration = 1 * time.Minute
teamNameByIdKey = "TeamName:team:%d"
scheduledQueriesForAgentsKey = "ScheduledQueriesAgents:team:%d"
)
// cloner represents any type that can clone itself. Used by types to provide a more efficient clone method.
@ -296,10 +298,12 @@ func (ds *cachedMysql) SaveTeam(ctx context.Context, team *fleet.Team) (*fleet.T
agentOptionsKey := fmt.Sprintf(teamAgentOptionsKey, team.ID)
featuresKey := fmt.Sprintf(teamFeaturesKey, team.ID)
mdmConfigKey := fmt.Sprintf(teamMDMConfigKey, team.ID)
teamNameKey := fmt.Sprintf(teamNameByIdKey, team.ID)
ds.c.Set(agentOptionsKey, team.Config.AgentOptions, ds.teamAgentOptionsExp)
ds.c.Set(featuresKey, &team.Config.Features, ds.teamFeaturesExp)
ds.c.Set(mdmConfigKey, &team.Config.MDM, ds.teamMDMConfigExp)
ds.c.Set(teamNameKey, &team.Name, ds.scheduledQueriesExp)
return team, nil
}
@ -313,10 +317,49 @@ func (ds *cachedMysql) DeleteTeam(ctx context.Context, teamID uint) error {
agentOptionsKey := fmt.Sprintf(teamAgentOptionsKey, teamID)
featuresKey := fmt.Sprintf(teamFeaturesKey, teamID)
mdmConfigKey := fmt.Sprintf(teamMDMConfigKey, teamID)
teamNameKey := fmt.Sprintf(teamNameByIdKey, teamID)
ds.c.Delete(agentOptionsKey)
ds.c.Delete(featuresKey)
ds.c.Delete(mdmConfigKey)
ds.c.Delete(teamNameKey)
return nil
}
func (ds *cachedMysql) GetTeamName(ctx context.Context, teamID uint) (*string, error) {
key := fmt.Sprintf(teamNameByIdKey, teamID)
if x, found := ds.c.Get(key); found {
if teamName, ok := x.(*string); ok {
return teamName, nil
}
}
teamName, err := ds.Datastore.GetTeamName(ctx, teamID)
if err != nil {
return nil, err
}
ds.c.Set(key, teamName, ds.scheduledQueriesExp)
return teamName, nil
}
func (ds *cachedMysql) ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*fleet.Query, error) {
var teamIDVal uint
if teamID != nil {
teamIDVal = *teamID
}
key := fmt.Sprintf(scheduledQueriesForAgentsKey, teamIDVal)
if x, found := ds.c.Get(key); found {
if queries, ok := x.([]*fleet.Query); ok {
return queries, nil
}
}
queries, err := ds.Datastore.ListScheduledQueriesForAgents(ctx, teamID)
if err != nil {
return nil, err
}
ds.c.Set(key, queries, ds.scheduledQueriesExp)
return queries, nil
}

89
server/datastore/cached_mysql/cached_mysql_test.go
View file

@ -12,6 +12,7 @@ import (
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/fleetdm/fleet/v4/server/mock"
"github.com/fleetdm/fleet/v4/server/ptr"
"github.com/fleetdm/fleet/v4/server/test"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@ -537,3 +538,91 @@ func TestCachedTeamMDMConfig(t *testing.T) {
_, err = ds.TeamMDMConfig(context.Background(), testTeam.ID)
require.Error(t, err)
}
func TestCachedGetTeamName(t *testing.T) {
t.Parallel()
ctx := context.Background()
mockedDS := new(mock.Store)
ds := New(mockedDS, WithScheduledQueriesExpiration(100*time.Millisecond))
team := fleet.Team{
ID: 1,
CreatedAt: time.Now(),
Name: "test",
}
deleted := false
mockedDS.GetTeamNameFunc = func(ctx context.Context, teamID uint) (*string, error) {
if deleted {
return nil, errors.New("not found")
}
return &team.Name, nil
}
mockedDS.SaveTeamFunc = func(ctx context.Context, team *fleet.Team) (*fleet.Team, error) {
return team, nil
}
mockedDS.DeleteTeamFunc = func(ctx context.Context, teamID uint) error {
deleted = true
return nil
}
// updating updates the cache
result, err := ds.GetTeamName(ctx, 1)
require.NoError(t, err)
require.Equal(t, team.Name, *result)
updatedTeam := &fleet.Team{
ID: team.ID,
CreatedAt: team.CreatedAt,
Name: "test II",
}
_, err = ds.SaveTeam(ctx, updatedTeam)
require.NoError(t, err)
result, err = ds.GetTeamName(ctx, team.ID)
require.NoError(t, err)
require.Equal(t, updatedTeam.Name, *result)
// deleting updates the cache
err = ds.DeleteTeam(ctx, team.ID)
require.NoError(t, err)
_, err = ds.GetTeamName(ctx, team.ID)
require.Error(t, err)
}
func TestCachedListScheduledQueriesForAgents(t *testing.T) {
t.Parallel()
ctx := context.Background()
mockedDS := new(mock.Store)
ds := New(mockedDS, WithScheduledQueriesExpiration(100*time.Millisecond))
teamID := ptr.Uint(1)
scheduledQueries := []*fleet.Query{
{
ID: 1,
Name: "test",
ScheduleInterval: 100,
AutomationsEnabled: true,
TeamID: teamID,
},
{
ID: 2,
Name: "test II",
ScheduleInterval: 100,
AutomationsEnabled: true,
TeamID: teamID,
},
}
mockedDS.ListScheduledQueriesForAgentsFunc = func(ctx context.Context, teamID *uint) ([]*fleet.Query, error) {
return scheduledQueries, nil
}
result, err := ds.ListScheduledQueriesForAgents(ctx, teamID)
require.NoError(t, err)
test.QueryElementsMatch(t, result, scheduledQueries)
}

33
server/datastore/mysql/hosts_test.go
View file

@ -564,10 +564,6 @@ func testHostsWithTeamPackStats(t *testing.T, ds *Datastore) {
})
require.NoError(t, err)
require.NoError(t, ds.AddHostsToTeam(context.Background(), &team.ID, []uint{host.ID}))
tp, err := ds.EnsureTeamPack(context.Background(), team.ID)
require.NoError(t, err)
tpQuery := test.NewQuery(t, ds, nil, "tp-time", "select * from time", 0, true)
tpSquery := test.NewScheduledQuery(t, ds, tp.ID, tpQuery.ID, 30, true, true, "time-scheduled")
// Create a new pack and target to the host.
// Pack and query must exist for stats to save successfully
@ -596,28 +592,8 @@ func testHostsWithTeamPackStats(t *testing.T, ds *Datastore) {
WallTime: 0,
},
}
stats2 := []fleet.ScheduledQueryStats{
{
ScheduledQueryName: tpSquery.Name,
ScheduledQueryID: tpSquery.ID,
QueryName: tpQuery.Name,
PackName: tp.Name,
PackID: tp.ID,
AverageMemory: 8000,
Denylisted: false,
Executions: 164,
Interval: 30,
LastExecuted: time.Unix(1620325191, 0).UTC(),
OutputSize: 1337,
SystemTime: 150,
UserTime: 180,
WallTime: 0,
},
}
packStats := []fleet.PackStats{
{PackID: pack1.ID, PackName: pack1.Name, QueryStats: stats1},
{PackID: tp.ID, PackName: teamScheduleName(team), QueryStats: stats2},
}
err = ds.SaveHostPackStats(context.Background(), host.ID, packStats)
require.NoError(t, err)
@ -625,14 +601,11 @@ func testHostsWithTeamPackStats(t *testing.T, ds *Datastore) {
host, err = ds.Host(context.Background(), host.ID)
require.NoError(t, err)
require.Len(t, host.PackStats, 2)
require.Len(t, host.PackStats, 1)
sort.Sort(packStatsSlice(host.PackStats))
assert.Equal(t, host.PackStats[0].PackName, teamScheduleName(team))
assert.ElementsMatch(t, host.PackStats[0].QueryStats, stats2)
assert.Equal(t, host.PackStats[1].PackName, pack1.Name)
assert.ElementsMatch(t, host.PackStats[1].QueryStats, stats1)
assert.Equal(t, host.PackStats[0].PackName, pack1.Name)
assert.ElementsMatch(t, host.PackStats[0].QueryStats, stats1)
}
type packStatsSlice []fleet.PackStats

36
server/datastore/mysql/queries.go
View file

@ -353,10 +353,10 @@ func (ds *Datastore) ListQueries(ctx context.Context, opt fleet.ListQueryOptions
FROM queries q
LEFT JOIN users u ON (q.author_id = u.id)
LEFT JOIN aggregated_stats ag ON (ag.id = q.id AND ag.global_stats = ? AND ag.type = ?)
WHERE saved = true`
`
args := []interface{}{false, aggregatedStatsTypeQuery}
whereClauses := ""
whereClauses := "WHERE saved = true"
if opt.OnlyObserverCanRun {
whereClauses += " AND q.observer_can_run=true"
@ -453,3 +453,35 @@ func (ds *Datastore) ObserverCanRunQuery(ctx context.Context, queryID uint) (boo
return observerCanRun, nil
}
func (ds *Datastore) ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*fleet.Query, error) {
sql := `
SELECT
q.name,
q.query,
q.team_id,
q.schedule_interval,
q.platform,
q.min_osquery_version,
q.automations_enabled,
q.logging_type
FROM queries q
WHERE q.saved = true
AND (q.schedule_interval > 0 AND q.automations_enabled = 1)
`
args := []interface{}{}
if teamID != nil {
args = append(args, *teamID)
sql += " AND team_id = ?"
} else {
sql += " AND team_id IS NULL"
}
results := []*fleet.Query{}
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &results, sql, args...); err != nil {
return nil, ctxerr.Wrap(ctx, err, "list scheduled queries for agents")
}
return results, nil
}

65
server/datastore/mysql/queries_test.go
View file

@ -30,8 +30,9 @@ func TestQueries(t *testing.T) {
{"DuplicateNew", testQueriesDuplicateNew},
{"ListFiltersObservers", testQueriesListFiltersObservers},
{"ObserverCanRunQuery", testObserverCanRunQuery},
{"ListFiltersByTeamID", testQueriesListFiltersByTeamID},
{"ListFiltersByIsScheduled", testQueriesListFiltersByIsScheduled},
{"ListQueriesFiltersByTeamID", testListQueriesFiltersByTeamID},
{"ListQueriesFiltersByIsScheduled", testListQueriesFiltersByIsScheduled},
{"ListScheduledQueriesForAgents", testListScheduledQueriesForAgents},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
@ -555,7 +556,7 @@ func testObserverCanRunQuery(t *testing.T, ds *Datastore) {
}
}
func testQueriesListFiltersByTeamID(t *testing.T, ds *Datastore) {
func testListQueriesFiltersByTeamID(t *testing.T, ds *Datastore) {
globalQ1, err := ds.NewQuery(context.Background(), &fleet.Query{
Name: "query1",
Query: "select 1;",
@ -617,7 +618,7 @@ func testQueriesListFiltersByTeamID(t *testing.T, ds *Datastore) {
test.QueryElementsMatch(t, queries, []*fleet.Query{teamQ1, teamQ2, teamQ3})
}
func testQueriesListFiltersByIsScheduled(t *testing.T, ds *Datastore) {
func testListQueriesFiltersByIsScheduled(t *testing.T, ds *Datastore) {
q1, err := ds.NewQuery(context.Background(), &fleet.Query{
Name: "query1",
Query: "select 1;",
@ -669,3 +670,59 @@ func testQueriesListFiltersByIsScheduled(t *testing.T, ds *Datastore) {
test.QueryElementsMatch(t, queries, tCase.expected, i)
}
}
func testListScheduledQueriesForAgents(t *testing.T, ds *Datastore) {
ctx := context.Background()
team, err := ds.NewTeam(context.Background(), &fleet.Team{
Name: "Team 1",
Description: "Team 1",
})
require.NoError(t, err)
for i, teamID := range []*uint{nil, &team.ID} {
var teamIDStr string
if teamID != nil {
teamIDStr = fmt.Sprintf("%d", *teamID)
}
_, err := ds.NewQuery(context.Background(), &fleet.Query{
Name: fmt.Sprintf("%s query1", teamIDStr),
Query: "select 1;",
Saved: true,
ScheduleInterval: 0,
TeamID: teamID,
})
require.NoError(t, err)
_, err = ds.NewQuery(context.Background(), &fleet.Query{
Name: fmt.Sprintf("%s query2", teamIDStr),
Query: "select 1;",
Saved: false,
ScheduleInterval: 10,
AutomationsEnabled: false,
TeamID: teamID,
})
require.NoError(t, err)
q3, err := ds.NewQuery(context.Background(), &fleet.Query{
Name: fmt.Sprintf("%s query3", teamIDStr),
Query: "select 1;",
Saved: true,
ScheduleInterval: 20,
AutomationsEnabled: true,
TeamID: teamID,
})
require.NoError(t, err)
_, err = ds.NewQuery(context.Background(), &fleet.Query{
Name: fmt.Sprintf("%s query4", teamIDStr),
Query: "select 1;",
Saved: true,
ScheduleInterval: 0,
AutomationsEnabled: true,
TeamID: teamID,
})
require.NoError(t, err)
result, err := ds.ListScheduledQueriesForAgents(ctx, teamID)
require.NoError(t, err)
test.QueryElementsMatch(t, result, []*fleet.Query{q3}, i)
}
}

14
server/datastore/mysql/teams.go
View file

@ -428,3 +428,17 @@ func (ds *Datastore) DeleteIntegrationsFromTeams(ctx context.Context, deletedInt
}
return rows.Err()
}
func (ds *Datastore) GetTeamName(ctx context.Context, teamID uint) (*string, error) {
stmt := `SELECT name FROM teams WHERE id = ?`
var teamName string
if err := sqlx.GetContext(ctx, ds.reader(ctx), &teamName, stmt, teamID); err != nil {
if err == sql.ErrNoRows {
return nil, ctxerr.Wrap(ctx, notFound("Team").WithID(teamID))
}
return nil, ctxerr.Wrap(ctx, err, "select team")
}
return &teamName, nil
}

22
server/datastore/mysql/teams_test.go
View file

@ -35,6 +35,7 @@ func TestTeams(t *testing.T) {
{"DeleteIntegrationsFromTeams", testTeamsDeleteIntegrationsFromTeams},
{"TeamsFeatures", testTeamsFeatures},
{"TeamsMDMConfig", testTeamsMDMConfig},
{"GetTeamByName", testGetTeamByName},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
@ -624,3 +625,24 @@ func testTeamsMDMConfig(t *testing.T, ds *Datastore) {
}, mdm)
})
}
func testGetTeamByName(t *testing.T, ds *Datastore) {
ctx := context.Background()
t.Run("team does not exists", func(t *testing.T) {
r, err := ds.GetTeamName(ctx, 123)
require.Nil(t, r)
require.Error(t, err)
})
t.Run("returns the team name", func(t *testing.T) {
team, err := ds.NewTeam(ctx, &fleet.Team{
Name: "team1",
})
require.NoError(t, err)
result, err := ds.GetTeamName(ctx, team.ID)
require.NoError(t, err)
require.Equal(t, team.Name, *result)
})
}

5
server/fleet/datastore.go
View file

@ -82,6 +82,9 @@ type Datastore interface {
// ListQueries returns a list of queries with the provided sorting and paging options. Associated packs should also
// be loaded.
ListQueries(ctx context.Context, opt ListQueryOptions) ([]*Query, error)
// ListScheduledQueriesForAgents returns a list of scheduled queries (without stats) for the
// given teamID. If teamID is nil, then all scheduled queries for the 'global' team are returned.
ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*Query, error)
// QueryByName looks up a query by name on a team. If teamID is nil, then the query is looked up in
// the 'global' team.
QueryByName(ctx context.Context, teamID *uint, name string, opts ...OptionalArg) (*Query, error)
@ -397,6 +400,8 @@ type Datastore interface {
SaveTeam(ctx context.Context, team *Team) (*Team, error)
// Team retrieves the Team by ID.
Team(ctx context.Context, tid uint) (*Team, error)
// GetTeamName retrieves the team name by their ID.
GetTeamName(ctx context.Context, teamID uint) (*string, error)
// Team deletes the Team by ID.
DeleteTeam(ctx context.Context, tid uint) error
// TeamByName retrieves the Team by Name.

24
server/mock/datastore_mock.go
View file

@ -70,6 +70,8 @@ type QueryFunc func(ctx context.Context, id uint) (*fleet.Query, error)
type ListQueriesFunc func(ctx context.Context, opt fleet.ListQueryOptions) ([]*fleet.Query, error)
type ListScheduledQueriesForAgentsFunc func(ctx context.Context, teamID *uint) ([]*fleet.Query, error)
type QueryByNameFunc func(ctx context.Context, teamID *uint, name string, opts ...fleet.OptionalArg) (*fleet.Query, error)
type ObserverCanRunQueryFunc func(ctx context.Context, queryID uint) (bool, error)
@ -300,6 +302,8 @@ type SaveTeamFunc func(ctx context.Context, team *fleet.Team) (*fleet.Team, erro
type TeamFunc func(ctx context.Context, tid uint) (*fleet.Team, error)
type GetTeamNameFunc func(ctx context.Context, teamID uint) (*string, error)
type DeleteTeamFunc func(ctx context.Context, tid uint) error
type TeamByNameFunc func(ctx context.Context, name string) (*fleet.Team, error)
@ -739,6 +743,9 @@ type DataStore struct {
ListQueriesFunc ListQueriesFunc
ListQueriesFuncInvoked bool
ListScheduledQueriesForAgentsFunc ListScheduledQueriesForAgentsFunc
ListScheduledQueriesForAgentsFuncInvoked bool
QueryByNameFunc QueryByNameFunc
QueryByNameFuncInvoked bool
@ -1084,6 +1091,9 @@ type DataStore struct {
TeamFunc TeamFunc
TeamFuncInvoked bool
GetTeamNameFunc GetTeamNameFunc
GetTeamNameFuncInvoked bool
DeleteTeamFunc DeleteTeamFunc
DeleteTeamFuncInvoked bool
@ -1809,6 +1819,13 @@ func (s *DataStore) ListQueries(ctx context.Context, opt fleet.ListQueryOptions)
return s.ListQueriesFunc(ctx, opt)
}
func (s *DataStore) ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*fleet.Query, error) {
s.mu.Lock()
s.ListScheduledQueriesForAgentsFuncInvoked = true
s.mu.Unlock()
return s.ListScheduledQueriesForAgentsFunc(ctx, teamID)
}
func (s *DataStore) QueryByName(ctx context.Context, teamID *uint, name string, opts ...fleet.OptionalArg) (*fleet.Query, error) {
s.mu.Lock()
s.QueryByNameFuncInvoked = true
@ -2614,6 +2631,13 @@ func (s *DataStore) Team(ctx context.Context, tid uint) (*fleet.Team, error) {
return s.TeamFunc(ctx, tid)
}
func (s *DataStore) GetTeamName(ctx context.Context, teamID uint) (*string, error) {
s.mu.Lock()
s.GetTeamNameFuncInvoked = true
s.mu.Unlock()
return s.GetTeamNameFunc(ctx, teamID)
}
func (s *DataStore) DeleteTeam(ctx context.Context, tid uint) error {
s.mu.Lock()
s.DeleteTeamFuncInvoked = true

9
server/service/osquery.go
View file

@ -347,8 +347,7 @@ func getClientConfigEndpoint(ctx context.Context, request interface{}, svc fleet
}
func (svc *Service) getScheduledQueries(ctx context.Context, teamID *uint) (fleet.Queries, error) {
opts := fleet.ListQueryOptions{IsScheduled: ptr.Bool(true), TeamID: teamID}
queries, err := svc.ds.ListQueries(ctx, opts)
queries, err := svc.ds.ListScheduledQueriesForAgents(ctx, teamID)
if err != nil {
return nil, err
}
@ -444,18 +443,18 @@ func (svc *Service) GetClientConfig(ctx context.Context) (map[string]interface{}
}
if host.TeamID != nil {
team, err := svc.ds.Team(ctx, *host.TeamID)
teamName, err := svc.ds.GetTeamName(ctx, *host.TeamID)
if err != nil {
return nil, newOsqueryError("database error: " + err.Error())
}
if team != nil {
if teamName != nil {
teamQueries, err := svc.getScheduledQueries(ctx, host.TeamID)
if err != nil {
return nil, newOsqueryError("database error: " + err.Error())
}
if len(teamQueries) > 0 {
packName := fmt.Sprintf("Team: %s", team.Name)
packName := fmt.Sprintf("Team: %s", *teamName)
packConfig[packName] = fleet.PackContent{
Queries: teamQueries,
}

20
server/service/osquery_test.go
View file

@ -40,11 +40,9 @@ import (
func TestGetClientConfig(t *testing.T) {
ds := new(mock.Store)
ds.TeamFunc = func(ctx context.Context, tid uint) (*fleet.Team, error) {
return &fleet.Team{
Name: "Alamo",
ID: 1,
}, nil
ds.GetTeamNameFunc = func(ctx context.Context, tid uint) (*string, error) {
teamName := "Alamo"
return &teamName, nil
}
ds.TeamAgentOptionsFunc = func(ctx context.Context, teamID uint) (*json.RawMessage, error) {
@ -71,8 +69,8 @@ func TestGetClientConfig(t *testing.T) {
return []*fleet.ScheduledQuery{}, nil
}
}
ds.ListQueriesFunc = func(ctx context.Context, opt fleet.ListQueryOptions) ([]*fleet.Query, error) {
if opt.TeamID == nil {
ds.ListScheduledQueriesForAgentsFunc = func(ctx context.Context, teamID *uint) ([]*fleet.Query, error) {
if teamID == nil {
return nil, nil
}
return []*fleet.Query{
@ -2006,6 +2004,14 @@ func TestUpdateHostIntervals(t *testing.T) {
svc, ctx := newTestService(t, ds, nil, nil)
ds.GetTeamNameFunc = func(ctx context.Context, tid uint) (*string, error) {
return nil, nil
}
ds.ListScheduledQueriesForAgentsFunc = func(ctx context.Context, teamID *uint) ([]*fleet.Query, error) {
return nil, nil
}
ds.ListPacksForHostFunc = func(ctx context.Context, hid uint) ([]*fleet.Pack, error) {
return []*fleet.Pack{}, nil
}