Discharge7443/fleet
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Support async saving of hosts' last seen time (#5640)

Browse Source
This commit is contained in:
Martin Angers
2022-05-10 11:29:17 -04:00
committed by GitHub
parent 421a0ec635
commit 1fa7bb7a19
32 changed files with 581 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
changes/issue-5536-async-host-last-seen
+1
View File
@@ -0,0 +1 @@
* Add experimental support for asynchronous saving of hosts' last seen timestamp.
cmd/fleet/serve.go
+12 -9
View File
@@ -305,6 +305,7 @@ the way that the Fleet server works.
task := &async.Task{
Datastore: ds,
Pool: redisPool,
Clock: clock.C,
AsyncEnabled: config.Osquery.EnableAsyncHostProcessing,
LockTimeout: config.Osquery.AsyncHostCollectLockTimeout,
LogStatsInterval: config.Osquery.AsyncHostCollectLogStatsInterval,
@@ -360,16 +361,18 @@ the way that the Fleet server works.
cancelBackground := runCrons(ds, task, kitlog.With(logger, "component", "crons"), config, license, failingPolicySet)
// Flush seen hosts every second
go func() {
for range time.Tick(time.Duration(rand.Intn(10)+1) * time.Second) {
if err := svc.FlushSeenHosts(context.Background()); err != nil {
level.Info(logger).Log(
"err", err,
"msg", "failed to update host seen times",
)
if !task.AsyncEnabled {
go func() {
for range time.Tick(time.Duration(rand.Intn(10)+1) * time.Second) {
if err := task.FlushHostsLastSeen(context.Background(), clock.C.Now()); err != nil {
level.Info(logger).Log(
"err", err,
"msg", "failed to update host seen times",
)
}
}
}
}()
}()
}
fieldKeys := []string{"method", "error"}
requestCount := kitprometheus.NewCounterFrom(prometheus.CounterOpts{
cmd/fleetctl/apply_test.go
+1 -1
View File
@@ -102,7 +102,7 @@ spec:
func TestApplyTeamSpecs(t *testing.T) {
license := &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}
_, ds := runServerWithMockedDS(t, service.TestServerOpts{License: license})
_, ds := runServerWithMockedDS(t, &service.TestServerOpts{License: license})
teamsByName := map[string]*fleet.Team{
"team1": {
cmd/fleetctl/get_test.go
+2 -2
View File
@@ -115,7 +115,7 @@ func TestGetTeams(t *testing.T) {
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
license := tt.license
_, ds := runServerWithMockedDS(t, service.TestServerOpts{License: license})
_, ds := runServerWithMockedDS(t, &service.TestServerOpts{License: license})
agentOpts := json.RawMessage(`{"config":{"foo":"bar"},"overrides":{"platforms":{"darwin":{"foo":"override"}}}}`)
ds.ListTeamsFunc = func(ctx context.Context, filter fleet.TeamFilter, opt fleet.ListOptions) ([]*fleet.Team, error) {
@@ -209,7 +209,7 @@ spec:
}
func TestGetTeamsByName(t *testing.T) {
_, ds := runServerWithMockedDS(t, service.TestServerOpts{License: &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}})
_, ds := runServerWithMockedDS(t, &service.TestServerOpts{License: &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}})
ds.ListTeamsFunc = func(ctx context.Context, filter fleet.TeamFilter, opt fleet.ListOptions) ([]*fleet.Team, error) {
require.Equal(t, "test1", opt.MatchQuery)
cmd/fleetctl/query_test.go
+1 -1
View File
@@ -23,7 +23,7 @@ func TestLiveQuery(t *testing.T) {
logger := kitlog.NewJSONLogger(os.Stdout)
logger = level.NewFilter(logger, level.AllowDebug())
_, ds := runServerWithMockedDS(t, service.TestServerOpts{
_, ds := runServerWithMockedDS(t, &service.TestServerOpts{
Rs: rs,
Lq: lq,
Logger: logger,
cmd/fleetctl/testing_utils.go
+1 -1
View File
@@ -18,7 +18,7 @@ import (
// runServerWithMockedDS runs the fleet server with several mocked DS methods.
//
// NOTE: Assumes the current session is always from the admin user (see ds.SessionByKeyFunc below).
func runServerWithMockedDS(t *testing.T, opts ...service.TestServerOpts) (*httptest.Server, *mock.Store) {
func runServerWithMockedDS(t *testing.T, opts ...*service.TestServerOpts) (*httptest.Server, *mock.Store) {
ds := new(mock.Store)
var users []*fleet.User
var admin *fleet.User
docs/Deploying/Configuration.md
+2 -2
View File
@@ -1067,7 +1067,7 @@ to the amount of time it takes for Fleet to give the host the label queries.
##### osquery_enable_async_host_processing
**Experimental feature**. Enable asynchronous processing of hosts query results. Currently, only supported for label query execution and policy membership results. This may improve performance and CPU usage of the Fleet instances and MySQL database servers for setups with a large number of hosts, while requiring more resources from Redis server(s).
**Experimental feature**. Enable asynchronous processing of hosts query results. Currently, only supported for label query execution, policy membership results and hosts' last seen timestamp. This may improve performance and CPU usage of the Fleet instances and MySQL database servers for setups with a large number of hosts, while requiring more resources from Redis server(s).
Note that currently, if both the failing policies webhook *and* this `osquery.enable_async_host_processing` option are set, some failing policies webhooks could be missing (some transitions from succeeding to failing or vice-versa could happen without triggering a webhook request).
@@ -1186,7 +1186,7 @@ Applies only when `osquery_enable_async_host_processing` is enabled. Maximum num
##### osquery_async_host_redis_scan_keys_count
Applies only when `osquery_enable_async_host_processing` is enabled. Order of magnitude (e.g. 10, 100, 1000, etc.) of keys to scan in a single SCAN request for keys to process when collecting host data into the database.
Applies only when `osquery_enable_async_host_processing` is enabled. Order of magnitude (e.g. 10, 100, 1000, etc.) of set members to scan in a single ZSCAN/SSCAN request for items to process when collecting host data into the database.
- Default value: 1000
- Environment variable: `FLEET_OSQUERY_ASYNC_HOST_REDIS_SCAN_KEYS_COUNT`
server/datastore/mysql/hosts.go
+1 -1
View File
@@ -383,7 +383,7 @@ func (ds *Datastore) ListHosts(ctx context.Context, filter fleet.TeamFilter, opt
`
if opt.DeviceMapping {
sql += `,
sql += `,
COALESCE(dm.device_mapping, 'null') as device_mapping
`
}
server/fleet/service.go
-2
View File
@@ -247,8 +247,6 @@ type Service interface {
HostByIdentifier(ctx context.Context, identifier string) (*HostDetail, error)
// RefetchHost requests a refetch of host details for the provided host.
RefetchHost(ctx context.Context, id uint) (err error)
FlushSeenHosts(ctx context.Context) error
// AddHostsToTeam adds hosts to an existing team, clearing their team settings if teamID is nil.
AddHostsToTeam(ctx context.Context, teamID *uint, hostIDs []uint) error
// AddHostsToTeamByFilter adds hosts to an existing team, clearing their team settings if teamID is nil. Hosts are
server/service/async/async.go
+16 -1
View File
@@ -5,6 +5,7 @@ import (
"fmt"
"time"
"github.com/WatchBeam/clock"
"github.com/fleetdm/fleet/v4/server/datastore/redis"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/getsentry/sentry-go"
@@ -18,6 +19,7 @@ const collectorLockKey = "locks:async_collector:{%s}"
type Task struct {
Datastore fleet.Datastore
Pool fleet.RedisPool
Clock clock.Clock
// AsyncEnabled indicates if async processing is enabled in the
// configuration. Note that Pool can be nil if this is false.
AsyncEnabled bool
@@ -30,6 +32,8 @@ type Task struct {
RedisPopCount int
RedisScanKeysCount int
CollectorInterval time.Duration
seenHostSet seenHostSet
}
// Collect runs the various collectors as distinct background goroutines if
@@ -69,7 +73,18 @@ func (t *Task) StartCollectors(ctx context.Context, jitterPct int, logger kitlog
errHandler: collectorErrHandler,
}
colls := []*collector{labelColl, policyColl}
lastSeenColl := &collector{
name: "collect_last_seen",
pool: t.Pool,
ds: t.Datastore,
execInterval: t.CollectorInterval,
jitterPct: jitterPct,
lockTimeout: t.LockTimeout,
handler: t.collectHostsLastSeen,
errHandler: collectorErrHandler,
}
colls := []*collector{labelColl, policyColl, lastSeenColl}
for _, coll := range colls {
go coll.Start(ctx)
}
server/service/async/async_host_seen.go
+190
View File
@@ -0,0 +1,190 @@
package async
import (
"context"
"sort"
"sync"
"time"
"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
"github.com/fleetdm/fleet/v4/server/datastore/redis"
"github.com/fleetdm/fleet/v4/server/fleet"
redigo "github.com/gomodule/redigo/redis"
)
const (
hostSeenRecordedHostIDsKey = "{host_seen:host_ids}" // the SET of current (pending) host ids
hostSeenProcessingHostIDsKey = "{host_seen:host_ids}:processing" // the SET of host ids in the process of being collected
hostSeenKeysMinTTL = 7 * 24 * time.Hour // 1 week
)
// RecordHostLastSeen records that the specified host ID was seen.
func (t *Task) RecordHostLastSeen(ctx context.Context, hostID uint) error {
if !t.AsyncEnabled {
t.seenHostSet.addHostID(hostID)
return nil
}
// set an expiration on the SET key, ensuring that if async processing is
// disabled, the set (eventually) does not use any redis space. Ensure that
// TTL is reasonably big to avoid deleting information that hasn't been
// collected yet - 1 week or 10 * the collector interval, whichever is
// biggest.
ttl := hostSeenKeysMinTTL
if maxTTL := 10 * t.CollectorInterval; maxTTL > ttl {
ttl = maxTTL
}
// keys and arguments passed to the script are:
// KEYS[1]: recorded set (hostSeenRecordedHostIDsKey)
// ARGV[1]: host id
// ARGV[2]: ttl for the key
script := redigo.NewScript(1, `
redis.call('SADD', KEYS[1], ARGV[1])
return redis.call('EXPIRE', KEYS[1], ARGV[2])
`)
conn := t.Pool.Get()
defer conn.Close()
if err := redis.BindConn(t.Pool, conn, hostSeenRecordedHostIDsKey); err != nil {
return ctxerr.Wrap(ctx, err, "bind redis connection")
}
if _, err := script.Do(conn, hostSeenRecordedHostIDsKey, hostID, int(ttl.Seconds())); err != nil {
return ctxerr.Wrap(ctx, err, "run redis script")
}
return nil
}
// FlushHostsLastSeen updates the last seen timestamp for the hosts that have
// been recorded since the last time FlushHostsLastSeen was called. It is a
// no-op if asychronous host processing is enabled, because then it is the
// task collector that will process the writes to mysql.
func (t *Task) FlushHostsLastSeen(ctx context.Context, now time.Time) error {
if !t.AsyncEnabled {
hostIDs := t.seenHostSet.getAndClearHostIDs()
return t.Datastore.MarkHostsSeen(ctx, hostIDs, now)
}
// no-op, flushing the hosts' last seen is done via the cron that runs the
// Task's collectors.
return nil
}
func (t *Task) collectHostsLastSeen(ctx context.Context, ds fleet.Datastore, pool fleet.RedisPool, stats *collectorExecStats) error {
hostIDs, err := t.loadSeenHostsIDs(ctx, pool)
if err != nil {
return err
}
stats.RedisCmds++ // the script to load seen hosts
stats.Keys = 2 // the reported and processing set keys
stats.Items = len(hostIDs)
// process in batches, as there could be many thousand host IDs
if len(hostIDs) > 0 {
// globally sort the host IDs so they are sent ordered as batches to MarkHostsSeen
sort.Slice(hostIDs, func(i, j int) bool { return hostIDs[i] < hostIDs[j] })
ts := t.Clock.Now()
batch := make([]uint, t.InsertBatch)
for {
n := copy(batch, hostIDs)
if n == 0 {
break
}
if err := ds.MarkHostsSeen(ctx, batch[:n], ts); err != nil {
return err
}
stats.Inserts++
hostIDs = hostIDs[n:]
}
}
conn := pool.Get()
defer conn.Close()
if _, err := conn.Do("DEL", hostSeenProcessingHostIDsKey); err != nil {
return ctxerr.Wrap(ctx, err, "delete processing set key")
}
return nil
}
func (t *Task) loadSeenHostsIDs(ctx context.Context, pool fleet.RedisPool) ([]uint, error) {
// compute the TTL for the processing key just as we do for the storage key,
// in case the collection fails before removing the working key, we don't
// want it to stick around forever.
ttl := hostSeenKeysMinTTL
if maxTTL := 10 * t.CollectorInterval; maxTTL > ttl {
ttl = maxTTL
}
// keys and arguments passed to the script are:
// KEYS[1]: recorded set (hostSeenRecordedHostIDsKey)
// KEYS[2]: processing set (hostSeenProcessingHostIDsKey)
// ARGV[1]: ttl for the processing key
script := redigo.NewScript(2, `
redis.call('SUNIONSTORE', KEYS[2], KEYS[1], KEYS[2])
redis.call('DEL', KEYS[1])
return redis.call('EXPIRE', KEYS[2], ARGV[1])
`)
conn := pool.Get()
defer conn.Close()
if err := redis.BindConn(pool, conn, hostSeenRecordedHostIDsKey, hostSeenProcessingHostIDsKey); err != nil {
return nil, ctxerr.Wrap(ctx, err, "bind redis connection")
}
if _, err := script.Do(conn, hostSeenRecordedHostIDsKey, hostSeenProcessingHostIDsKey, int(ttl.Seconds())); err != nil {
return nil, ctxerr.Wrap(ctx, err, "run redis script")
}
var ids []uint
cursor := 0
for {
res, err := redigo.Values(conn.Do("SSCAN", hostSeenProcessingHostIDsKey, cursor, "COUNT", t.RedisScanKeysCount))
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "scan seen host ids")
}
var scanIDs []uint
if _, err := redigo.Scan(res, &cursor, &scanIDs); err != nil {
return nil, ctxerr.Wrap(ctx, err, "convert scan results")
}
ids = append(ids, scanIDs...)
if cursor == 0 {
// iteration completed
return ids, nil
}
}
}
// seenHostSet implements synchronized storage for the set of seen hosts.
type seenHostSet struct {
mutex sync.Mutex
hostIDs map[uint]bool
}
// addHostID adds the host identified by ID to the set
func (m *seenHostSet) addHostID(id uint) {
m.mutex.Lock()
defer m.mutex.Unlock()
if m.hostIDs == nil {
m.hostIDs = make(map[uint]bool)
}
m.hostIDs[id] = true
}
// getAndClearHostIDs gets the list of unique host IDs from the set and empties
// the set.
func (m *seenHostSet) getAndClearHostIDs() []uint {
m.mutex.Lock()
defer m.mutex.Unlock()
var ids []uint
for id := range m.hostIDs {
ids = append(ids, id)
}
m.hostIDs = make(map[uint]bool)
return ids
}
server/service/async/async_host_seen_test.go
+274
View File
@@ -0,0 +1,274 @@
package async
import (
"context"
"database/sql"
"math"
"testing"
"time"
"github.com/WatchBeam/clock"
"github.com/fleetdm/fleet/v4/server/datastore/mysql"
"github.com/fleetdm/fleet/v4/server/datastore/redis"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/fleetdm/fleet/v4/server/mock"
redigo "github.com/gomodule/redigo/redis"
"github.com/jmoiron/sqlx"
"github.com/stretchr/testify/require"
)
func testCollectHostsLastSeen(t *testing.T, ds *mysql.Datastore, pool fleet.RedisPool) {
ctx := context.Background()
type hostLastSeen struct {
HostID int `db:"host_id"`
SeenTime sql.NullTime `db:"seen_time"`
}
mockTime := clock.NewMockClock()
startTime := mockTime.Now()
hostIDs := createHosts(t, ds, 4, startTime)
t.Logf("real host IDs: %v", hostIDs)
hid := func(id int) int {
return int(hostIDs[id-1])
}
// note that cases cannot be run in isolation, each case builds on the
// previous one's state, so they are not run as distinct sub-tests. An
// hour is added to mockTime after each test case.
cases := []struct {
name string
hostIDs []int
want []hostLastSeen
}{
{
"no key",
nil,
[]hostLastSeen{
// createHosts (called above) stores the initial host seen time
{HostID: hid(1), SeenTime: sql.NullTime{Time: startTime}},
{HostID: hid(2), SeenTime: sql.NullTime{Time: startTime}},
{HostID: hid(3), SeenTime: sql.NullTime{Time: startTime}},
{HostID: hid(4), SeenTime: sql.NullTime{Time: startTime}},
},
},
{
"report host 1",
[]int{hid(1)},
[]hostLastSeen{
{HostID: hid(1), SeenTime: sql.NullTime{Time: startTime.Add(time.Hour)}},
{HostID: hid(2), SeenTime: sql.NullTime{Time: startTime}},
{HostID: hid(3), SeenTime: sql.NullTime{Time: startTime}},
{HostID: hid(4), SeenTime: sql.NullTime{Time: startTime}},
},
},
{
"report hosts 2, 3",
[]int{hid(2), hid(3)},
[]hostLastSeen{
{HostID: hid(1), SeenTime: sql.NullTime{Time: startTime.Add(time.Hour)}},
{HostID: hid(2), SeenTime: sql.NullTime{Time: startTime.Add(2 * time.Hour)}},
{HostID: hid(3), SeenTime: sql.NullTime{Time: startTime.Add(2 * time.Hour)}},
{HostID: hid(4), SeenTime: sql.NullTime{Time: startTime}},
},
},
{
"report hosts 1, 2, 3, 4",
[]int{hid(1), hid(2), hid(3), hid(4)},
[]hostLastSeen{
{HostID: hid(1), SeenTime: sql.NullTime{Time: startTime.Add(3 * time.Hour)}},
{HostID: hid(2), SeenTime: sql.NullTime{Time: startTime.Add(3 * time.Hour)}},
{HostID: hid(3), SeenTime: sql.NullTime{Time: startTime.Add(3 * time.Hour)}},
{HostID: hid(4), SeenTime: sql.NullTime{Time: startTime.Add(3 * time.Hour)}},
},
},
{
"report hosts 2, 3, 4",
[]int{hid(2), hid(3), hid(4)},
[]hostLastSeen{
{HostID: hid(1), SeenTime: sql.NullTime{Time: startTime.Add(3 * time.Hour)}},
{HostID: hid(2), SeenTime: sql.NullTime{Time: startTime.Add(4 * time.Hour)}},
{HostID: hid(3), SeenTime: sql.NullTime{Time: startTime.Add(4 * time.Hour)}},
{HostID: hid(4), SeenTime: sql.NullTime{Time: startTime.Add(4 * time.Hour)}},
},
},
{
"report no new hosts",
[]int{},
[]hostLastSeen{
{HostID: hid(1), SeenTime: sql.NullTime{Time: startTime.Add(3 * time.Hour)}},
{HostID: hid(2), SeenTime: sql.NullTime{Time: startTime.Add(4 * time.Hour)}},
{HostID: hid(3), SeenTime: sql.NullTime{Time: startTime.Add(4 * time.Hour)}},
{HostID: hid(4), SeenTime: sql.NullTime{Time: startTime.Add(4 * time.Hour)}},
},
},
}
const batchSizes = 3
setupTest := func(t *testing.T, ids []int) collectorExecStats {
conn := redis.ConfigureDoer(pool, pool.Get())
defer conn.Close()
// store the host memberships and prepare the expected stats
var wantStats collectorExecStats
wantStats.Keys = 2
wantStats.RedisCmds = 1
if len(ids) > 0 {
args := redigo.Args{hostSeenRecordedHostIDsKey}
args = args.AddFlat(ids)
_, err := conn.Do("SADD", args...)
require.NoError(t, err)
wantStats.Items = len(ids)
wantStats.Inserts = int(math.Ceil(float64(len(ids)) / float64(batchSizes)))
}
return wantStats
}
selectRows := func(t *testing.T) []hostLastSeen {
var rows []hostLastSeen
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
return sqlx.SelectContext(ctx, tx, &rows, `SELECT host_id, seen_time FROM host_seen_times ORDER BY 1`)
})
return rows
}
for _, c := range cases {
func() {
t.Log("test name: ", c.name)
wantStats := setupTest(t, c.hostIDs)
// run the collection
var stats collectorExecStats
task := Task{
Clock: mockTime,
InsertBatch: batchSizes,
UpdateBatch: batchSizes,
DeleteBatch: batchSizes,
RedisPopCount: batchSizes,
RedisScanKeysCount: 10,
}
err := task.collectHostsLastSeen(ctx, ds, pool, &stats)
require.NoError(t, err)
require.Equal(t, wantStats, stats)
// check that the table contains the expected rows
rows := selectRows(t)
require.Equal(t, len(c.want), len(rows))
for i := range c.want {
want, got := c.want[i], rows[i]
require.Equal(t, want.HostID, got.HostID)
require.WithinDuration(t, want.SeenTime.Time, got.SeenTime.Time, time.Second)
}
mockTime.AddTime(time.Hour)
}()
}
}
func testRecordHostLastSeenSync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
ctx := context.Background()
var calledWithHostIDs []uint
ds.MarkHostsSeenFunc = func(ctx context.Context, hostIDs []uint, ts time.Time) error {
calledWithHostIDs = append(calledWithHostIDs, hostIDs...)
return nil
}
task := Task{
Datastore: ds,
Pool: pool,
Clock: clock.C,
AsyncEnabled: false,
}
err := task.RecordHostLastSeen(ctx, 1)
require.NoError(t, err)
err = task.RecordHostLastSeen(ctx, 2)
require.NoError(t, err)
err = task.RecordHostLastSeen(ctx, 3)
require.NoError(t, err)
require.False(t, ds.MarkHostsSeenFuncInvoked)
err = task.FlushHostsLastSeen(ctx, time.Now())
require.NoError(t, err)
require.True(t, ds.MarkHostsSeenFuncInvoked)
require.ElementsMatch(t, []uint{1, 2, 3}, calledWithHostIDs)
ds.MarkHostsSeenFuncInvoked = false
conn := redis.ConfigureDoer(pool, pool.Get())
defer conn.Close()
defer conn.Do("DEL", hostSeenRecordedHostIDsKey, hostSeenProcessingHostIDsKey)
n, err := redigo.Int(conn.Do("EXISTS", hostSeenRecordedHostIDsKey))
require.NoError(t, err)
require.Equal(t, 0, n)
n, err = redigo.Int(conn.Do("EXISTS", hostSeenProcessingHostIDsKey))
require.NoError(t, err)
require.Equal(t, 0, n)
}
func testRecordHostLastSeenAsync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
ctx := context.Background()
var calledWithHostIDs []uint
ds.MarkHostsSeenFunc = func(ctx context.Context, hostIDs []uint, ts time.Time) error {
calledWithHostIDs = append(calledWithHostIDs, hostIDs...)
return nil
}
task := Task{
Datastore: ds,
Pool: pool,
AsyncEnabled: true,
Clock: clock.C,
InsertBatch: 2,
RedisScanKeysCount: 10,
}
err := task.RecordHostLastSeen(ctx, 1)
require.NoError(t, err)
err = task.RecordHostLastSeen(ctx, 2)
require.NoError(t, err)
err = task.RecordHostLastSeen(ctx, 3)
require.NoError(t, err)
require.False(t, ds.MarkHostsSeenFuncInvoked)
err = task.FlushHostsLastSeen(ctx, time.Now())
require.NoError(t, err)
require.False(t, ds.MarkHostsSeenFuncInvoked)
conn := redis.ConfigureDoer(pool, pool.Get())
defer conn.Close()
defer conn.Do("DEL", hostSeenRecordedHostIDsKey, hostSeenProcessingHostIDsKey)
n, err := redigo.Int(conn.Do("SCARD", hostSeenRecordedHostIDsKey))
require.NoError(t, err)
require.Equal(t, 3, n)
n, err = redigo.Int(conn.Do("EXISTS", hostSeenProcessingHostIDsKey))
require.NoError(t, err)
require.Equal(t, 0, n)
// running the collector removes the recorded host IDs key
var stats collectorExecStats
err = task.collectHostsLastSeen(ctx, ds, pool, &stats)
require.NoError(t, err)
require.Equal(t, 2, stats.Keys)
require.Equal(t, 3, stats.Items)
require.False(t, stats.Failed)
require.True(t, ds.MarkHostsSeenFuncInvoked)
require.ElementsMatch(t, []uint{1, 2, 3}, calledWithHostIDs)
ds.MarkHostsSeenFuncInvoked = false
n, err = redigo.Int(conn.Do("EXISTS", hostSeenRecordedHostIDsKey))
require.NoError(t, err)
require.Equal(t, 0, n)
n, err = redigo.Int(conn.Do("EXISTS", hostSeenProcessingHostIDsKey))
require.NoError(t, err)
require.Equal(t, 0, n)
}
server/service/async/async_label.go
+1 -1
View File
@@ -200,7 +200,7 @@ func (t *Task) collectLabelQueryExecutions(ctx context.Context, ds fleet.Datasto
hostIDs[i] = host.HostID
}
ts := time.Now()
ts := t.Clock.Now()
updateBatch := make([]uint, t.UpdateBatch)
for {
n := copy(updateBatch, hostIDs)
server/service/async/async_label_test.go
+5
View File
@@ -6,6 +6,7 @@ import (
"testing"
"time"
"github.com/WatchBeam/clock"
"github.com/fleetdm/fleet/v4/server/datastore/mysql"
"github.com/fleetdm/fleet/v4/server/datastore/redis"
"github.com/fleetdm/fleet/v4/server/fleet"
@@ -204,6 +205,7 @@ func testCollectLabelQueryExecutions(t *testing.T, ds *mysql.Datastore, pool fle
// run the collection
var stats collectorExecStats
task := Task{
Clock: clock.C,
InsertBatch: batchSizes,
UpdateBatch: batchSizes,
DeleteBatch: batchSizes,
@@ -251,6 +253,7 @@ func testCollectLabelQueryExecutions(t *testing.T, ds *mysql.Datastore, pool fle
setupTest(t, map[int]map[int]bool{1: {1: true}})
var stats collectorExecStats
task := Task{
Clock: clock.C,
InsertBatch: batchSizes,
UpdateBatch: batchSizes,
DeleteBatch: batchSizes,
@@ -288,6 +291,7 @@ func testRecordLabelQueryExecutionsSync(t *testing.T, ds *mock.Store, pool fleet
task := Task{
Datastore: ds,
Pool: pool,
Clock: clock.C,
AsyncEnabled: false,
}
@@ -335,6 +339,7 @@ func testRecordLabelQueryExecutionsAsync(t *testing.T, ds *mock.Store, pool flee
task := Task{
Datastore: ds,
Pool: pool,
Clock: clock.C,
AsyncEnabled: true,
InsertBatch: 3,
server/service/async/async_policy.go
+1 -1
View File
@@ -195,7 +195,7 @@ func (t *Task) collectPolicyQueryExecutions(ctx context.Context, ds fleet.Datast
hostIDs[i] = host.HostID
}
ts := time.Now()
ts := t.Clock.Now()
updateBatch := make([]uint, t.UpdateBatch)
for {
n := copy(updateBatch, hostIDs)
server/service/async/async_policy_test.go
+5
View File
@@ -7,6 +7,7 @@ import (
"testing"
"time"
"github.com/WatchBeam/clock"
"github.com/fleetdm/fleet/v4/server/datastore/mysql"
"github.com/fleetdm/fleet/v4/server/datastore/redis"
"github.com/fleetdm/fleet/v4/server/fleet"
@@ -227,6 +228,7 @@ func testCollectPolicyQueryExecutions(t *testing.T, ds *mysql.Datastore, pool fl
// run the collection
var stats collectorExecStats
task := Task{
Clock: clock.C,
InsertBatch: batchSizes,
UpdateBatch: batchSizes,
DeleteBatch: batchSizes,
@@ -275,6 +277,7 @@ func testCollectPolicyQueryExecutions(t *testing.T, ds *mysql.Datastore, pool fl
setupTest(t, map[int]map[int]*bool{1: {1: nil}})
var stats collectorExecStats
task := Task{
Clock: clock.C,
InsertBatch: batchSizes,
UpdateBatch: batchSizes,
DeleteBatch: batchSizes,
@@ -312,6 +315,7 @@ func testRecordPolicyQueryExecutionsSync(t *testing.T, ds *mock.Store, pool flee
task := Task{
Datastore: ds,
Pool: pool,
Clock: clock.C,
AsyncEnabled: false,
}
@@ -359,6 +363,7 @@ func testRecordPolicyQueryExecutionsAsync(t *testing.T, ds *mock.Store, pool fle
task := Task{
Datastore: ds,
Pool: pool,
Clock: clock.C,
AsyncEnabled: true,
InsertBatch: 3,
server/service/async/async_test.go
+30 -2
View File
@@ -13,7 +13,7 @@ import (
"github.com/stretchr/testify/require"
)
func TestCollectQueryExecutions(t *testing.T) {
func TestCollect(t *testing.T) {
ds := mysql.CreateMySQLDS(t)
oldMaxPolicy := maxRedisPolicyResultsPerHost
@@ -49,9 +49,23 @@ func TestCollectQueryExecutions(t *testing.T) {
testCollectPolicyQueryExecutions(t, ds, pool)
})
})
t.Run("Host Last Seen", func(t *testing.T) {
t.Run("standalone", func(t *testing.T) {
defer mysql.TruncateTables(t, ds)
pool := redistest.SetupRedis(t, "host_last_seen", false, false, false)
testCollectHostsLastSeen(t, ds, pool)
})
t.Run("cluster", func(t *testing.T) {
defer mysql.TruncateTables(t, ds)
pool := redistest.SetupRedis(t, "host_last_seen", true, true, false)
testCollectHostsLastSeen(t, ds, pool)
})
})
}
func TestRecordQueryExecutions(t *testing.T) {
func TestRecord(t *testing.T) {
ds := new(mock.Store)
ds.RecordLabelQueryExecutionsFunc = func(ctx context.Context, host *fleet.Host, results map[uint]*bool, ts time.Time, deferred bool) error {
return nil
@@ -96,6 +110,20 @@ func TestRecordQueryExecutions(t *testing.T) {
t.Run("async", func(t *testing.T) { testRecordPolicyQueryExecutionsAsync(t, ds, pool) })
})
})
t.Run("Host Last Seen", func(t *testing.T) {
t.Run("standalone", func(t *testing.T) {
pool := redistest.SetupRedis(t, "host_last_seen", false, false, false)
t.Run("sync", func(t *testing.T) { testRecordHostLastSeenSync(t, ds, pool) })
t.Run("async", func(t *testing.T) { testRecordHostLastSeenAsync(t, ds, pool) })
})
t.Run("cluster", func(t *testing.T) {
pool := redistest.SetupRedis(t, "host_last_seen", true, true, false)
t.Run("sync", func(t *testing.T) { testRecordHostLastSeenSync(t, ds, pool) })
t.Run("async", func(t *testing.T) { testRecordHostLastSeenAsync(t, ds, pool) })
})
})
}
func TestActiveHostIDsSet(t *testing.T) {
server/service/async/collect.go
+1 -1
View File
@@ -71,7 +71,7 @@ type collectorExecStats struct {
Inserts int
Updates int
Deletes int
RedisCmds int // does not include scan keys iteration commands
RedisCmds int // script counts as 1, does not include ZSCAN of active hosts iteration commands
Failed bool
}
server/service/async/doc.go
-4
View File
@@ -5,8 +5,4 @@
// setups with lots of hosts, by collecting the data in fast storage (Redis)
// and then using a background task to store it down to persistent storage
// (mysql) in a controlled manner.
//
// TODO: for now, only label queries are done like this - if it proves to be
// efficient and manageable, policy queries and saving of hosts could use
// this approach too.
package async
server/service/hosts.go
-6
View File
@@ -52,12 +52,6 @@ func hostDetailResponseForHost(ctx context.Context, svc fleet.Service, host *fle
}, nil
}
func (svc *Service) FlushSeenHosts(ctx context.Context) error {
// No authorization check because this is used only internally.
hostIDs := svc.seenHostSet.getAndClearHostIDs()
return svc.ds.MarkHostsSeen(ctx, hostIDs, svc.clock.Now())
}
////////////////////////////////////////////////////////////////////////////////
// List Hosts
////////////////////////////////////////////////////////////////////////////////
server/service/hosts_test.go
+1 -1
View File
@@ -218,7 +218,7 @@ func TestHostAuth(t *testing.T) {
})
}
// List, GetHostSummary, FlushSeenHost work for all
// List, GetHostSummary work for all
}
func TestListHosts(t *testing.T) {
server/service/integration_ds_only_test.go
+1 -1
View File
@@ -44,7 +44,7 @@ func (s *integrationDSTestSuite) TestLicenseExpiration() {
t := s.T()
license := &fleet.LicenseInfo{Tier: tt.tier, Expiration: tt.expiration}
_, server := RunServerForTestsWithDS(t, s.ds, TestServerOpts{License: license, SkipCreateTestUsers: true})
_, server := RunServerForTestsWithDS(t, s.ds, &TestServerOpts{License: license, SkipCreateTestUsers: true})
ts := withServer{server: server}
ts.s = &s.Suite
server/service/integration_enterprise_test.go
+1 -1
View File
@@ -31,7 +31,7 @@ func (s *integrationEnterpriseTestSuite) SetupSuite() {
s.withDS.SetupSuite("integrationEnterpriseTestSuite")
users, server := RunServerForTestsWithDS(
s.T(), s.ds, TestServerOpts{License: &fleet.LicenseInfo{Tier: fleet.TierPremium}})
s.T(), s.ds, &TestServerOpts{License: &fleet.LicenseInfo{Tier: fleet.TierPremium}})
s.server = server
s.users = users
s.token = s.getTestAdminToken()
server/service/integration_live_queries_test.go
+1 -1
View File
@@ -44,7 +44,7 @@ func (s *liveQueriesTestSuite) SetupSuite() {
lq := new(live_query.MockLiveQuery)
s.lq = lq
users, server := RunServerForTestsWithDS(s.T(), s.ds, TestServerOpts{Lq: lq, Rs: rs})
users, server := RunServerForTestsWithDS(s.T(), s.ds, &TestServerOpts{Lq: lq, Rs: rs})
s.server = server
s.users = users
s.token = getTestAdminToken(s.T(), s.server)
server/service/integration_logger_test.go
+1 -1
View File
@@ -42,7 +42,7 @@ func (s *integrationLoggerTestSuite) SetupSuite() {
logger := log.NewJSONLogger(s.buf)
logger = level.NewFilter(logger, level.AllowDebug())
users, server := RunServerForTestsWithDS(s.T(), s.ds, TestServerOpts{Logger: logger})
users, server := RunServerForTestsWithDS(s.T(), s.ds, &TestServerOpts{Logger: logger})
s.server = server
s.users = users
}
server/service/integration_sso_test.go
+1 -1
View File
@@ -28,7 +28,7 @@ func (s *integrationSSOTestSuite) SetupSuite() {
s.withDS.SetupSuite("integrationSSOTestSuite")
pool := redistest.SetupRedis(s.T(), "zz", false, false, false)
users, server := RunServerForTestsWithDS(s.T(), s.ds, TestServerOpts{Pool: pool})
users, server := RunServerForTestsWithDS(s.T(), s.ds, &TestServerOpts{Pool: pool})
s.server = server
s.users = users
s.token = s.getTestAdminToken()
server/service/osquery.go
+3 -1
View File
@@ -69,7 +69,9 @@ func (svc *Service) AuthenticateHost(ctx context.Context, nodeKey string) (*flee
// updating the seen time for these hosts. This seems to be an acceptable
// tradeoff as an online host will continue to check in and quickly be
// marked online again.
svc.seenHostSet.addHostID(host.ID)
if err := svc.task.RecordHostLastSeen(ctx, host.ID); err != nil {
logging.WithErr(ctx, ctxerr.Wrap(ctx, err, "record host last seen"))
}
host.SeenTime = svc.clock.Now()
return host, svc.debugEnabledForHost(ctx, host.ID), nil
server/service/osquery_test.go
+8 -5
View File
@@ -28,6 +28,7 @@ import (
"github.com/fleetdm/fleet/v4/server/mock"
"github.com/fleetdm/fleet/v4/server/ptr"
"github.com/fleetdm/fleet/v4/server/pubsub"
"github.com/fleetdm/fleet/v4/server/service/async"
"github.com/fleetdm/fleet/v4/server/service/osquery_utils"
"github.com/fleetdm/fleet/v4/server/service/redis_policy_set"
"github.com/go-kit/kit/log"
@@ -278,7 +279,8 @@ func TestEnrollAgentDetails(t *testing.T) {
func TestAuthenticateHost(t *testing.T) {
ds := new(mock.Store)
svc := newTestService(t, ds, nil, nil)
task := &async.Task{Datastore: ds, AsyncEnabled: false}
svc := newTestService(t, ds, nil, nil, &TestServerOpts{Task: task})
var gotKey string
host := fleet.Host{ID: 1, Hostname: "foobar"}
@@ -312,12 +314,13 @@ func TestAuthenticateHost(t *testing.T) {
assert.Equal(t, "floobar", gotKey)
assert.False(t, ds.MarkHostsSeenFuncInvoked)
err = svc.FlushSeenHosts(context.Background())
err = task.FlushHostsLastSeen(context.Background(), time.Now())
require.NoError(t, err)
assert.True(t, ds.MarkHostsSeenFuncInvoked)
ds.MarkHostsSeenFuncInvoked = false
assert.ElementsMatch(t, []uint{1, 7}, gotHostIDs)
err = svc.FlushSeenHosts(context.Background())
err = task.FlushHostsLastSeen(context.Background(), time.Now())
require.NoError(t, err)
assert.True(t, ds.MarkHostsSeenFuncInvoked)
require.Len(t, gotHostIDs, 0)
@@ -2260,7 +2263,7 @@ func TestPolicyWebhooks(t *testing.T) {
pool := redistest.SetupRedis(t, t.Name(), false, false, false)
failingPolicySet := redis_policy_set.NewFailingTest(t, pool)
testConfig := config.TestConfig()
svc := newTestServiceWithConfig(t, ds, testConfig, nil, lq, TestServerOpts{
svc := newTestServiceWithConfig(t, ds, testConfig, nil, lq, &TestServerOpts{
FailingPolicySet: failingPolicySet,
Clock: mockClock,
})
@@ -2527,7 +2530,7 @@ func TestLiveQueriesFailing(t *testing.T) {
cfg := config.TestConfig()
buf := new(bytes.Buffer)
logger := log.NewLogfmtLogger(buf)
svc := newTestServiceWithConfig(t, ds, cfg, nil, lq, TestServerOpts{
svc := newTestServiceWithConfig(t, ds, cfg, nil, lq, &TestServerOpts{
Logger: logger,
})
server/service/service.go
-36
View File
@@ -38,8 +38,6 @@ type Service struct {
mailService fleet.MailService
ssoSessionStore sso.SessionStore
seenHostSet *seenHostSet
failingPolicySet fleet.FailingPolicySet
authz *authz.Authorizer
@@ -89,7 +87,6 @@ func NewService(
osqueryLogWriter: osqueryLogger,
mailService: mailService,
ssoSessionStore: sso,
seenHostSet: newSeenHostSet(),
license: license,
failingPolicySet: failingPolicySet,
authz: authorizer,
@@ -114,36 +111,3 @@ type validationMiddleware struct {
func getAssetURL() template.URL {
return template.URL("https://fleetdm.com/images/permanent")
}
// seenHostSet implements synchronized storage for the set of seen hosts.
type seenHostSet struct {
mutex sync.Mutex
hostIDs map[uint]bool
}
func newSeenHostSet() *seenHostSet {
return &seenHostSet{
mutex: sync.Mutex{},
hostIDs: make(map[uint]bool),
}
}
// addHostID adds the host identified by ID to the set
func (m *seenHostSet) addHostID(id uint) {
m.mutex.Lock()
defer m.mutex.Unlock()
m.hostIDs[id] = true
}
// getAndClearHostIDs gets the list of unique host IDs from the set and empties
// the set.
func (m *seenHostSet) getAndClearHostIDs() []uint {
m.mutex.Lock()
defer m.mutex.Unlock()
var ids []uint
for id := range m.hostIDs {
ids = append(ids, id)
}
m.hostIDs = make(map[uint]bool)
return ids
}
server/service/teams_test.go
+1 -1
View File
@@ -14,7 +14,7 @@ import (
func TestTeamAuth(t *testing.T) {
ds := new(mock.Store)
license := &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}
svc := newTestService(t, ds, nil, nil, TestServerOpts{License: license, SkipCreateTestUsers: true})
svc := newTestService(t, ds, nil, nil, &TestServerOpts{License: license, SkipCreateTestUsers: true})
ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
return &fleet.AppConfig{}, nil
server/service/testing_client.go
+1 -1
View File
@@ -46,7 +46,7 @@ func (ts *withServer) SetupSuite(dbName string) {
ts.withDS.SetupSuite(dbName)
rs := pubsub.NewInmemQueryResults()
users, server := RunServerForTestsWithDS(ts.s.T(), ts.ds, TestServerOpts{Rs: rs})
users, server := RunServerForTestsWithDS(ts.s.T(), ts.ds, &TestServerOpts{Rs: rs})
ts.server = server
ts.users = users
ts.token = ts.getTestAdminToken()
server/service/testing_utils.go
+18 -8
View File
@@ -26,11 +26,11 @@ import (
"github.com/throttled/throttled/v2/store/memstore"
)
func newTestService(t *testing.T, ds fleet.Datastore, rs fleet.QueryResultStore, lq fleet.LiveQueryStore, opts ...TestServerOpts) fleet.Service {
func newTestService(t *testing.T, ds fleet.Datastore, rs fleet.QueryResultStore, lq fleet.LiveQueryStore, opts ...*TestServerOpts) fleet.Service {
return newTestServiceWithConfig(t, ds, config.TestConfig(), rs, lq, opts...)
}
func newTestServiceWithConfig(t *testing.T, ds fleet.Datastore, fleetConfig config.FleetConfig, rs fleet.QueryResultStore, lq fleet.LiveQueryStore, opts ...TestServerOpts) fleet.Service {
func newTestServiceWithConfig(t *testing.T, ds fleet.Datastore, fleetConfig config.FleetConfig, rs fleet.QueryResultStore, lq fleet.LiveQueryStore, opts ...*TestServerOpts) fleet.Service {
mailer := &mockMailService{SendEmailFn: func(e fleet.Email) error { return nil }}
license := &fleet.LicenseInfo{Tier: fleet.TierFree}
writer, err := logging.NewFilesystemLogWriter(
@@ -49,6 +49,12 @@ func newTestServiceWithConfig(t *testing.T, ds fleet.Datastore, fleetConfig conf
var failingPolicySet fleet.FailingPolicySet = NewMemFailingPolicySet()
var c clock.Clock = clock.C
task := &async.Task{
Datastore: ds,
AsyncEnabled: false,
}
if len(opts) > 0 {
if opts[0].Logger != nil {
logger = opts[0].Logger
@@ -65,11 +71,14 @@ func newTestServiceWithConfig(t *testing.T, ds fleet.Datastore, fleetConfig conf
if opts[0].Clock != nil {
c = opts[0].Clock
}
if opts[0].Task != nil {
task = opts[0].Task
} else {
opts[0].Task = task
}
}
task := &async.Task{
Datastore: ds,
AsyncEnabled: false,
}
task.Clock = c
svc, err := NewService(context.Background(), ds, task, rs, logger, osqlogger, fleetConfig, mailer, c, ssoStore, lq, ds, *license, failingPolicySet, &fleet.NoOpGeoIP{})
if err != nil {
panic(err)
@@ -85,7 +94,7 @@ func newTestServiceWithConfig(t *testing.T, ds fleet.Datastore, fleetConfig conf
func newTestServiceWithClock(t *testing.T, ds fleet.Datastore, rs fleet.QueryResultStore, lq fleet.LiveQueryStore, c clock.Clock) fleet.Service {
testConfig := config.TestConfig()
svc := newTestServiceWithConfig(t, ds, testConfig, rs, lq, TestServerOpts{
svc := newTestServiceWithConfig(t, ds, testConfig, rs, lq, &TestServerOpts{
Clock: c,
})
return svc
@@ -156,9 +165,10 @@ type TestServerOpts struct {
Pool fleet.RedisPool
FailingPolicySet fleet.FailingPolicySet
Clock clock.Clock
Task *async.Task
}
func RunServerForTestsWithDS(t *testing.T, ds fleet.Datastore, opts ...TestServerOpts) (map[string]fleet.User, *httptest.Server) {
func RunServerForTestsWithDS(t *testing.T, ds fleet.Datastore, opts ...*TestServerOpts) (map[string]fleet.User, *httptest.Server) {
var rs fleet.QueryResultStore
if len(opts) > 0 && opts[0].Rs != nil {
rs = opts[0].Rs