[RC] Optimize retrieval of raw target files in ClientGetConfigs (#31987)

Author: ameske

pkg/config/remote/service/service.go

@@ -111,50 +111,21 @@ func (s *Service) getNewDirectorRoots(uptane uptaneClient, currentVersion uint64
 	return roots, nil
 }
 
-func (s *Service) getTargetFiles(uptane uptaneClient, products []rdata.Product, cachedTargetFiles []*pbgo.TargetFileMeta) ([]*pbgo.File, error) {
-	productSet := make(map[rdata.Product]struct{})
-	for _, product := range products {
-		productSet[product] = struct{}{}
-	}
-	targets, err := uptane.Targets()
+func (s *Service) getTargetFiles(uptaneClient uptaneClient, targetFilePaths []string) ([]*pbgo.File, error) {
+	files, err := uptaneClient.TargetFiles(targetFilePaths)
 	if err != nil {
 		return nil, err
 	}
-	cachedTargets := make(map[string]data.FileMeta)
-	for _, cachedTarget := range cachedTargetFiles {
-		hashes := make(data.Hashes)
-		for _, hash := range cachedTarget.Hashes {
-			h, err := hex.DecodeString(hash.Hash)
-			if err != nil {
-				return nil, err
-			}
-			hashes[hash.Algorithm] = h
-		}
-		cachedTargets[cachedTarget.Path] = data.FileMeta{
-			Hashes: hashes,
-			Length: cachedTarget.Length,
-		}
-	}
+
 	var configFiles []*pbgo.File
-	for targetPath, targetMeta := range targets {
-		configPathMeta, err := rdata.ParseConfigPath(targetPath)
-		if err != nil {
-			return nil, err
-		}
-		if _, inClientProducts := productSet[rdata.Product(configPathMeta.Product)]; inClientProducts {
-			if notEqualErr := tufutil.FileMetaEqual(cachedTargets[targetPath], targetMeta.FileMeta); notEqualErr == nil {
-				continue
-			}
-			fileContents, err := uptane.TargetFile(targetPath)
-			if err != nil {
-				return nil, err
-			}
-			configFiles = append(configFiles, &pbgo.File{
-				Path: targetPath,
-				Raw:  fileContents,
-			})
-		}
+	for path, contents := range files {
+		// Note: This unconditionally succeeds as long as we don't change bufferDestination earlier
+		configFiles = append(configFiles, &pbgo.File{
+			Path: path,
+			Raw:  contents,
+		})
 	}
+
 	return configFiles, nil
 }
 
@@ -211,6 +182,7 @@ type uptaneClient interface {
 	StoredOrgUUID() (string, error)
 	Targets() (data.TargetFiles, error)
 	TargetFile(path string) ([]byte, error)
+	TargetFiles(files []string) (map[string][]byte, error)
 	TargetsMeta() ([]byte, error)
 	TargetsCustom() ([]byte, error)
 	TUFVersionState() (uptane.TUFVersions, error)
@@ -828,36 +800,30 @@ func (s *CoreAgentService) ClientGetConfigs(_ context.Context, request *pbgo.Cli
 	if err != nil {
 		return nil, err
 	}
-	targetsRaw, err := s.uptane.TargetsMeta()
+
+	directorTargets, err := s.uptane.Targets()
 	if err != nil {
 		return nil, err
 	}
-	targetFiles, err := s.getTargetFiles(s.uptane, rdata.StringListToProduct(request.Client.Products), request.CachedTargetFiles)
+	matchedClientConfigs, err := executeTracerPredicates(request.Client, directorTargets)
 	if err != nil {
 		return nil, err
 	}
 
-	directorTargets, err := s.uptane.Targets()
+	neededFiles, err := filterNeededTargetFiles(matchedClientConfigs, request.CachedTargetFiles, directorTargets)
 	if err != nil {
 		return nil, err
 	}
-	matchedClientConfigs, err := executeTracerPredicates(request.Client, directorTargets)
+
+	targetFiles, err := s.getTargetFiles(s.uptane, neededFiles)
 	if err != nil {
 		return nil, err
 	}
 
-	// filter files to only return the ones that predicates marked for this client
-	matchedConfigsMap := make(map[string]interface{})
-	for _, configPointer := range matchedClientConfigs {
-		matchedConfigsMap[configPointer] = struct{}{}
-	}
-	filteredFiles := make([]*pbgo.File, 0, len(matchedClientConfigs))
-	for _, targetFile := range targetFiles {
-		if _, ok := matchedConfigsMap[targetFile.Path]; ok {
-			filteredFiles = append(filteredFiles, targetFile)
-		}
+	targetsRaw, err := s.uptane.TargetsMeta()
+	if err != nil {
+		return nil, err
 	}
-
 	canonicalTargets, err := enforceCanonicalJSON(targetsRaw)
 	if err != nil {
 		return nil, err
@@ -866,11 +832,42 @@ func (s *CoreAgentService) ClientGetConfigs(_ context.Context, request *pbgo.Cli
 	return &pbgo.ClientGetConfigsResponse{
 		Roots:         roots,
 		Targets:       canonicalTargets,
-		TargetFiles:   filteredFiles,
+		TargetFiles:   targetFiles,
 		ClientConfigs: matchedClientConfigs,
 	}, nil
 }
 
+func filterNeededTargetFiles(neededConfigs []string, cachedTargetFiles []*pbgo.TargetFileMeta, tufTargets data.TargetFiles) ([]string, error) {
+	// Build an O(1) lookup of cached target files
+	cachedTargetsMap := make(map[string]data.FileMeta)
+	for _, cachedTarget := range cachedTargetFiles {
+		hashes := make(data.Hashes)
+		for _, hash := range cachedTarget.Hashes {
+			h, err := hex.DecodeString(hash.Hash)
+			if err != nil {
+				return nil, err
+			}
+			hashes[hash.Algorithm] = h
+		}
+		cachedTargetsMap[cachedTarget.Path] = data.FileMeta{
+			Hashes: hashes,
+			Length: cachedTarget.Length,
+		}
+	}
+
+	// We don't need to pull the raw contents if the client already has the exact version of the file cached
+	filteredList := make([]string, 0, len(neededConfigs))
+	for _, path := range neededConfigs {
+		if notEqualErr := tufutil.FileMetaEqual(cachedTargetsMap[path], tufTargets[path].FileMeta); notEqualErr == nil {
+			continue
+		}
+
+		filteredList = append(filteredList, path)
+	}
+
+	return filteredList, nil
+}
+
 // ConfigGetState returns the state of the configuration and the director repos in the local store
 func (s *CoreAgentService) ConfigGetState() (*pbgo.GetStateConfigResponse, error) {
 	state, err := s.uptane.State()
@@ -1117,29 +1114,14 @@ func (c *HTTPClient) getUpdate(
 	if tufVersions.DirectorTargets == currentTargetsVersion {
 		return nil, nil
 	}
-	roots, err := c.getNewDirectorRoots(c.uptane, currentRootVersion, tufVersions.DirectorRoot)
-	if err != nil {
-		return nil, err
-	}
-	targetsRaw, err := c.uptane.TargetsMeta()
-	if err != nil {
-		return nil, err
-	}
-	targetFiles, err := c.getTargetFiles(c.uptane, rdata.StringListToProduct(products), cachedTargetFiles)
-	if err != nil {
-		return nil, err
-	}
-
-	canonicalTargets, err := enforceCanonicalJSON(targetsRaw)
-	if err != nil {
-		return nil, err
-	}
 
+	// Filter out files that either:
+	//	- don't correspond to the product list the client is requesting
+	//	- have expired
 	directorTargets, err := c.uptane.Targets()
 	if err != nil {
 		return nil, err
 	}
-
 	productsMap := make(map[string]struct{})
 	for _, product := range products {
 		productsMap[product] = struct{}{}
@@ -1165,14 +1147,33 @@ func (c *HTTPClient) getUpdate(
 
 		configs = append(configs, path)
 	}
+	span.SetTag("configs.returned", configs)
+	span.SetTag("configs.expired", expiredConfigs)
 
+	// Gather the files and map-ify them for the state data structure
+	targetFiles, err := c.getTargetFiles(c.uptane, configs)
+	if err != nil {
+		return nil, err
+	}
 	fileMap := make(map[string][]byte, len(targetFiles))
 	for _, f := range targetFiles {
 		fileMap[f.Path] = f.Raw
 	}
 
-	span.SetTag("configs.returned", configs)
-	span.SetTag("configs.expired", expiredConfigs)
+	// Gather some TUF metadata files we need to send down
+	roots, err := c.getNewDirectorRoots(c.uptane, currentRootVersion, tufVersions.DirectorRoot)
+	if err != nil {
+		return nil, err
+	}
+	targetsRaw, err := c.uptane.TargetsMeta()
+	if err != nil {
+		return nil, err
+	}
+	canonicalTargets, err := enforceCanonicalJSON(targetsRaw)
+	if err != nil {
+		return nil, err
+	}
+
 	return &state.Update{
 		TUFRoots:      roots,
 		TUFTargets:    canonicalTargets,

pkg/config/remote/service/service_test.go

@@ -116,6 +116,11 @@ func (m *mockUptane) TargetFile(path string) ([]byte, error) {
 	return args.Get(0).([]byte), args.Error(1)
 }
 
+func (m *mockUptane) TargetFiles(files []string) (map[string][]byte, error) {
+	args := m.Called(files)
+	return args.Get(0).(map[string][]byte), args.Error(1)
+}
+
 func (m *mockUptane) TargetsMeta() ([]byte, error) {
 	args := m.Called()
 	return args.Get(0).([]byte), args.Error(1)
@@ -486,8 +491,8 @@ func TestService(t *testing.T) {
 	}, nil)
 	uptaneClient.On("DirectorRoot", uint64(3)).Return(root3, nil)
 	uptaneClient.On("DirectorRoot", uint64(4)).Return(root4, nil)
-	uptaneClient.On("TargetFile", "datadog/2/APM_SAMPLING/id/1").Return(fileAPM1, nil)
-	uptaneClient.On("TargetFile", "datadog/2/APM_SAMPLING/id/2").Return(fileAPM2, nil)
+
+	uptaneClient.On("TargetFiles", mock.MatchedBy(listsEqual([]string{"datadog/2/APM_SAMPLING/id/1", "datadog/2/APM_SAMPLING/id/2"}))).Return(map[string][]byte{"datadog/2/APM_SAMPLING/id/1": fileAPM1, "datadog/2/APM_SAMPLING/id/2": fileAPM2}, nil)
 	uptaneClient.On("Update", lastConfigResponse).Return(nil)
 	api.On("Fetch", mock.Anything, &pbgo.LatestConfigsRequest{
 		Hostname:                     service.hostname,
@@ -513,7 +518,7 @@ func TestService(t *testing.T) {
 	configResponse, err := service.ClientGetConfigs(context.Background(), &pbgo.ClientGetConfigsRequest{Client: client})
 	assert.NoError(t, err)
 	assert.ElementsMatch(t, [][]byte{canonicalRoot3, canonicalRoot4}, configResponse.Roots)
-	assert.ElementsMatch(t, []*pbgo.File{{Path: "datadog/2/APM_SAMPLING/id/1", Raw: fileAPM1}, {Path: "datadog/2/APM_SAMPLING/id/2", Raw: fileAPM2}}, configResponse.TargetFiles)
+	assert.ElementsMatch(t, []*pbgo.File{{Path: "datadog/2/APM_SAMPLING/id/1", Raw: fileAPM1}, {Path: "datadog/2/APM_SAMPLING/id/2", Raw: fileAPM2}}, configResponse.TargetFiles, nil)
 	assert.Equal(t, canonicalTargets, configResponse.Targets)
 	assert.ElementsMatch(t,
 		configResponse.ClientConfigs,
@@ -597,8 +602,7 @@ func TestServiceClientPredicates(t *testing.T) {
 		DirectorRoot:    1,
 		DirectorTargets: 5,
 	}, nil)
-	uptaneClient.On("TargetFile", "datadog/2/APM_SAMPLING/id/1").Return([]byte(``), nil)
-	uptaneClient.On("TargetFile", "datadog/2/APM_SAMPLING/id/2").Return([]byte(``), nil)
+	uptaneClient.On("TargetFiles", mock.MatchedBy(listsEqual([]string{"datadog/2/APM_SAMPLING/id/1", "datadog/2/APM_SAMPLING/id/2"}))).Return(map[string][]byte{"datadog/2/APM_SAMPLING/id/1": []byte(``), "datadog/2/APM_SAMPLING/id/2": []byte(``)}, nil)
 	uptaneClient.On("Update", lastConfigResponse).Return(nil)
 	api.On("Fetch", mock.Anything, &pbgo.LatestConfigsRequest{
 		Hostname:                     service.hostname,
@@ -887,8 +891,7 @@ func TestConfigExpiration(t *testing.T) {
 		DirectorRoot:    1,
 		DirectorTargets: 5,
 	}, nil)
-	uptaneClient.On("TargetFile", "datadog/2/APM_SAMPLING/id/1").Return([]byte(``), nil)
-	uptaneClient.On("TargetFile", "datadog/2/APM_SAMPLING/id/2").Return([]byte(``), nil)
+	uptaneClient.On("TargetFiles", []string{"datadog/2/APM_SAMPLING/id/1"}).Return(map[string][]byte{"datadog/2/APM_SAMPLING/id/1": []byte(``), "datadog/2/APM_SAMPLING/id/2": []byte(``)}, nil)
 	uptaneClient.On("Update", lastConfigResponse).Return(nil)
 	api.On("Fetch", mock.Anything, &pbgo.LatestConfigsRequest{
 		Hostname:                     service.hostname,
@@ -1190,7 +1193,7 @@ func TestHTTPClientRecentUpdate(t *testing.T) {
 		},
 		nil,
 	)
-	uptaneClient.On("TargetFile", "datadog/2/TESTING1/id/1").Return([]byte(`testing_1`), nil)
+	uptaneClient.On("TargetFiles", []string{"datadog/2/TESTING1/id/1"}).Return(map[string][]byte{"datadog/2/TESTING1/id/1": []byte(`testing_1`)}, nil)
 
 	client := setupCDNClient(t, uptaneClient)
 	defer client.Close()
@@ -1236,7 +1239,7 @@ func TestHTTPClientUpdateSuccess(t *testing.T) {
 				},
 				nil,
 			)
-			uptaneClient.On("TargetFile", "datadog/2/TESTING1/id/1").Return([]byte(`testing_1`), nil)
+			uptaneClient.On("TargetFiles", []string{"datadog/2/TESTING1/id/1"}).Return(map[string][]byte{"datadog/2/TESTING1/id/1": []byte(`testing_1`)}, nil)
 
 			updateErr := fmt.Errorf("uh oh")
 			if tt.updateSucceeds {
@@ -1261,3 +1264,24 @@ func TestHTTPClientUpdateSuccess(t *testing.T) {
 		})
 	}
 }
+
+func listsEqual(mustMatch []string) func(candidate []string) bool {
+	return func(candidate []string) bool {
+		if len(candidate) != len(mustMatch) {
+			return false
+		}
+
+		candidateSet := make(map[string]struct{})
+		for _, item := range candidate {
+			candidateSet[item] = struct{}{}
+		}
+
+		for _, item := range mustMatch {
+			if _, ok := candidateSet[item]; !ok {
+				return false
+			}
+		}
+
+		return true
+	}
+}

pkg/config/remote/uptane/client.go

@@ -342,6 +342,37 @@ func (c *Client) TargetFile(path string) ([]byte, error) {
 	return c.unsafeTargetFile(path)
 }
 
+// TargetFiles returns the content of various multiple target files if the repository is in a
+// verified state.
+func (c *Client) TargetFiles(targetFiles []string) (map[string][]byte, error) {
+	c.Lock()
+	defer c.Unlock()
+
+	err := c.verify()
+	if err != nil {
+		return nil, err
+	}
+
+	// Build the storage space
+	destinations := make(map[string]client.Destination)
+	for _, path := range targetFiles {
+		destinations[path] = &bufferDestination{}
+	}
+
+	err = c.directorTUFClient.DownloadBatch(destinations)
+	if err != nil {
+		return nil, err
+	}
+
+	// Build the return type
+	files := make(map[string][]byte)
+	for path, contents := range destinations {
+		files[path] = contents.(*bufferDestination).Bytes()
+	}
+
+	return files, nil
+}
+
 // TargetsMeta returns the current raw targets.json meta of this uptane client
 func (c *Client) TargetsMeta() ([]byte, error) {
 	c.Lock()