From 4c671e7a159c19ed255c95ad27da54985d94c680 Mon Sep 17 00:00:00 2001
From: Lucas Manuel Rodriguez <lucas@fleetdm.com>
Date: Thu, 14 Mar 2024 12:59:57 -0300
Subject: [PATCH] Fix scheduled query results in osquery-perf for frequencies >
 1m (#17576)

This is a bug in the generation of results of scheduled queries in
osquery-perf.
It seems the bug has been around since we added scheduled query result
support in osquery-perf.

PS: In my Fleet downtime tests for
https://github.com/fleetdm/fleet/issues/16423 I probably missed this
because I was using a high frequency queries (150s intervals IIRC).
---
 cmd/osquery-perf/agent.go | 34 ++++++++++++++++++++--------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/cmd/osquery-perf/agent.go b/cmd/osquery-perf/agent.go
index 0d8d9c77e9..1e76edd1bb 100644
--- a/cmd/osquery-perf/agent.go
+++ b/cmd/osquery-perf/agent.go
@@ -387,8 +387,7 @@ type agent struct {
 	DiskEncryptionEnabled bool
 
 	scheduledQueriesMu sync.Mutex // protects the below members
-	scheduledQueries   []string
-	scheduledQueryData []scheduledQuery
+	scheduledQueryData map[string]scheduledQuery
 	// bufferedResults contains result logs that are buffered when
 	// /api/v1/osquery/log requests to the Fleet server fail.
 	//
@@ -528,9 +527,10 @@ type scheduledQuery struct {
 	Platform         string  `json:"platform"`
 	Version          string  `json:"version"`
 	Snapshot         bool    `json:"snapshot"`
-	nextRun          float64
-	numRows          uint
-	packName         string
+
+	lastRun  int64
+	numRows  uint
+	packName string
 }
 
 func (a *agent) isOrbit() bool {
@@ -620,14 +620,17 @@ func (a *agent) runLoop(i int, onlyAlreadyEnrolled bool) {
 		now := time.Now().Unix()
 		a.scheduledQueriesMu.Lock()
 		prevCount := a.countBuffered()
-		for i, query := range a.scheduledQueryData {
-			if query.nextRun == 0 || now >= int64(query.nextRun) {
+		for queryName, query := range a.scheduledQueryData {
+			if query.lastRun == 0 || now >= (query.lastRun+int64(query.ScheduleInterval)) {
 				results = append(results, resultLog{
 					packName:  query.packName,
 					queryName: query.Name,
 					numRows:   int(query.numRows),
 				})
-				a.scheduledQueryData[i].nextRun = float64(now + int64(query.ScheduleInterval))
+				// Update lastRun
+				v := a.scheduledQueryData[queryName]
+				v.lastRun = now
+				a.scheduledQueryData[queryName] = v
 			}
 		}
 		if prevCount+len(results) < 1_000_000 { // osquery buffered_log_max is 1M
@@ -1122,11 +1125,9 @@ func (a *agent) config() error {
 		return fmt.Errorf("json parse at config: %w", err)
 	}
 
-	var scheduledQueries []string
-	var scheduledQueryData []scheduledQuery
+	scheduledQueryData := make(map[string]scheduledQuery)
 	for packName, pack := range parsedResp.Packs {
 		for queryName, query := range pack.Queries {
-			scheduledQueries = append(scheduledQueries, packName+"_"+queryName)
 			m, ok := query.(map[string]interface{})
 			if !ok {
 				return fmt.Errorf("processing scheduled query failed: %v", query)
@@ -1146,12 +1147,17 @@ func (a *agent) config() error {
 			}
 			q.ScheduleInterval = m["interval"].(float64)
 			q.Query = m["query"].(string)
-			scheduledQueryData = append(scheduledQueryData, q)
+
+			scheduledQueryName := packName + "_" + queryName
+			if existingEntry, ok := a.scheduledQueryData[scheduledQueryName]; ok {
+				// Keep lastRun if the query is already scheduled.
+				q.lastRun = existingEntry.lastRun
+			}
+			scheduledQueryData[scheduledQueryName] = q
 		}
 	}
 
 	a.scheduledQueriesMu.Lock()
-	a.scheduledQueries = scheduledQueries
 	a.scheduledQueryData = scheduledQueryData
 	a.scheduledQueriesMu.Unlock()
 
@@ -1337,7 +1343,7 @@ func (a *agent) randomQueryStats() []map[string]string {
 	defer a.scheduledQueriesMu.Unlock()
 
 	var stats []map[string]string
-	for _, scheduledQuery := range a.scheduledQueries {
+	for scheduledQuery := range a.scheduledQueryData {
 		stats = append(stats, map[string]string{
 			"name":           scheduledQuery,
 			"delimiter":      "_",