Discharge7443/fleet
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Update NewLabel method to use more efficient update mechanism (#25777)

Browse Source 
For #25555 

# Checklist for submitter

- [X] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/Committing-Changes.md#changes-files)
for more information.
- [X] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)

This PR updates the `NewLabel` service to use the
`UpdateLabelMembershipByHostIDs` method previously added by
@jacobshandling rather than using `ApplyLabels`. The latter method has
performance issues when adding large numbers of hosts at once to a
manual label (see #25555) because it does an expensive lookup of host
names before transforming those into Fleet host IDs. The new code skips
the middleman and transforms host identifiers directly to Fleet host
IDs, and does so using a batching strategy to ensure the queries don't
get too large.

This PR does update `UpdateLabelMembershipByHostIDs` slightly to return
an updated Label object and host IDs array, as this is the expected
return value for `NewLabel`. I update the method's tests accordingly. I
don't think any new tests for `NewLabel` are needed as it should have
the same functionality and return values.

## Manual Testing

On the main branch, I launched my local MySQL with the thread stack size
set to the minimal allowed, and used the API to try and create a new
label with 5,000 hosts attached, and received a 422 response from the
server. Server logs showed:
```
level=error ts=2025-01-28T15:08:20.465401Z component=http user=scott@fleetdm.com method=POST 
uri=/api/latest/fleet/labels took=16.610292ms err="get hostnames by identifiers: Error 1436 (HY000): Thread stack 
overrun:  111136 bytes used of a 131072 byte stack, and 20000 bytes needed.  Use 'mysqld --thread_stack=#' to specify 
a bigger stack."
```

On this branch, I kept the same MySQL settings and tried my API request
again and it was successful:
<img width="776" alt="image"
src="https://github.com/user-attachments/assets/c4f0f52b-4d09-457b-8096-4dd3a747b1f4"
/>

## QA

The script I used to create a new manual label with lots of hosts is at:
https://gist.github.com/sgress454/84f12064c437da456c456e25c26d9069

To run it, first grab a bearer token from any API request by opening the
network tab, clicking a Fleet API request, and in the headers tab
scrolling down to Authorization:
<img width="892" alt="image"
src="https://github.com/user-attachments/assets/5680f3bf-8db8-469a-9f03-000b86622c04"
/>
(only take the part _after_ "Bearer")

Then download the script from that gist and in its folder run:
```
NODE_TLS_REJECT_UNAUTHORIZED=0 node ./add_hosts_to_label.js <the bearer token> "<a label name>"
```
e.g.
```
NODE_TLS_REJECT_UNAUTHORIZED=0 node ./add_hosts_to_label.js U3HpbdtadmJXGKYSB0U/PbwfOpHbBt7FpkWmGKKYolOO1moLNZA6XxP+QO5LVukvAotZ7d+JbNUEEhYHZtxoqg== "some test label"
```
This will invoke the API on https://localhost:8080 and try to add 5000
hosts a new label "some test label".

If you need to change the # of hosts or the url of the server, there are
additional arguments:
```
NODE_TLS_REJECT_UNAUTHORIZED=0 node ./add_hosts_to_label.js <the bearer token> "<a label name>" <number of hosts> <url>
```
e.g.
```
NODE_TLS_REJECT_UNAUTHORIZED=0 node ./add_hosts_to_label.js U3HpbdtadmJXGKYSB0U/PbwfOpHbBt7FpkWmGKKYolOO1moLNZA6XxP+QO5LVukvAotZ7d+JbNUEEhYHZtxoqg== "some test label" 10000 https://foo.bar
```
This commit is contained in:
Scott Gress
2025-01-31 09:19:36 -06:00
committed by GitHub
parent 49fe510ab0
commit 1cd37ef966
7 changed files with 58 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/datastore/mysql/hosts.go
+4 -2
View File
@@ -2681,12 +2681,14 @@ func (ds *Datastore) HostIDsByIdentifier(ctx context.Context, filter fleet.TeamF
WHERE
(hostname IN (?)
OR uuid IN (?)
OR hardware_serial IN (?))
OR hardware_serial IN (?)
OR node_key IN (?)
OR osquery_host_id IN (?))
AND %s
`, ds.whereFilterHostsByTeams(filter, "hosts"),
)
sql, args, err := sqlx.In(sqlStatement, hostIdentifiers, hostIdentifiers, hostIdentifiers)
sql, args, err := sqlx.In(sqlStatement, hostIdentifiers, hostIdentifiers, hostIdentifiers, hostIdentifiers, hostIdentifiers)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "building query to get host IDs by identifier")
}
server/datastore/mysql/labels.go
+6 -3
View File
@@ -123,8 +123,8 @@ func batchHostnames(hostnames []string) [][]string {
return batches
}
func (ds *Datastore) UpdateLabelMembershipByHostIDs(ctx context.Context, labelID uint, hostIds []uint) (err error) {
err = ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
func (ds *Datastore) UpdateLabelMembershipByHostIDs(ctx context.Context, labelID uint, hostIds []uint, teamFilter fleet.TeamFilter) (*fleet.Label, []uint, error) {
err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
// delete all label membership
sql := `
DELETE FROM label_membership WHERE label_id = ?
@@ -165,8 +165,11 @@ VALUES ` + strings.Join(placeholders, ", ")
}
return nil
})
if err != nil {
return nil, nil, ctxerr.Wrap(ctx, err, "UpdateLabelMembershipByHostIDs transaction")
}
return ctxerr.Wrap(ctx, err, "UpdateLabelMembershipByHostIDs transaction")
return ds.labelDB(ctx, labelID, teamFilter, ds.writer(ctx))
}
func batchHostIds(hostIds []uint) [][]uint {
server/datastore/mysql/labels_test.go
+38 -14
View File
@@ -1769,6 +1769,8 @@ func labelIDFromName(t *testing.T, ds fleet.Datastore, name string) uint {
func testUpdateLabelMembershipByHostIDs(t *testing.T, ds *Datastore) {
ctx := context.Background()
filter := fleet.TeamFilter{User: test.UserAdmin}
host1, err := ds.NewHost(ctx, &fleet.Host{
OsqueryHostID: ptr.String("1"),
NodeKey: ptr.String("1"),
@@ -1804,17 +1806,24 @@ func testUpdateLabelMembershipByHostIDs(t *testing.T, ds *Datastore) {
require.NoError(t, err)
// add hosts 1 and 2 to the label
err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host1.ID, host2.ID})
label, hostIDs, err := ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host1.ID, host2.ID}, filter)
require.NoError(t, err)
require.Equal(t, label.HostCount, 2)
// expect hosts 1 and 2 to be in the label, but not 3
require.NoError(t, err)
// correct hosts were added to label
require.Len(t, hostIDs, 2)
require.Equal(t, host1.ID, hostIDs[0])
require.Equal(t, host2.ID, hostIDs[1])
label, err := ds.GetLabelSpec(ctx, label1.Name)
labelSpec, err := ds.GetLabelSpec(ctx, label1.Name)
require.NoError(t, err)
// label.Hosts contains hostnames
require.Len(t, label.Hosts, 2)
require.Equal(t, host1.Hostname, label.Hosts[0])
require.Equal(t, host2.Hostname, label.Hosts[1])
require.Len(t, labelSpec.Hosts, 2)
require.Equal(t, host1.Hostname, labelSpec.Hosts[0])
require.Equal(t, host2.Hostname, labelSpec.Hosts[1])
labels, err := ds.ListLabelsForHost(ctx, host1.ID)
require.NoError(t, err)
@@ -1831,9 +1840,11 @@ func testUpdateLabelMembershipByHostIDs(t *testing.T, ds *Datastore) {
require.Len(t, labels, 0)
// modify the label to contain hosts 1 and 3, confirm
err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host1.ID, host3.ID})
label, _, err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host1.ID, host3.ID}, filter)
require.NoError(t, err)
require.Equal(t, label.HostCount, 2)
labels, err = ds.ListLabelsForHost(ctx, host1.ID)
require.NoError(t, err)
require.Len(t, labels, 1)
@@ -1849,9 +1860,11 @@ func testUpdateLabelMembershipByHostIDs(t *testing.T, ds *Datastore) {
require.Equal(t, "label1", labels[0].Name)
// modify the label to contain hosts 2 and 3, confirm
err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host2.ID, host3.ID})
label, _, err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host2.ID, host3.ID}, filter)
require.NoError(t, err)
require.Equal(t, label.HostCount, 2)
labels, err = ds.ListLabelsForHost(ctx, host1.ID)
require.NoError(t, err)
require.Len(t, labels, 0)
@@ -1867,8 +1880,9 @@ func testUpdateLabelMembershipByHostIDs(t *testing.T, ds *Datastore) {
require.Equal(t, "label1", labels[0].Name)
// modify the label to contain no hosts, confirm
err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{})
label, _, err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{}, filter)
require.NoError(t, err)
require.Equal(t, label.HostCount, 0)
labels, err = ds.ListLabelsForHost(ctx, host1.ID)
require.NoError(t, err)
@@ -1883,9 +1897,11 @@ func testUpdateLabelMembershipByHostIDs(t *testing.T, ds *Datastore) {
require.Len(t, labels, 0)
// modify the label to contain all 3 hosts, confirm
err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host1.ID, host2.ID, host3.ID})
label, hostIDs, err = ds.UpdateLabelMembershipByHostIDs(ctx, label1.ID, []uint{host1.ID, host2.ID, host3.ID}, filter)
require.NoError(t, err)
require.Equal(t, label.HostCount, 3)
labels, err = ds.ListLabelsForHost(ctx, host1.ID)
require.NoError(t, err)
require.Len(t, labels, 1)
@@ -1901,11 +1917,19 @@ func testUpdateLabelMembershipByHostIDs(t *testing.T, ds *Datastore) {
require.Len(t, labels, 1)
require.Equal(t, "label1", labels[0].Name)
label, err = ds.GetLabelSpec(ctx, label1.Name)
require.NoError(t, err)
require.Len(t, label.Hosts, 3)
require.Equal(t, host1.Hostname, label.Hosts[0])
require.Len(t, hostIDs, 3)
require.Equal(t, host1.ID, hostIDs[0])
// 2 and 3 have same name
require.Equal(t, host2.Hostname, label.Hosts[1])
require.Equal(t, host3.Hostname, label.Hosts[2])
require.Equal(t, host2.ID, hostIDs[1])
require.Equal(t, host3.ID, hostIDs[2])
labelSpec, err = ds.GetLabelSpec(ctx, label1.Name)
require.NoError(t, err)
// label.Hosts contains hostnames
require.Len(t, labelSpec.Hosts, 3)
require.Equal(t, host1.Hostname, labelSpec.Hosts[0])
require.Equal(t, host2.Hostname, labelSpec.Hosts[1])
require.Equal(t, host3.Hostname, labelSpec.Hosts[2])
}
server/fleet/datastore.go
+2 -2
View File
@@ -191,7 +191,7 @@ type Datastore interface {
// UpdateLabelMembershipByHostIDs updates the label membership for the given label ID with host
// IDs, applied in batches
UpdateLabelMembershipByHostIDs(ctx context.Context, labelID uint, hostIds []uint) (err error)
UpdateLabelMembershipByHostIDs(ctx context.Context, labelID uint, hostIds []uint, teamFilter TeamFilter) (*Label, []uint, error)
NewLabel(ctx context.Context, Label *Label, opts ...OptionalArg) (*Label, error)
// SaveLabel updates the label and returns the label and an array of host IDs
@@ -274,7 +274,7 @@ type Datastore interface {
CleanupIncomingHosts(ctx context.Context, now time.Time) ([]uint, error)
// GenerateHostStatusStatistics retrieves the count of online, offline, MIA and new hosts.
GenerateHostStatusStatistics(ctx context.Context, filter TeamFilter, now time.Time, platform *string, lowDiskSpace *int) (*HostSummary, error)
// HostIDsByIdentifier retrieves the IDs associated with the given hostnames, UUIDs, or hardware serials.
// HostIDsByIdentifier retrieves the IDs associated with the given hostnames, UUIDs, hardware serials, node keys or osquery host IDs.
HostIDsByIdentifier(ctx context.Context, filter TeamFilter, hostnames []string) ([]uint, error)
// HostIDsByOSID retrieves the IDs of all host for the given OS ID
server/mock/datastore_mock.go
+3 -3
View File
@@ -129,7 +129,7 @@ type ListPacksForHostFunc func(ctx context.Context, hid uint) (packs []*fleet.Pa
type ApplyLabelSpecsFunc func(ctx context.Context, specs []*fleet.LabelSpec) error
type UpdateLabelMembershipByHostIDsFunc func(ctx context.Context, labelID uint, hostIDs []uint) (err error)
type UpdateLabelMembershipByHostIDsFunc func(ctx context.Context, labelID uint, hostIDs []uint, teamFilter fleet.TeamFilter) (*fleet.Label, []uint, error)
type GetLabelSpecsFunc func(ctx context.Context) ([]*fleet.LabelSpec, error)
@@ -3378,11 +3378,11 @@ func (s *DataStore) ApplyLabelSpecs(ctx context.Context, specs []*fleet.LabelSpe
return s.ApplyLabelSpecsFunc(ctx, specs)
}
func (s *DataStore) UpdateLabelMembershipByHostIDs(ctx context.Context, labelID uint, hostIDs []uint) (err error) {
func (s *DataStore) UpdateLabelMembershipByHostIDs(ctx context.Context, labelID uint, hostIDs []uint, teamFilter fleet.TeamFilter) (*fleet.Label, []uint, error) {
s.mu.Lock()
s.UpdateLabelMembershipByHostIDsFuncInvoked = true
s.mu.Unlock()
return s.UpdateLabelMembershipByHostIDsFunc(ctx, labelID, hostIDs)
return s.UpdateLabelMembershipByHostIDsFunc(ctx, labelID, hostIDs, teamFilter)
}
func (s *DataStore) GetLabelSpecs(ctx context.Context) ([]*fleet.LabelSpec, error) {
server/service/labels.go
+4 -26
View File
@@ -7,7 +7,6 @@ import (
"github.com/fleetdm/fleet/v4/server"
authz_ctx "github.com/fleetdm/fleet/v4/server/contexts/authz"
"github.com/fleetdm/fleet/v4/server/contexts/ctxdb"
"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
"github.com/fleetdm/fleet/v4/server/contexts/license"
"github.com/fleetdm/fleet/v4/server/contexts/viewer"
@@ -89,36 +88,15 @@ func (svc *Service) NewLabel(ctx context.Context, p fleet.LabelPayload) (*fleet.
return nil, nil, err
}
// Next, if membership type is manual, use ApplyLabelSpecs to create label
// memberships. Must resolve the host identifiers to hostname so that
// ApplySpecs can be used.
var hostIDs []uint
if label.LabelMembershipType == fleet.LabelMembershipTypeManual {
spec := fleet.LabelSpec{
Name: label.Name,
Description: label.Description,
Query: label.Query,
Platform: label.Platform,
LabelType: label.LabelType,
LabelMembershipType: label.LabelMembershipType,
}
hostnames, err := svc.ds.HostnamesByIdentifiers(ctx, p.Hosts)
if err != nil {
return nil, nil, err
}
spec.Hosts = hostnames
if err := svc.ds.ApplyLabelSpecs(ctx, []*fleet.LabelSpec{&spec}); err != nil {
return nil, nil, err
}
// must reload it to get the host IDs, refresh its count
ctx = ctxdb.RequirePrimary(ctx, true)
label, hostIDs, err = svc.ds.Label(ctx, label.ID, filter)
hostIDs, err = svc.ds.HostIDsByIdentifier(ctx, filter, p.Hosts)
if err != nil {
return nil, nil, err
}
return svc.ds.UpdateLabelMembershipByHostIDs(ctx, label.ID, hostIDs, filter)
}
return label, hostIDs, nil
return label, nil, nil
}
////////////////////////////////////////////////////////////////////////////////
@@ -196,7 +174,7 @@ func (svc *Service) ModifyLabel(ctx context.Context, id uint, payload fleet.Modi
if err != nil {
return nil, nil, err
}
if err := svc.ds.UpdateLabelMembershipByHostIDs(ctx, label.ID, hostIds); err != nil {
if _, _, err := svc.ds.UpdateLabelMembershipByHostIDs(ctx, label.ID, hostIds, filter); err != nil {
return nil, nil, err
}
}