Speed up manage packs page loading (#1429)

Push the calculation of target counts into the SQL query, rather than loading
all of the targets and then counting them. This provides a dramatic (>100x)
speedup in loading of the manage packs page when large numbers of hosts are
present.

Closes #1426

server/datastore/datastore_targets_test.go

@@ -0,0 +1,162 @@
+package datastore
+
+import (
+	"testing"
+	"time"
+
+	"github.com/WatchBeam/clock"
+	"github.com/kolide/kolide/server/kolide"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func testCountHostsInTargets(t *testing.T, ds kolide.Datastore) {
+	if ds.Name() == "inmem" {
+		t.Skip("inmem is being deprecated, test skipped")
+	}
+
+	mockClock := clock.NewMockClock()
+
+	h1, err := ds.NewHost(&kolide.Host{
+		OsqueryHostID:    "1",
+		DetailUpdateTime: time.Now(),
+		SeenTime:         time.Now(),
+		HostName:         "foo.local",
+		NodeKey:          "1",
+		UUID:             "1",
+	})
+	require.Nil(t, err)
+	require.Nil(t, ds.MarkHostSeen(h1, mockClock.Now()))
+
+	h2, err := ds.NewHost(&kolide.Host{
+		OsqueryHostID:    "2",
+		DetailUpdateTime: time.Now(),
+		SeenTime:         time.Now(),
+		HostName:         "bar.local",
+		NodeKey:          "2",
+		UUID:             "2",
+	})
+	require.Nil(t, err)
+	// make this host "offline"
+	require.Nil(t, ds.MarkHostSeen(h2, mockClock.Now().Add(-1*time.Hour)))
+
+	h3, err := ds.NewHost(&kolide.Host{
+		OsqueryHostID:    "3",
+		DetailUpdateTime: time.Now(),
+		SeenTime:         time.Now(),
+		HostName:         "baz.local",
+		NodeKey:          "3",
+		UUID:             "3",
+	})
+	require.Nil(t, err)
+	require.Nil(t, ds.MarkHostSeen(h3, mockClock.Now().Add(-5*time.Second)))
+
+	h4, err := ds.NewHost(&kolide.Host{
+		OsqueryHostID:    "4",
+		DetailUpdateTime: time.Now(),
+		SeenTime:         time.Now(),
+		HostName:         "xxx.local",
+		NodeKey:          "4",
+		UUID:             "4",
+	})
+	require.Nil(t, err)
+	require.Nil(t, ds.MarkHostSeen(h4, mockClock.Now()))
+
+	h5, err := ds.NewHost(&kolide.Host{
+		OsqueryHostID:    "5",
+		DetailUpdateTime: time.Now(),
+		SeenTime:         time.Now(),
+		HostName:         "yyy.local",
+		NodeKey:          "5",
+		UUID:             "5",
+	})
+	require.Nil(t, err)
+	require.Nil(t, ds.MarkHostSeen(h5, mockClock.Now()))
+
+	h6, err := ds.NewHost(&kolide.Host{
+		OsqueryHostID:    "6",
+		DetailUpdateTime: time.Now(),
+		SeenTime:         time.Now(),
+		HostName:         "zzz.local",
+		NodeKey:          "6",
+		UUID:             "6",
+	})
+	require.Nil(t, err)
+	const thirtyDaysAndAMinuteAgo = -1 * (30*24*60 + 1)
+	require.Nil(t, ds.MarkHostSeen(h6, mockClock.Now().Add(thirtyDaysAndAMinuteAgo*time.Minute)))
+
+	l1, err := ds.NewLabel(&kolide.Label{
+		Name:  "label foo",
+		Query: "query foo",
+	})
+	require.Nil(t, err)
+	require.NotZero(t, l1.ID)
+
+	l2, err := ds.NewLabel(&kolide.Label{
+		Name:  "label bar",
+		Query: "query foo",
+	})
+	require.Nil(t, err)
+	require.NotZero(t, l2.ID)
+
+	for _, h := range []*kolide.Host{h1, h2, h3, h6} {
+		err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, mockClock.Now())
+		assert.Nil(t, err)
+	}
+
+	for _, h := range []*kolide.Host{h3, h4, h5} {
+		err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l2.ID: true}, mockClock.Now())
+		assert.Nil(t, err)
+	}
+
+	metrics, err := ds.CountHostsInTargets(nil, []uint{l1.ID, l2.ID}, mockClock.Now(), 30*time.Minute)
+	require.Nil(t, err)
+	require.NotNil(t, metrics)
+	assert.Equal(t, uint(6), metrics.TotalHosts)
+	assert.Equal(t, uint(1), metrics.OfflineHosts)
+	assert.Equal(t, uint(4), metrics.OnlineHosts)
+	assert.Equal(t, uint(1), metrics.MissingInActionHosts)
+
+	metrics, err = ds.CountHostsInTargets([]uint{h1.ID, h2.ID}, []uint{l1.ID, l2.ID}, mockClock.Now(), 30*time.Minute)
+	require.Nil(t, err)
+	require.NotNil(t, metrics)
+	assert.Equal(t, uint(6), metrics.TotalHosts)
+	assert.Equal(t, uint(1), metrics.OfflineHosts)
+	assert.Equal(t, uint(4), metrics.OnlineHosts)
+	assert.Equal(t, uint(1), metrics.MissingInActionHosts)
+
+	metrics, err = ds.CountHostsInTargets([]uint{h1.ID, h2.ID}, nil, mockClock.Now(), 30*time.Minute)
+	require.Nil(t, err)
+	require.NotNil(t, metrics)
+	assert.Equal(t, uint(2), metrics.TotalHosts)
+	assert.Equal(t, uint(1), metrics.OnlineHosts)
+	assert.Equal(t, uint(1), metrics.OfflineHosts)
+	assert.Equal(t, uint(0), metrics.MissingInActionHosts)
+
+	metrics, err = ds.CountHostsInTargets([]uint{h1.ID}, []uint{l2.ID}, mockClock.Now(), 30*time.Minute)
+	require.Nil(t, err)
+	require.NotNil(t, metrics)
+	assert.Equal(t, uint(4), metrics.TotalHosts)
+	assert.Equal(t, uint(4), metrics.OnlineHosts)
+	assert.Equal(t, uint(0), metrics.OfflineHosts)
+	assert.Equal(t, uint(0), metrics.MissingInActionHosts)
+
+	metrics, err = ds.CountHostsInTargets(nil, nil, mockClock.Now(), 30*time.Minute)
+	require.Nil(t, err)
+	require.NotNil(t, metrics)
+	assert.Equal(t, uint(0), metrics.TotalHosts)
+	assert.Equal(t, uint(0), metrics.OnlineHosts)
+	assert.Equal(t, uint(0), metrics.OfflineHosts)
+	assert.Equal(t, uint(0), metrics.MissingInActionHosts)
+
+	// Advance clock so all hosts are offline
+	mockClock.AddTime(1 * time.Hour)
+	metrics, err = ds.CountHostsInTargets(nil, []uint{l1.ID, l2.ID}, mockClock.Now(), 30*time.Minute)
+	require.Nil(t, err)
+	require.NotNil(t, metrics)
+	assert.Equal(t, uint(6), metrics.TotalHosts)
+	assert.Equal(t, uint(0), metrics.OnlineHosts)
+	assert.Equal(t, uint(5), metrics.OfflineHosts)
+	assert.Equal(t, uint(1), metrics.MissingInActionHosts)
+
+}

