diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3c5f860f7c..64b1bb5404 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -174,7 +174,7 @@ jobs: - name: Run Go Tests run: | - MYSQL_TEST=1 REDIS_TEST=1 make test-go + MYSQL_TEST=1 make test-go lint-go: diff --git a/changes/issue-1531-thread-safe-redis b/changes/issue-1531-thread-safe-redis new file mode 100644 index 0000000000..974ce7b0a3 --- /dev/null +++ b/changes/issue-1531-thread-safe-redis @@ -0,0 +1 @@ +* Reads live query results from redis in a thread safe manner. diff --git a/docs/3-Contributing/2-Testing.md b/docs/3-Contributing/2-Testing.md index 03f86b5585..561dc1d0e9 100644 --- a/docs/3-Contributing/2-Testing.md +++ b/docs/3-Contributing/2-Testing.md @@ -34,7 +34,7 @@ Check out [the instructions in the `/tools/osquery` directory](../../tools/osque To execute the basic unit and integration tests, run the following from the root of the repository: ``` -MYSQL_TEST=1 REDIS_TEST=1 make test +MYSQL_TEST=1 make test ``` It is a good idea to run `make test` before submitting a Pull Request. @@ -91,14 +91,6 @@ To run MySQL integration tests set environment variables as follows: MYSQL_TEST=1 make test-go ``` -#### Redis tests - -To run Redis integration tests set environment variables as follows: - -``` -REDIS_TEST=1 make test-go -``` - #### Email tests To run email related integration tests using MailHog set environment as follows: diff --git a/server/live_query/redis_live_query_test.go b/server/live_query/redis_live_query_test.go index 4e6830074a..fd2fd5242b 100644 --- a/server/live_query/redis_live_query_test.go +++ b/server/live_query/redis_live_query_test.go @@ -1,7 +1,6 @@ package live_query import ( - "os" "testing" "github.com/fleetdm/fleet/v4/server/pubsub" @@ -11,10 +10,6 @@ import ( ) func TestRedisLiveQuery(t *testing.T) { - if _, ok := os.LookupEnv("REDIS_TEST"); !ok { - t.Skip("Redis tests not requested. Skipping.") - } - for _, f := range testFunctions { t.Run(test.FunctionName(f), func(t *testing.T) { store, teardown := setupRedisLiveQuery(t) diff --git a/server/pubsub/query_results_test.go b/server/pubsub/query_results_test.go index 74d1e2ff5f..24a84042ac 100644 --- a/server/pubsub/query_results_test.go +++ b/server/pubsub/query_results_test.go @@ -2,13 +2,11 @@ package pubsub import ( "context" - "os" "sync" "testing" "time" "github.com/fleetdm/fleet/v4/server/fleet" - "github.com/fleetdm/fleet/v4/server/test" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -29,36 +27,6 @@ func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool { } } -var testFunctions = [...]func(*testing.T, fleet.QueryResultStore){ - testQueryResultsStore, - testQueryResultsStoreErrors, -} - -func TestRedis(t *testing.T) { - if _, ok := os.LookupEnv("REDIS_TEST"); !ok { - t.SkipNow() - } - - store, teardown := setupRedis(t) - defer teardown() - - for _, f := range testFunctions { - t.Run(test.FunctionName(f), func(t *testing.T) { - f(t, store) - }) - } -} - -func TestInmem(t *testing.T) { - for _, f := range testFunctions { - t.Run(test.FunctionName(f), func(t *testing.T) { - t.Parallel() - store := NewInmemQueryResults() - f(t, store) - }) - } -} - func setupRedis(t *testing.T) (store *redisQueryResults, teardown func()) { var ( addr = "127.0.0.1:6379" @@ -82,7 +50,10 @@ func setupRedis(t *testing.T) (store *redisQueryResults, teardown func()) { return store, teardown } -func testQueryResultsStoreErrors(t *testing.T, store fleet.QueryResultStore) { +func TestQueryResultsStoreErrors(t *testing.T) { + store, teardown := setupRedis(t) + defer teardown() + // Write with no subscriber err := store.WriteResult( fleet.DistributedQueryResult{ @@ -106,7 +77,10 @@ func testQueryResultsStoreErrors(t *testing.T, store fleet.QueryResultStore) { } } -func testQueryResultsStore(t *testing.T, store fleet.QueryResultStore) { +func TestQueryResultsStore(t *testing.T) { + store, teardown := setupRedis(t) + defer teardown() + // Test handling results for two campaigns in parallel campaign1 := fleet.DistributedQueryCampaign{ID: 1} @@ -115,7 +89,7 @@ func testQueryResultsStore(t *testing.T, store fleet.QueryResultStore) { assert.Nil(t, err) expected1 := []fleet.DistributedQueryResult{ - fleet.DistributedQueryResult{ + { DistributedQueryCampaignID: 1, Rows: []map[string]string{{"foo": "bar"}}, Host: fleet.Host{ @@ -136,7 +110,7 @@ func testQueryResultsStore(t *testing.T, store fleet.QueryResultStore) { SeenTime: time.Now().UTC(), }, }, - fleet.DistributedQueryResult{ + { DistributedQueryCampaignID: 1, Rows: []map[string]string{{"whoo": "wahh"}}, Host: fleet.Host{ @@ -154,7 +128,7 @@ func testQueryResultsStore(t *testing.T, store fleet.QueryResultStore) { SeenTime: time.Now().UTC(), }, }, - fleet.DistributedQueryResult{ + { DistributedQueryCampaignID: 1, Rows: []map[string]string{{"bing": "fds"}}, Host: fleet.Host{ diff --git a/server/pubsub/redis_query_results.go b/server/pubsub/redis_query_results.go index 7c94044ac0..de1c6157b5 100644 --- a/server/pubsub/redis_query_results.go +++ b/server/pubsub/redis_query_results.go @@ -125,46 +125,55 @@ func (r *redisQueryResults) WriteResult(result fleet.DistributedQueryResult) err // connection over the provided channel. This effectively allows a select // statement to run on conn.Receive() (by running on the channel that is being // fed by this function) -func receiveMessages(conn *redis.PubSubConn, outChan chan<- interface{}) { +func receiveMessages(ctx context.Context, pool *redisc.Cluster, query fleet.DistributedQueryCampaign, outChan chan<- interface{}) { + conn := redis.PubSubConn{Conn: pool.Get()} + defer conn.Close() + + pubSubName := pubSubForID(query.ID) + err := conn.Subscribe(pubSubName) + if err != nil { + outChan <- errors.Wrap(err, "subscribe to channel") + } + defer conn.Unsubscribe(pubSubName) + defer func() { close(outChan) }() for { msg := conn.Receive() - outChan <- msg - switch msg := msg.(type) { - case error: - // If an error occurred (i.e. connection was closed), - // then we should exit - return - case redis.Subscription: - // If the subscription count is 0, the ReadChannel call - // that invoked this goroutine has unsubscribed, and we - // can exit - if msg.Count == 0 { + select { + case outChan <- msg: + switch msg := msg.(type) { + case error: + // If an error occurred (i.e. connection was closed), + // then we should exit return + case redis.Subscription: + // If the subscription count is 0, the ReadChannel call + // that invoked this goroutine has unsubscribed, and we + // can exit + if msg.Count == 0 { + return + } } + case <-ctx.Done(): + conn.Unsubscribe(pubSubName) + return } } } func (r *redisQueryResults) ReadChannel(ctx context.Context, query fleet.DistributedQueryCampaign) (<-chan interface{}, error) { outChannel := make(chan interface{}) - - conn := redis.PubSubConn{Conn: r.pool.Get()} - - pubSubName := pubSubForID(query.ID) - conn.Subscribe(pubSubName) - msgChannel := make(chan interface{}) + // Run a separate goroutine feeding redis messages into // msgChannel - go receiveMessages(&conn, msgChannel) + go receiveMessages(ctx, r.pool, query, msgChannel) go func() { defer close(outChannel) - defer conn.Close() for { // Loop reading messages from conn.Receive() (via @@ -185,13 +194,11 @@ func (r *redisQueryResults) ReadChannel(ctx context.Context, query fleet.Distrib case error: outChannel <- errors.Wrap(msg, "reading from redis") } - case <-ctx.Done(): - conn.Unsubscribe() - + return } - } + } }() return outChannel, nil } diff --git a/server/service/service_campaigns.go b/server/service/service_campaigns.go index a41b67eaa1..a9b50c1574 100644 --- a/server/service/service_campaigns.go +++ b/server/service/service_campaigns.go @@ -221,7 +221,10 @@ func (svc Service) StreamCampaignResults(ctx context.Context, conn *websocket.Co // Open the channel from which we will receive incoming query results // (probably from the redis pubsub implementation) - readChan, err := svc.resultStore.ReadChannel(context.Background(), *campaign) + cancelCtx, cancelFunc := context.WithCancel(ctx) + defer cancelFunc() + + readChan, err := svc.resultStore.ReadChannel(cancelCtx, *campaign) if err != nil { conn.WriteJSONError(fmt.Sprintf("cannot open read channel for campaign %d ", campaignID)) return