diff --git a/changes/42441-live-query-redis-scaling b/changes/42441-live-query-redis-scaling new file mode 100644 index 0000000000..0da1baa19d --- /dev/null +++ b/changes/42441-live-query-redis-scaling @@ -0,0 +1 @@ +- Fixed a bug where running many concurrent live queries that each target a small number of hosts could overload Redis and slow down host check-ins. diff --git a/cmd/fleet/serve.go b/cmd/fleet/serve.go index 12ba7843e7..c76fadaad2 100644 --- a/cmd/fleet/serve.go +++ b/cmd/fleet/serve.go @@ -278,7 +278,8 @@ func runServeCmd(cmd *cobra.Command, configManager configpkg.Manager, debug, dev resultStore := pubsub.NewRedisQueryResults(redisPool, config.Redis.DuplicateResults, logger.With("component", "query-results"), ) - liveQueryStore := live_query.NewRedisLiveQuery(redisPool, logger, liveQueryMemCacheDuration) + liveQueryStore := live_query.NewRedisLiveQuery(redisPool, logger, liveQueryMemCacheDuration, + config.Redis.LiveQuerySmallTargetThreshold) ssoSessionStore := sso.NewSessionStore(redisPool) osquerydStatusLogger, osquerydResultLogger, auditLogger := initOsqueryLogging(cmd.Context(), config, license, logger, initFatal) diff --git a/server/config/config.go b/server/config/config.go index fc33269ea5..cc4bac5f35 100644 --- a/server/config/config.go +++ b/server/config/config.go @@ -101,6 +101,13 @@ type RedisConfig struct { // per-entry TTL is jittered by ±10% to avoid synchronized expiry waves. // Only meaningful when HostCacheEnabled is true. Hidden from --help. HostCacheTTL time.Duration `yaml:"host_cache_ttl"` + // LiveQuerySmallTargetThreshold is the maximum number of targeted hosts for a + // live query to use the per-host reverse index instead of a fleet-wide + // bitfield. Storing small-target queries as a per-host set means a host + // checkin no longer issues one GETBIT per such query. Set to 0 to disable the + // reverse index entirely (kill-switch) and use the bitfield for all live + // queries. + LiveQuerySmallTargetThreshold int `yaml:"live_query_small_target_threshold"` } const ( @@ -1355,6 +1362,10 @@ func (man Manager) addConfigs() { "Base TTL for Redis-backed host lookup cache entries. Actual per-entry TTL is jittered by ±10% to avoid "+ "synchronized expiry waves. Must be > 0 when redis.host_cache_enabled is true; set "+ "redis.host_cache_enabled=false to disable the cache.") + man.addConfigInt("redis.live_query_small_target_threshold", 1000, + "Maximum number of targeted hosts for a live query to use the per-host reverse index instead of a "+ + "fleet-wide bitfield, avoiding one GETBIT per query on every host check-in. Set to 0 to disable "+ + "the reverse index and use the bitfield for all live queries.") // Server man.addConfigString("server.address", "0.0.0.0:8080", @@ -1836,35 +1847,36 @@ func (man Manager) LoadConfig() FleetConfig { Mysql: loadMysqlConfig("mysql"), MysqlReadReplica: loadMysqlConfig("mysql_read_replica"), Redis: RedisConfig{ - Address: man.getConfigString("redis.address"), - Username: man.getConfigString("redis.username"), - Password: man.getConfigString("redis.password"), - Database: man.getConfigInt("redis.database"), - Region: man.getConfigString("redis.region"), - CacheName: man.getConfigString("redis.cache_name"), - UseTLS: man.getConfigBool("redis.use_tls"), - DuplicateResults: man.getConfigBool("redis.duplicate_results"), - ConnectTimeout: man.getConfigDuration("redis.connect_timeout"), - KeepAlive: man.getConfigDuration("redis.keep_alive"), - ConnectRetryAttempts: man.getConfigInt("redis.connect_retry_attempts"), - ClusterFollowRedirections: man.getConfigBool("redis.cluster_follow_redirections"), - ClusterReadFromReplica: man.getConfigBool("redis.cluster_read_from_replica"), - TLSCert: man.getConfigString("redis.tls_cert"), - TLSKey: man.getConfigString("redis.tls_key"), - TLSCA: man.getConfigString("redis.tls_ca"), - TLSServerName: man.getConfigString("redis.tls_server_name"), - TLSHandshakeTimeout: man.getConfigDuration("redis.tls_handshake_timeout"), - MaxIdleConns: man.getConfigInt("redis.max_idle_conns"), - MaxOpenConns: man.getConfigInt("redis.max_open_conns"), - ConnMaxLifetime: man.getConfigDuration("redis.conn_max_lifetime"), - IdleTimeout: man.getConfigDuration("redis.idle_timeout"), - ConnWaitTimeout: man.getConfigDuration("redis.conn_wait_timeout"), - WriteTimeout: man.getConfigDuration("redis.write_timeout"), - ReadTimeout: man.getConfigDuration("redis.read_timeout"), - StsAssumeRoleArn: man.getConfigString("redis.sts_assume_role_arn"), - StsExternalID: man.getConfigString("redis.sts_external_id"), - HostCacheEnabled: man.getConfigBool("redis.host_cache_enabled"), - HostCacheTTL: man.getConfigDuration("redis.host_cache_ttl"), + Address: man.getConfigString("redis.address"), + Username: man.getConfigString("redis.username"), + Password: man.getConfigString("redis.password"), + Database: man.getConfigInt("redis.database"), + Region: man.getConfigString("redis.region"), + CacheName: man.getConfigString("redis.cache_name"), + UseTLS: man.getConfigBool("redis.use_tls"), + DuplicateResults: man.getConfigBool("redis.duplicate_results"), + ConnectTimeout: man.getConfigDuration("redis.connect_timeout"), + KeepAlive: man.getConfigDuration("redis.keep_alive"), + ConnectRetryAttempts: man.getConfigInt("redis.connect_retry_attempts"), + ClusterFollowRedirections: man.getConfigBool("redis.cluster_follow_redirections"), + ClusterReadFromReplica: man.getConfigBool("redis.cluster_read_from_replica"), + TLSCert: man.getConfigString("redis.tls_cert"), + TLSKey: man.getConfigString("redis.tls_key"), + TLSCA: man.getConfigString("redis.tls_ca"), + TLSServerName: man.getConfigString("redis.tls_server_name"), + TLSHandshakeTimeout: man.getConfigDuration("redis.tls_handshake_timeout"), + MaxIdleConns: man.getConfigInt("redis.max_idle_conns"), + MaxOpenConns: man.getConfigInt("redis.max_open_conns"), + ConnMaxLifetime: man.getConfigDuration("redis.conn_max_lifetime"), + IdleTimeout: man.getConfigDuration("redis.idle_timeout"), + ConnWaitTimeout: man.getConfigDuration("redis.conn_wait_timeout"), + WriteTimeout: man.getConfigDuration("redis.write_timeout"), + ReadTimeout: man.getConfigDuration("redis.read_timeout"), + StsAssumeRoleArn: man.getConfigString("redis.sts_assume_role_arn"), + StsExternalID: man.getConfigString("redis.sts_external_id"), + HostCacheEnabled: man.getConfigBool("redis.host_cache_enabled"), + HostCacheTTL: man.getConfigDuration("redis.host_cache_ttl"), + LiveQuerySmallTargetThreshold: man.getConfigInt("redis.live_query_small_target_threshold"), }, Server: ServerConfig{ Address: man.getConfigString("server.address"), diff --git a/server/live_query/redis_live_query.go b/server/live_query/redis_live_query.go index 3911b12aca..f794868e13 100644 --- a/server/live_query/redis_live_query.go +++ b/server/live_query/redis_live_query.go @@ -17,15 +17,19 @@ // into each host's set. This model has many potential writes for LQ // creation, but a host checkin has very few. // -// We believe that normal fleet usage has many hosts, and a small -// number of live queries targeting all of them. This was a big -// factor in choosing this implementation. +// The bitfield model fits "many hosts, few queries", but it scales poorly when +// many queries run concurrently: a host checkin must probe (GETBIT) every active +// query's bitfield, so the per-checkin cost grows with the number of queries - +// even queries that target a single host force a probe on every host. To handle +// that case, this package uses a hybrid: queries targeting at most +// smallTargetThreshold hosts are stored using the per-host set model above +// (the "reverse index"), while larger ("broadcast") queries keep the bitfield. // // # Implementation // -// As mentioned in the Design section, there are three keys for each -// live query: the bitfield, the SQL of the query and the set containing -// the IDs of all active live queries: +// There are three keys for each bitfield (broadcast) live query: the bitfield, +// the SQL of the query and the set containing the IDs of all active live +// queries: // // livequery: is the bitfield that indicates the hosts // sql:livequery: is the SQL of the query. @@ -38,6 +42,20 @@ // are always stored on the same node (as they hash to the same cluster slot). // See https://redis.io/topics/cluster-spec#keys-hash-tags for details. // +// Small-target queries instead use the reverse index. There is no bitfield; +// the campaign ID is added to a per-host set for each targeted host, and the +// campaign ID is also added to a set of reverse-model queries: +// +// livequery:host: is the set of campaign IDs targeting that host +// livequery:active:reverse is the set of campaign IDs using the reverse model +// +// The sql:livequery: and livequery:active keys are used by both models. A +// host checkin reads its own livequery:host: set once (instead of one +// GETBIT per small-target query) and probes the bitfield only for the remaining +// broadcast queries. The per-host sets have a TTL and stale entries (campaigns +// no longer active) are filtered against the active set at read time, so they do +// not need to be removed on StopQuery. +// // It is a noted downside that the active live queries set will necessarily // live on a single node in cluster mode (a "hot key"), and that node will see // increased activity due to that. Should that become a significant problem, an @@ -65,6 +83,8 @@ const ( queryKeyPrefix = "livequery:" sqlKeyPrefix = "sql:" activeQueriesKey = "livequery:active" + activeReverseQueriesKey = "livequery:active:reverse" + reverseHostKeyPrefix = "livequery:host:" queryExpiration = 7 * 24 * time.Hour queryResultsCountPrefix = "query_results_count:" ) @@ -77,6 +97,11 @@ type redisLiveQuery struct { // in memory cache expiration cacheExpiration time.Duration + // smallTargetThreshold is the maximum number of targeted hosts for a query to + // use the per-host reverse index instead of the bitfield. A value of 0 + // disables the reverse index entirely (all queries use the bitfield). + smallTargetThreshold int + logger *slog.Logger } @@ -86,6 +111,10 @@ type redisLiveQuery struct { type memCache struct { sqlCache map[string]string activeQueriesCache []string + // reverseActiveCache holds the campaign IDs (among the active queries) that + // use the reverse per-host index. It is used by the read path to exclude + // those queries from the per-host bitfield (GETBIT) probes. + reverseActiveCache map[string]struct{} cacheExp time.Time mu sync.RWMutex } @@ -106,14 +135,29 @@ func (r *redisLiveQuery) getSQLByCampaignID(campaignID string) (string, bool) { return sql, found } -// NewRedisQueryResults creates a new Redis implementation of the -// QueryResultStore interface using the provided Redis connection pool. -func NewRedisLiveQuery(pool fleet.RedisPool, logger *slog.Logger, memCacheExp time.Duration) *redisLiveQuery { +// isReverse is a thread-safe method that reports whether the given active +// campaign ID is stored using the reverse per-host index (rather than a +// bitfield). +func (r *redisLiveQuery) isReverse(campaignID string) bool { + r.cache.mu.RLock() + defer r.cache.mu.RUnlock() + _, found := r.cache.reverseActiveCache[campaignID] + return found +} + +// NewRedisLiveQuery creates a new Redis implementation of the live query store +// using the provided Redis connection pool. +// +// smallTargetThreshold is the maximum number of targeted hosts for a query to +// use the reverse per-host index instead of the bitfield; a value of 0 disables +// the reverse index entirely (kill-switch), so all queries use the bitfield. +func NewRedisLiveQuery(pool fleet.RedisPool, logger *slog.Logger, memCacheExp time.Duration, smallTargetThreshold int) *redisLiveQuery { return &redisLiveQuery{ - pool: pool, - cache: newMemCache(), - cacheExpiration: memCacheExp, - logger: logger, + pool: pool, + cache: newMemCache(), + cacheExpiration: memCacheExp, + smallTargetThreshold: smallTargetThreshold, + logger: logger, } } @@ -121,6 +165,7 @@ func newMemCache() memCache { return memCache{ sqlCache: make(map[string]string), activeQueriesCache: make([]string, 0), + reverseActiveCache: make(map[string]struct{}), } } @@ -148,6 +193,14 @@ func extractTargetKeyName(key string) string { return name } +// reverseHostKey returns the key of the per-host set that stores the campaign +// IDs of the small-target live queries targeting the given host. The host ID is +// used as the cluster hash tag so that a host's set always lives on a single +// node (the set is read on every checkin for that host). +func reverseHostKey(hostID uint) string { + return reverseHostKeyPrefix + "{" + strconv.FormatUint(uint64(hostID), 10) + "}" +} + // RunQuery stores the live query information in ephemeral storage for the // duration of the query or its TTL. Note that hostIDs *must* be sorted // in ascending order. The name is the campaign ID as a string. @@ -156,12 +209,27 @@ func (r *redisLiveQuery) RunQuery(name, sql string, hostIDs []uint) error { return errors.New("no hosts targeted") } - // store the sql and targeted hosts information - if err := r.storeQueryInfo(name, sql, hostIDs); err != nil { - return fmt.Errorf("store query info: %w", err) + // Small-target queries use the per-host reverse index so that a host checkin + // does not have to probe this query's bitfield (one GETBIT per query). Large + // (broadcast) queries keep the bitfield, which is compact relative to a large + // target set and cheap to create/stop. A threshold of 0 disables the reverse + // index (no query has <= 0 targets), so all queries use the bitfield. + if len(hostIDs) <= r.smallTargetThreshold { + if err := r.storeQueryInfoReverse(name, sql, hostIDs); err != nil { + return fmt.Errorf("store reverse query info: %w", err) + } + // mark the campaign id as using the reverse model + if err := r.storeReverseQueryName(name); err != nil { + return fmt.Errorf("store reverse query name: %w", err) + } + } else { + // store the sql and targeted hosts information (bitfield) + if err := r.storeQueryInfo(name, sql, hostIDs); err != nil { + return fmt.Errorf("store query info: %w", err) + } } - // store name (campaign id) into the active live queries set + // store name (campaign id) into the active live queries set (both models) if err := r.storeQueryNames(name); err != nil { return fmt.Errorf("store query name: %w", err) } @@ -170,7 +238,8 @@ func (r *redisLiveQuery) RunQuery(name, sql string, hostIDs []uint) error { } func (r *redisLiveQuery) StopQuery(name string) error { - // remove the sql and targeted hosts keys + // remove the sql and targeted hosts keys (DEL of the bitfield key is a no-op + // for reverse queries, which don't have one) if err := r.removeQueryInfo(name); err != nil { return fmt.Errorf("remove query info: %w", err) } @@ -180,6 +249,16 @@ func (r *redisLiveQuery) StopQuery(name string) error { return fmt.Errorf("remove query name: %w", err) } + // remove from the reverse model set. The per-host sets cannot be enumerated + // by campaign, so they are left to expire via their TTL and are filtered out + // at read time against the active set. This is safe only because campaign IDs + // are monotonic (MySQL auto-increment) and never reused: a stale per-host + // membership can therefore never collide with a different, newly-active + // campaign that happens to share the same ID. + if err := r.removeReverseQueryNames(name); err != nil { + return fmt.Errorf("remove reverse query name: %w", err) + } + return nil } @@ -187,30 +266,75 @@ func (r *redisLiveQuery) StopQuery(name string) error { var cleanupExpiredQueriesModulo int64 = 10 func (r *redisLiveQuery) QueriesForHost(hostID uint) (map[string]string, error) { - // Get keys for active queries + // Get keys for active queries (this also (re)loads the in-memory cache, which + // is what isReverse below relies on). names, err := r.LoadActiveQueryNames() if err != nil { return nil, fmt.Errorf("load active queries: %w", err) } - // convert the query name (campaign id) to the key name + queries := make(map[string]string) + + // Broadcast queries: probe this host's bit in each query's bitfield. Reverse + // (small-target) queries are excluded here - probing them is the per-checkin + // command storm this whole change is meant to avoid. keyNames := make([]string, 0, len(names)) for _, name := range names { + if r.isReverse(name) { + continue + } tkey, _ := generateKeys(name) keyNames = append(keyNames, tkey) } keysBySlot := redis.SplitKeysBySlot(r.pool, keyNames...) - queries := make(map[string]string) for _, qkeys := range keysBySlot { if err := r.collectBatchQueriesForHost(hostID, qkeys, queries); err != nil { return nil, err } } + // Reverse (small-target) queries: a single read of this host's own set. + if err := r.collectReverseQueriesForHost(hostID, queries); err != nil { + return nil, err + } + return queries, nil } +// collectReverseQueriesForHost reads the per-host reverse-index set and adds any +// still-active small-target queries targeting this host to queriesByHost. Stale +// campaign IDs (lingering in the per-host set after the query was stopped) are +// filtered out because their SQL is no longer in the cache. +func (r *redisLiveQuery) collectReverseQueriesForHost(hostID uint, queriesByHost map[string]string) error { + conn := redis.ReadOnlyConn(r.pool, r.pool.Get()) + defer conn.Close() + + // Stale-entry filtering below relies on the SQL cache holding only active + // queries. Refresh it on expiry here so this path stays correct on its own, + // independent of any cache (re)load done by the caller or the bitfield path + // (which is skipped when every active query is small-target). + if r.cacheIsExpired() { + if err := r.loadCache(); err != nil { + return fmt.Errorf("load cache: %w", err) + } + } + + names, err := redigo.Strings(conn.Do("SMEMBERS", reverseHostKey(hostID))) + if err != nil && err != redigo.ErrNil { + return fmt.Errorf("smembers reverse host key: %w", err) + } + + for _, name := range names { + // The SQL cache only holds active queries, so a missing entry means the + // campaign is no longer active (stale entry) and is skipped. + if sql, found := r.getSQLByCampaignID(name); found { + queriesByHost[name] = sql + } + } + return nil +} + func (r *redisLiveQuery) collectBatchQueriesForHost(hostID uint, queryKeys []string, queriesByHost map[string]string) error { conn := redis.ReadOnlyConn(r.pool, r.pool.Get()) defer conn.Close() @@ -261,6 +385,15 @@ func (r *redisLiveQuery) QueryCompletedByHost(name string, hostID uint) error { conn := redis.ConfigureDoer(r.pool, r.pool.Get()) defer conn.Close() + // Clear completion in both models without depending on which one this query + // uses: exactly one of these has an effect, the other is a harmless no-op + // (SREM on an absent member, and the guarded SETBIT below on an absent key). + // This avoids relying on a possibly-stale cache to pick the model, where a + // wrong guess would leave the host still receiving the query. + if _, err := conn.Do("SREM", reverseHostKey(hostID), name); err != nil { + return fmt.Errorf("srem reverse host key: %w", err) + } + targetKey, _ := generateKeys(name) // Update the bitfield for this host only if the key exists. @@ -308,6 +441,82 @@ func (r *redisLiveQuery) storeQueryInfo(name, sql string, hostIDs []uint) error return nil } +// storeQueryInfoReverse stores the SQL of the query and adds the campaign id to +// the per-host set of every targeted host (the reverse index). The per-host +// sets are given a TTL so that orphaned entries (e.g. if StopQuery is missed) +// eventually expire; they are also filtered against the active set at read time. +func (r *redisLiveQuery) storeQueryInfoReverse(name, sql string, hostIDs []uint) error { + // Store the SQL first, so a host never sees the query as targeted before its + // SQL can be looked up (same ordering guarantee as the bitfield path). + _, sqlKey := generateKeys(name) + conn := redis.ConfigureDoer(r.pool, r.pool.Get()) + if _, err := conn.Do("SET", sqlKey, sql, "EX", queryExpiration.Seconds()); err != nil { + conn.Close() + return fmt.Errorf("set sql: %w", err) + } + conn.Close() + + // Add the campaign id to each targeted host's set, pipelined per cluster slot. + hostKeys := make([]string, len(hostIDs)) + for i, hostID := range hostIDs { + hostKeys[i] = reverseHostKey(hostID) + } + + keysBySlot := redis.SplitKeysBySlot(r.pool, hostKeys...) + for _, keys := range keysBySlot { + if err := r.storeBatchReverseHostKeys(name, keys); err != nil { + return err + } + } + return nil +} + +func (r *redisLiveQuery) storeBatchReverseHostKeys(name string, hostKeys []string) error { + conn := r.pool.Get() + defer conn.Close() + + for _, hostKey := range hostKeys { + if err := conn.Send("SADD", hostKey, name); err != nil { + return fmt.Errorf("sadd reverse host key: %w", err) + } + if err := conn.Send("EXPIRE", hostKey, int(queryExpiration.Seconds())); err != nil { + return fmt.Errorf("expire reverse host key: %w", err) + } + } + if err := conn.Flush(); err != nil { + return fmt.Errorf("flush pipeline: %w", err) + } + // drain replies (2 per host key) to complete the pipeline + for range hostKeys { + if _, err := conn.Receive(); err != nil { + return fmt.Errorf("receive sadd reply: %w", err) + } + if _, err := conn.Receive(); err != nil { + return fmt.Errorf("receive expire reply: %w", err) + } + } + return nil +} + +func (r *redisLiveQuery) storeReverseQueryName(name string) error { + conn := redis.ConfigureDoer(r.pool, r.pool.Get()) + defer conn.Close() + + _, err := conn.Do("SADD", activeReverseQueriesKey, name) + return err +} + +func (r *redisLiveQuery) removeReverseQueryNames(names ...string) error { + conn := redis.ConfigureDoer(r.pool, r.pool.Get()) + defer conn.Close() + + var args redigo.Args + args = args.Add(activeReverseQueriesKey) + args = args.AddFlat(names) + _, err := conn.Do("SREM", args...) + return err +} + func (r *redisLiveQuery) storeQueryNames(names ...string) error { conn := redis.ConfigureDoer(r.pool, r.pool.Get()) defer conn.Close() @@ -376,6 +585,17 @@ func (r *redisLiveQuery) loadCache() error { return fmt.Errorf("get active queries: %w", err) } + // Load which active campaigns use the reverse per-host index, so the read + // path can exclude them from the per-host bitfield (GETBIT) probes. + reverseIDs, err := redigo.Strings(conn.Do("SMEMBERS", activeReverseQueriesKey)) + if err != nil && err != redigo.ErrNil { + return fmt.Errorf("get reverse active queries: %w", err) + } + reverseActive := make(map[string]struct{}, len(reverseIDs)) + for _, id := range reverseIDs { + reverseActive[id] = struct{}{} + } + for _, id := range activeIDs { _, sqlKey := generateKeys(id) @@ -409,6 +629,7 @@ func (r *redisLiveQuery) loadCache() error { r.cache.mu.Lock() r.cache.sqlCache = sqlCache r.cache.activeQueriesCache = activeIDs + r.cache.reverseActiveCache = reverseActive r.cache.cacheExp = time.Now().Add(r.cacheExpiration) r.cache.mu.Unlock() @@ -487,6 +708,13 @@ func (r *redisLiveQuery) removeInactiveQueries(ctx context.Context, inactiveCamp if _, err := conn.Do("SREM", args...); err != nil { return ctxerr.Wrap(ctx, err, "remove inactive campaign IDs") } + + // Also remove from the reverse model set. The per-host sets are left to expire + // via their TTL and are filtered against the active set at read time. + reverseArgs := redigo.Args{}.Add(activeReverseQueriesKey).AddFlat(inactiveCampaignIDs) + if _, err := conn.Do("SREM", reverseArgs...); err != nil { + return ctxerr.Wrap(ctx, err, "remove inactive reverse campaign IDs") + } return nil } diff --git a/server/live_query/redis_live_query_test.go b/server/live_query/redis_live_query_test.go index eb85b8b773..b335f39754 100644 --- a/server/live_query/redis_live_query_test.go +++ b/server/live_query/redis_live_query_test.go @@ -4,30 +4,59 @@ import ( "log/slog" "testing" + "github.com/fleetdm/fleet/v4/server/datastore/redis" "github.com/fleetdm/fleet/v4/server/datastore/redis/redistest" "github.com/fleetdm/fleet/v4/server/test" + redigo "github.com/gomodule/redigo/redis" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestRedisLiveQuery(t *testing.T) { + // Run every interface-contract test against both storage models: the legacy + // bitfield and the reverse per-host index. The reverse index uses a large + // threshold so the small target sets used in these tests are all stored as + // reverse queries. + models := []struct { + name string + reverse bool + }{ + {"bitfield", false}, + {"reverse", true}, + } for _, f := range testFunctions { t.Run(test.FunctionName(f), func(t *testing.T) { - t.Run("standalone", func(t *testing.T) { - store := setupRedisLiveQuery(t, false) - f(t, store) - }) + for _, m := range models { + t.Run(m.name, func(t *testing.T) { + t.Run("standalone", func(t *testing.T) { + store := setupRedisLiveQuery(t, false, m.reverse) + f(t, store) + }) - t.Run("cluster", func(t *testing.T) { - store := setupRedisLiveQuery(t, true) - f(t, store) - }) + t.Run("cluster", func(t *testing.T) { + store := setupRedisLiveQuery(t, true, m.reverse) + f(t, store) + }) + }) + } }) } } -func setupRedisLiveQuery(t *testing.T, cluster bool) *redisLiveQuery { +func setupRedisLiveQuery(t *testing.T, cluster, reverseEnabled bool) *redisLiveQuery { + // A 0 threshold disables the reverse index (bitfield model); a large threshold + // ensures the small target sets used in the contract tests all qualify for the + // reverse index. + threshold := 0 + if reverseEnabled { + threshold = 1 << 30 + } + return setupRedisLiveQueryThreshold(t, cluster, threshold) +} + +func setupRedisLiveQueryThreshold(t *testing.T, cluster bool, threshold int) *redisLiveQuery { pool := redistest.SetupRedis(t, "*livequery", cluster, true, true) - return NewRedisLiveQuery(pool, slog.New(slog.DiscardHandler), 0) + return NewRedisLiveQuery(pool, slog.New(slog.DiscardHandler), 0, threshold) } func TestMapBitfield(t *testing.T) { @@ -65,3 +94,207 @@ func TestMapBitfield(t *testing.T) { mapBitfield([]uint{79}), ) } + +// TestReverseIndexThreshold verifies that queries at or below the small-target +// threshold are stored in the per-host reverse index (no bitfield) while larger +// queries keep the bitfield, and that the read path returns the union of both. +func TestReverseIndexThreshold(t *testing.T) { + for _, cluster := range []bool{false, true} { + clusterName := "standalone" + if cluster { + clusterName = "cluster" + } + t.Run(clusterName, func(t *testing.T) { + // threshold 2: up to 2 targeted hosts use the reverse index. + store := setupRedisLiveQueryThreshold(t, cluster, 2) + conn := redis.ConfigureDoer(store.pool, store.pool.Get()) + defer conn.Close() + + // Small-target query (2 hosts == threshold) -> reverse index. + require.NoError(t, store.RunQuery("small", "SELECT 1", []uint{1, 2})) + + for _, h := range []uint{1, 2} { + isMember, err := redigo.Bool(conn.Do("SISMEMBER", reverseHostKey(h), "small")) + require.NoError(t, err) + assert.True(t, isMember, "host %d should be in the reverse set", h) + } + isReverse, err := redigo.Bool(conn.Do("SISMEMBER", activeReverseQueriesKey, "small")) + require.NoError(t, err) + assert.True(t, isReverse, "small-target query should be marked reverse") + bitfieldExists, err := redigo.Int(conn.Do("EXISTS", queryKeyPrefix+"{small}")) + require.NoError(t, err) + assert.Zero(t, bitfieldExists, "small-target query must not create a bitfield") + + // Broadcast query (3 hosts > threshold) -> bitfield. + require.NoError(t, store.RunQuery("big", "SELECT 2", []uint{1, 2, 3})) + + bitfieldExists, err = redigo.Int(conn.Do("EXISTS", queryKeyPrefix+"{big}")) + require.NoError(t, err) + assert.Equal(t, 1, bitfieldExists, "broadcast query must create a bitfield") + isReverse, err = redigo.Bool(conn.Do("SISMEMBER", activeReverseQueriesKey, "big")) + require.NoError(t, err) + assert.False(t, isReverse, "broadcast query should not be marked reverse") + isMember, err := redigo.Bool(conn.Do("SISMEMBER", reverseHostKey(1), "big")) + require.NoError(t, err) + assert.False(t, isMember, "broadcast query must not be added to per-host sets") + + // Read path returns the union for a host targeted by both models. + queries, err := store.QueriesForHost(1) + require.NoError(t, err) + assert.Equal(t, map[string]string{"small": "SELECT 1", "big": "SELECT 2"}, queries) + + // Host targeted only by the broadcast query. + queries, err = store.QueriesForHost(3) + require.NoError(t, err) + assert.Equal(t, map[string]string{"big": "SELECT 2"}, queries) + }) + } +} + +// TestReverseIndexStaleEntryFiltering verifies that after StopQuery, a reverse +// query is no longer returned to a targeted host even though its campaign ID is +// intentionally left lingering in that host's per-host set (StopQuery cannot +// enumerate the per-host sets). The read-time filter against the active set is +// what keeps this correct. +func TestReverseIndexStaleEntryFiltering(t *testing.T) { + for _, cluster := range []bool{false, true} { + clusterName := "standalone" + if cluster { + clusterName = "cluster" + } + t.Run(clusterName, func(t *testing.T) { + store := setupRedisLiveQueryThreshold(t, cluster, 2) + conn := redis.ConfigureDoer(store.pool, store.pool.Get()) + defer conn.Close() + + require.NoError(t, store.RunQuery("small", "SELECT 1", []uint{1})) + + queries, err := store.QueriesForHost(1) + require.NoError(t, err) + require.Equal(t, map[string]string{"small": "SELECT 1"}, queries) + + require.NoError(t, store.StopQuery("small")) + + // The per-host set still contains the (now stale) campaign ID: StopQuery + // deliberately does not clean it up. + isMember, err := redigo.Bool(conn.Do("SISMEMBER", reverseHostKey(1), "small")) + require.NoError(t, err) + require.True(t, isMember, "stale entry should remain in the per-host set after StopQuery") + + // Despite the lingering membership, the query must not be delivered: it is + // filtered out because it is no longer in the active set / SQL cache. + queries, err = store.QueriesForHost(1) + require.NoError(t, err) + require.Empty(t, queries, "stopped reverse query must not be returned despite stale per-host membership") + }) + } +} + +// TestReverseIndexQueryCompletedByHost verifies that QueryCompletedByHost on a +// reverse query removes only the completing host's per-host membership, so that +// host stops receiving the query while other targeted hosts still do. +func TestReverseIndexQueryCompletedByHost(t *testing.T) { + for _, cluster := range []bool{false, true} { + clusterName := "standalone" + if cluster { + clusterName = "cluster" + } + t.Run(clusterName, func(t *testing.T) { + store := setupRedisLiveQueryThreshold(t, cluster, 2) + conn := redis.ConfigureDoer(store.pool, store.pool.Get()) + defer conn.Close() + + require.NoError(t, store.RunQuery("small", "SELECT 1", []uint{1, 2})) + + // Host 1 completes the query. + require.NoError(t, store.QueryCompletedByHost("small", 1)) + + // Host 1's per-host membership is removed, host 2's remains. + isMember, err := redigo.Bool(conn.Do("SISMEMBER", reverseHostKey(1), "small")) + require.NoError(t, err) + require.False(t, isMember, "completing host's membership should be removed") + isMember, err = redigo.Bool(conn.Do("SISMEMBER", reverseHostKey(2), "small")) + require.NoError(t, err) + require.True(t, isMember, "other targeted host's membership should remain") + + // The query is still active, so the bitfield no-op SETBIT in + // QueryCompletedByHost must not have created a lingering bitfield key. + bitfieldExists, err := redigo.Int(conn.Do("EXISTS", queryKeyPrefix+"{small}")) + require.NoError(t, err) + require.Zero(t, bitfieldExists, "reverse query must not gain a bitfield from completion") + + // Host 1 no longer receives the query; host 2 still does. + queries, err := store.QueriesForHost(1) + require.NoError(t, err) + require.Empty(t, queries) + queries, err = store.QueriesForHost(2) + require.NoError(t, err) + require.Equal(t, map[string]string{"small": "SELECT 1"}, queries) + }) + } +} + +// TestReverseIndexCleanupInactiveQueries verifies that CleanupInactiveQueries +// removes a reverse campaign from both the active set and the reverse-model set. +func TestReverseIndexCleanupInactiveQueries(t *testing.T) { + for _, cluster := range []bool{false, true} { + clusterName := "standalone" + if cluster { + clusterName = "cluster" + } + t.Run(clusterName, func(t *testing.T) { + store := setupRedisLiveQueryThreshold(t, cluster, 2) + conn := redis.ConfigureDoer(store.pool, store.pool.Get()) + defer conn.Close() + + // Campaign IDs must be numeric so they match the uint IDs passed to + // CleanupInactiveQueries. + require.NoError(t, store.RunQuery("5", "SELECT 1", []uint{1})) + + isReverse, err := redigo.Bool(conn.Do("SISMEMBER", activeReverseQueriesKey, "5")) + require.NoError(t, err) + require.True(t, isReverse) + + require.NoError(t, store.CleanupInactiveQueries(t.Context(), []uint{5})) + + isActive, err := redigo.Bool(conn.Do("SISMEMBER", activeQueriesKey, "5")) + require.NoError(t, err) + require.False(t, isActive, "inactive campaign should be removed from the active set") + isReverse, err = redigo.Bool(conn.Do("SISMEMBER", activeReverseQueriesKey, "5")) + require.NoError(t, err) + require.False(t, isReverse, "inactive campaign should be removed from the reverse-model set") + }) + } +} + +// TestReverseIndexKillSwitch verifies that a threshold of 0 disables the reverse +// index and forces the legacy bitfield model even for single-host queries. +func TestReverseIndexKillSwitch(t *testing.T) { + for _, cluster := range []bool{false, true} { + clusterName := "standalone" + if cluster { + clusterName = "cluster" + } + t.Run(clusterName, func(t *testing.T) { + store := setupRedisLiveQueryThreshold(t, cluster, 0) + conn := redis.ConfigureDoer(store.pool, store.pool.Get()) + defer conn.Close() + + require.NoError(t, store.RunQuery("q", "SELECT 1", []uint{1})) + + bitfieldExists, err := redigo.Int(conn.Do("EXISTS", queryKeyPrefix+"{q}")) + require.NoError(t, err) + assert.Equal(t, 1, bitfieldExists, "threshold 0 should force the bitfield model") + isReverse, err := redigo.Bool(conn.Do("SISMEMBER", activeReverseQueriesKey, "q")) + require.NoError(t, err) + assert.False(t, isReverse) + isMember, err := redigo.Bool(conn.Do("SISMEMBER", reverseHostKey(1), "q")) + require.NoError(t, err) + assert.False(t, isMember) + + queries, err := store.QueriesForHost(1) + require.NoError(t, err) + assert.Equal(t, map[string]string{"q": "SELECT 1"}, queries) + }) + } +}