server/datastore/datastore_test.go

@@ -76,4 +76,5 @@ var testFunctions = [...]func(*testing.T, kolide.Datastore){
 	testReplaceDeletedLabel,
 	testMigrationStatus,
 	testUnicode,
+	testCountHostsInTargets,
 }

server/datastore/inmem/inmem.go

@@ -39,6 +39,7 @@ type Datastore struct {
 	yaraSignatureGroups             map[uint]*kolide.YARASignatureGroup
 	appConfig                       *kolide.AppConfig
 	config                          *config.KolideConfig
+	kolide.TargetStore
 }
 
 func New(config config.KolideConfig) (*Datastore, error) {

server/datastore/mysql/targets.go

@@ -0,0 +1,54 @@
+package mysql
+
+import (
+	"time"
+
+	"github.com/jmoiron/sqlx"
+	"github.com/kolide/kolide/server/kolide"
+	"github.com/pkg/errors"
+)
+
+func (d *Datastore) CountHostsInTargets(hostIDs []uint, labelIDs []uint, now time.Time, onlineInterval time.Duration) (kolide.TargetMetrics, error) {
+	if len(hostIDs) == 0 && len(labelIDs) == 0 {
+		// No need to query if no targets selected
+		return kolide.TargetMetrics{}, nil
+	}
+
+	sql := `
+		SELECT
+			COUNT(*) total,
+			COALESCE(SUM(CASE WHEN DATE_ADD(seen_time, INTERVAL 30 DAY) <= ? THEN 1 ELSE 0 END), 0) mia,
+			COALESCE(SUM(CASE WHEN DATE_ADD(seen_time, INTERVAL ? SECOND) <= ? AND DATE_ADD(seen_time, INTERVAL 30 DAY) >= ? THEN 1 ELSE 0 END), 0) offline,
+			COALESCE(SUM(CASE WHEN DATE_ADD(seen_time, INTERVAL ? SECOND) > ? THEN 1 ELSE 0 END), 0) online,
+			COALESCE(SUM(CASE WHEN DATE_ADD(created_at, INTERVAL 1 DAY) >= ? THEN 1 ELSE 0 END), 0) new
+		FROM hosts h
+		WHERE (id IN (?) OR (id IN (SELECT DISTINCT host_id FROM label_query_executions WHERE label_id IN (?) AND matches = 1)))
+		AND NOT deleted
+`
+
+	// Using -1 in the ID slices for the IN clause allows us to include the
+	// IN clause even if we have no IDs to use. -1 will not match the
+	// auto-increment IDs, and will also allow us to use the same query in
+	// all situations (no need to remove the clause when there are no values)
+	queryLabelIDs := []int{-1}
+	for _, id := range labelIDs {
+		queryLabelIDs = append(queryLabelIDs, int(id))
+	}
+	queryHostIDs := []int{-1}
+	for _, id := range hostIDs {
+		queryHostIDs = append(queryHostIDs, int(id))
+	}
+
+	query, args, err := sqlx.In(sql, now, onlineInterval.Seconds(), now, now, onlineInterval.Seconds(), now, now, queryHostIDs, queryLabelIDs)
+	if err != nil {
+		return kolide.TargetMetrics{}, errors.Wrap(err, "sqlx.In CountHostsInTargets")
+	}
+
+	res := kolide.TargetMetrics{}
+	err = d.db.Get(&res, query, args...)
+	if err != nil {
+		return kolide.TargetMetrics{}, errors.Wrap(err, "sqlx.Get CountHostsInTargets")
+	}
+
+	return res, nil
+}

server/kolide/datastore.go

@@ -8,6 +8,7 @@ type Datastore interface {
 	PackStore
 	LabelStore
 	HostStore
+	TargetStore
 	PasswordResetStore
 	SessionStore
 	AppConfigStore

server/kolide/targets.go

@@ -2,6 +2,7 @@ package kolide
 
 import (
 	"context"
+	"time"
 )
 
 type TargetSearchResults struct {
@@ -12,14 +13,15 @@ type TargetSearchResults struct {
 // TargetMetrics contains information about the state
 // of hosts that are tracked by the app
 type TargetMetrics struct {
-	TotalHosts uint
+	TotalHosts uint `db:"total"`
 	// OnlineHosts have updated within the last 30 minutes
-	OnlineHosts uint
+	OnlineHosts uint `db:"online"`
 	// OfflineHosts are hosts that haven't updated in 30 minutes
-	OfflineHosts uint
+	OfflineHosts uint `db:"offline"`
 	// MissingInActionHosts are hosts that haven't had an update for more
 	// than thirty days
-	MissingInActionHosts uint
+	MissingInActionHosts uint `db:"mia"`
+	NewHosts             uint `db:"new"`
 }
 
 type TargetService interface {
@@ -36,6 +38,10 @@ type TargetService interface {
 	CountHostsInTargets(ctx context.Context, hostIDs []uint, labelIDs []uint) (*TargetMetrics, error)
 }
 
+type TargetStore interface {
+	CountHostsInTargets(hostIDs []uint, labelIDs []uint, now time.Time, onlineInterval time.Duration) (TargetMetrics, error)
+}
+
 type TargetType int
 
 const (

server/mock/datastore.go

@@ -22,6 +22,7 @@ type Store struct {
 	kolide.ScheduledQueryStore
 	kolide.FileIntegrityMonitoringStore
 	kolide.YARAStore
+	kolide.TargetStore
 	LicenseStore
 	InviteStore
 	UserStore

server/service/endpoint_packs.go

@@ -27,19 +27,11 @@ func packResponseForPack(ctx context.Context, svc kolide.Service, pack kolide.Pa
 		return nil, err
 	}
 
-	// ListHostsInPack returns hosts which were explicitly set +
-	// the hosts which are part of a packs labels. We want both for the totals,
-	// but only the explicit host ids for the host_id field.
-	allHosts, err := svc.ListHostsInPack(ctx, pack.ID, opts)
-	if err != nil {
-		return nil, err
-	}
-	totalHostCount := uint(len(allHosts))
-
 	hosts, err := svc.ListExplicitHostsInPack(ctx, pack.ID, opts)
 	if err != nil {
 		return nil, err
 	}
+
 	labels, err := svc.ListLabelsForPack(ctx, pack.ID)
 	labelIDs := make([]uint, len(labels), len(labels))
 	for i, label := range labels {
@@ -48,10 +40,16 @@ func packResponseForPack(ctx context.Context, svc kolide.Service, pack kolide.Pa
 	if err != nil {
 		return nil, err
 	}
+
+	hostMetrics, err := svc.CountHostsInTargets(ctx, hosts, labelIDs)
+	if err != nil {
+		return nil, err
+	}
+
 	return &packResponse{
 		Pack:            pack,
 		QueryCount:      uint(len(queries)),
-		TotalHostsCount: totalHostCount,
+		TotalHostsCount: hostMetrics.TotalHosts,
 		HostIDs:         hosts,
 		LabelIDs:        labelIDs,
 	}, nil

server/service/service_targets.go

@@ -29,43 +29,15 @@ func (svc service) SearchTargets(ctx context.Context, query string, selectedHost
 }
 
 func (svc service) CountHostsInTargets(ctx context.Context, hostIDs []uint, labelIDs []uint) (*kolide.TargetMetrics, error) {
-	hosts, err := svc.ds.ListUniqueHostsInLabels(labelIDs)
-	if err != nil {
-		return nil, err
-	}
-
-	for _, id := range hostIDs {
-		h, err := svc.ds.Host(id)
-		if err != nil {
-			return nil, err
-		}
-		hosts = append(hosts, *h)
-	}
-
-	hostLookup := map[uint]bool{}
-
-	result := &kolide.TargetMetrics{}
-
 	onlineInterval, err := svc.ExpectedCheckinInterval(ctx)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting expected check-in interval")
 	}
 
-	for _, host := range hosts {
-		if !hostLookup[host.ID] {
-			hostLookup[host.ID] = true
-			switch host.Status(svc.clock.Now().UTC(), onlineInterval) {
-			case kolide.StatusOnline:
-				result.OnlineHosts++
-			case kolide.StatusOffline:
-				result.OfflineHosts++
-			case kolide.StatusMIA:
-				result.MissingInActionHosts++
-			}
-		}
+	metrics, err := svc.ds.CountHostsInTargets(hostIDs, labelIDs, svc.clock.Now(), onlineInterval)
+	if err != nil {
+		return nil, err
 	}
 
-	result.TotalHosts = uint(len(hostLookup))
-
-	return result, nil
+	return &metrics, nil
 }

server/service/service_targets_test.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/WatchBeam/clock"
 	"github.com/kolide/kolide/server/config"
 	"github.com/kolide/kolide/server/datastore/inmem"
 	"github.com/kolide/kolide/server/kolide"
@@ -44,143 +43,6 @@ func TestSearchTargets(t *testing.T) {
 	assert.Equal(t, l1.Name, results.Labels[0].Name)
 }
 
-func TestCountHostsInTargets(t *testing.T) {
-	ds, err := inmem.New(config.TestConfig())
-	require.Nil(t, err)
-
-	mockClock := clock.NewMockClock()
-
-	svc, err := newTestServiceWithClock(ds, nil, mockClock)
-	require.Nil(t, err)
-
-	ctx := context.Background()
-
-	h1, err := ds.NewHost(&kolide.Host{
-		HostName: "foo.local",
-		NodeKey:  "1",
-		UUID:     "1",
-	})
-	require.Nil(t, err)
-	require.Nil(t, ds.MarkHostSeen(h1, mockClock.Now()))
-
-	h2, err := ds.NewHost(&kolide.Host{
-		HostName: "bar.local",
-		NodeKey:  "2",
-		UUID:     "2",
-	})
-	require.Nil(t, err)
-	// make this host "offline"
-	require.Nil(t, ds.MarkHostSeen(h2, mockClock.Now().Add(-1*time.Hour)))
-
-	h3, err := ds.NewHost(&kolide.Host{
-		HostName: "baz.local",
-		NodeKey:  "3",
-		UUID:     "3",
-	})
-	require.Nil(t, err)
-	require.Nil(t, ds.MarkHostSeen(h3, mockClock.Now().Add(-5*time.Second)))
-
-	h4, err := ds.NewHost(&kolide.Host{
-		HostName: "xxx.local",
-		NodeKey:  "4",
-		UUID:     "4",
-	})
-	require.Nil(t, err)
-	require.Nil(t, ds.MarkHostSeen(h4, mockClock.Now()))
-
-	h5, err := ds.NewHost(&kolide.Host{
-		HostName: "yyy.local",
-		NodeKey:  "5",
-		UUID:     "5",
-	})
-	require.Nil(t, err)
-	require.Nil(t, ds.MarkHostSeen(h5, mockClock.Now()))
-
-	h6, err := ds.NewHost(&kolide.Host{
-		HostName: "zzz.local",
-		NodeKey:  "6",
-		UUID:     "6",
-	})
-	require.Nil(t, err)
-	const thirtyDaysAndAMinuteAgo = -1 * (30*24*60 + 1)
-	require.Nil(t, ds.MarkHostSeen(h6, mockClock.Now().Add(thirtyDaysAndAMinuteAgo*time.Minute)))
-
-	l1, err := ds.NewLabel(&kolide.Label{
-		Name:  "label foo",
-		Query: "query foo",
-	})
-	require.Nil(t, err)
-	require.NotZero(t, l1.ID)
-
-	l2, err := ds.NewLabel(&kolide.Label{
-		Name:  "label bar",
-		Query: "query foo",
-	})
-	require.Nil(t, err)
-	require.NotZero(t, l2.ID)
-
-	for _, h := range []*kolide.Host{h1, h2, h3, h6} {
-		err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, mockClock.Now())
-		assert.Nil(t, err)
-	}
-
-	for _, h := range []*kolide.Host{h3, h4, h5} {
-		err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l2.ID: true}, mockClock.Now())
-		assert.Nil(t, err)
-	}
-
-	metrics, err := svc.CountHostsInTargets(ctx, nil, []uint{l1.ID, l2.ID})
-	require.Nil(t, err)
-	require.NotNil(t, metrics)
-	assert.Equal(t, uint(6), metrics.TotalHosts)
-	assert.Equal(t, uint(1), metrics.OfflineHosts)
-	assert.Equal(t, uint(4), metrics.OnlineHosts)
-	assert.Equal(t, uint(1), metrics.MissingInActionHosts)
-
-	metrics, err = svc.CountHostsInTargets(ctx, []uint{h1.ID, h2.ID}, []uint{l1.ID, l2.ID})
-	require.Nil(t, err)
-	require.NotNil(t, metrics)
-	assert.Equal(t, uint(6), metrics.TotalHosts)
-	assert.Equal(t, uint(1), metrics.OfflineHosts)
-	assert.Equal(t, uint(4), metrics.OnlineHosts)
-	assert.Equal(t, uint(1), metrics.MissingInActionHosts)
-
-	metrics, err = svc.CountHostsInTargets(ctx, []uint{h1.ID, h2.ID}, nil)
-	require.Nil(t, err)
-	require.NotNil(t, metrics)
-	assert.Equal(t, uint(2), metrics.TotalHosts)
-	assert.Equal(t, uint(1), metrics.OnlineHosts)
-	assert.Equal(t, uint(1), metrics.OfflineHosts)
-	assert.Equal(t, uint(0), metrics.MissingInActionHosts)
-
-	metrics, err = svc.CountHostsInTargets(ctx, []uint{h1.ID}, []uint{l2.ID})
-	require.Nil(t, err)
-	require.NotNil(t, metrics)
-	assert.Equal(t, uint(4), metrics.TotalHosts)
-	assert.Equal(t, uint(4), metrics.OnlineHosts)
-	assert.Equal(t, uint(0), metrics.OfflineHosts)
-	assert.Equal(t, uint(0), metrics.MissingInActionHosts)
-
-	metrics, err = svc.CountHostsInTargets(ctx, nil, nil)
-	require.Nil(t, err)
-	require.NotNil(t, metrics)
-	assert.Equal(t, uint(0), metrics.TotalHosts)
-	assert.Equal(t, uint(0), metrics.OnlineHosts)
-	assert.Equal(t, uint(0), metrics.OfflineHosts)
-	assert.Equal(t, uint(0), metrics.MissingInActionHosts)
-
-	// Advance clock so all hosts are offline
-	mockClock.AddTime(1 * time.Hour)
-	metrics, err = svc.CountHostsInTargets(ctx, nil, []uint{l1.ID, l2.ID})
-	require.Nil(t, err)
-	require.NotNil(t, metrics)
-	assert.Equal(t, uint(6), metrics.TotalHosts)
-	assert.Equal(t, uint(0), metrics.OnlineHosts)
-	assert.Equal(t, uint(5), metrics.OfflineHosts)
-	assert.Equal(t, uint(1), metrics.MissingInActionHosts)
-
-}
-
 func TestSearchWithOmit(t *testing.T) {
 	ds, err := inmem.New(config.TestConfig())
 	require.Nil(t, err)