Refactor async host processing to avoid redis SCAN keys (for policies) (#3657)
This commit is contained in:
Martin Angers
@@ -432,19 +432,19 @@ The following options are available when configuring SMTP authentication:
|
||||
|
||||
##### Host Status
|
||||
|
||||
The following options allow the configuration of a webhook that will be triggered if the specified percentage of hosts
|
||||
The following options allow the configuration of a webhook that will be triggered if the specified percentage of hosts
|
||||
are offline for the specified amount of time.
|
||||
|
||||
- `webhook_settings.host_status_webhook.enable_host_status_webhook`: true or false. Defines whether the check for host status will run or not.
|
||||
- `webhook_settings.host_status_webhook.destination_url`: the URL to POST to when the condition for the webhook triggers.
|
||||
- `webhook_settings.host_status_webhook.host_percentage`: the percentage of hosts that need to be offline
|
||||
- `webhook_settings.host_status_webhook.host_percentage`: the percentage of hosts that need to be offline
|
||||
- `webhook_settings.host_status_webhook.days_count`: amount of days that hosts need to be offline for to count as part of the percentage.
|
||||
|
||||
##### Failing Policies
|
||||
|
||||
The following options allow the configuration of a webhook that will be triggered if selected policies are not passing for some hosts.
|
||||
|
||||
- `webhook_settings.failing_policies_webhook.enable_failing_policies_webhook`: true or false. Defines whether to enable the failing policies webhook.
|
||||
- `webhook_settings.failing_policies_webhook.enable_failing_policies_webhook`: true or false. Defines whether to enable the failing policies webhook. Note that currently, if the failing policies webhook *and* the `osquery.enable_async_host_processing` options 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).
|
||||
- `webhook_settings.failing_policies_webhook.destination_url`: the URL to POST to when the condition for the webhook triggers.
|
||||
- `webhook_settings.failing_policies_webhook.policy_ids`: the IDs of the policies for which the webhook will be enabled.
|
||||
- `webhook_settings.failing_policies_webhook.host_batch_size`: Maximum number of hosts to batch on POST requests. A value of `0`, the default, means no batching, all hosts failing a policy will be sent on one POST request.
|
||||
@@ -454,7 +454,7 @@ The following options allow the configuration of a webhook that will be triggere
|
||||
There's a lot of information coming from hosts, but it's sometimes useful to see exactly what a host is returning in order
|
||||
to debug different scenarios.
|
||||
|
||||
So for example, let's say the hosts with ids 342 and 98 are not behaving as you expect in Fleet, you can enable verbose
|
||||
So for example, let's say the hosts with ids 342 and 98 are not behaving as you expect in Fleet, you can enable verbose
|
||||
logging with the following configuration:
|
||||
|
||||
```yaml
|
||||
|
||||
@@ -965,7 +965,9 @@ 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 results. This may improve performance and CPU usage of the fleet instances and MySQL database servers for setups with a large number of hosts (100 000+), while requiring more resources from Redis server(s). Using Redis Cluster is recommended to enable this mode.
|
||||
**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).
|
||||
|
||||
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).
|
||||
|
||||
- Default value: false
|
||||
- Environment variable: `FLEET_OSQUERY_ENABLE_ASYNC_HOST_PROCESSING`
|
||||
|
||||
@@ -363,6 +363,14 @@ func (d *Datastore) RecordLabelQueryExecutions(ctx context.Context, host *fleet.
|
||||
}
|
||||
}
|
||||
|
||||
// NOTE: the insert/delete of label membership that follows must be kept in
|
||||
// sync with the async implementations in
|
||||
// AsyncBatch{Insert,Delete}LabelMembership, and the update of the
|
||||
// label_updated_at timestamp in sync with the
|
||||
// AsyncBatchUpdateLabelTimestamp method (that is, their processing must be
|
||||
// semantically equivalent, even though here it processes a single host and
|
||||
// in async mode it processes a batch of hosts).
|
||||
|
||||
err := d.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
|
||||
// Complete inserts if necessary
|
||||
if len(vals) > 0 {
|
||||
@@ -755,7 +763,7 @@ func (d *Datastore) AsyncBatchDeleteLabelMembership(ctx context.Context, batch [
|
||||
})
|
||||
}
|
||||
|
||||
// AsyncBatchUpdateLabelTimestamp updates the table the hosts' label_updated_at timestamp
|
||||
// AsyncBatchUpdateLabelTimestamp updates the hosts' label_updated_at timestamp
|
||||
// for the batch of host ids provided.
|
||||
func (d *Datastore) AsyncBatchUpdateLabelTimestamp(ctx context.Context, ids []uint, ts time.Time) error {
|
||||
// NOTE: this is tested via the server/service/async package tests.
|
||||
|
||||
@@ -203,6 +203,13 @@ func (ds *Datastore) RecordPolicyQueryExecutions(ctx context.Context, host *flee
|
||||
vals = append(vals, updated, policyID, host.ID, matches)
|
||||
}
|
||||
|
||||
// NOTE: the insert of policy membership that follows must be kept in sync
|
||||
// with the async implementation in AsyncBatchInsertPolicyMembership, and the
|
||||
// update of the policy_updated_at timestamp in sync with the
|
||||
// AsyncBatchUpdatePolicyTimestamp method (that is, their processing must be
|
||||
// semantically equivalent, even though here it processes a single host and
|
||||
// in async mode it processes a batch of hosts).
|
||||
|
||||
query := fmt.Sprintf(
|
||||
`INSERT INTO policy_membership (updated_at, policy_id, host_id, passes)
|
||||
VALUES %s ON DUPLICATE KEY UPDATE updated_at=VALUES(updated_at), passes=VALUES(passes)`,
|
||||
@@ -435,3 +442,48 @@ func amountPoliciesDB(db sqlx.Queryer) (int, error) {
|
||||
}
|
||||
return amount, nil
|
||||
}
|
||||
|
||||
// AsyncBatchInsertPolicyMembership inserts into the policy_membership table
|
||||
// the batch of policy membership results.
|
||||
func (ds *Datastore) AsyncBatchInsertPolicyMembership(ctx context.Context, batch []fleet.PolicyMembershipResult) error {
|
||||
// NOTE: this is tested via the server/service/async package tests.
|
||||
|
||||
// INSERT IGNORE, to avoid failing if policy / host does not exist (as this
|
||||
// runs asynchronously, they could get deleted in between the data being
|
||||
// received and being upserted).
|
||||
sql := `INSERT IGNORE INTO policy_membership (policy_id, host_id, passes) VALUES `
|
||||
sql += strings.Repeat(`(?, ?, ?),`, len(batch))
|
||||
sql = strings.TrimSuffix(sql, ",")
|
||||
sql += ` ON DUPLICATE KEY UPDATE updated_at = VALUES(updated_at), passes = VALUES(passes)`
|
||||
|
||||
vals := make([]interface{}, 0, len(batch)*3)
|
||||
for _, tup := range batch {
|
||||
vals = append(vals, tup.PolicyID, tup.HostID, tup.Passes)
|
||||
}
|
||||
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
|
||||
_, err := tx.ExecContext(ctx, sql, vals...)
|
||||
return ctxerr.Wrap(ctx, err, "insert into policy_membership")
|
||||
})
|
||||
}
|
||||
|
||||
// AsyncBatchUpdatePolicyTimestamp updates the hosts' policy_updated_at timestamp
|
||||
// for the batch of host ids provided.
|
||||
func (ds *Datastore) AsyncBatchUpdatePolicyTimestamp(ctx context.Context, ids []uint, ts time.Time) error {
|
||||
// NOTE: this is tested via the server/service/async package tests.
|
||||
|
||||
sql := `
|
||||
UPDATE
|
||||
hosts
|
||||
SET
|
||||
policy_updated_at = ?
|
||||
WHERE
|
||||
id IN (?)`
|
||||
query, args, err := sqlx.In(sql, ts, ids)
|
||||
if err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "building query to update hosts.policy_updated_at")
|
||||
}
|
||||
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
|
||||
_, err := tx.ExecContext(ctx, query, args...)
|
||||
return ctxerr.Wrap(ctx, err, "update hosts.policy_updated_at")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -355,6 +355,10 @@ type Datastore interface {
|
||||
PolicyQueriesForHost(ctx context.Context, host *Host) (map[string]string, error)
|
||||
ApplyPolicySpecs(ctx context.Context, authorID uint, specs []*PolicySpec) error
|
||||
|
||||
// Methods used for async processing of host policy query results.
|
||||
AsyncBatchInsertPolicyMembership(ctx context.Context, batch []PolicyMembershipResult) error
|
||||
AsyncBatchUpdatePolicyTimestamp(ctx context.Context, ids []uint, ts time.Time) error
|
||||
|
||||
// MigrateTables creates and migrates the table schemas
|
||||
MigrateTables(ctx context.Context) error
|
||||
// MigrateData populates built-in data
|
||||
|
||||
@@ -231,3 +231,9 @@ type PolicySetHost struct {
|
||||
// Hostname is the host's name.
|
||||
Hostname string
|
||||
}
|
||||
|
||||
type PolicyMembershipResult struct {
|
||||
HostID uint
|
||||
PolicyID uint
|
||||
Passes *bool
|
||||
}
|
||||
|
||||
@@ -288,6 +288,10 @@ type PolicyQueriesForHostFunc func(ctx context.Context, host *fleet.Host) (map[s
|
||||
|
||||
type ApplyPolicySpecsFunc func(ctx context.Context, authorID uint, specs []*fleet.PolicySpec) error
|
||||
|
||||
type AsyncBatchInsertPolicyMembershipFunc func(ctx context.Context, batch []fleet.PolicyMembershipResult) error
|
||||
|
||||
type AsyncBatchUpdatePolicyTimestampFunc func(ctx context.Context, ids []uint, ts time.Time) error
|
||||
|
||||
type MigrateTablesFunc func(ctx context.Context) error
|
||||
|
||||
type MigrateDataFunc func(ctx context.Context) error
|
||||
@@ -771,6 +775,12 @@ type DataStore struct {
|
||||
ApplyPolicySpecsFunc ApplyPolicySpecsFunc
|
||||
ApplyPolicySpecsFuncInvoked bool
|
||||
|
||||
AsyncBatchInsertPolicyMembershipFunc AsyncBatchInsertPolicyMembershipFunc
|
||||
AsyncBatchInsertPolicyMembershipFuncInvoked bool
|
||||
|
||||
AsyncBatchUpdatePolicyTimestampFunc AsyncBatchUpdatePolicyTimestampFunc
|
||||
AsyncBatchUpdatePolicyTimestampFuncInvoked bool
|
||||
|
||||
MigrateTablesFunc MigrateTablesFunc
|
||||
MigrateTablesFuncInvoked bool
|
||||
|
||||
@@ -1564,6 +1574,16 @@ func (s *DataStore) ApplyPolicySpecs(ctx context.Context, authorID uint, specs [
|
||||
return s.ApplyPolicySpecsFunc(ctx, authorID, specs)
|
||||
}
|
||||
|
||||
func (s *DataStore) AsyncBatchInsertPolicyMembership(ctx context.Context, batch []fleet.PolicyMembershipResult) error {
|
||||
s.AsyncBatchInsertPolicyMembershipFuncInvoked = true
|
||||
return s.AsyncBatchInsertPolicyMembershipFunc(ctx, batch)
|
||||
}
|
||||
|
||||
func (s *DataStore) AsyncBatchUpdatePolicyTimestamp(ctx context.Context, ids []uint, ts time.Time) error {
|
||||
s.AsyncBatchUpdatePolicyTimestampFuncInvoked = true
|
||||
return s.AsyncBatchUpdatePolicyTimestampFunc(ctx, ids, ts)
|
||||
}
|
||||
|
||||
func (s *DataStore) MigrateTables(ctx context.Context) error {
|
||||
s.MigrateTablesFuncInvoked = true
|
||||
return s.MigrateTablesFunc(ctx)
|
||||
|
||||
+93
-299
@@ -5,7 +5,6 @@ import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
|
||||
"github.com/fleetdm/fleet/v4/server/datastore/redis"
|
||||
"github.com/fleetdm/fleet/v4/server/fleet"
|
||||
kitlog "github.com/go-kit/kit/log"
|
||||
@@ -13,13 +12,7 @@ import (
|
||||
redigo "github.com/gomodule/redigo/redis"
|
||||
)
|
||||
|
||||
const (
|
||||
labelMembershipActiveHostIDsKey = "label_membership:active_host_ids"
|
||||
labelMembershipHostKey = "label_membership:{%d}"
|
||||
labelMembershipReportedKey = "label_membership_reported:{%d}"
|
||||
labelMembershipKeysMinTTL = 7 * 24 * time.Hour // 1 week
|
||||
collectorLockKey = "locks:async_collector:{%s}"
|
||||
)
|
||||
const collectorLockKey = "locks:async_collector:{%s}"
|
||||
|
||||
type Task struct {
|
||||
Datastore fleet.Datastore
|
||||
@@ -48,6 +41,10 @@ func (t *Task) StartCollectors(ctx context.Context, jitterPct int, logger kitlog
|
||||
}
|
||||
level.Debug(logger).Log("task", "async enabled, starting collectors", "interval", t.CollectorInterval, "jitter", jitterPct)
|
||||
|
||||
collectorErrHandler := func(name string, err error) {
|
||||
level.Error(logger).Log("err", fmt.Sprintf("%s collector", name), "details", err)
|
||||
}
|
||||
|
||||
labelColl := &collector{
|
||||
name: "collect_labels",
|
||||
pool: t.Pool,
|
||||
@@ -56,11 +53,24 @@ func (t *Task) StartCollectors(ctx context.Context, jitterPct int, logger kitlog
|
||||
jitterPct: jitterPct,
|
||||
lockTimeout: t.LockTimeout,
|
||||
handler: t.collectLabelQueryExecutions,
|
||||
errHandler: func(name string, err error) {
|
||||
level.Error(logger).Log("err", fmt.Sprintf("%s collector", name), "details", err)
|
||||
},
|
||||
errHandler: collectorErrHandler,
|
||||
}
|
||||
|
||||
policyColl := &collector{
|
||||
name: "collect_policies",
|
||||
pool: t.Pool,
|
||||
ds: t.Datastore,
|
||||
execInterval: t.CollectorInterval,
|
||||
jitterPct: jitterPct,
|
||||
lockTimeout: t.LockTimeout,
|
||||
handler: t.collectPolicyQueryExecutions,
|
||||
errHandler: collectorErrHandler,
|
||||
}
|
||||
|
||||
colls := []*collector{labelColl, policyColl}
|
||||
for _, coll := range colls {
|
||||
go coll.Start(ctx)
|
||||
}
|
||||
go labelColl.Start(ctx)
|
||||
|
||||
// log stats at regular intervals
|
||||
if t.LogStatsInterval > 0 {
|
||||
@@ -69,8 +79,10 @@ func (t *Task) StartCollectors(ctx context.Context, jitterPct int, logger kitlog
|
||||
for {
|
||||
select {
|
||||
case <-tick:
|
||||
stats := labelColl.ReadStats()
|
||||
level.Debug(logger).Log("stats", fmt.Sprintf("%#v", stats), "name", labelColl.name)
|
||||
for _, coll := range colls {
|
||||
stats := coll.ReadStats()
|
||||
level.Debug(logger).Log("stats", fmt.Sprintf("%#v", stats), "name", coll.name)
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
@@ -79,234 +91,13 @@ func (t *Task) StartCollectors(ctx context.Context, jitterPct int, logger kitlog
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task) RecordLabelQueryExecutions(ctx context.Context, host *fleet.Host, results map[uint]*bool, ts time.Time) error {
|
||||
if !t.AsyncEnabled {
|
||||
host.LabelUpdatedAt = ts
|
||||
return t.Datastore.RecordLabelQueryExecutions(ctx, host, results, ts, false)
|
||||
}
|
||||
|
||||
keySet := fmt.Sprintf(labelMembershipHostKey, host.ID)
|
||||
keyTs := fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
|
||||
// set an expiration on both keys (set and ts), ensuring that a deleted host
|
||||
// (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.
|
||||
//
|
||||
// This means that it will only expire if that host hasn't reported labels
|
||||
// during that (TTL) time (each time it does report, the TTL is reset), and
|
||||
// the collector will have plenty of time to run (multiple times) to try to
|
||||
// persist all the data in mysql.
|
||||
ttl := labelMembershipKeysMinTTL
|
||||
if maxTTL := 10 * t.CollectorInterval; maxTTL > ttl {
|
||||
ttl = maxTTL
|
||||
}
|
||||
|
||||
// keys and arguments passed to the script are:
|
||||
// KEYS[1]: keySet (labelMembershipHostKey)
|
||||
// KEYS[2]: keyTs (labelMembershipReportedKey)
|
||||
// ARGV[1]: timestamp for "reported at"
|
||||
// ARGV[2]: ttl for both keys
|
||||
// ARGV[3..]: the arguments to ZADD to keySet
|
||||
script := redigo.NewScript(2, `
|
||||
redis.call('ZADD', KEYS[1], unpack(ARGV, 3))
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||||
redis.call('SET', KEYS[2], ARGV[1])
|
||||
return redis.call('EXPIRE', KEYS[2], ARGV[2])
|
||||
`)
|
||||
|
||||
// convert results to ZADD arguments, store as -1 for delete, +1 for insert
|
||||
args := make(redigo.Args, 0, 4+(len(results)*2))
|
||||
args = args.Add(keySet, keyTs, ts.Unix(), int(ttl.Seconds()))
|
||||
for k, v := range results {
|
||||
score := -1
|
||||
if v != nil && *v {
|
||||
score = 1
|
||||
}
|
||||
args = args.Add(score, k)
|
||||
}
|
||||
|
||||
conn := t.Pool.Get()
|
||||
defer conn.Close()
|
||||
if err := redis.BindConn(t.Pool, conn, keySet, keyTs); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "bind redis connection")
|
||||
}
|
||||
|
||||
if _, err := script.Do(conn, args...); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "run redis script")
|
||||
}
|
||||
|
||||
// Storing the host id in the set of active host IDs for label membership
|
||||
// outside of the redis script because in Redis Cluster mode the key may not
|
||||
// live on the same node as the host's keys. At the same time, purge any
|
||||
// entry in the set that is older than now - TTL.
|
||||
if err := storePurgeActiveHostID(t.Pool, host.ID, ts, ts.Add(-ttl)); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "store active host id")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Task) collectLabelQueryExecutions(ctx context.Context, ds fleet.Datastore, pool fleet.RedisPool, stats *collectorExecStats) error {
|
||||
hosts, err := loadActiveHostIDs(pool, t.RedisScanKeysCount)
|
||||
if err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "load active host ids")
|
||||
}
|
||||
stats.Keys = len(hosts)
|
||||
|
||||
getKeyTuples := func(hostID uint) (inserts, deletes [][2]uint, err error) {
|
||||
keySet := fmt.Sprintf(labelMembershipHostKey, hostID)
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
for {
|
||||
stats.RedisCmds++
|
||||
|
||||
vals, err := redigo.Ints(conn.Do("ZPOPMIN", keySet, t.RedisPopCount))
|
||||
if err != nil {
|
||||
return nil, nil, ctxerr.Wrap(ctx, err, "redis ZPOPMIN")
|
||||
}
|
||||
items := len(vals) / 2 // each item has the label id and the score (-1=delete, +1=insert)
|
||||
stats.Items += items
|
||||
|
||||
for i := 0; i < len(vals); i += 2 {
|
||||
labelID := vals[i]
|
||||
|
||||
var score int
|
||||
if i+1 < len(vals) { // just to be safe we received all pairs
|
||||
score = vals[i+1]
|
||||
}
|
||||
|
||||
switch score {
|
||||
case 1:
|
||||
inserts = append(inserts, [2]uint{uint(labelID), hostID})
|
||||
case -1:
|
||||
deletes = append(deletes, [2]uint{uint(labelID), hostID})
|
||||
}
|
||||
}
|
||||
if items < t.RedisPopCount {
|
||||
return inserts, deletes, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Based on those pages, the best approach appears to be INSERT with multiple
|
||||
// rows in the VALUES section (short of doing LOAD FILE, which we can't):
|
||||
// https://www.databasejournal.com/features/mysql/optimize-mysql-inserts-using-batch-processing.html
|
||||
// https://dev.mysql.com/doc/refman/5.7/en/insert-optimization.html
|
||||
// https://dev.mysql.com/doc/refman/5.7/en/optimizing-innodb-bulk-data-loading.html
|
||||
//
|
||||
// Given that there are no UNIQUE constraints in label_membership (well,
|
||||
// apart from the primary key columns), no AUTO_INC column and no FOREIGN
|
||||
// KEY, there is no obvious setting to tweak (based on the recommendations of
|
||||
// the third link above).
|
||||
//
|
||||
// However, in label_membership, updated_at defaults to the current timestamp
|
||||
// both on INSERT and when UPDATEd, so it does not need to be provided.
|
||||
|
||||
runInsertBatch := func(batch [][2]uint) error {
|
||||
stats.Inserts++
|
||||
return ds.AsyncBatchInsertLabelMembership(ctx, batch)
|
||||
}
|
||||
|
||||
runDeleteBatch := func(batch [][2]uint) error {
|
||||
stats.Deletes++
|
||||
return ds.AsyncBatchDeleteLabelMembership(ctx, batch)
|
||||
}
|
||||
|
||||
runUpdateBatch := func(ids []uint, ts time.Time) error {
|
||||
stats.Updates++
|
||||
return ds.AsyncBatchUpdateLabelTimestamp(ctx, ids, ts)
|
||||
}
|
||||
|
||||
insertBatch := make([][2]uint, 0, t.InsertBatch)
|
||||
deleteBatch := make([][2]uint, 0, t.DeleteBatch)
|
||||
for _, host := range hosts {
|
||||
hid := host.HostID
|
||||
ins, del, err := getKeyTuples(hid)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
insertBatch = append(insertBatch, ins...)
|
||||
deleteBatch = append(deleteBatch, del...)
|
||||
|
||||
if len(insertBatch) >= t.InsertBatch {
|
||||
if err := runInsertBatch(insertBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
insertBatch = insertBatch[:0]
|
||||
}
|
||||
if len(deleteBatch) >= t.DeleteBatch {
|
||||
if err := runDeleteBatch(deleteBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
deleteBatch = deleteBatch[:0]
|
||||
}
|
||||
}
|
||||
|
||||
// process any remaining batch that did not reach the batchSize limit in the
|
||||
// loop.
|
||||
if len(insertBatch) > 0 {
|
||||
if err := runInsertBatch(insertBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if len(deleteBatch) > 0 {
|
||||
if err := runDeleteBatch(deleteBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if len(hosts) > 0 {
|
||||
hostIDs := make([]uint, len(hosts))
|
||||
for i, host := range hosts {
|
||||
hostIDs[i] = host.HostID
|
||||
}
|
||||
|
||||
ts := time.Now()
|
||||
updateBatch := make([]uint, t.UpdateBatch)
|
||||
for {
|
||||
n := copy(updateBatch, hostIDs)
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
if err := runUpdateBatch(updateBatch[:n], ts); err != nil {
|
||||
return err
|
||||
}
|
||||
hostIDs = hostIDs[n:]
|
||||
}
|
||||
|
||||
// batch-remove any host ID from the active set that still has its score to
|
||||
// the initial value, so that the active set does not keep all (potentially
|
||||
// 100K+) host IDs to process at all times - only those with reported
|
||||
// results to process.
|
||||
if err := removeProcessedHostIDs(pool, hosts); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "remove processed host ids")
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Task) GetHostLabelReportedAt(ctx context.Context, host *fleet.Host) time.Time {
|
||||
if t.AsyncEnabled {
|
||||
conn := redis.ConfigureDoer(t.Pool, t.Pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
key := fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
epoch, err := redigo.Int64(conn.Do("GET", key))
|
||||
if err == nil {
|
||||
if reported := time.Unix(epoch, 0); reported.After(host.LabelUpdatedAt) {
|
||||
return reported
|
||||
}
|
||||
}
|
||||
}
|
||||
return host.LabelUpdatedAt
|
||||
}
|
||||
|
||||
func storePurgeActiveHostID(pool fleet.RedisPool, hid uint, reportedAt, purgeOlder time.Time) error {
|
||||
// KEYS[1]: labelMembershipActiveHostIDsKey
|
||||
func storePurgeActiveHostID(pool fleet.RedisPool, zsetKey string, hid uint, reportedAt, purgeOlder time.Time) (int, error) {
|
||||
// KEYS[1]: the zsetKey
|
||||
// ARGV[1]: the host ID to add
|
||||
// ARGV[2]: the added host's reported-at timestamp
|
||||
// ARGV[3]: purge any entry with score older than this (purgeOlder timestamp)
|
||||
//
|
||||
// returns how many hosts were removed
|
||||
script := redigo.NewScript(1, `
|
||||
redis.call('ZADD', KEYS[1], ARGV[2], ARGV[1])
|
||||
return redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[3])
|
||||
@@ -315,67 +106,15 @@ func storePurgeActiveHostID(pool fleet.RedisPool, hid uint, reportedAt, purgeOld
|
||||
conn := pool.Get()
|
||||
defer conn.Close()
|
||||
|
||||
if err := redis.BindConn(pool, conn, labelMembershipActiveHostIDsKey); err != nil {
|
||||
return fmt.Errorf("bind redis connection: %w", err)
|
||||
if err := redis.BindConn(pool, conn, zsetKey); err != nil {
|
||||
return 0, fmt.Errorf("bind redis connection: %w", err)
|
||||
}
|
||||
|
||||
if _, err := script.Do(conn, labelMembershipActiveHostIDsKey, hid, reportedAt.Unix(), purgeOlder.Unix()); err != nil {
|
||||
return fmt.Errorf("run redis script: %w", err)
|
||||
count, err := redigo.Int(script.Do(conn, zsetKey, hid, reportedAt.Unix(), purgeOlder.Unix()))
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("run redis script: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func removeProcessedHostIDs(pool fleet.RedisPool, batch []hostIDLastReported) error {
|
||||
// This script removes from the set of active hosts for label membership all
|
||||
// those that still have the same score as when the batch was read (via
|
||||
// loadActiveHostIDs). This is so that any host that would've reported new
|
||||
// data since the call to loadActiveHostIDs would *not* get deleted (as the
|
||||
// score would change if that was the case).
|
||||
//
|
||||
// Note that this approach is correct - in that it is safe and won't delete
|
||||
// any host that has unsaved reported data - but it is potentially slow, as
|
||||
// it needs to check the score of each member before deleting it. Should that
|
||||
// become too slow, we have some options:
|
||||
//
|
||||
// * split the batch in smaller, capped ones (that would be if the redis
|
||||
// server gets blocked for too long processing a single batch)
|
||||
// * use ZREMRANGEBYSCORE to remove in one command all members with a score
|
||||
// (reported-at timestamp) lower than the maximum timestamp in batch.
|
||||
// While this would be almost certainly faster, it might be incorrect as
|
||||
// new data could be reported with timestamps older than the maximum one,
|
||||
// e.g. if the clocks are not exactly in sync between fleet instances, or
|
||||
// if hosts report new data while the ZSCAN is going on and don't get picked
|
||||
// up by the SCAN (this is possible, as part of the guarantees of SCAN).
|
||||
|
||||
// KEYS[1]: labelMembershipActiveHostIDsKey
|
||||
// ARGV...: the list of host ID-last reported timestamp pairs
|
||||
script := redigo.NewScript(1, `
|
||||
local count = 0
|
||||
for i = 1, #ARGV, 2 do
|
||||
local member, ts = ARGV[i], ARGV[i+1]
|
||||
if redis.call('ZSCORE', KEYS[1], member) == ts then
|
||||
count = count + 1
|
||||
redis.call('ZREM', KEYS[1], member)
|
||||
end
|
||||
end
|
||||
return count
|
||||
`)
|
||||
|
||||
conn := pool.Get()
|
||||
defer conn.Close()
|
||||
|
||||
if err := redis.BindConn(pool, conn, labelMembershipActiveHostIDsKey); err != nil {
|
||||
return fmt.Errorf("bind redis connection: %w", err)
|
||||
}
|
||||
|
||||
args := redigo.Args{labelMembershipActiveHostIDsKey}
|
||||
for _, host := range batch {
|
||||
args = args.Add(host.HostID, host.LastReported)
|
||||
}
|
||||
if _, err := script.Do(conn, args...); err != nil {
|
||||
return fmt.Errorf("run redis script: %w", err)
|
||||
}
|
||||
return nil
|
||||
return count, nil
|
||||
}
|
||||
|
||||
type hostIDLastReported struct {
|
||||
@@ -383,7 +122,7 @@ type hostIDLastReported struct {
|
||||
LastReported int64 // timestamp in unix epoch
|
||||
}
|
||||
|
||||
func loadActiveHostIDs(pool fleet.RedisPool, scanCount int) ([]hostIDLastReported, error) {
|
||||
func loadActiveHostIDs(pool fleet.RedisPool, zsetKey string, scanCount int) ([]hostIDLastReported, error) {
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
@@ -392,7 +131,7 @@ func loadActiveHostIDs(pool fleet.RedisPool, scanCount int) ([]hostIDLastReporte
|
||||
var hosts []hostIDLastReported
|
||||
cursor := 0
|
||||
for {
|
||||
res, err := redigo.Values(conn.Do("ZSCAN", labelMembershipActiveHostIDsKey, cursor, "COUNT", scanCount))
|
||||
res, err := redigo.Values(conn.Do("ZSCAN", zsetKey, cursor, "COUNT", scanCount))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("scan active host ids: %w", err)
|
||||
}
|
||||
@@ -410,3 +149,58 @@ func loadActiveHostIDs(pool fleet.RedisPool, scanCount int) ([]hostIDLastReporte
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func removeProcessedHostIDs(pool fleet.RedisPool, zsetKey string, batch []hostIDLastReported) (int, error) {
|
||||
// This script removes from the set of active hosts all those that still have
|
||||
// the same score as when the batch was read (via loadActiveHostIDs). This is
|
||||
// so that any host that would've reported new data since the call to
|
||||
// loadActiveHostIDs would *not* get deleted (as the score would change if
|
||||
// that was the case).
|
||||
//
|
||||
// Note that this approach is correct - in that it is safe and won't delete
|
||||
// any host that has unsaved reported data - but it is potentially slow, as
|
||||
// it needs to check the score of each member before deleting it. Should that
|
||||
// become too slow, we have some options:
|
||||
//
|
||||
// * split the batch in smaller, capped ones (that would be if the redis
|
||||
// server gets blocked for too long processing a single batch)
|
||||
// * use ZREMRANGEBYSCORE to remove in one command all members with a score
|
||||
// (reported-at timestamp) lower than the maximum timestamp in batch.
|
||||
// While this would be almost certainly faster, it might be incorrect as
|
||||
// new data could be reported with timestamps older than the maximum one,
|
||||
// e.g. if the clocks are not exactly in sync between fleet instances, or
|
||||
// if hosts report new data while the ZSCAN is going on and don't get picked
|
||||
// up by the SCAN (this is possible, as part of the guarantees of SCAN).
|
||||
|
||||
// KEYS[1]: zsetKey
|
||||
// ARGV...: the list of host ID-last reported timestamp pairs
|
||||
// returns the count of hosts removed
|
||||
script := redigo.NewScript(1, `
|
||||
local count = 0
|
||||
for i = 1, #ARGV, 2 do
|
||||
local member, ts = ARGV[i], ARGV[i+1]
|
||||
if redis.call('ZSCORE', KEYS[1], member) == ts then
|
||||
count = count + 1
|
||||
redis.call('ZREM', KEYS[1], member)
|
||||
end
|
||||
end
|
||||
return count
|
||||
`)
|
||||
|
||||
conn := pool.Get()
|
||||
defer conn.Close()
|
||||
|
||||
if err := redis.BindConn(pool, conn, zsetKey); err != nil {
|
||||
return 0, fmt.Errorf("bind redis connection: %w", err)
|
||||
}
|
||||
|
||||
args := redigo.Args{zsetKey}
|
||||
for _, host := range batch {
|
||||
args = args.Add(host.HostID, host.LastReported)
|
||||
}
|
||||
count, err := redigo.Int(script.Do(conn, args...))
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("run redis script: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
package async
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"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 (
|
||||
labelMembershipActiveHostIDsKey = "label_membership:active_host_ids"
|
||||
labelMembershipHostKey = "label_membership:{%d}"
|
||||
labelMembershipReportedKey = "label_membership_reported:{%d}"
|
||||
labelMembershipKeysMinTTL = 7 * 24 * time.Hour // 1 week
|
||||
)
|
||||
|
||||
func (t *Task) RecordLabelQueryExecutions(ctx context.Context, host *fleet.Host, results map[uint]*bool, ts time.Time, deferred bool) error {
|
||||
if !t.AsyncEnabled {
|
||||
host.LabelUpdatedAt = ts
|
||||
return t.Datastore.RecordLabelQueryExecutions(ctx, host, results, ts, deferred)
|
||||
}
|
||||
|
||||
keySet := fmt.Sprintf(labelMembershipHostKey, host.ID)
|
||||
keyTs := fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
|
||||
// set an expiration on both keys (set and ts), ensuring that a deleted host
|
||||
// (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.
|
||||
//
|
||||
// This means that it will only expire if that host hasn't reported labels
|
||||
// during that (TTL) time (each time it does report, the TTL is reset), and
|
||||
// the collector will have plenty of time to run (multiple times) to try to
|
||||
// persist all the data in mysql.
|
||||
ttl := labelMembershipKeysMinTTL
|
||||
if maxTTL := 10 * t.CollectorInterval; maxTTL > ttl {
|
||||
ttl = maxTTL
|
||||
}
|
||||
|
||||
// keys and arguments passed to the script are:
|
||||
// KEYS[1]: keySet (labelMembershipHostKey)
|
||||
// KEYS[2]: keyTs (labelMembershipReportedKey)
|
||||
// ARGV[1]: timestamp for "reported at"
|
||||
// ARGV[2]: ttl for both keys
|
||||
// ARGV[3..]: the arguments to ZADD to keySet
|
||||
script := redigo.NewScript(2, `
|
||||
redis.call('ZADD', KEYS[1], unpack(ARGV, 3))
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||||
redis.call('SET', KEYS[2], ARGV[1])
|
||||
return redis.call('EXPIRE', KEYS[2], ARGV[2])
|
||||
`)
|
||||
|
||||
// convert results to ZADD arguments, store as -1 for delete, +1 for insert
|
||||
args := make(redigo.Args, 0, 4+(len(results)*2))
|
||||
args = args.Add(keySet, keyTs, ts.Unix(), int(ttl.Seconds()))
|
||||
for k, v := range results {
|
||||
score := -1
|
||||
if v != nil && *v {
|
||||
score = 1
|
||||
}
|
||||
args = args.Add(score, k)
|
||||
}
|
||||
|
||||
conn := t.Pool.Get()
|
||||
defer conn.Close()
|
||||
if err := redis.BindConn(t.Pool, conn, keySet, keyTs); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "bind redis connection")
|
||||
}
|
||||
|
||||
if _, err := script.Do(conn, args...); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "run redis script")
|
||||
}
|
||||
|
||||
// Storing the host id in the set of active host IDs for label membership
|
||||
// outside of the redis script because in Redis Cluster mode the key may not
|
||||
// live on the same node as the host's keys. At the same time, purge any
|
||||
// entry in the set that is older than now - TTL.
|
||||
if _, err := storePurgeActiveHostID(t.Pool, labelMembershipActiveHostIDsKey, host.ID, ts, ts.Add(-ttl)); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "store active host id")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Task) collectLabelQueryExecutions(ctx context.Context, ds fleet.Datastore, pool fleet.RedisPool, stats *collectorExecStats) error {
|
||||
hosts, err := loadActiveHostIDs(pool, labelMembershipActiveHostIDsKey, t.RedisScanKeysCount)
|
||||
if err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "load active host ids")
|
||||
}
|
||||
stats.Keys = len(hosts)
|
||||
|
||||
getKeyTuples := func(hostID uint) (inserts, deletes [][2]uint, err error) {
|
||||
keySet := fmt.Sprintf(labelMembershipHostKey, hostID)
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
for {
|
||||
stats.RedisCmds++
|
||||
|
||||
vals, err := redigo.Ints(conn.Do("ZPOPMIN", keySet, t.RedisPopCount))
|
||||
if err != nil {
|
||||
return nil, nil, ctxerr.Wrap(ctx, err, "redis ZPOPMIN")
|
||||
}
|
||||
items := len(vals) / 2 // each item has the label id and the score (-1=delete, +1=insert)
|
||||
stats.Items += items
|
||||
|
||||
for i := 0; i < len(vals); i += 2 {
|
||||
labelID := vals[i]
|
||||
|
||||
var score int
|
||||
if i+1 < len(vals) { // just to be safe we received all pairs
|
||||
score = vals[i+1]
|
||||
}
|
||||
|
||||
switch score {
|
||||
case 1:
|
||||
inserts = append(inserts, [2]uint{uint(labelID), hostID})
|
||||
case -1:
|
||||
deletes = append(deletes, [2]uint{uint(labelID), hostID})
|
||||
}
|
||||
}
|
||||
if items < t.RedisPopCount {
|
||||
return inserts, deletes, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Based on those pages, the best approach appears to be INSERT with multiple
|
||||
// rows in the VALUES section (short of doing LOAD FILE, which we can't):
|
||||
// https://www.databasejournal.com/features/mysql/optimize-mysql-inserts-using-batch-processing.html
|
||||
// https://dev.mysql.com/doc/refman/5.7/en/insert-optimization.html
|
||||
// https://dev.mysql.com/doc/refman/5.7/en/optimizing-innodb-bulk-data-loading.html
|
||||
//
|
||||
// Given that there are no UNIQUE constraints in label_membership (well,
|
||||
// apart from the primary key columns), no AUTO_INC column and no FOREIGN
|
||||
// KEY, there is no obvious setting to tweak (based on the recommendations of
|
||||
// the third link above).
|
||||
//
|
||||
// However, in label_membership, updated_at defaults to the current timestamp
|
||||
// both on INSERT and when UPDATEd, so it does not need to be provided.
|
||||
|
||||
runInsertBatch := func(batch [][2]uint) error {
|
||||
stats.Inserts++
|
||||
return ds.AsyncBatchInsertLabelMembership(ctx, batch)
|
||||
}
|
||||
|
||||
runDeleteBatch := func(batch [][2]uint) error {
|
||||
stats.Deletes++
|
||||
return ds.AsyncBatchDeleteLabelMembership(ctx, batch)
|
||||
}
|
||||
|
||||
runUpdateBatch := func(ids []uint, ts time.Time) error {
|
||||
stats.Updates++
|
||||
return ds.AsyncBatchUpdateLabelTimestamp(ctx, ids, ts)
|
||||
}
|
||||
|
||||
insertBatch := make([][2]uint, 0, t.InsertBatch)
|
||||
deleteBatch := make([][2]uint, 0, t.DeleteBatch)
|
||||
for _, host := range hosts {
|
||||
hid := host.HostID
|
||||
ins, del, err := getKeyTuples(hid)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
insertBatch = append(insertBatch, ins...)
|
||||
deleteBatch = append(deleteBatch, del...)
|
||||
|
||||
if len(insertBatch) >= t.InsertBatch {
|
||||
if err := runInsertBatch(insertBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
insertBatch = insertBatch[:0]
|
||||
}
|
||||
if len(deleteBatch) >= t.DeleteBatch {
|
||||
if err := runDeleteBatch(deleteBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
deleteBatch = deleteBatch[:0]
|
||||
}
|
||||
}
|
||||
|
||||
// process any remaining batch that did not reach the batchSize limit in the
|
||||
// loop.
|
||||
if len(insertBatch) > 0 {
|
||||
if err := runInsertBatch(insertBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if len(deleteBatch) > 0 {
|
||||
if err := runDeleteBatch(deleteBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if len(hosts) > 0 {
|
||||
hostIDs := make([]uint, len(hosts))
|
||||
for i, host := range hosts {
|
||||
hostIDs[i] = host.HostID
|
||||
}
|
||||
|
||||
ts := time.Now()
|
||||
updateBatch := make([]uint, t.UpdateBatch)
|
||||
for {
|
||||
n := copy(updateBatch, hostIDs)
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
if err := runUpdateBatch(updateBatch[:n], ts); err != nil {
|
||||
return err
|
||||
}
|
||||
hostIDs = hostIDs[n:]
|
||||
}
|
||||
|
||||
// batch-remove any host ID from the active set that still has its score to
|
||||
// the initial value, so that the active set does not keep all (potentially
|
||||
// 100K+) host IDs to process at all times - only those with reported
|
||||
// results to process.
|
||||
if _, err := removeProcessedHostIDs(pool, labelMembershipActiveHostIDsKey, hosts); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "remove processed host ids")
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Task) GetHostLabelReportedAt(ctx context.Context, host *fleet.Host) time.Time {
|
||||
if t.AsyncEnabled {
|
||||
conn := redis.ConfigureDoer(t.Pool, t.Pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
key := fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
epoch, err := redigo.Int64(conn.Do("GET", key))
|
||||
if err == nil {
|
||||
if reported := time.Unix(epoch, 0); reported.After(host.LabelUpdatedAt) {
|
||||
return reported
|
||||
}
|
||||
}
|
||||
}
|
||||
return host.LabelUpdatedAt
|
||||
}
|
||||
@@ -0,0 +1,392 @@
|
||||
package async
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"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 testCollectLabelQueryExecutions(t *testing.T, ds *mysql.Datastore, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
|
||||
type labelMembership struct {
|
||||
HostID int `db:"host_id"`
|
||||
LabelID uint `db:"label_id"`
|
||||
UpdatedAt time.Time `db:"updated_at"`
|
||||
}
|
||||
|
||||
hostIDs := createHosts(t, ds, 4, time.Now().Add(-24*time.Hour))
|
||||
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.
|
||||
cases := []struct {
|
||||
name string
|
||||
// map of host ID to label IDs to insert (true) or delete (false)
|
||||
reported map[int]map[int]bool
|
||||
want []labelMembership
|
||||
}{
|
||||
{"no key", nil, nil},
|
||||
{
|
||||
"report host 1 label 1",
|
||||
map[int]map[int]bool{hid(1): {1: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 1, 2",
|
||||
map[int]map[int]bool{hid(1): {1: true, 2: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 1, 2, 3",
|
||||
map[int]map[int]bool{1: {1: true, 2: true, 3: true}},
|
||||
[]labelMembership{
|
||||
{HostID: 1, LabelID: 1},
|
||||
{HostID: 1, LabelID: 2},
|
||||
{HostID: 1, LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -1",
|
||||
map[int]map[int]bool{hid(1): {1: false}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(1), LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -2, -3",
|
||||
map[int]map[int]bool{hid(1): {2: false, 3: false}},
|
||||
[]labelMembership{},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 1, 2, 3, 4",
|
||||
map[int]map[int]bool{hid(1): {1: true, 2: true, 3: true, 4: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(1), LabelID: 3},
|
||||
{HostID: hid(1), LabelID: 4},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -2, -3, -4, -5",
|
||||
map[int]map[int]bool{hid(1): {2: false, 3: false, 4: false, 5: false}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 2, host 2 labels 2, 3",
|
||||
map[int]map[int]bool{hid(1): {2: true}, hid(2): {2: true, 3: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -99, non-existing",
|
||||
map[int]map[int]bool{hid(1): {99: false}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report hosts 1, 2, 3, 4 labels 1, 2, -3, 4",
|
||||
map[int]map[int]bool{
|
||||
hid(1): {1: true, 2: true, 3: false, 4: true},
|
||||
hid(2): {1: true, 2: true, 3: false, 4: true},
|
||||
hid(3): {1: true, 2: true, 3: false, 4: true},
|
||||
hid(4): {1: true, 2: true, 3: false, 4: true},
|
||||
},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(1), LabelID: 4},
|
||||
{HostID: hid(2), LabelID: 1},
|
||||
{HostID: hid(2), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 4},
|
||||
{HostID: hid(3), LabelID: 1},
|
||||
{HostID: hid(3), LabelID: 2},
|
||||
{HostID: hid(3), LabelID: 4},
|
||||
{HostID: hid(4), LabelID: 1},
|
||||
{HostID: hid(4), LabelID: 2},
|
||||
{HostID: hid(4), LabelID: 4},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
const batchSizes = 3
|
||||
|
||||
setupTest := func(t *testing.T, data map[int]map[int]bool) collectorExecStats {
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
// store the host memberships and prepare the expected stats
|
||||
var wantStats collectorExecStats
|
||||
for hostID, res := range data {
|
||||
if len(res) > 0 {
|
||||
key := fmt.Sprintf(labelMembershipHostKey, hostID)
|
||||
args := make(redigo.Args, 0, 1+(len(res)*2))
|
||||
args = args.Add(key)
|
||||
for lblID, ins := range res {
|
||||
score := -1
|
||||
if ins {
|
||||
score = 1
|
||||
}
|
||||
args = args.Add(score, lblID)
|
||||
}
|
||||
_, err := conn.Do("ZADD", args...)
|
||||
require.NoError(t, err)
|
||||
_, err = conn.Do("ZADD", labelMembershipActiveHostIDsKey, time.Now().Unix(), hostID)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
cnt, err := redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
wantStats.Keys = cnt
|
||||
wantStats.Items += len(res)
|
||||
wantStats.RedisCmds++
|
||||
wantStats.RedisCmds += len(res) / batchSizes
|
||||
}
|
||||
return wantStats
|
||||
}
|
||||
|
||||
selectRows := func(t *testing.T) ([]labelMembership, map[int]time.Time) {
|
||||
var rows []labelMembership
|
||||
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
|
||||
return sqlx.SelectContext(ctx, tx, &rows, `SELECT host_id, label_id, updated_at FROM label_membership ORDER BY 1, 2`)
|
||||
})
|
||||
|
||||
var hosts []struct {
|
||||
ID int `db:"id"`
|
||||
LabelUpdatedAt time.Time `db:"label_updated_at"`
|
||||
}
|
||||
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
|
||||
return sqlx.SelectContext(ctx, tx, &hosts, `SELECT id, label_updated_at FROM hosts`)
|
||||
})
|
||||
|
||||
hostsUpdated := make(map[int]time.Time, len(hosts))
|
||||
for _, h := range hosts {
|
||||
hostsUpdated[h.ID] = h.LabelUpdatedAt
|
||||
}
|
||||
return rows, hostsUpdated
|
||||
}
|
||||
|
||||
minUpdatedAt := time.Now()
|
||||
for _, c := range cases {
|
||||
func() {
|
||||
t.Log("test name: ", c.name)
|
||||
wantStats := setupTest(t, c.reported)
|
||||
|
||||
// run the collection
|
||||
var stats collectorExecStats
|
||||
task := Task{
|
||||
InsertBatch: batchSizes,
|
||||
UpdateBatch: batchSizes,
|
||||
DeleteBatch: batchSizes,
|
||||
RedisPopCount: batchSizes,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
err := task.collectLabelQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
// inserts, updates and deletes are a bit tricky to track automatically,
|
||||
// just ignore them when comparing stats.
|
||||
stats.Inserts, stats.Updates, stats.Deletes = 0, 0, 0
|
||||
require.Equal(t, wantStats, stats)
|
||||
|
||||
// check that the table contains the expected rows
|
||||
rows, hostsUpdated := 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.Equal(t, want.LabelID, got.LabelID)
|
||||
require.WithinDuration(t, minUpdatedAt, got.UpdatedAt, 10*time.Second)
|
||||
|
||||
ts, ok := hostsUpdated[want.HostID]
|
||||
require.True(t, ok)
|
||||
require.WithinDuration(t, minUpdatedAt, ts, 10*time.Second)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// after all cases, run one last upsert (an update) to make sure that the
|
||||
// updated at column is properly updated. First we need to ensure that this
|
||||
// runs in a distinct second, because the mysql resolution is not precise.
|
||||
time.Sleep(time.Second)
|
||||
|
||||
var h1l1Before labelMembership
|
||||
beforeRows, _ := selectRows(t)
|
||||
for _, row := range beforeRows {
|
||||
if row.HostID == 1 && row.LabelID == 1 {
|
||||
h1l1Before = row
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// update host 1, label 1, already existing
|
||||
setupTest(t, map[int]map[int]bool{1: {1: true}})
|
||||
var stats collectorExecStats
|
||||
task := Task{
|
||||
InsertBatch: batchSizes,
|
||||
UpdateBatch: batchSizes,
|
||||
DeleteBatch: batchSizes,
|
||||
RedisPopCount: batchSizes,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
err := task.collectLabelQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
|
||||
var h1l1After labelMembership
|
||||
afterRows, _ := selectRows(t)
|
||||
for _, row := range afterRows {
|
||||
if row.HostID == 1 && row.LabelID == 1 {
|
||||
h1l1After = row
|
||||
break
|
||||
}
|
||||
}
|
||||
require.True(t, h1l1Before.UpdatedAt.Before(h1l1After.UpdatedAt))
|
||||
}
|
||||
|
||||
func testRecordLabelQueryExecutionsSync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
lastYear := now.Add(-365 * 24 * time.Hour)
|
||||
host := &fleet.Host{
|
||||
ID: 1,
|
||||
Platform: "linux",
|
||||
LabelUpdatedAt: lastYear,
|
||||
}
|
||||
|
||||
var yes, no = true, false
|
||||
results := map[uint]*bool{1: &yes, 2: &yes, 3: &no, 4: nil}
|
||||
keySet, keyTs := fmt.Sprintf(labelMembershipHostKey, host.ID), fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
|
||||
task := Task{
|
||||
Datastore: ds,
|
||||
Pool: pool,
|
||||
AsyncEnabled: false,
|
||||
}
|
||||
|
||||
labelReportedAt := task.GetHostLabelReportedAt(ctx, host)
|
||||
require.True(t, labelReportedAt.Equal(lastYear))
|
||||
|
||||
err := task.RecordLabelQueryExecutions(ctx, host, results, now, false)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ds.RecordLabelQueryExecutionsFuncInvoked)
|
||||
ds.RecordLabelQueryExecutionsFuncInvoked = false
|
||||
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
defer conn.Do("DEL", keySet, keyTs)
|
||||
|
||||
n, err := redigo.Int(conn.Do("EXISTS", keySet))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
n, err = redigo.Int(conn.Do("EXISTS", keyTs))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
n, err = redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
labelReportedAt = task.GetHostLabelReportedAt(ctx, host)
|
||||
require.True(t, labelReportedAt.Equal(now))
|
||||
}
|
||||
|
||||
func testRecordLabelQueryExecutionsAsync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
lastYear := now.Add(-365 * 24 * time.Hour)
|
||||
host := &fleet.Host{
|
||||
ID: 1,
|
||||
Platform: "linux",
|
||||
LabelUpdatedAt: lastYear,
|
||||
}
|
||||
var yes, no = true, false
|
||||
results := map[uint]*bool{1: &yes, 2: &yes, 3: &no, 4: nil}
|
||||
keySet, keyTs := fmt.Sprintf(labelMembershipHostKey, host.ID), fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
|
||||
task := Task{
|
||||
Datastore: ds,
|
||||
Pool: pool,
|
||||
AsyncEnabled: true,
|
||||
|
||||
InsertBatch: 3,
|
||||
UpdateBatch: 3,
|
||||
DeleteBatch: 3,
|
||||
RedisPopCount: 3,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
|
||||
labelReportedAt := task.GetHostLabelReportedAt(ctx, host)
|
||||
require.True(t, labelReportedAt.Equal(lastYear))
|
||||
|
||||
err := task.RecordLabelQueryExecutions(ctx, host, results, now, false)
|
||||
require.NoError(t, err)
|
||||
require.False(t, ds.RecordLabelQueryExecutionsFuncInvoked)
|
||||
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
defer conn.Do("DEL", keySet, keyTs)
|
||||
|
||||
res, err := redigo.IntMap(conn.Do("ZPOPMIN", keySet, 10))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 4, len(res))
|
||||
require.Equal(t, map[string]int{"1": 1, "2": 1, "3": -1, "4": -1}, res)
|
||||
|
||||
ts, err := redigo.Int64(conn.Do("GET", keyTs))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, now.Unix(), ts)
|
||||
|
||||
count, err := redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count)
|
||||
tsActive, err := redigo.Int64(conn.Do("ZSCORE", labelMembershipActiveHostIDsKey, host.ID))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tsActive, ts)
|
||||
|
||||
labelReportedAt = task.GetHostLabelReportedAt(ctx, host)
|
||||
// because we transition via unix epoch (seconds), not exactly equal
|
||||
require.WithinDuration(t, now, labelReportedAt, time.Second)
|
||||
// host's LabelUpdatedAt field hasn't been updated yet, because the label
|
||||
// results are in redis, not in mysql yet.
|
||||
require.True(t, host.LabelUpdatedAt.Equal(lastYear))
|
||||
|
||||
// running the collector removes the host from the active set
|
||||
var stats collectorExecStats
|
||||
err = task.collectLabelQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, stats.Keys)
|
||||
require.Equal(t, 0, stats.Items) // zero because we cleared the host's set with ZPOPMIN above
|
||||
require.False(t, stats.Failed)
|
||||
|
||||
count, err = redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, count)
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
package async
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
|
||||
"github.com/fleetdm/fleet/v4/server/datastore/redis"
|
||||
"github.com/fleetdm/fleet/v4/server/fleet"
|
||||
"github.com/fleetdm/fleet/v4/server/ptr"
|
||||
redigo "github.com/gomodule/redigo/redis"
|
||||
)
|
||||
|
||||
const (
|
||||
policyPassHostIDsKey = "policy_pass:active_host_ids"
|
||||
policyPassHostKey = "policy_pass:{%d}"
|
||||
policyPassReportedKey = "policy_pass_reported:{%d}"
|
||||
policyPassKeysMinTTL = 7 * 24 * time.Hour // 1 week
|
||||
)
|
||||
|
||||
var (
|
||||
// redis list will be LTRIM'd if there are more policy IDs than this.
|
||||
maxRedisPolicyResultsPerHost = 1000
|
||||
)
|
||||
|
||||
func (t *Task) RecordPolicyQueryExecutions(ctx context.Context, host *fleet.Host, results map[uint]*bool, ts time.Time, deferred bool) error {
|
||||
if !t.AsyncEnabled {
|
||||
host.PolicyUpdatedAt = ts
|
||||
return t.Datastore.RecordPolicyQueryExecutions(ctx, host, results, ts, deferred)
|
||||
}
|
||||
|
||||
keyList := fmt.Sprintf(policyPassHostKey, host.ID)
|
||||
keyTs := fmt.Sprintf(policyPassReportedKey, host.ID)
|
||||
|
||||
// set an expiration on both keys (list and ts), ensuring that a deleted host
|
||||
// (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.
|
||||
//
|
||||
// This means that it will only expire if that host hasn't reported policies
|
||||
// during that (TTL) time (each time it does report, the TTL is reset), and
|
||||
// the collector will have plenty of time to run (multiple times) to try to
|
||||
// persist all the data in mysql.
|
||||
ttl := policyPassKeysMinTTL
|
||||
if maxTTL := 10 * t.CollectorInterval; maxTTL > ttl {
|
||||
ttl = maxTTL
|
||||
}
|
||||
|
||||
// KEYS[1]: keyList (policyPassHostKey)
|
||||
// KEYS[2]: keyTs (policyPassReportedKey)
|
||||
// ARGV[1]: timestamp for "reported at"
|
||||
// ARGV[2]: max policy results to keep per host (list is trimmed to that size)
|
||||
// ARGV[3]: ttl for both keys
|
||||
// ARGV[4..]: policy_id=pass entries to LPUSH to the list
|
||||
script := redigo.NewScript(2, `
|
||||
redis.call('LPUSH', KEYS[1], unpack(ARGV, 4))
|
||||
redis.call('LTRIM', KEYS[1], 0, ARGV[2])
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[3])
|
||||
redis.call('SET', KEYS[2], ARGV[1])
|
||||
return redis.call('EXPIRE', KEYS[2], ARGV[3])
|
||||
`)
|
||||
|
||||
// convert results to LPUSH arguments, store as policy_id=1 for pass,
|
||||
// policy_id=-1 for fail, policy_id=0 for null result.
|
||||
args := make(redigo.Args, 0, 5+len(results))
|
||||
args = args.Add(keyList, keyTs, ts.Unix(), maxRedisPolicyResultsPerHost, int(ttl.Seconds()))
|
||||
for k, v := range results {
|
||||
pass := 0
|
||||
if v != nil {
|
||||
if *v {
|
||||
pass = 1
|
||||
} else {
|
||||
pass = -1
|
||||
}
|
||||
}
|
||||
args = args.Add(fmt.Sprintf("%d=%d", k, pass))
|
||||
}
|
||||
|
||||
conn := t.Pool.Get()
|
||||
defer conn.Close()
|
||||
if err := redis.BindConn(t.Pool, conn, keyList, keyTs); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "bind redis connection")
|
||||
}
|
||||
|
||||
if _, err := script.Do(conn, args...); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "run redis script")
|
||||
}
|
||||
|
||||
// Storing the host id in the set of active host IDs for policy membership
|
||||
// outside of the redis script because in Redis Cluster mode the key may not
|
||||
// live on the same node as the host's keys. At the same time, purge any
|
||||
// entry in the set that is older than now - TTL.
|
||||
if _, err := storePurgeActiveHostID(t.Pool, policyPassHostIDsKey, host.ID, ts, ts.Add(-ttl)); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "store active host id")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Task) collectPolicyQueryExecutions(ctx context.Context, ds fleet.Datastore, pool fleet.RedisPool, stats *collectorExecStats) error {
|
||||
hosts, err := loadActiveHostIDs(pool, policyPassHostIDsKey, t.RedisScanKeysCount)
|
||||
if err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "load active host ids")
|
||||
}
|
||||
stats.Keys = len(hosts)
|
||||
|
||||
// need to use a script as the RPOP command only supports a COUNT since
|
||||
// 6.2. Because we use LTRIM when inserting, we know the total number
|
||||
// of results is at most maxRedisPolicyResultsPerHost, so it is capped
|
||||
// and can be returned in one go.
|
||||
script := redigo.NewScript(1, `
|
||||
local res = redis.call('LRANGE', KEYS[1], 0, -1)
|
||||
redis.call('DEL', KEYS[1])
|
||||
return res
|
||||
`)
|
||||
|
||||
getKeyTuples := func(hostID uint) (inserts []fleet.PolicyMembershipResult, err error) {
|
||||
keyList := fmt.Sprintf(policyPassHostKey, hostID)
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
stats.RedisCmds++
|
||||
res, err := redigo.Strings(script.Do(conn, keyList))
|
||||
if err != nil {
|
||||
return nil, ctxerr.Wrap(ctx, err, "redis LRANGE script")
|
||||
}
|
||||
|
||||
inserts = make([]fleet.PolicyMembershipResult, 0, len(res))
|
||||
stats.Items += len(res)
|
||||
for _, item := range res {
|
||||
parts := strings.Split(item, "=")
|
||||
if len(parts) != 2 {
|
||||
continue
|
||||
}
|
||||
|
||||
var tup fleet.PolicyMembershipResult
|
||||
if id, _ := strconv.ParseUint(parts[0], 10, 32); id > 0 {
|
||||
tup.HostID = hostID
|
||||
tup.PolicyID = uint(id)
|
||||
switch parts[1] {
|
||||
case "1":
|
||||
tup.Passes = ptr.Bool(true)
|
||||
case "-1":
|
||||
tup.Passes = ptr.Bool(false)
|
||||
case "0":
|
||||
tup.Passes = nil
|
||||
default:
|
||||
continue
|
||||
}
|
||||
inserts = append(inserts, tup)
|
||||
}
|
||||
}
|
||||
return inserts, nil
|
||||
}
|
||||
|
||||
runInsertBatch := func(batch []fleet.PolicyMembershipResult) error {
|
||||
stats.Inserts++
|
||||
return ds.AsyncBatchInsertPolicyMembership(ctx, batch)
|
||||
}
|
||||
|
||||
runUpdateBatch := func(ids []uint, ts time.Time) error {
|
||||
stats.Updates++
|
||||
return ds.AsyncBatchUpdatePolicyTimestamp(ctx, ids, ts)
|
||||
}
|
||||
|
||||
insertBatch := make([]fleet.PolicyMembershipResult, 0, t.InsertBatch)
|
||||
for _, host := range hosts {
|
||||
hid := host.HostID
|
||||
ins, err := getKeyTuples(hid)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
insertBatch = append(insertBatch, ins...)
|
||||
|
||||
if len(insertBatch) >= t.InsertBatch {
|
||||
if err := runInsertBatch(insertBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
insertBatch = insertBatch[:0]
|
||||
}
|
||||
}
|
||||
|
||||
// process any remaining batch that did not reach the batchSize limit in the
|
||||
// loop.
|
||||
if len(insertBatch) > 0 {
|
||||
if err := runInsertBatch(insertBatch); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if len(hosts) > 0 {
|
||||
hostIDs := make([]uint, len(hosts))
|
||||
for i, host := range hosts {
|
||||
hostIDs[i] = host.HostID
|
||||
}
|
||||
|
||||
ts := time.Now()
|
||||
updateBatch := make([]uint, t.UpdateBatch)
|
||||
for {
|
||||
n := copy(updateBatch, hostIDs)
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
if err := runUpdateBatch(updateBatch[:n], ts); err != nil {
|
||||
return err
|
||||
}
|
||||
hostIDs = hostIDs[n:]
|
||||
}
|
||||
|
||||
// batch-remove any host ID from the active set that still has its score to
|
||||
// the initial value, so that the active set does not keep all (potentially
|
||||
// 100K+) host IDs to process at all times - only those with reported
|
||||
// results to process.
|
||||
if _, err := removeProcessedHostIDs(pool, policyPassHostIDsKey, hosts); err != nil {
|
||||
return ctxerr.Wrap(ctx, err, "remove processed host ids")
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Task) GetHostPolicyReportedAt(ctx context.Context, host *fleet.Host) time.Time {
|
||||
if t.AsyncEnabled {
|
||||
conn := redis.ConfigureDoer(t.Pool, t.Pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
key := fmt.Sprintf(policyPassReportedKey, host.ID)
|
||||
epoch, err := redigo.Int64(conn.Do("GET", key))
|
||||
if err == nil {
|
||||
if reported := time.Unix(epoch, 0); reported.After(host.PolicyUpdatedAt) {
|
||||
return reported
|
||||
}
|
||||
}
|
||||
}
|
||||
return host.PolicyUpdatedAt
|
||||
}
|
||||
@@ -0,0 +1,434 @@
|
||||
package async
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"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"
|
||||
"github.com/fleetdm/fleet/v4/server/ptr"
|
||||
redigo "github.com/gomodule/redigo/redis"
|
||||
"github.com/jmoiron/sqlx"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func testCollectPolicyQueryExecutions(t *testing.T, ds *mysql.Datastore, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
|
||||
type policyMembership struct {
|
||||
HostID int `db:"host_id"`
|
||||
PolicyID int `db:"policy_id"`
|
||||
Passes sql.NullBool `db:"passes"`
|
||||
UpdatedAt time.Time `db:"updated_at"`
|
||||
}
|
||||
|
||||
hostIDs := createHosts(t, ds, 4, time.Now().Add(-24*time.Hour))
|
||||
policyIDs := createPolicies(t, ds, 4)
|
||||
t.Logf("real host IDs: %v", hostIDs)
|
||||
t.Logf("real policy IDs: %v", policyIDs)
|
||||
hid := func(id int) int {
|
||||
return int(hostIDs[id-1])
|
||||
}
|
||||
pid := func(id int) int {
|
||||
if id < 0 || id >= len(policyIDs) {
|
||||
return id
|
||||
}
|
||||
return int(policyIDs[id-1])
|
||||
}
|
||||
|
||||
nbTrue := sql.NullBool{Valid: true, Bool: true}
|
||||
nbFalse := sql.NullBool{Valid: true, Bool: false}
|
||||
nbNull := sql.NullBool{Valid: false}
|
||||
|
||||
// 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.
|
||||
cases := []struct {
|
||||
name string
|
||||
// map of host ID to policy IDs to insert with passes set to the bool.
|
||||
reported map[int]map[int]*bool
|
||||
want []policyMembership
|
||||
}{
|
||||
{"no key", nil, nil},
|
||||
{
|
||||
"report host 1 policy 1",
|
||||
map[int]map[int]*bool{hid(1): {pid(1): ptr.Bool(true)}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbTrue},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 policies 1, 2",
|
||||
map[int]map[int]*bool{hid(1): {pid(1): ptr.Bool(true), pid(2): ptr.Bool(true)}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbTrue},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 policies 1, 2, 3",
|
||||
map[int]map[int]*bool{hid(1): {pid(1): ptr.Bool(true), pid(2): ptr.Bool(true), pid(3): ptr.Bool(true)}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(3), Passes: nbTrue},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 policy -1",
|
||||
map[int]map[int]*bool{hid(1): {pid(1): ptr.Bool(false)}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbFalse},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(3), Passes: nbTrue},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 policies -2, -3",
|
||||
map[int]map[int]*bool{hid(1): {pid(2): ptr.Bool(false), pid(3): ptr.Bool(false)}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbFalse},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbFalse},
|
||||
{HostID: hid(1), PolicyID: pid(3), Passes: nbFalse},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 policies 1, 2, (3), 4",
|
||||
map[int]map[int]*bool{hid(1): {pid(1): ptr.Bool(true), pid(2): ptr.Bool(true), pid(3): nil, pid(4): ptr.Bool(true)}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(3), Passes: nbNull},
|
||||
{HostID: hid(1), PolicyID: pid(4), Passes: nbTrue},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 policies -2, -3, -4, 1",
|
||||
map[int]map[int]*bool{hid(1): {pid(2): ptr.Bool(false), pid(3): ptr.Bool(false), pid(4): ptr.Bool(false), pid(1): ptr.Bool(true)}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbFalse},
|
||||
{HostID: hid(1), PolicyID: pid(3), Passes: nbFalse},
|
||||
{HostID: hid(1), PolicyID: pid(4), Passes: nbFalse},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 policy 2, host 2 policies 2, 3",
|
||||
map[int]map[int]*bool{hid(1): {pid(2): ptr.Bool(true)}, hid(2): {pid(2): ptr.Bool(true), pid(3): nil}},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(3), Passes: nbFalse},
|
||||
{HostID: hid(1), PolicyID: pid(4), Passes: nbFalse},
|
||||
{HostID: hid(2), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(2), PolicyID: pid(3), Passes: nbNull},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report hosts 1, 2, 3, 4 policies 1, 2, -3, (4)",
|
||||
map[int]map[int]*bool{
|
||||
hid(1): {pid(1): ptr.Bool(true), pid(2): ptr.Bool(true), pid(3): ptr.Bool(false), pid(4): nil},
|
||||
hid(2): {pid(1): ptr.Bool(true), pid(2): ptr.Bool(true), pid(3): ptr.Bool(false), pid(4): nil},
|
||||
hid(3): {pid(1): ptr.Bool(true), pid(2): ptr.Bool(true), pid(3): ptr.Bool(false), pid(4): nil},
|
||||
hid(4): {pid(1): ptr.Bool(true), pid(2): ptr.Bool(true), pid(3): ptr.Bool(false), pid(4): nil},
|
||||
},
|
||||
[]policyMembership{
|
||||
{HostID: hid(1), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(1), PolicyID: pid(3), Passes: nbFalse},
|
||||
{HostID: hid(1), PolicyID: pid(4), Passes: nbNull},
|
||||
{HostID: hid(2), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(2), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(2), PolicyID: pid(3), Passes: nbFalse},
|
||||
{HostID: hid(2), PolicyID: pid(4), Passes: nbNull},
|
||||
{HostID: hid(3), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(3), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(3), PolicyID: pid(3), Passes: nbFalse},
|
||||
{HostID: hid(3), PolicyID: pid(4), Passes: nbNull},
|
||||
{HostID: hid(4), PolicyID: pid(1), Passes: nbTrue},
|
||||
{HostID: hid(4), PolicyID: pid(2), Passes: nbTrue},
|
||||
{HostID: hid(4), PolicyID: pid(3), Passes: nbFalse},
|
||||
{HostID: hid(4), PolicyID: pid(4), Passes: nbNull},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
const batchSizes = 3
|
||||
|
||||
setupTest := func(t *testing.T, data map[int]map[int]*bool) collectorExecStats {
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
// store the host memberships and prepare the expected stats
|
||||
var wantStats collectorExecStats
|
||||
for hostID, res := range data {
|
||||
if len(res) > 0 {
|
||||
key := fmt.Sprintf(policyPassHostKey, hostID)
|
||||
args := make(redigo.Args, 0, 1+(len(res)))
|
||||
args = args.Add(key)
|
||||
for polID, pass := range res {
|
||||
score := 0
|
||||
if pass != nil {
|
||||
if *pass {
|
||||
score = 1
|
||||
} else {
|
||||
score = -1
|
||||
}
|
||||
}
|
||||
args = args.Add(fmt.Sprintf("%d=%d", polID, score))
|
||||
}
|
||||
_, err := conn.Do("LPUSH", args...)
|
||||
require.NoError(t, err)
|
||||
_, err = conn.Do("ZADD", policyPassHostIDsKey, time.Now().Unix(), hostID)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
cnt, err := redigo.Int(conn.Do("ZCARD", policyPassHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
wantStats.Keys = cnt
|
||||
wantStats.RedisCmds++
|
||||
wantStats.Items += len(res)
|
||||
}
|
||||
return wantStats
|
||||
}
|
||||
|
||||
selectRows := func(t *testing.T) ([]policyMembership, map[int]time.Time) {
|
||||
var rows []policyMembership
|
||||
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
|
||||
return sqlx.SelectContext(ctx, tx, &rows, `SELECT host_id, policy_id, passes, updated_at
|
||||
FROM policy_membership
|
||||
ORDER BY host_id, policy_id`)
|
||||
})
|
||||
|
||||
var hosts []struct {
|
||||
ID int `db:"id"`
|
||||
PolicyUpdatedAt time.Time `db:"policy_updated_at"`
|
||||
}
|
||||
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
|
||||
return sqlx.SelectContext(ctx, tx, &hosts, `SELECT id, policy_updated_at FROM hosts`)
|
||||
})
|
||||
|
||||
hostsUpdated := make(map[int]time.Time, len(hosts))
|
||||
for _, h := range hosts {
|
||||
hostsUpdated[h.ID] = h.PolicyUpdatedAt
|
||||
}
|
||||
return rows, hostsUpdated
|
||||
}
|
||||
|
||||
minUpdatedAt := time.Now()
|
||||
for _, c := range cases {
|
||||
func() {
|
||||
t.Log("test name: ", c.name)
|
||||
wantStats := setupTest(t, c.reported)
|
||||
|
||||
// run the collection
|
||||
var stats collectorExecStats
|
||||
task := Task{
|
||||
InsertBatch: batchSizes,
|
||||
UpdateBatch: batchSizes,
|
||||
DeleteBatch: batchSizes,
|
||||
RedisPopCount: batchSizes,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
err := task.collectPolicyQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
// inserts, updates and deletes are a bit tricky to track automatically,
|
||||
// just ignore them when comparing stats.
|
||||
stats.Inserts, stats.Updates, stats.Deletes = 0, 0, 0
|
||||
require.Equal(t, wantStats, stats)
|
||||
|
||||
// check that the table contains the expected rows
|
||||
rows, hostsUpdated := 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, "[%d] host id", i)
|
||||
require.Equal(t, want.PolicyID, got.PolicyID, "[%d] policy id", i)
|
||||
require.Equal(t, want.Passes, got.Passes, "[%d] passes", i)
|
||||
require.WithinDuration(t, minUpdatedAt, got.UpdatedAt, 10*time.Second, "[%d] membership updated at", i)
|
||||
|
||||
ts, ok := hostsUpdated[want.HostID]
|
||||
require.True(t, ok)
|
||||
require.WithinDuration(t, minUpdatedAt, ts, 10*time.Second, "[%d] host updated at", i)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// after all cases, run one last upsert (an update) to make sure that the
|
||||
// updated at column is properly updated. First we need to ensure that this
|
||||
// runs in a distinct second, because the mysql resolution is not precise.
|
||||
time.Sleep(time.Second)
|
||||
|
||||
var h1p1Before policyMembership
|
||||
beforeRows, _ := selectRows(t)
|
||||
for _, row := range beforeRows {
|
||||
if row.HostID == 1 && row.PolicyID == 1 {
|
||||
h1p1Before = row
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// update host 1, policy 1, already existing
|
||||
setupTest(t, map[int]map[int]*bool{1: {1: nil}})
|
||||
var stats collectorExecStats
|
||||
task := Task{
|
||||
InsertBatch: batchSizes,
|
||||
UpdateBatch: batchSizes,
|
||||
DeleteBatch: batchSizes,
|
||||
RedisPopCount: batchSizes,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
err := task.collectPolicyQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
|
||||
var h1p1After policyMembership
|
||||
afterRows, _ := selectRows(t)
|
||||
for _, row := range afterRows {
|
||||
if row.HostID == 1 && row.PolicyID == 1 {
|
||||
h1p1After = row
|
||||
break
|
||||
}
|
||||
}
|
||||
require.True(t, h1p1Before.UpdatedAt.Before(h1p1After.UpdatedAt))
|
||||
}
|
||||
|
||||
func testRecordPolicyQueryExecutionsSync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
lastYear := now.Add(-365 * 24 * time.Hour)
|
||||
host := &fleet.Host{
|
||||
ID: 1,
|
||||
Platform: "linux",
|
||||
PolicyUpdatedAt: lastYear,
|
||||
}
|
||||
|
||||
var yes, no = true, false
|
||||
results := map[uint]*bool{1: &yes, 2: &yes, 3: &no, 4: nil}
|
||||
keyList, keyTs := fmt.Sprintf(policyPassHostKey, host.ID), fmt.Sprintf(policyPassReportedKey, host.ID)
|
||||
|
||||
task := Task{
|
||||
Datastore: ds,
|
||||
Pool: pool,
|
||||
AsyncEnabled: false,
|
||||
}
|
||||
|
||||
policyReportedAt := task.GetHostPolicyReportedAt(ctx, host)
|
||||
require.True(t, policyReportedAt.Equal(lastYear))
|
||||
|
||||
err := task.RecordPolicyQueryExecutions(ctx, host, results, now, false)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ds.RecordPolicyQueryExecutionsFuncInvoked)
|
||||
ds.RecordPolicyQueryExecutionsFuncInvoked = false
|
||||
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
defer conn.Do("DEL", keyList, keyTs)
|
||||
|
||||
n, err := redigo.Int(conn.Do("EXISTS", keyList))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
n, err = redigo.Int(conn.Do("EXISTS", keyTs))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
n, err = redigo.Int(conn.Do("ZCARD", policyPassHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
policyReportedAt = task.GetHostPolicyReportedAt(ctx, host)
|
||||
require.True(t, policyReportedAt.Equal(now))
|
||||
}
|
||||
|
||||
func testRecordPolicyQueryExecutionsAsync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
lastYear := now.Add(-365 * 24 * time.Hour)
|
||||
host := &fleet.Host{
|
||||
ID: 1,
|
||||
Platform: "linux",
|
||||
PolicyUpdatedAt: lastYear,
|
||||
}
|
||||
var yes, no = true, false
|
||||
results := map[uint]*bool{1: &yes, 2: &yes, 3: &no, 4: nil}
|
||||
keyList, keyTs := fmt.Sprintf(policyPassHostKey, host.ID), fmt.Sprintf(policyPassReportedKey, host.ID)
|
||||
|
||||
task := Task{
|
||||
Datastore: ds,
|
||||
Pool: pool,
|
||||
AsyncEnabled: true,
|
||||
|
||||
InsertBatch: 3,
|
||||
UpdateBatch: 3,
|
||||
DeleteBatch: 3,
|
||||
RedisPopCount: 3,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
|
||||
policyReportedAt := task.GetHostPolicyReportedAt(ctx, host)
|
||||
require.True(t, policyReportedAt.Equal(lastYear))
|
||||
|
||||
err := task.RecordPolicyQueryExecutions(ctx, host, results, now, false)
|
||||
require.NoError(t, err)
|
||||
require.False(t, ds.RecordPolicyQueryExecutionsFuncInvoked)
|
||||
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
defer conn.Do("DEL", keyList, keyTs)
|
||||
|
||||
res, err := redigo.Strings(conn.Do("LRANGE", keyList, 0, -1))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 4, len(res))
|
||||
require.ElementsMatch(t, []string{"1=1", "2=1", "3=-1", "4=0"}, res)
|
||||
|
||||
ts, err := redigo.Int64(conn.Do("GET", keyTs))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, now.Unix(), ts)
|
||||
|
||||
count, err := redigo.Int(conn.Do("ZCARD", policyPassHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count)
|
||||
tsActive, err := redigo.Int64(conn.Do("ZSCORE", policyPassHostIDsKey, host.ID))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tsActive, ts)
|
||||
|
||||
policyReportedAt = task.GetHostPolicyReportedAt(ctx, host)
|
||||
// because we transition via unix epoch (seconds), not exactly equal
|
||||
require.WithinDuration(t, now, policyReportedAt, time.Second)
|
||||
// host's PolicyUpdatedAt field hasn't been updated yet, because the label
|
||||
// results are in redis, not in mysql yet.
|
||||
require.True(t, host.PolicyUpdatedAt.Equal(lastYear))
|
||||
|
||||
// running the collector removes the host from the active set
|
||||
var stats collectorExecStats
|
||||
err = task.collectPolicyQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, stats.Keys)
|
||||
require.Equal(t, 4, stats.Items)
|
||||
require.False(t, stats.Failed)
|
||||
|
||||
count, err = redigo.Int(conn.Do("ZCARD", policyPassHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, count)
|
||||
}
|
||||
|
||||
func createPolicies(t *testing.T, ds *mysql.Datastore, count int) []uint {
|
||||
ctx := context.Background()
|
||||
|
||||
ids := make([]uint, count)
|
||||
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
|
||||
for i := 0; i < count; i++ {
|
||||
res, err := tx.ExecContext(ctx, `INSERT INTO policies (name, description, query) VALUES (?, ?, ?)`, fmt.Sprintf("%s-%d", t.Name(), i), t.Name(), "SELECT 1")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
pid, _ := res.LastInsertId()
|
||||
ids[i] = uint(pid)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return ids
|
||||
}
|
||||
+152
-392
@@ -7,287 +7,51 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/fleetdm/fleet/v4/server/datastore/mysql"
|
||||
"github.com/fleetdm/fleet/v4/server/datastore/redis"
|
||||
"github.com/fleetdm/fleet/v4/server/datastore/redis/redistest"
|
||||
"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 TestCollectLabelQueryExecutions(t *testing.T) {
|
||||
func TestCollectQueryExecutions(t *testing.T) {
|
||||
ds := mysql.CreateMySQLDS(t)
|
||||
|
||||
t.Run("standalone", func(t *testing.T) {
|
||||
defer mysql.TruncateTables(t, ds)
|
||||
pool := redistest.SetupRedis(t, false, false, false)
|
||||
testCollectLabelQueryExecutions(t, ds, pool)
|
||||
oldMaxPolicy := maxRedisPolicyResultsPerHost
|
||||
maxRedisPolicyResultsPerHost = 3
|
||||
t.Cleanup(func() {
|
||||
maxRedisPolicyResultsPerHost = oldMaxPolicy
|
||||
})
|
||||
|
||||
t.Run("cluster", func(t *testing.T) {
|
||||
defer mysql.TruncateTables(t, ds)
|
||||
pool := redistest.SetupRedis(t, true, true, false)
|
||||
testCollectLabelQueryExecutions(t, ds, pool)
|
||||
t.Run("Label", func(t *testing.T) {
|
||||
t.Run("standalone", func(t *testing.T) {
|
||||
defer mysql.TruncateTables(t, ds)
|
||||
pool := redistest.SetupRedis(t, false, false, false)
|
||||
testCollectLabelQueryExecutions(t, ds, pool)
|
||||
})
|
||||
|
||||
t.Run("cluster", func(t *testing.T) {
|
||||
defer mysql.TruncateTables(t, ds)
|
||||
pool := redistest.SetupRedis(t, true, true, false)
|
||||
testCollectLabelQueryExecutions(t, ds, pool)
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("Policy", func(t *testing.T) {
|
||||
t.Run("standalone", func(t *testing.T) {
|
||||
defer mysql.TruncateTables(t, ds)
|
||||
pool := redistest.SetupRedis(t, false, false, false)
|
||||
testCollectPolicyQueryExecutions(t, ds, pool)
|
||||
})
|
||||
|
||||
t.Run("cluster", func(t *testing.T) {
|
||||
defer mysql.TruncateTables(t, ds)
|
||||
pool := redistest.SetupRedis(t, true, true, false)
|
||||
testCollectPolicyQueryExecutions(t, ds, pool)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func testCollectLabelQueryExecutions(t *testing.T, ds *mysql.Datastore, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
|
||||
type labelMembership struct {
|
||||
HostID int `db:"host_id"`
|
||||
LabelID uint `db:"label_id"`
|
||||
UpdatedAt time.Time `db:"updated_at"`
|
||||
}
|
||||
|
||||
hostIDs := createHosts(t, ds, 4, time.Now().Add(-24*time.Hour))
|
||||
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.
|
||||
cases := []struct {
|
||||
name string
|
||||
// map of host ID to label IDs to insert (true) or delete (false)
|
||||
reported map[int]map[int]bool
|
||||
want []labelMembership
|
||||
}{
|
||||
{"no key", nil, nil},
|
||||
{
|
||||
"report host 1 label 1",
|
||||
map[int]map[int]bool{hid(1): {1: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 1, 2",
|
||||
map[int]map[int]bool{hid(1): {1: true, 2: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 1, 2, 3",
|
||||
map[int]map[int]bool{1: {1: true, 2: true, 3: true}},
|
||||
[]labelMembership{
|
||||
{HostID: 1, LabelID: 1},
|
||||
{HostID: 1, LabelID: 2},
|
||||
{HostID: 1, LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -1",
|
||||
map[int]map[int]bool{hid(1): {1: false}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(1), LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -2, -3",
|
||||
map[int]map[int]bool{hid(1): {2: false, 3: false}},
|
||||
[]labelMembership{},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 1, 2, 3, 4",
|
||||
map[int]map[int]bool{hid(1): {1: true, 2: true, 3: true, 4: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(1), LabelID: 3},
|
||||
{HostID: hid(1), LabelID: 4},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -2, -3, -4, -5",
|
||||
map[int]map[int]bool{hid(1): {2: false, 3: false, 4: false, 5: false}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels 2, host 2 labels 2, 3",
|
||||
map[int]map[int]bool{hid(1): {2: true}, hid(2): {2: true, 3: true}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report host 1 labels -99, non-existing",
|
||||
map[int]map[int]bool{hid(1): {99: false}},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"report hosts 1, 2, 3, 4 labels 1, 2, -3, 4",
|
||||
map[int]map[int]bool{
|
||||
hid(1): {1: true, 2: true, 3: false, 4: true},
|
||||
hid(2): {1: true, 2: true, 3: false, 4: true},
|
||||
hid(3): {1: true, 2: true, 3: false, 4: true},
|
||||
hid(4): {1: true, 2: true, 3: false, 4: true},
|
||||
},
|
||||
[]labelMembership{
|
||||
{HostID: hid(1), LabelID: 1},
|
||||
{HostID: hid(1), LabelID: 2},
|
||||
{HostID: hid(1), LabelID: 4},
|
||||
{HostID: hid(2), LabelID: 1},
|
||||
{HostID: hid(2), LabelID: 2},
|
||||
{HostID: hid(2), LabelID: 4},
|
||||
{HostID: hid(3), LabelID: 1},
|
||||
{HostID: hid(3), LabelID: 2},
|
||||
{HostID: hid(3), LabelID: 4},
|
||||
{HostID: hid(4), LabelID: 1},
|
||||
{HostID: hid(4), LabelID: 2},
|
||||
{HostID: hid(4), LabelID: 4},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
const batchSizes = 3
|
||||
|
||||
setupTest := func(t *testing.T, data map[int]map[int]bool) collectorExecStats {
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
|
||||
// store the host memberships and prepare the expected stats
|
||||
var wantStats collectorExecStats
|
||||
for hostID, res := range data {
|
||||
if len(res) > 0 {
|
||||
key := fmt.Sprintf(labelMembershipHostKey, hostID)
|
||||
args := make(redigo.Args, 0, 1+(len(res)*2))
|
||||
args = args.Add(key)
|
||||
for lblID, ins := range res {
|
||||
score := -1
|
||||
if ins {
|
||||
score = 1
|
||||
}
|
||||
args = args.Add(score, lblID)
|
||||
}
|
||||
_, err := conn.Do("ZADD", args...)
|
||||
require.NoError(t, err)
|
||||
_, err = conn.Do("ZADD", labelMembershipActiveHostIDsKey, time.Now().Unix(), hostID)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
cnt, err := redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
wantStats.Keys = cnt
|
||||
wantStats.Items += len(res)
|
||||
wantStats.RedisCmds++
|
||||
wantStats.RedisCmds += len(res) / batchSizes
|
||||
}
|
||||
return wantStats
|
||||
}
|
||||
|
||||
selectRows := func(t *testing.T) ([]labelMembership, map[int]time.Time) {
|
||||
var rows []labelMembership
|
||||
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
|
||||
return sqlx.SelectContext(ctx, tx, &rows, `SELECT host_id, label_id, updated_at FROM label_membership ORDER BY 1, 2`)
|
||||
})
|
||||
|
||||
var hosts []struct {
|
||||
ID int `db:"id"`
|
||||
LabelUpdatedAt time.Time `db:"label_updated_at"`
|
||||
}
|
||||
mysql.ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error {
|
||||
return sqlx.SelectContext(ctx, tx, &hosts, `SELECT id, label_updated_at FROM hosts`)
|
||||
})
|
||||
|
||||
hostsUpdated := make(map[int]time.Time, len(hosts))
|
||||
for _, h := range hosts {
|
||||
hostsUpdated[h.ID] = h.LabelUpdatedAt
|
||||
}
|
||||
return rows, hostsUpdated
|
||||
}
|
||||
|
||||
minUpdatedAt := time.Now()
|
||||
for _, c := range cases {
|
||||
func() {
|
||||
t.Log("test name: ", c.name)
|
||||
wantStats := setupTest(t, c.reported)
|
||||
|
||||
// run the collection
|
||||
var stats collectorExecStats
|
||||
task := Task{
|
||||
InsertBatch: batchSizes,
|
||||
UpdateBatch: batchSizes,
|
||||
DeleteBatch: batchSizes,
|
||||
RedisPopCount: batchSizes,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
err := task.collectLabelQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
// inserts, updates and deletes are a bit tricky to track automatically,
|
||||
// just ignore them when comparing stats.
|
||||
stats.Inserts, stats.Updates, stats.Deletes = 0, 0, 0
|
||||
require.Equal(t, wantStats, stats)
|
||||
|
||||
// check that the table contains the expected rows
|
||||
rows, hostsUpdated := 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.Equal(t, want.LabelID, got.LabelID)
|
||||
require.WithinDuration(t, minUpdatedAt, got.UpdatedAt, 10*time.Second)
|
||||
|
||||
ts, ok := hostsUpdated[want.HostID]
|
||||
require.True(t, ok)
|
||||
require.WithinDuration(t, minUpdatedAt, ts, 10*time.Second)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// after all cases, run one last upsert (an update) to make sure that the
|
||||
// updated at column is properly updated. First we need to ensure that this
|
||||
// runs in a distinct second, because the mysql resolution is not precise.
|
||||
time.Sleep(time.Second)
|
||||
|
||||
var h1l1Before labelMembership
|
||||
beforeRows, _ := selectRows(t)
|
||||
for _, row := range beforeRows {
|
||||
if row.HostID == 1 && row.LabelID == 1 {
|
||||
h1l1Before = row
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// update host 1, label 1, already existing
|
||||
setupTest(t, map[int]map[int]bool{1: {1: true}})
|
||||
var stats collectorExecStats
|
||||
task := Task{
|
||||
InsertBatch: batchSizes,
|
||||
UpdateBatch: batchSizes,
|
||||
DeleteBatch: batchSizes,
|
||||
RedisPopCount: batchSizes,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
err := task.collectLabelQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
|
||||
var h1l1After labelMembership
|
||||
afterRows, _ := selectRows(t)
|
||||
for _, row := range afterRows {
|
||||
if row.HostID == 1 && row.LabelID == 1 {
|
||||
h1l1After = row
|
||||
break
|
||||
}
|
||||
}
|
||||
require.True(t, h1l1Before.UpdatedAt.Before(h1l1After.UpdatedAt))
|
||||
}
|
||||
|
||||
func TestRecordLabelQueryExecutions(t *testing.T) {
|
||||
func TestRecordQueryExecutions(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
|
||||
@@ -295,140 +59,136 @@ func TestRecordLabelQueryExecutions(t *testing.T) {
|
||||
ds.AsyncBatchUpdateLabelTimestampFunc = func(ctx context.Context, ids []uint, ts time.Time) error {
|
||||
return nil
|
||||
}
|
||||
ds.RecordPolicyQueryExecutionsFunc = func(ctx context.Context, host *fleet.Host, results map[uint]*bool, ts time.Time, deferred bool) error {
|
||||
return nil
|
||||
}
|
||||
ds.AsyncBatchInsertPolicyMembershipFunc = func(ctx context.Context, batch []fleet.PolicyMembershipResult) error {
|
||||
return nil
|
||||
}
|
||||
ds.AsyncBatchUpdatePolicyTimestampFunc = func(ctx context.Context, ids []uint, ts time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
t.Run("Label", func(t *testing.T) {
|
||||
t.Run("standalone", func(t *testing.T) {
|
||||
pool := redistest.SetupRedis(t, false, false, false)
|
||||
t.Run("sync", func(t *testing.T) { testRecordLabelQueryExecutionsSync(t, ds, pool) })
|
||||
t.Run("async", func(t *testing.T) { testRecordLabelQueryExecutionsAsync(t, ds, pool) })
|
||||
})
|
||||
|
||||
t.Run("cluster", func(t *testing.T) {
|
||||
pool := redistest.SetupRedis(t, true, true, false)
|
||||
t.Run("sync", func(t *testing.T) { testRecordLabelQueryExecutionsSync(t, ds, pool) })
|
||||
t.Run("async", func(t *testing.T) { testRecordLabelQueryExecutionsAsync(t, ds, pool) })
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("Policy", func(t *testing.T) {
|
||||
t.Run("standalone", func(t *testing.T) {
|
||||
pool := redistest.SetupRedis(t, false, false, false)
|
||||
t.Run("sync", func(t *testing.T) { testRecordPolicyQueryExecutionsSync(t, ds, pool) })
|
||||
t.Run("async", func(t *testing.T) { testRecordPolicyQueryExecutionsAsync(t, ds, pool) })
|
||||
})
|
||||
|
||||
t.Run("cluster", func(t *testing.T) {
|
||||
pool := redistest.SetupRedis(t, true, true, false)
|
||||
t.Run("sync", func(t *testing.T) { testRecordPolicyQueryExecutionsSync(t, ds, pool) })
|
||||
t.Run("async", func(t *testing.T) { testRecordPolicyQueryExecutionsAsync(t, ds, pool) })
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestActiveHostIDsSet(t *testing.T) {
|
||||
runTest := func(t *testing.T, pool fleet.RedisPool) {
|
||||
const zkey = "testActiveHostIDsSet"
|
||||
|
||||
activeHosts, err := loadActiveHostIDs(pool, zkey, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, activeHosts, 0)
|
||||
|
||||
// add a few hosts with a timestamp that increases by a second for each
|
||||
// note that host IDs will be 1..10 (t[0] == host 1, t[1] == host 2, etc.)
|
||||
tpurgeNone := time.Now()
|
||||
ts := make([]int64, 10)
|
||||
for i := range ts {
|
||||
if i > 0 {
|
||||
ts[i] = time.Unix(ts[i-1], 0).Add(time.Second).Unix()
|
||||
} else {
|
||||
ts[i] = tpurgeNone.Add(time.Second).Unix()
|
||||
}
|
||||
|
||||
// none ever get deleted, all are after tpurgeNone
|
||||
n, err := storePurgeActiveHostID(pool, zkey, uint(i+1), time.Unix(ts[i], 0), tpurgeNone)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
}
|
||||
|
||||
activeHosts, err = loadActiveHostIDs(pool, zkey, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, activeHosts, len(ts))
|
||||
for i, host := range activeHosts {
|
||||
require.Equal(t, ts[i], host.LastReported)
|
||||
}
|
||||
|
||||
// store a new one but now use t[1] as purge date - will remove two
|
||||
ts = append(ts, time.Unix(ts[len(ts)-1], 0).Add(time.Second).Unix())
|
||||
n, err := storePurgeActiveHostID(pool, zkey, uint(len(ts)), time.Unix(ts[len(ts)-1], 0), time.Unix(ts[1], 0))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, n)
|
||||
|
||||
// report t[3] and t[5] (hosts 4 and 6) as processed
|
||||
batch := []hostIDLastReported{
|
||||
{HostID: 4, LastReported: ts[3]},
|
||||
{HostID: 6, LastReported: ts[5]},
|
||||
}
|
||||
n, err = removeProcessedHostIDs(pool, zkey, batch)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, n)
|
||||
|
||||
// update t[6] of host 7, as if it had reported new data since the load
|
||||
newT6 := time.Unix(ts[len(ts)-1], 0).Add(time.Second)
|
||||
n, err = storePurgeActiveHostID(pool, zkey, 7, newT6, tpurgeNone)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
// report t[6] and t[7] (hosts 7 and 8) as processed, but only host 8
|
||||
// will get deleted, because the timestamp of host 7 has changed (we pass
|
||||
// its old timestamp, to simluate that it changed since loading the
|
||||
// information)
|
||||
batch = []hostIDLastReported{
|
||||
{HostID: 7, LastReported: ts[6]},
|
||||
{HostID: 8, LastReported: ts[7]},
|
||||
}
|
||||
n, err = removeProcessedHostIDs(pool, zkey, batch)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, n)
|
||||
|
||||
// check the remaining active hosts (only 6 remain)
|
||||
activeHosts, err = loadActiveHostIDs(pool, zkey, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, activeHosts, 6)
|
||||
want := []hostIDLastReported{
|
||||
{HostID: 3, LastReported: ts[2]},
|
||||
{HostID: 5, LastReported: ts[4]},
|
||||
{HostID: 7, LastReported: newT6.Unix()},
|
||||
{HostID: 9, LastReported: ts[8]},
|
||||
{HostID: 10, LastReported: ts[9]},
|
||||
{HostID: 11, LastReported: ts[10]},
|
||||
}
|
||||
require.ElementsMatch(t, want, activeHosts)
|
||||
}
|
||||
|
||||
t.Run("standalone", func(t *testing.T) {
|
||||
pool := redistest.SetupRedis(t, false, false, false)
|
||||
t.Run("sync", func(t *testing.T) { testRecordLabelQueryExecutionsSync(t, ds, pool) })
|
||||
t.Run("async", func(t *testing.T) { testRecordLabelQueryExecutionsAsync(t, ds, pool) })
|
||||
t.Run("sync", func(t *testing.T) { runTest(t, pool) })
|
||||
})
|
||||
|
||||
t.Run("cluster", func(t *testing.T) {
|
||||
pool := redistest.SetupRedis(t, true, true, false)
|
||||
t.Run("sync", func(t *testing.T) { testRecordLabelQueryExecutionsSync(t, ds, pool) })
|
||||
t.Run("async", func(t *testing.T) { testRecordLabelQueryExecutionsAsync(t, ds, pool) })
|
||||
t.Run("sync", func(t *testing.T) { runTest(t, pool) })
|
||||
})
|
||||
}
|
||||
|
||||
func testRecordLabelQueryExecutionsSync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
lastYear := now.Add(-365 * 24 * time.Hour)
|
||||
host := &fleet.Host{
|
||||
ID: 1,
|
||||
Platform: "linux",
|
||||
LabelUpdatedAt: lastYear,
|
||||
}
|
||||
|
||||
var yes, no = true, false
|
||||
results := map[uint]*bool{1: &yes, 2: &yes, 3: &no, 4: nil}
|
||||
keySet, keyTs := fmt.Sprintf(labelMembershipHostKey, host.ID), fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
|
||||
task := Task{
|
||||
Datastore: ds,
|
||||
Pool: pool,
|
||||
AsyncEnabled: false,
|
||||
}
|
||||
|
||||
labelReportedAt := task.GetHostLabelReportedAt(ctx, host)
|
||||
require.True(t, labelReportedAt.Equal(lastYear))
|
||||
|
||||
err := task.RecordLabelQueryExecutions(ctx, host, results, now)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ds.RecordLabelQueryExecutionsFuncInvoked)
|
||||
ds.RecordLabelQueryExecutionsFuncInvoked = false
|
||||
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
defer conn.Do("DEL", keySet, keyTs)
|
||||
|
||||
n, err := redigo.Int(conn.Do("EXISTS", keySet))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
n, err = redigo.Int(conn.Do("EXISTS", keyTs))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
n, err = redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, n)
|
||||
|
||||
labelReportedAt = task.GetHostLabelReportedAt(ctx, host)
|
||||
require.True(t, labelReportedAt.Equal(now))
|
||||
}
|
||||
|
||||
func testRecordLabelQueryExecutionsAsync(t *testing.T, ds *mock.Store, pool fleet.RedisPool) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
lastYear := now.Add(-365 * 24 * time.Hour)
|
||||
host := &fleet.Host{
|
||||
ID: 1,
|
||||
Platform: "linux",
|
||||
LabelUpdatedAt: lastYear,
|
||||
}
|
||||
var yes, no = true, false
|
||||
results := map[uint]*bool{1: &yes, 2: &yes, 3: &no, 4: nil}
|
||||
keySet, keyTs := fmt.Sprintf(labelMembershipHostKey, host.ID), fmt.Sprintf(labelMembershipReportedKey, host.ID)
|
||||
|
||||
task := Task{
|
||||
Datastore: ds,
|
||||
Pool: pool,
|
||||
AsyncEnabled: true,
|
||||
|
||||
InsertBatch: 3,
|
||||
UpdateBatch: 3,
|
||||
DeleteBatch: 3,
|
||||
RedisPopCount: 3,
|
||||
RedisScanKeysCount: 10,
|
||||
}
|
||||
|
||||
labelReportedAt := task.GetHostLabelReportedAt(ctx, host)
|
||||
require.True(t, labelReportedAt.Equal(lastYear))
|
||||
|
||||
err := task.RecordLabelQueryExecutions(ctx, host, results, now)
|
||||
require.NoError(t, err)
|
||||
require.False(t, ds.RecordLabelQueryExecutionsFuncInvoked)
|
||||
|
||||
conn := redis.ConfigureDoer(pool, pool.Get())
|
||||
defer conn.Close()
|
||||
defer conn.Do("DEL", keySet, keyTs)
|
||||
|
||||
res, err := redigo.IntMap(conn.Do("ZPOPMIN", keySet, 10))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 4, len(res))
|
||||
require.Equal(t, map[string]int{"1": 1, "2": 1, "3": -1, "4": -1}, res)
|
||||
|
||||
ts, err := redigo.Int64(conn.Do("GET", keyTs))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, now.Unix(), ts)
|
||||
|
||||
count, err := redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count)
|
||||
tsActive, err := redigo.Int64(conn.Do("ZSCORE", labelMembershipActiveHostIDsKey, host.ID))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tsActive, ts)
|
||||
|
||||
labelReportedAt = task.GetHostLabelReportedAt(ctx, host)
|
||||
// because we transition via unix epoch (seconds), not exactly equal
|
||||
require.WithinDuration(t, now, labelReportedAt, time.Second)
|
||||
// host's LabelUpdatedAt field hasn't been updated yet, because the label
|
||||
// results are in redis, not in mysql yet.
|
||||
require.True(t, host.LabelUpdatedAt.Equal(lastYear))
|
||||
|
||||
// running the collector removes the host from the active set
|
||||
var stats collectorExecStats
|
||||
err = task.collectLabelQueryExecutions(ctx, ds, pool, &stats)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, stats.Keys)
|
||||
require.Equal(t, 0, stats.Items) // zero because we cleared the host's set with ZPOPMIN above
|
||||
require.False(t, stats.Failed)
|
||||
|
||||
count, err = redigo.Int(conn.Do("ZCARD", labelMembershipActiveHostIDsKey))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, count)
|
||||
}
|
||||
|
||||
func createHosts(t *testing.T, ds fleet.Datastore, count int, ts time.Time) []uint {
|
||||
ids := make([]uint, count)
|
||||
for i := 0; i < count; i++ {
|
||||
|
||||
@@ -468,7 +468,8 @@ func (svc *Service) labelQueriesForHost(ctx context.Context, host *fleet.Host) (
|
||||
}
|
||||
|
||||
func (svc *Service) policyQueriesForHost(ctx context.Context, host *fleet.Host) (map[string]string, error) {
|
||||
if !svc.shouldUpdate(host.PolicyUpdatedAt, svc.config.Osquery.PolicyUpdateInterval, host.ID) && !host.RefetchRequested {
|
||||
policyReportedAt := svc.task.GetHostPolicyReportedAt(ctx, host)
|
||||
if !svc.shouldUpdate(policyReportedAt, svc.config.Osquery.PolicyUpdateInterval, host.ID) && !host.RefetchRequested {
|
||||
return nil, nil
|
||||
}
|
||||
policyQueries, err := svc.ds.PolicyQueriesForHost(ctx, host)
|
||||
@@ -740,14 +741,8 @@ func (svc *Service) SubmitDistributedQueryResults(
|
||||
}
|
||||
|
||||
if len(labelResults) > 0 {
|
||||
if ac.ServerSettings.DeferredSaveHost {
|
||||
if err := svc.ds.RecordLabelQueryExecutions(ctx, host, labelResults, svc.clock.Now(), true); err != nil {
|
||||
logging.WithErr(ctx, err)
|
||||
}
|
||||
} else {
|
||||
if err := svc.task.RecordLabelQueryExecutions(ctx, host, labelResults, svc.clock.Now()); err != nil {
|
||||
logging.WithErr(ctx, err)
|
||||
}
|
||||
if err := svc.task.RecordLabelQueryExecutions(ctx, host, labelResults, svc.clock.Now(), ac.ServerSettings.DeferredSaveHost); err != nil {
|
||||
logging.WithErr(ctx, err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -765,10 +760,14 @@ func (svc *Service) SubmitDistributedQueryResults(
|
||||
}()
|
||||
}
|
||||
}
|
||||
// NOTE(mna): currently, failing policies webhook wouldn't see the new
|
||||
// flipped policies on the next run if async processing is enabled and the
|
||||
// collection has not been done yet (not persisted in mysql). Should
|
||||
// FlippingPoliciesForHost take pending redis data into consideration, or
|
||||
// maybe we should impose restrictions between async collection interval
|
||||
// and policy update interval?
|
||||
|
||||
host.PolicyUpdatedAt = svc.clock.Now()
|
||||
err = svc.ds.RecordPolicyQueryExecutions(ctx, host, policyResults, svc.clock.Now(), ac.ServerSettings.DeferredSaveHost)
|
||||
if err != nil {
|
||||
if err := svc.task.RecordPolicyQueryExecutions(ctx, host, policyResults, svc.clock.Now(), ac.ServerSettings.DeferredSaveHost); err != nil {
|
||||
logging.WithErr(ctx, err)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user