feat: add headers for tenantID and clientRequestID (#1228)

Author: jwtty

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/Azure/karpenter-provider-azure v1.5.1
 	github.com/crossplane/crossplane-runtime v1.17.0
 	github.com/evanphx/json-patch/v5 v5.9.11
+	github.com/gofrs/uuid v4.4.0+incompatible
 	github.com/google/go-cmp v0.7.0
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0
 	github.com/onsi/ginkgo/v2 v2.23.4

go.sum

@@ -157,6 +157,8 @@ github.com/go-playground/validator/v10 v10.26.0 h1:SP05Nqhjcvz81uJaRfEV0YBSSSGMc
 github.com/go-playground/validator/v10 v10.26.0/go.mod h1:I5QpIEbmr8On7W0TktmJAumgzX4CA1XNl4ZmDuVHKKo=
 github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI=
 github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
+github.com/gofrs/uuid v4.4.0+incompatible h1:3qXRTX8/NbyulANqlc0lchS1gqAVxRgsuW1YrTJupqA=
+github.com/gofrs/uuid v4.4.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI=

pkg/clients/azure/compute/vmsizerecommenderclient.go

@@ -13,23 +13,31 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"os"
+	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/gofrs/uuid"
 	"google.golang.org/protobuf/encoding/protojson"
+	"k8s.io/klog/v2"
 
 	computev1 "go.goms.io/fleet/apis/protos/azure/compute/v1"
 	"go.goms.io/fleet/pkg/clients/httputil"
 	"go.goms.io/fleet/pkg/utils/controller"
 )
 
 const (
+	tenantIDEnvVarName = "AZURE_TENANT_ID"
 	// recommendationsPathTemplate is the URL path template for VM size recommendations API.
 	recommendationsPathTemplate = "/subscriptions/%s/providers/Microsoft.Compute/locations/%s/vmSizeRecommendations/vmAttributeBased/generate"
 )
 
 // AttributeBasedVMSizeRecommenderClient accesses Azure Attribute-Based VM Size Recommender API
 // to provide VM size recommendations based on specified attributes.
 type AttributeBasedVMSizeRecommenderClient struct {
+	// tenantID is the ID of the Azure fleet's tenant.
+	// At the moment, Azure fleet is single-tenant, the fleet and all its members must be in the same tenant.
+	tenantID string
 	// baseURL is the base URL of the http(s) requests to the attribute-based VM size recommender service endpoint.
 	baseURL string
 	// httpClient is the HTTP client used for making requests.
@@ -43,13 +51,18 @@ func NewAttributeBasedVMSizeRecommenderClient(
 	serverAddress string,
 	httpClient *http.Client,
 ) (*AttributeBasedVMSizeRecommenderClient, error) {
+	tenantID := os.Getenv(tenantIDEnvVarName)
+	if tenantID == "" {
+		return nil, fmt.Errorf("failed to get tenantID: environment variable %s is not set", tenantIDEnvVarName)
+	}
 	if len(serverAddress) == 0 {
 		return nil, fmt.Errorf("serverAddress cannot be empty")
 	}
 	if httpClient == nil {
 		return nil, fmt.Errorf("httpClient cannot be nil")
 	}
 	return &AttributeBasedVMSizeRecommenderClient{
+		tenantID:   tenantID,
 		baseURL:    serverAddress,
 		httpClient: httpClient,
 	}, nil
@@ -59,7 +72,7 @@ func NewAttributeBasedVMSizeRecommenderClient(
 func (c *AttributeBasedVMSizeRecommenderClient) GenerateAttributeBasedRecommendations(
 	ctx context.Context,
 	req *computev1.GenerateAttributeBasedRecommendationsRequest,
-) (*computev1.GenerateAttributeBasedRecommendationsResponse, error) {
+) (response *computev1.GenerateAttributeBasedRecommendationsResponse, err error) {
 	if req == nil {
 		return nil, controller.NewUnexpectedBehaviorError(errors.New("request cannot be nil"))
 	}
@@ -94,10 +107,23 @@ func (c *AttributeBasedVMSizeRecommenderClient) GenerateAttributeBasedRecommenda
 	}
 
 	// Set headers
+	clientRequestID := uuid.Must(uuid.NewV4()).String()
 	httpReq.Header.Set(httputil.HeaderContentTypeKey, httputil.HeaderContentTypeJSON)
 	httpReq.Header.Set(httputil.HeaderAcceptKey, httputil.HeaderContentTypeJSON)
+	httpReq.Header.Set(httputil.HeaderAzureSubscriptionTenantIDKey, c.tenantID)
+	httpReq.Header.Set(httputil.HeaderAzureClientRequestIDKey, clientRequestID)
 
 	// Execute the request
+	startTime := time.Now()
+	klog.V(2).InfoS("Generating VM size recommendations", "subscriptionID", req.SubscriptionId, "location", req.Location, "clientRequestID", clientRequestID)
+	defer func() {
+		latency := time.Since(startTime).Milliseconds()
+		if err != nil {
+			klog.ErrorS(err, "Failed to generate VM size recommendations", "subscriptionID", req.SubscriptionId, "location", req.Location, "clientRequestID", clientRequestID, "latency", latency)
+		}
+		klog.V(2).InfoS("Generated VM size recommendations", "subscriptionID", req.SubscriptionId, "location", req.Location, "clientRequestID", clientRequestID, "latency", latency)
+	}()
+
 	resp, err := c.httpClient.Do(httpReq)
 	if err != nil {
 		return nil, fmt.Errorf("failed to execute request: %w", err)
@@ -116,13 +142,13 @@ func (c *AttributeBasedVMSizeRecommenderClient) GenerateAttributeBasedRecommenda
 	}
 
 	// Unmarshal response using protojson for proper proto3 support
-	var response computev1.GenerateAttributeBasedRecommendationsResponse
+	response = &computev1.GenerateAttributeBasedRecommendationsResponse{}
 	unmarshaler := protojson.UnmarshalOptions{
 		DiscardUnknown: true,
 	}
-	if err := unmarshaler.Unmarshal(respBody, &response); err != nil {
+	if err := unmarshaler.Unmarshal(respBody, response); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
 	}
 
-	return &response, nil
+	return response, nil
 }

pkg/clients/azure/compute/vmsizerecommenderclient_test.go

@@ -21,33 +21,49 @@ import (
 	computev1 "go.goms.io/fleet/apis/protos/azure/compute/v1"
 )
 
+const (
+	testTenantID = "test-tenant-id"
+)
+
 func TestNewAttributeBasedVMSizeRecommenderClient(t *testing.T) {
 	tests := []struct {
 		name          string
+		tenantID      string
 		serverAddress string
 		httpClient    *http.Client
 		wantClient    *AttributeBasedVMSizeRecommenderClient
 		wantErr       bool
 	}{
+		{
+			name:          "with missing tenant ID environment variable",
+			serverAddress: "https://example.com",
+			httpClient:    http.DefaultClient,
+			wantClient:    nil,
+			wantErr:       true,
+		},
 		{
 			name:          "with empty server address",
+			tenantID:      testTenantID,
 			serverAddress: "",
 			httpClient:    http.DefaultClient,
 			wantClient:    nil,
 			wantErr:       true,
 		},
 		{
 			name:          "with nil HTTP client",
+			tenantID:      testTenantID,
 			serverAddress: "http://localhost:8080",
 			httpClient:    nil,
 			wantClient:    nil,
 			wantErr:       true,
 		},
 		{
-			name:          "with both server address and HTTP client",
+			name:          "with all fields properly set",
+			tenantID:      testTenantID,
 			serverAddress: "https://example.com",
 			httpClient:    http.DefaultClient,
 			wantClient: &AttributeBasedVMSizeRecommenderClient{
+				tenantID:   testTenantID,
 				baseURL:    "https://example.com",
 				httpClient: http.DefaultClient,
 			},
@@ -57,6 +73,7 @@ func TestNewAttributeBasedVMSizeRecommenderClient(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			t.Setenv(tenantIDEnvVarName, tt.tenantID)
 			got, gotErr := NewAttributeBasedVMSizeRecommenderClient(tt.serverAddress, tt.httpClient)
 			if (gotErr != nil) != tt.wantErr {
 				t.Errorf("NewAttributeBasedVMSizeRecommenderClient() error = %v, wantErr %v", gotErr, tt.wantErr)
@@ -204,29 +221,38 @@ func TestClient_GenerateAttributeBasedRecommendations(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			// Create mock server
+			// Set tenant ID environment variable to create client.
+			t.Setenv(tenantIDEnvVarName, testTenantID)
+			// Create mock server.
 			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				// Verify request method
+				// Verify request method.
 				if r.Method != http.MethodPost {
 					t.Errorf("got %s, want POST request", r.Method)
 				}
 
-				// Verify headers
+				// Verify headers.
 				if r.Header.Get("Content-Type") != "application/json" {
 					t.Errorf("got %s, want Content-Type: application/json", r.Header.Get("Content-Type"))
 				}
 				if r.Header.Get("Accept") != "application/json" {
 					t.Errorf("got %s, want Accept: application/json", r.Header.Get("Accept"))
 				}
+				if r.Header.Get("Grpc-Metadata-subscriptionTenantID") != testTenantID {
+					t.Errorf("got %s, want Grpc-Metadata-subscriptionTenantID: %s",
+						r.Header.Get("Grpc-Metadata-subscriptionTenantID"), testTenantID)
+				}
+				if r.Header.Get("Grpc-Metadata-clientRequestID") == "" {
+					t.Error("Grpc-Metadata-clientRequestID header is missing")
+				}
 
-				// Verify URL path if request is not nil
+				// Verify URL path if request is not nil.
 				if tt.request != nil && tt.request.SubscriptionId != "" && tt.request.Location != "" {
 					wantPath := fmt.Sprintf(recommendationsPathTemplate, tt.request.SubscriptionId, tt.request.Location)
 					if r.URL.Path != wantPath {
 						t.Errorf("got %s, want path %s", r.URL.Path, wantPath)
 					}
 
-					// Verify request body using protojson for proper proto3 oneof support
+					// Verify request body using protojson for proper proto3 oneof support.
 					body, err := io.ReadAll(r.Body)
 					if err != nil {
 						t.Fatalf("failed to read request body: %v", err)
@@ -243,24 +269,24 @@ func TestClient_GenerateAttributeBasedRecommendations(t *testing.T) {
 					}
 				}
 
-				// Write mock response
+				// Write mock response.
 				w.WriteHeader(tt.mockStatusCode)
 				if _, err := w.Write([]byte(tt.mockResponse)); err != nil {
 					t.Fatalf("failed to write response: %v", err)
 				}
 			}))
 			defer server.Close()
 
-			// Create client
+			// Create client.
 			client, err := NewAttributeBasedVMSizeRecommenderClient(server.URL, http.DefaultClient)
 			if err != nil {
 				t.Errorf("failed to create client: %v", err)
 			}
 
-			// Execute request
+			// Execute request.
 			got, err := client.GenerateAttributeBasedRecommendations(context.Background(), tt.request)
 
-			// Check error
+			// Check error.
 			if (err != nil) != tt.wantErr {
 				t.Errorf("GenerateAttributeBasedRecommendations() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -271,7 +297,7 @@ func TestClient_GenerateAttributeBasedRecommendations(t *testing.T) {
 				return
 			}
 
-			// Compare response
+			// Compare response.
 			if !proto.Equal(tt.wantResponse, got) {
 				t.Errorf("GenerateAttributeBasedRecommendations() = %+v, want %+v", got, tt.wantResponse)
 			}

pkg/clients/httputil/httputil.go

@@ -9,15 +9,12 @@ package httputil
 import (
 	"net/http"
 	"time"
+
+	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 )
 
 // Common HTTP constants.
 const (
-	// HTTPTimeoutAzure is the timeout for HTTP requests to Azure services.
-	// Setting to 60 seconds, following ARM client request timeout conventions:
-	// https://github.com/Azure/azure-resource-manager-rpc/blob/master/v1.0/common-api-details.md#client-request-timeout.
-	HTTPTimeoutAzure = 60 * time.Second
-
 	// HeaderContentTypeKey is the HTTP header key for Content-Type.
 	HeaderContentTypeKey = "Content-Type"
 	// HeaderAcceptKey is the HTTP header key for Accept.
@@ -26,6 +23,20 @@ const (
 	HeaderContentTypeJSON = "application/json"
 )
 
+const (
+	// HTTPTimeoutAzure is the timeout for HTTP requests to Azure services.
+	// Setting to 60 seconds, following ARM client request timeout conventions:
+	// https://github.com/Azure/azure-resource-manager-rpc/blob/master/v1.0/common-api-details.md#client-request-timeout.
+	HTTPTimeoutAzure = 60 * time.Second
+
+	// HeaderAzureSubscriptionTenantIDKey is the HTTP header key for the tenantID of the requested Azure Subscription.
+	// grpc-gateway maps headers with Grpc-Metadata- prefix to grpc metadata after removing it.
+	// See: https://github.com/grpc-ecosystem/grpc-gateway.
+	HeaderAzureSubscriptionTenantIDKey = runtime.MetadataHeaderPrefix + "subscriptionTenantID"
+	// HeaderAzureClientRequestIDKey is the HTTP header key for Azure Client Request ID.
+	HeaderAzureClientRequestIDKey = runtime.MetadataHeaderPrefix + "clientRequestID"
+)
+
 var (
 	// DefaultClientForAzure is the default HTTP client to access Azure services.
 	DefaultClientForAzure = &http.Client{Timeout: HTTPTimeoutAzure}

pkg/propertychecker/azure/checker_test.go

@@ -430,6 +430,7 @@ func TestCheckIfMeetSKUCapacityRequirement(t *testing.T) {
 			server := createMockAttributeBasedVMSizeRecommenderServer(t, tt.mockStatusCode)
 			defer server.Close()
 
+			t.Setenv("AZURE_TENANT_ID", "test-tenant-id")
 			client, err := compute.NewAttributeBasedVMSizeRecommenderClient(server.URL, http.DefaultClient)
 			if err != nil {
 				t.Fatalf("failed to create VM size recommender client: %v", err)