Make receive calls to redis conn thread safe (#1641)

* Make receive calls to redis conn thread safe

Also removes REDIS_TEST env var. Redis is lightweight and fast, no need
to skip these tests.

* No need to increase the wait

.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:

changes/issue-1531-thread-safe-redis

@@ -0,0 +1 @@
+* Reads live query results from redis in a thread safe manner.

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:

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)

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{

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
 }

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