Batch extension label-membership checks in GetOrbitConfig

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.

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 {

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(),

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)

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

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)
+				}
 			}
 		}
 	}