diff --git a/changes/batch-extension-label-check b/changes/batch-extension-label-check new file mode 100644 index 00000000000..4bc7af2bc90 --- /dev/null +++ b/changes/batch-extension-label-check @@ -0,0 +1 @@ +- Improved performance of Orbit config endpoint by batching extension label-membership checks into a single database query. diff --git a/server/datastore/mysql/labels.go b/server/datastore/mysql/labels.go index 4bfa8aadb57..b423c5aa906 100644 --- a/server/datastore/mysql/labels.go +++ b/server/datastore/mysql/labels.go @@ -1743,6 +1743,35 @@ func (ds *Datastore) HostMemberOfAllLabels(ctx context.Context, hostID uint, lab return ok, nil } +// HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. +// Labels that do not exist are not included in the result. +func (ds *Datastore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) { + if len(labelNames) == 0 { + return nil, nil + } + + sqlStatement := ` + SELECT l.name FROM labels l + JOIN label_membership lm ON l.id = lm.label_id + WHERE lm.host_id = ? AND l.name IN (?) + ` + sql, args, err := sqlx.In(sqlStatement, hostID, labelNames) + if err != nil { + return nil, ctxerr.Wrap(ctx, err, "building query for host label membership") + } + + var names []string + if err := sqlx.SelectContext(ctx, ds.reader(ctx), &names, sql, args...); err != nil { + return nil, ctxerr.Wrap(ctx, err, "get host label membership") + } + + result := make(map[string]struct{}, len(names)) + for _, n := range names { + result[n] = struct{}{} + } + return result, nil +} + // AddLabelsToHost skips auth as it's only used in tests, and where label teams have already been validated. func (ds *Datastore) AddLabelsToHost(ctx context.Context, hostID uint, labelIDs []uint) error { if len(labelIDs) == 0 { diff --git a/server/datastore/mysql/labels_test.go b/server/datastore/mysql/labels_test.go index 623e8c0e282..5500554f736 100644 --- a/server/datastore/mysql/labels_test.go +++ b/server/datastore/mysql/labels_test.go @@ -96,6 +96,7 @@ func TestLabels(t *testing.T) { {"ListHostsInLabelIssues", testListHostsInLabelIssues}, {"ListHostsInLabelDiskEncryptionStatus", testListHostsInLabelDiskEncryptionStatus}, {"HostMemberOfAllLabels", testHostMemberOfAllLabels}, + {"HostMembershipForLabels", testHostMembershipForLabels}, {"ListHostsInLabelOSSettings", testLabelsListHostsInLabelOSSettings}, {"AddDeleteLabelsToFromHost", testAddDeleteLabelsToFromHost}, {"ApplyLabelSpecSerialUUID", testApplyLabelSpecsForSerialUUID}, @@ -2106,6 +2107,126 @@ func testHostMemberOfAllLabels(t *testing.T, ds *Datastore) { } } +func testHostMembershipForLabels(t *testing.T, ds *Datastore) { + ctx := context.Background() + + // + // Setup test + // - h1 member of 'All hosts', 'Foobar' and 'Zoobar' + // - h2 member of 'All hosts' and 'Foobar' + // - h3 member of no labels + // + + allHostsLabel, err := ds.NewLabel(ctx, + &fleet.Label{ + Name: "All hosts", + Query: "SELECT 1", + LabelType: fleet.LabelTypeBuiltIn, + LabelMembershipType: fleet.LabelMembershipTypeDynamic, + }, + ) + require.NoError(t, err) + foobarLabel, err := ds.NewLabel(ctx, &fleet.Label{ + Name: "Foobar", + Query: "SELECT 1;", + LabelType: fleet.LabelTypeRegular, + LabelMembershipType: fleet.LabelMembershipTypeDynamic, + }) + require.NoError(t, err) + zoobarLabel, err := ds.NewLabel(ctx, &fleet.Label{ + Name: "Zoobar", + Query: "SELECT 2;", + LabelType: fleet.LabelTypeRegular, + LabelMembershipType: fleet.LabelMembershipTypeDynamic, + }) + require.NoError(t, err) + + newHostFunc := func(name string) *fleet.Host { + h, err := ds.NewHost(ctx, &fleet.Host{ + DetailUpdatedAt: time.Now(), + LabelUpdatedAt: time.Now(), + PolicyUpdatedAt: time.Now(), + SeenTime: time.Now(), + OsqueryHostID: new(name), + NodeKey: new(name), + UUID: name, + Hostname: "foo.local" + name, + }) + require.NoError(t, err) + return h + } + + h1 := newHostFunc("h1") + h2 := newHostFunc("h2") + h3 := newHostFunc("h3") + + err = ds.RecordLabelQueryExecutions(ctx, h1, map[uint]*bool{ + allHostsLabel.ID: new(true), + foobarLabel.ID: new(true), + zoobarLabel.ID: new(true), + }, time.Now(), false) + require.NoError(t, err) + err = ds.RecordLabelQueryExecutions(ctx, h2, map[uint]*bool{ + allHostsLabel.ID: new(true), + foobarLabel.ID: new(true), + }, time.Now(), false) + require.NoError(t, err) + + for _, tc := range []struct { + name string + hostID uint + labelNames []string + expectedResult map[string]struct{} + }{ + { + name: "empty label names returns nil", + hostID: h1.ID, + labelNames: nil, + expectedResult: nil, + }, + { + name: "h1 is member of all three labels", + hostID: h1.ID, + labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, + expectedResult: map[string]struct{}{allHostsLabel.Name: {}, foobarLabel.Name: {}, zoobarLabel.Name: {}}, + }, + { + name: "h2 is member of some but not all labels", + hostID: h2.ID, + labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, + expectedResult: map[string]struct{}{allHostsLabel.Name: {}, foobarLabel.Name: {}}, + }, + { + name: "nonexistent labels are not included", + hostID: h1.ID, + labelNames: []string{allHostsLabel.Name, "nonexistent-label"}, + expectedResult: map[string]struct{}{allHostsLabel.Name: {}}, + }, + { + name: "nonexistent host returns empty result", + hostID: 999, + labelNames: []string{allHostsLabel.Name}, + expectedResult: map[string]struct{}{}, + }, + { + name: "h3 is member of no labels", + hostID: h3.ID, + labelNames: []string{allHostsLabel.Name, foobarLabel.Name}, + expectedResult: map[string]struct{}{}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + result, err := ds.HostMembershipForLabels(ctx, tc.hostID, tc.labelNames) + require.NoError(t, err) + if tc.expectedResult == nil { + require.Nil(t, result) + } else { + require.Equal(t, tc.expectedResult, result) + } + }) + } +} + func testLabelsListHostsInLabelOSSettings(t *testing.T, db *Datastore) { h1, err := db.NewHost(context.Background(), &fleet.Host{ DetailUpdatedAt: time.Now(), diff --git a/server/fleet/datastore.go b/server/fleet/datastore.go index 97e7bb4c122..376b737f1c2 100644 --- a/server/fleet/datastore.go +++ b/server/fleet/datastore.go @@ -337,6 +337,10 @@ type Datastore interface { // if labelNames is empty. HostMemberOfAllLabels(ctx context.Context, hostID uint, labelNames []string) (bool, error) + // HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. + // Labels that do not exist are not included in the result. + HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) + // TODO JUAN: Refactor this to use the Operating System type instead. // HostIDsByOSVersion retrieves the IDs of all host matching osVersion HostIDsByOSVersion(ctx context.Context, osVersion OSVersion, offset int, limit int) ([]uint, error) diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go index 6ecb4d300d9..7d841ac6238 100644 --- a/server/mock/datastore_mock.go +++ b/server/mock/datastore_mock.go @@ -256,6 +256,8 @@ type HostIDsByOSIDFunc func(ctx context.Context, osID uint, offset int, limit in type HostMemberOfAllLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (bool, error) +type HostMembershipForLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) + type HostIDsByOSVersionFunc func(ctx context.Context, osVersion fleet.OSVersion, offset int, limit int) ([]uint, error) type HostByIdentifierFunc func(ctx context.Context, identifier string) (*fleet.Host, error) @@ -2507,6 +2509,9 @@ type DataStore struct { HostMemberOfAllLabelsFunc HostMemberOfAllLabelsFunc HostMemberOfAllLabelsFuncInvoked bool + HostMembershipForLabelsFunc HostMembershipForLabelsFunc + HostMembershipForLabelsFuncInvoked bool + HostIDsByOSVersionFunc HostIDsByOSVersionFunc HostIDsByOSVersionFuncInvoked bool @@ -6175,6 +6180,13 @@ func (s *DataStore) HostMemberOfAllLabels(ctx context.Context, hostID uint, labe return s.HostMemberOfAllLabelsFunc(ctx, hostID, labelNames) } +func (s *DataStore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) { + s.mu.Lock() + s.HostMembershipForLabelsFuncInvoked = true + s.mu.Unlock() + return s.HostMembershipForLabelsFunc(ctx, hostID, labelNames) +} + func (s *DataStore) HostIDsByOSVersion(ctx context.Context, osVersion fleet.OSVersion, offset int, limit int) ([]uint, error) { s.mu.Lock() s.HostIDsByOSVersionFuncInvoked = true diff --git a/server/service/orbit.go b/server/service/orbit.go index 791e33f52e7..ce3faf16979 100644 --- a/server/service/orbit.go +++ b/server/service/orbit.go @@ -971,17 +971,39 @@ func (svc *Service) filterExtensionsForHost(ctx context.Context, extensions json // Filter the extensions by labels (premium only feature). if license, _ := license.FromContext(ctx); license != nil && license.IsPremium() { - for extensionName, extensionInfo := range extensionsInfo { - hostIsMemberOfAllLabels, err := svc.ds.HostMemberOfAllLabels(ctx, host.ID, extensionInfo.Labels) + // Collect all unique label names across all extensions. + allLabels := make(map[string]struct{}) + for _, extInfo := range extensionsInfo { + for _, l := range extInfo.Labels { + allLabels[l] = struct{}{} + } + } + + if len(allLabels) > 0 { + labelNames := make([]string, 0, len(allLabels)) + for l := range allLabels { + labelNames = append(labelNames, l) + } + + memberOf, err := svc.ds.HostMembershipForLabels(ctx, host.ID, labelNames) if err != nil { - return nil, ctxerr.Wrap(ctx, err, "check host labels") + return nil, ctxerr.Wrap(ctx, err, "check host label membership") } - if hostIsMemberOfAllLabels { - // Do not filter out, but there's no need to send the label names to the devices. - extensionInfo.Labels = nil - extensionsInfo[extensionName] = extensionInfo - } else { - delete(extensionsInfo, extensionName) + + for extensionName, extensionInfo := range extensionsInfo { + allMatch := true + for _, l := range extensionInfo.Labels { + if _, ok := memberOf[l]; !ok { + allMatch = false + break + } + } + if allMatch { + extensionInfo.Labels = nil + extensionsInfo[extensionName] = extensionInfo + } else { + delete(extensionsInfo, extensionName) + } } } }