From 96c0244c04eef7e8840fa9770822cbf00b0a1885 Mon Sep 17 00:00:00 2001 From: Tomas Touceda Date: Fri, 13 Aug 2021 11:16:11 -0300 Subject: [PATCH] Fix redis locking receive (#1655) * Receive redis data with timeout to not hold the connection * Address review comments --- changes/fix-redis-locking-receive | 1 + go.mod | 2 +- go.sum | 2 + server/pubsub/query_results_test.go | 28 +--- server/pubsub/redis_query_results.go | 55 +++++--- server/pubsub/testing_utils.go | 30 +++++ server/service/service_campaign_test.go | 163 ++++++++++++++++++++++++ 7 files changed, 239 insertions(+), 42 deletions(-) create mode 100644 changes/fix-redis-locking-receive create mode 100644 server/pubsub/testing_utils.go create mode 100644 server/service/service_campaign_test.go diff --git a/changes/fix-redis-locking-receive b/changes/fix-redis-locking-receive new file mode 100644 index 0000000000..8bf0149259 --- /dev/null +++ b/changes/fix-redis-locking-receive @@ -0,0 +1 @@ +* When a connection from a live query websocket is closed, Fleet now timeouts the receive and handles the different cases correctly to not hold the connection to Redis. diff --git a/go.mod b/go.mod index 3c9a7c300d..91eae30f4e 100644 --- a/go.mod +++ b/go.mod @@ -30,7 +30,7 @@ require ( github.com/go-logfmt/logfmt v0.5.0 // indirect github.com/go-sql-driver/mysql v1.5.0 github.com/golang-jwt/jwt/v4 v4.0.0 - github.com/gomodule/redigo v1.8.4 + github.com/gomodule/redigo v1.8.5 github.com/google/go-cmp v0.5.6 github.com/google/go-github/v37 v37.0.0 github.com/google/uuid v1.1.2 diff --git a/go.sum b/go.sum index eb34ade776..1b99aff915 100644 --- a/go.sum +++ b/go.sum @@ -350,6 +350,8 @@ github.com/golangci/unconvert v0.0.0-20180507085042-28b1c447d1f4 h1:zwtduBRr5SSW github.com/golangci/unconvert v0.0.0-20180507085042-28b1c447d1f4/go.mod h1:Izgrg8RkN3rCIMLGE9CyYmU9pY2Jer6DgANEnZ/L/cQ= github.com/gomodule/redigo v1.8.4 h1:Z5JUg94HMTR1XpwBaSH4vq3+PNSIykBLxMdglbw10gg= github.com/gomodule/redigo v1.8.4/go.mod h1:P9dn9mFrCBvWhGE1wpxx6fgq7BAeLBk+UUUzlpkBYO0= +github.com/gomodule/redigo v1.8.5 h1:nRAxCa+SVsyjSBrtZmG/cqb6VbTmuRzpg/PoTFlpumc= +github.com/gomodule/redigo v1.8.5/go.mod h1:P9dn9mFrCBvWhGE1wpxx6fgq7BAeLBk+UUUzlpkBYO0= github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= diff --git a/server/pubsub/query_results_test.go b/server/pubsub/query_results_test.go index 24a84042ac..e3b280c1ab 100644 --- a/server/pubsub/query_results_test.go +++ b/server/pubsub/query_results_test.go @@ -8,7 +8,6 @@ import ( "github.com/fleetdm/fleet/v4/server/fleet" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" ) // waitTimeout waits for the waitgroup for the specified max timeout. @@ -27,31 +26,8 @@ func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool { } } -func setupRedis(t *testing.T) (store *redisQueryResults, teardown func()) { - var ( - addr = "127.0.0.1:6379" - password = "" - database = 0 - useTLS = false - dupResults = false - ) - - pool, err := NewRedisPool(addr, password, database, useTLS) - require.NoError(t, err) - store = NewRedisQueryResults(pool, dupResults) - - _, err = store.pool.Get().Do("PING") - require.Nil(t, err) - - teardown = func() { - store.pool.Close() - } - - return store, teardown -} - func TestQueryResultsStoreErrors(t *testing.T) { - store, teardown := setupRedis(t) + store, teardown := SetupRedisForTest(t) defer teardown() // Write with no subscriber @@ -78,7 +54,7 @@ func TestQueryResultsStoreErrors(t *testing.T) { } func TestQueryResultsStore(t *testing.T) { - store, teardown := setupRedis(t) + store, teardown := SetupRedisForTest(t) defer teardown() // Test handling results for two campaigns in parallel diff --git a/server/pubsub/redis_query_results.go b/server/pubsub/redis_query_results.go index de1c6157b5..74804d110a 100644 --- a/server/pubsub/redis_query_results.go +++ b/server/pubsub/redis_query_results.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "net" "strings" "time" @@ -94,6 +95,10 @@ func pubSubForID(id uint) string { return fmt.Sprintf("results_%d", id) } +func (r *redisQueryResults) Pool() *redisc.Cluster { + return r.pool +} + func (r *redisQueryResults) WriteResult(result fleet.DistributedQueryResult) error { conn := r.pool.Get() defer conn.Close() @@ -121,6 +126,17 @@ func (r *redisQueryResults) WriteResult(result fleet.DistributedQueryResult) err return nil } +// writeOrDone tries to write the item into the channel taking into account context.Done(). If context is done, returns +// true, otherwise false +func writeOrDone(ctx context.Context, ch chan<- interface{}, item interface{}) bool { + select { + case ch <- item: + case <-ctx.Done(): + return true + } + return false +} + // receiveMessages runs in a goroutine, forwarding messages from the Pub/Sub // connection over the provided channel. This effectively allows a select // statement to run on conn.Receive() (by running on the channel that is being @@ -131,8 +147,8 @@ func receiveMessages(ctx context.Context, pool *redisc.Cluster, query fleet.Dist pubSubName := pubSubForID(query.ID) err := conn.Subscribe(pubSubName) - if err != nil { - outChan <- errors.Wrap(err, "subscribe to channel") + if err != nil && writeOrDone(ctx, outChan, errors.Wrap(err, "subscribe to channel")) { + return } defer conn.Unsubscribe(pubSubName) @@ -141,24 +157,28 @@ func receiveMessages(ctx context.Context, pool *redisc.Cluster, query fleet.Dist }() for { - msg := conn.Receive() + // This Receive needs to be with timeout, otherwise we might block on it forever + msg := conn.ReceiveWithTimeout(5 * time.Second) + select { case outChan <- msg: switch msg := msg.(type) { case error: - // If an error occurred (i.e. connection was closed), - // then we should exit - return + if err, ok := msg.(net.Error); ok && err.Timeout() { + // We ignore timeouts, we just want them there to make sure we don't hang on Receiving + continue + } else { + // 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 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 } } @@ -176,8 +196,7 @@ func (r *redisQueryResults) ReadChannel(ctx context.Context, query fleet.Distrib defer close(outChannel) for { - // Loop reading messages from conn.Receive() (via - // msgChannel) until the context is cancelled. + // Loop reading messages from conn.Receive() (via msgChannel) until the context is cancelled. select { case msg, ok := <-msgChannel: if !ok { @@ -188,11 +207,17 @@ func (r *redisQueryResults) ReadChannel(ctx context.Context, query fleet.Distrib var res fleet.DistributedQueryResult err := json.Unmarshal(msg.Data, &res) if err != nil { - outChannel <- err + if writeOrDone(ctx, outChannel, err) { + return + } + } + if writeOrDone(ctx, outChannel, res) { + return } - outChannel <- res case error: - outChannel <- errors.Wrap(msg, "reading from redis") + if writeOrDone(ctx, outChannel, errors.Wrap(msg, "reading from redis")) { + return + } } case <-ctx.Done(): return diff --git a/server/pubsub/testing_utils.go b/server/pubsub/testing_utils.go new file mode 100644 index 0000000000..e6088a0907 --- /dev/null +++ b/server/pubsub/testing_utils.go @@ -0,0 +1,30 @@ +package pubsub + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func SetupRedisForTest(t *testing.T) (store *redisQueryResults, teardown func()) { + var ( + addr = "127.0.0.1:6379" + password = "" + database = 0 + useTLS = false + dupResults = false + ) + + pool, err := NewRedisPool(addr, password, database, useTLS) + require.NoError(t, err) + store = NewRedisQueryResults(pool, dupResults) + + _, err = store.pool.Get().Do("PING") + require.Nil(t, err) + + teardown = func() { + store.pool.Close() + } + + return store, teardown +} diff --git a/server/service/service_campaign_test.go b/server/service/service_campaign_test.go new file mode 100644 index 0000000000..e2682ba2af --- /dev/null +++ b/server/service/service_campaign_test.go @@ -0,0 +1,163 @@ +package service + +import ( + "context" + "crypto/tls" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "github.com/WatchBeam/clock" + "github.com/fleetdm/fleet/v4/server/contexts/viewer" + "github.com/fleetdm/fleet/v4/server/fleet" + "github.com/fleetdm/fleet/v4/server/live_query" + "github.com/fleetdm/fleet/v4/server/mock" + "github.com/fleetdm/fleet/v4/server/ptr" + "github.com/fleetdm/fleet/v4/server/pubsub" + ws "github.com/fleetdm/fleet/v4/server/websocket" + kitlog "github.com/go-kit/kit/log" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +func TestStreamCampaignResultsClosesReditOnWSClose(t *testing.T) { + store, teardown := pubsub.SetupRedisForTest(t) + defer teardown() + + mockClock := clock.NewMockClock() + ds := new(mock.Store) + lq := new(live_query.MockLiveQuery) + svc := newTestServiceWithClock(ds, store, lq, mockClock) + + campaign := &fleet.DistributedQueryCampaign{ID: 42} + + ds.LabelQueriesForHostFunc = func(host *fleet.Host, cutoff time.Time) (map[string]string, error) { + return map[string]string{}, nil + } + ds.SaveHostFunc = func(host *fleet.Host) error { + return nil + } + ds.AppConfigFunc = func() (*fleet.AppConfig, error) { + return &fleet.AppConfig{EnableHostUsers: true}, nil + } + ds.NewQueryFunc = func(query *fleet.Query, opts ...fleet.OptionalArg) (*fleet.Query, error) { + return query, nil + } + ds.NewDistributedQueryCampaignFunc = func(camp *fleet.DistributedQueryCampaign) (*fleet.DistributedQueryCampaign, error) { + return camp, nil + } + ds.NewDistributedQueryCampaignTargetFunc = func(target *fleet.DistributedQueryCampaignTarget) (*fleet.DistributedQueryCampaignTarget, error) { + return target, nil + } + ds.HostIDsInTargetsFunc = func(filter fleet.TeamFilter, targets fleet.HostTargets) ([]uint, error) { + return []uint{1}, nil + } + ds.CountHostsInTargetsFunc = func(filter fleet.TeamFilter, targets fleet.HostTargets, now time.Time) (fleet.TargetMetrics, error) { + return fleet.TargetMetrics{TotalHosts: 1}, nil + } + ds.NewActivityFunc = func(user *fleet.User, activityType string, details *map[string]interface{}) error { + return nil + } + ds.SessionByKeyFunc = func(key string) (*fleet.Session, error) { + return &fleet.Session{ + CreateTimestamp: fleet.CreateTimestamp{CreatedAt: time.Now()}, + ID: 42, + AccessedAt: time.Now(), + UserID: 999, + Key: "asd", + }, nil + } + + host := &fleet.Host{ID: 1, Platform: "windows"} + + lq.On("QueriesForHost", uint(1)).Return( + map[string]string{ + strconv.Itoa(int(campaign.ID)): "select * from time", + }, + nil, + ) + lq.On("QueryCompletedByHost", strconv.Itoa(int(campaign.ID)), host.ID).Return(nil) + lq.On("RunQuery", "0", "select year, month, day, hour, minutes, seconds from time", []uint{1}).Return(nil) + viewerCtx := viewer.NewContext(context.Background(), viewer.Viewer{ + User: &fleet.User{ + ID: 0, + GlobalRole: ptr.String(fleet.RoleAdmin), + }, + }) + q := "select year, month, day, hour, minutes, seconds from time" + _, err := svc.NewDistributedQueryCampaign(viewerCtx, q, nil, fleet.HostTargets{HostIDs: []uint{2}, LabelIDs: []uint{1}}) + require.NoError(t, err) + + s := httptest.NewServer(makeStreamDistributedQueryCampaignResultsHandler(svc, kitlog.NewNopLogger())) + defer s.Close() + // Convert http://127.0.0.1 to ws://127.0.0.1 + u := "ws" + strings.TrimPrefix(s.URL, "http") + "/api/v1/fleet/results/websocket" + + // Connect to the server + dialer := &websocket.Dialer{ + Proxy: http.ProxyFromEnvironment, + HandshakeTimeout: 45 * time.Second, + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + } + + conn, _, err := dialer.Dial(u, nil) + if err != nil { + t.Fatalf("%v", err) + } + defer conn.Close() + + err = conn.WriteJSON(ws.JSONMessage{ + Type: "auth", + Data: map[string]interface{}{"token": "asd"}, + }) + require.NoError(t, err) + + err = conn.WriteJSON(ws.JSONMessage{ + Type: "select_campaign", + Data: map[string]interface{}{"campaign_id": campaign.ID}, + }) + require.NoError(t, err) + + ds.MarkSessionAccessedFunc = func(*fleet.Session) error { + return nil + } + ds.UserByIDFunc = func(id uint) (*fleet.User, error) { + return &fleet.User{GlobalRole: ptr.String(fleet.RoleAdmin)}, nil + } + ds.DistributedQueryCampaignFunc = func(id uint) (*fleet.DistributedQueryCampaign, error) { + return campaign, nil + } + ds.SaveDistributedQueryCampaignFunc = func(camp *fleet.DistributedQueryCampaign) error { + return nil + } + ds.DistributedQueryCampaignTargetIDsFunc = func(id uint) (targets *fleet.HostTargets, err error) { + return &fleet.HostTargets{HostIDs: []uint{1}}, nil + } + ds.QueryFunc = func(id uint) (*fleet.Query, error) { + return &fleet.Query{}, nil + } + + /*****************************************************************************************/ + /* THE ACTUAL TEST BEGINS HERE */ + /*****************************************************************************************/ + prevActiveConn := 0 + for prevActiveConn < 3 { + time.Sleep(2 * time.Second) + + for _, s := range store.Pool().Stats() { + prevActiveConn = s.ActiveCount + } + } + + conn.Close() + time.Sleep(10 * time.Second) + + newActiveConn := prevActiveConn + for _, s := range store.Pool().Stats() { + newActiveConn = s.ActiveCount + } + require.Equal(t, prevActiveConn-1, newActiveConn) +}