diff --git a/changes/33254-prevent-lingering-redis-keys b/changes/33254-prevent-lingering-redis-keys new file mode 100644 index 0000000000..1dac13540d --- /dev/null +++ b/changes/33254-prevent-lingering-redis-keys @@ -0,0 +1 @@ +* Fixed a bug in live queries that caused `livequery:{$CAMPAIGN_ID}` Redis keys to not be cleaned up or expire. diff --git a/server/live_query/live_query_test.go b/server/live_query/live_query_test.go index 73f64593de..aa57bcca83 100644 --- a/server/live_query/live_query_test.go +++ b/server/live_query/live_query_test.go @@ -19,6 +19,7 @@ var testFunctions = [...]func(*testing.T, fleet.LiveQueryStore){ testLiveQueryExpiredQuery, testLiveQueryOnlyExpired, testLiveQueryCleanupInactive, + testLiveQuerySetBitOnlyIfKeyExists, } func testLiveQuery(t *testing.T, store fleet.LiveQueryStore) { @@ -221,3 +222,45 @@ func testLiveQueryCleanupInactive(t *testing.T, store fleet.LiveQueryStore) { require.NoError(t, err) require.Empty(t, m) } + +func testLiveQuerySetBitOnlyIfKeyExists(t *testing.T, store fleet.LiveQueryStore) { + // Create a live query campaign. + err := store.RunQuery("test", "SELECT 1;", []uint{1}) + require.NoError(t, err) + + // Get the query for the host. + queries, err := store.QueriesForHost(1) + require.NoError(t, err) + require.Equal(t, + map[string]string{ + "test": "SELECT 1;", + }, + queries, + ) + + // Mark query as completed by host. + err = store.QueryCompletedByHost("test", 1) + require.NoError(t, err) + + // Query should not be returned anymore as it was marked as completed for this host. + queries, err = store.QueriesForHost(1) + require.NoError(t, err) + require.Empty(t, queries) + + // A host could be attempting to write a result for a query that was already deleted. + err = store.QueryCompletedByHost("test-2", 1) + require.NoError(t, err) + + // Let's test that such key was not created. + + // get a raw Redis connection to make direct checks + pool := store.(*redisLiveQuery).pool + conn := redis.ConfigureDoer(pool, pool.Get()) + t.Cleanup(func() { + conn.Close() + }) + + n, err := redigo.Int(conn.Do("EXISTS", queryKeyPrefix+"{test-2}")) + require.NoError(t, err) + require.Zero(t, n) +} diff --git a/server/live_query/redis_live_query.go b/server/live_query/redis_live_query.go index b54c1c1f3a..c5538ac9b1 100644 --- a/server/live_query/redis_live_query.go +++ b/server/live_query/redis_live_query.go @@ -263,8 +263,17 @@ func (r *redisLiveQuery) QueryCompletedByHost(name string, hostID uint) error { targetKey, _ := generateKeys(name) - // Update the bitfield for this host. - if _, err := conn.Do("SETBIT", targetKey, hostID, 0); err != nil { + // Update the bitfield for this host only if the key exists. + // If the key doesn't exist (e.g. query marked as completed or cancelled) + // then we don't want to call SETBIT because it will create a new + // key (that won't expire and linger "forever"). + const setBitScript = ` + if redis.call('EXISTS', KEYS[1]) == 1 then + return redis.call('SETBIT', KEYS[1], ARGV[1], ARGV[2]) + else + return nil + end` + if _, err := conn.Do("EVAL", setBitScript, 1, targetKey, hostID, 0); err != nil { return fmt.Errorf("setbit query key: %w", err) }