Broke apart the hourly host_software count query to reduce the individual query runtime (#18773)

#18221
Broke apart the hourly host_software count query to reduce the
individual query runtime. This fixes timeouts seen when host_software
table has over 25 million records.

I recommend hiding whitespace during review:
<img width="240" alt="image"
src="https://github.com/fleetdm/fleet/assets/2685025/6da9b643-8582-4d2f-bf32-8a1cc38f1032">


# Checklist for submitter

<!-- Note that API documentation changes are now addressed by the
product design team. -->

- [x] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://fleetdm.com/docs/contributing/committing-changes#changes-files)
for more information.
- [x] Added/updated tests
- [x] Manual QA for all new/changed functionality
This commit is contained in:
Victor Lyuboslavsky
2024-05-08 09:27:17 -05:00
committed by GitHub
parent a453979202
commit abf4837eff
3 changed files with 104 additions and 44 deletions
+1
View File
@@ -0,0 +1 @@
Broke apart the hourly host_software count query to reduce the individual query runtime. This fixes timeouts seen when host_software table has over 25 million records.
+64 -41
View File
@@ -18,6 +18,10 @@ import (
"github.com/jmoiron/sqlx"
)
// Since DB may have millions of software items, we need to batch the aggregation counts to avoid long SQL query times.
// This is a variable so it can be adjusted during unit testing.
var countHostSoftwareBatchSize = uint64(100000)
func softwareSliceToMap(softwares []fleet.Software) map[string]fleet.Software {
result := make(map[string]fleet.Software)
for _, s := range softwares {
@@ -1252,7 +1256,7 @@ func (ds *Datastore) SyncHostsSoftware(ctx context.Context, updatedAt time.Time)
globalCountsStmt = `
SELECT count(*), 0 as team_id, software_id
FROM host_software
WHERE software_id > 0
WHERE software_id > ? AND software_id <= ?
GROUP BY software_id`
teamCountsStmt = `
@@ -1260,7 +1264,7 @@ func (ds *Datastore) SyncHostsSoftware(ctx context.Context, updatedAt time.Time)
FROM host_software hs
INNER JOIN hosts h
ON hs.host_id = h.id
WHERE h.team_id IS NOT NULL AND hs.software_id > 0
WHERE h.team_id IS NOT NULL AND hs.software_id > ? AND hs.software_id <= ?
GROUP BY hs.software_id, h.team_id`
insertStmt = `
@@ -1307,55 +1311,74 @@ func (ds *Datastore) SyncHostsSoftware(ctx context.Context, updatedAt time.Time)
return ctxerr.Wrap(ctx, err, "reset all software_host_counts to 0")
}
// next get a cursor for the global and team counts for each software
stmtLabel := []string{"global", "team"}
for i, countStmt := range []string{globalCountsStmt, teamCountsStmt} {
rows, err := ds.reader(ctx).QueryContext(ctx, countStmt)
if err != nil {
return ctxerr.Wrapf(ctx, err, "read %s counts from host_software", stmtLabel[i])
}
defer rows.Close()
db := ds.reader(ctx)
// use a loop to iterate to prevent loading all in one go in memory, as it
// could get pretty big at >100K hosts with 1000+ software each. Use a write
// batch to prevent making too many single-row inserts.
const batchSize = 100
var batchCount int
args := make([]interface{}, 0, batchSize*4)
for rows.Next() {
var (
count int
teamID uint
sid uint
)
// Figure out how many software items we need to count.
type minMaxIDs struct {
Min uint64 `db:"min"`
Max uint64 `db:"max"`
}
minMax := minMaxIDs{}
err := sqlx.GetContext(
ctx, db, &minMax, "SELECT COALESCE(MIN(software_id),1) as min, COALESCE(MAX(software_id),0) as max FROM host_software",
)
if err != nil {
return ctxerr.Wrap(ctx, err, "get min/max software_id")
}
if err := rows.Scan(&count, &teamID, &sid); err != nil {
return ctxerr.Wrapf(ctx, err, "scan %s row into variables", stmtLabel[i])
for minSoftwareID, maxSoftwareID := minMax.Min-1, minMax.Min-1+countHostSoftwareBatchSize; minSoftwareID < minMax.Max; minSoftwareID, maxSoftwareID = maxSoftwareID, maxSoftwareID+countHostSoftwareBatchSize {
// next get a cursor for the global and team counts for each software
stmtLabel := []string{"global", "team"}
for i, countStmt := range []string{globalCountsStmt, teamCountsStmt} {
rows, err := db.QueryContext(ctx, countStmt, minSoftwareID, maxSoftwareID)
if err != nil {
return ctxerr.Wrapf(ctx, err, "read %s counts from host_software", stmtLabel[i])
}
defer rows.Close()
args = append(args, sid, count, teamID, updatedAt)
batchCount++
// use a loop to iterate to prevent loading all in one go in memory, as it
// could get pretty big at >100K hosts with 1000+ software each. Use a write
// batch to prevent making too many single-row inserts.
const batchSize = 100
var batchCount int
args := make([]interface{}, 0, batchSize*4)
for rows.Next() {
var (
count int
teamID uint
sid uint
)
if batchCount == batchSize {
values := strings.TrimSuffix(strings.Repeat(valuesPart, batchCount), ",")
if _, err := ds.writer(ctx).ExecContext(ctx, fmt.Sprintf(insertStmt, values), args...); err != nil {
return ctxerr.Wrapf(ctx, err, "insert %s batch into software_host_counts", stmtLabel[i])
if err := rows.Scan(&count, &teamID, &sid); err != nil {
return ctxerr.Wrapf(ctx, err, "scan %s row into variables", stmtLabel[i])
}
args = args[:0]
batchCount = 0
args = append(args, sid, count, teamID, updatedAt)
batchCount++
if batchCount == batchSize {
values := strings.TrimSuffix(strings.Repeat(valuesPart, batchCount), ",")
if _, err := ds.writer(ctx).ExecContext(ctx, fmt.Sprintf(insertStmt, values), args...); err != nil {
return ctxerr.Wrapf(ctx, err, "insert %s batch into software_host_counts", stmtLabel[i])
}
args = args[:0]
batchCount = 0
}
}
}
if batchCount > 0 {
values := strings.TrimSuffix(strings.Repeat(valuesPart, batchCount), ",")
if _, err := ds.writer(ctx).ExecContext(ctx, fmt.Sprintf(insertStmt, values), args...); err != nil {
return ctxerr.Wrapf(ctx, err, "insert last %s batch into software_host_counts", stmtLabel[i])
if batchCount > 0 {
values := strings.TrimSuffix(strings.Repeat(valuesPart, batchCount), ",")
if _, err := ds.writer(ctx).ExecContext(ctx, fmt.Sprintf(insertStmt, values), args...); err != nil {
return ctxerr.Wrapf(ctx, err, "insert last %s batch into software_host_counts", stmtLabel[i])
}
}
if err := rows.Err(); err != nil {
return ctxerr.Wrapf(ctx, err, "iterate over %s host_software counts", stmtLabel[i])
}
rows.Close()
}
if err := rows.Err(); err != nil {
return ctxerr.Wrapf(ctx, err, "iterate over %s host_software counts", stmtLabel[i])
}
rows.Close()
}
// remove any unused software (global counts = 0)
+39 -3
View File
@@ -927,6 +927,14 @@ func listSoftwareCheckCount(t *testing.T, ds *Datastore, expectedListCount int,
}
func testSoftwareSyncHostsSoftware(t *testing.T, ds *Datastore) {
countHostSoftwareBatchSizeOrig := countHostSoftwareBatchSize
t.Cleanup(
func() {
countHostSoftwareBatchSize = countHostSoftwareBatchSizeOrig
},
)
countHostSoftwareBatchSize = 2
ctx := context.Background()
cmpNameVersionCount := func(want, got []fleet.Software) {
@@ -947,9 +955,21 @@ func testSoftwareSyncHostsSoftware(t *testing.T, ds *Datastore) {
require.Equal(t, want, tableCount)
}
host0 := test.NewHost(t, ds, "host0", "", "host0key", "host0uuid", time.Now())
host1 := test.NewHost(t, ds, "host1", "", "host1key", "host1uuid", time.Now())
host2 := test.NewHost(t, ds, "host2", "", "host2key", "host2uuid", time.Now())
hostTemp := test.NewHost(t, ds, "hostTemp", "", "hostTempKey", "hostTempUuid", time.Now())
// Get counts without any software.
globalOpts := fleet.SoftwareListOptions{
WithHostCounts: true, ListOptions: fleet.ListOptions{OrderKey: "hosts_count", OrderDirection: fleet.OrderDescending},
}
_ = listSoftwareCheckCount(t, ds, 0, 0, globalOpts, false)
software0 := []fleet.Software{
{Name: "abc", Version: "0.0.1", Source: "apps"},
{Name: "def", Version: "0.0.1", Source: "apps"},
}
software1 := []fleet.Software{
{Name: "foo", Version: "0.0.1", Source: "chrome_extensions"},
{Name: "foo", Version: "0.0.3", Source: "chrome_extensions"},
@@ -959,17 +979,33 @@ func testSoftwareSyncHostsSoftware(t *testing.T, ds *Datastore) {
{Name: "foo", Version: "0.0.3", Source: "chrome_extensions"},
{Name: "bar", Version: "0.0.3", Source: "deb_packages"},
}
softwareTemp := make([]fleet.Software, 0, 10)
for i := 0; i < 10; i++ {
softwareTemp = append(
softwareTemp, fleet.Software{Name: fmt.Sprintf("foo%d", i), Version: fmt.Sprintf("%d.0.1", i), Source: "deb_packages"},
)
}
_, err := ds.UpdateHostSoftware(ctx, host1.ID, software1)
_, err := ds.UpdateHostSoftware(ctx, host0.ID, software0)
require.NoError(t, err)
_, err = ds.UpdateHostSoftware(ctx, host1.ID, software1)
require.NoError(t, err)
_, err = ds.UpdateHostSoftware(ctx, hostTemp.ID, softwareTemp)
require.NoError(t, err)
_, err = ds.UpdateHostSoftware(ctx, host2.ID, software2)
require.NoError(t, err)
require.NoError(t, ds.SyncHostsSoftware(ctx, time.Now()))
globalOpts := fleet.SoftwareListOptions{WithHostCounts: true, ListOptions: fleet.ListOptions{OrderKey: "hosts_count", OrderDirection: fleet.OrderDescending}}
globalCounts := listSoftwareCheckCount(t, ds, 4, 4, globalOpts, false)
_ = listSoftwareCheckCount(t, ds, 16, 16, globalOpts, false)
checkTableTotalCount(16)
// Now, delete 2 hosts. Software with the lowest ID is removed, and there should be a chunk with missing software IDs from the deleted hostTemp software.
require.NoError(t, ds.DeleteHost(ctx, host0.ID))
require.NoError(t, ds.DeleteHost(ctx, hostTemp.ID))
require.NoError(t, ds.SyncHostsSoftware(ctx, time.Now()))
globalCounts := listSoftwareCheckCount(t, ds, 4, 4, globalOpts, false)
want := []fleet.Software{
{Name: "foo", Version: "0.0.3", HostsCount: 2},
{Name: "foo", Version: "0.0.1", HostsCount: 1},