mgmt compute, support beginCreate with context for VirtualMachine (#45597)

XiaofeiCao · web-flow · commit abb34595d65f · 2025-06-06T10:26:00.000+08:00
* code

* add context to co-related resources

* changelog.md

* fix

* assets.json

* checkstyle

* address comments

* verify PUT count

* recording
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/CHANGELOG.md b/sdk/resourcemanager/azure-resourcemanager-compute/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Supported `beginCreate(Context)` in `VirtualMachine`.
+
 ### Breaking Changes
 
 ### Bugs Fixed
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/assets.json b/sdk/resourcemanager/azure-resourcemanager-compute/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "java",
   "TagPrefix": "java/resourcemanager/azure-resourcemanager-compute",
-  "Tag": "java/resourcemanager/azure-resourcemanager-compute_db05ab8f85"
+  "Tag": "java/resourcemanager/azure-resourcemanager-compute_6e783cad01"
 }
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/pom.xml b/sdk/resourcemanager/azure-resourcemanager-compute/pom.xml
@@ -53,6 +53,8 @@
       --add-opens com.azure.resourcemanager.resources/com.azure.resourcemanager.resources=ALL-UNNAMED
       --add-opens com.azure.resourcemanager.resources/com.azure.resourcemanager.resources.fluentcore.arm=ALL-UNNAMED
       --add-opens com.azure.resourcemanager.storage/com.azure.resourcemanager.storage=ALL-UNNAMED
+      --add-opens com.azure.core/com.azure.core.http=com.azure.resourcemanager.compute
+      --add-opens com.azure.core/com.azure.core.http.policy=com.azure.resourcemanager.compute
 
       --add-opens com.azure.core/com.azure.core.implementation.util=ALL-UNNAMED
     </javaModulesSurefireArgLine>
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/implementation/VirtualMachineImpl.java b/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/implementation/VirtualMachineImpl.java
@@ -10,6 +10,7 @@
 import com.azure.core.management.serializer.SerializerFactory;
 import com.azure.core.util.Context;
 import com.azure.core.util.CoreUtils;
+import com.azure.core.util.FluxUtil;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.serializer.SerializerAdapter;
 import com.azure.core.util.serializer.SerializerEncoding;
@@ -2089,23 +2090,29 @@ private Mono<VirtualMachine> prepareCreateResourceAsync() {
     }
 
     public Accepted<VirtualMachine> beginCreate() {
+        return beginCreate(Context.NONE);
+    }
+
+    public Accepted<VirtualMachine> beginCreate(Context context) {
         return AcceptedImpl.<VirtualMachine, VirtualMachineInner>newAccepted(logger,
             this.manager().serviceClient().getHttpPipeline(), this.manager().serviceClient().getDefaultPollInterval(),
             () -> this.manager()
                 .serviceClient()
                 .getVirtualMachines()
                 .createOrUpdateWithResponseAsync(resourceGroupName(), vmName, innerModel(), null, null)
+                .contextWrite(c -> c.putAll(FluxUtil.toReactorContext(context).readOnly()))
                 .block(),
             inner -> new VirtualMachineImpl(inner.name(), inner, this.manager(), this.storageManager,
                 this.networkManager, this.authorizationManager),
             VirtualMachineInner.class, () -> {
                 Flux<Indexable> dependencyTasksAsync
-                    = taskGroup().invokeDependencyAsync(taskGroup().newInvocationContext());
+                    = taskGroup().invokeDependencyAsync(taskGroup().newInvocationContext())
+                        .contextWrite(c -> c.putAll(FluxUtil.toReactorContext(context).readOnly()));
                 dependencyTasksAsync.blockLast();
 
                 // same as createResourceAsync
                 prepareCreateResourceAsync().block();
-            }, this::reset, Context.NONE);
+            }, this::reset, context);
     }
 
     @Override
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/models/VirtualMachine.java b/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/models/VirtualMachine.java
@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Fluent;
 import com.azure.core.http.rest.PagedIterable;
 import com.azure.core.management.Region;
+import com.azure.core.util.Context;
 import com.azure.resourcemanager.authorization.models.BuiltInRole;
 import com.azure.resourcemanager.compute.ComputeManager;
 import com.azure.resourcemanager.compute.fluent.models.VirtualMachineInner;
@@ -2190,6 +2191,18 @@ interface WithCreate extends Creatable<VirtualMachine>, Resource.DefinitionWithT
              * @return the accepted create operation
              */
             Accepted<VirtualMachine> beginCreate();
+
+            /**
+             * Begins creating the virtual machine resource.
+             *
+             * Virtual machine extensions can only be created after the completion of virtual machine.
+             * Therefore, the configuration of virtual machine extensions is not compatible with this operation.
+             * Please use {@link WithCreate#create(Context)} if virtual machine extensions is configured.
+             *
+             * @param context the {@link Context} of the request
+             * @return the accepted create operation
+             */
+            Accepted<VirtualMachine> beginCreate(Context context);
         }
     }
 
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/src/test/java/com/azure/resourcemanager/compute/VirtualMachineOperationsTests.java b/sdk/resourcemanager/azure-resourcemanager-compute/src/test/java/com/azure/resourcemanager/compute/VirtualMachineOperationsTests.java
@@ -3,7 +3,10 @@
 
 package com.azure.resourcemanager.compute;
 
+import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.policy.HttpPipelinePolicy;
 import com.azure.core.http.rest.PagedIterable;
 import com.azure.core.http.rest.PagedResponse;
 import com.azure.core.http.rest.Response;
@@ -40,8 +43,8 @@
 import com.azure.resourcemanager.compute.models.RunCommandInputParameter;
 import com.azure.resourcemanager.compute.models.RunCommandResult;
 import com.azure.resourcemanager.compute.models.SecurityTypes;
-import com.azure.resourcemanager.compute.models.StorageAccountTypes;
 import com.azure.resourcemanager.compute.models.Sku;
+import com.azure.resourcemanager.compute.models.StorageAccountTypes;
 import com.azure.resourcemanager.compute.models.UpgradeMode;
 import com.azure.resourcemanager.compute.models.VirtualMachine;
 import com.azure.resourcemanager.compute.models.VirtualMachineDiskOptions;
@@ -74,6 +77,7 @@
 import com.azure.resourcemanager.resources.models.ResourceGroup;
 import com.azure.resourcemanager.storage.models.StorageAccount;
 import com.azure.resourcemanager.storage.models.StorageAccountSkuType;
+import com.azure.resourcemanager.test.utils.TestIdentifierProvider;
 import com.azure.security.keyvault.keys.models.KeyType;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
@@ -89,6 +93,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class VirtualMachineOperationsTests extends ComputeManagementTest {
@@ -105,11 +110,13 @@ public class VirtualMachineOperationsTests extends ComputeManagementTest {
     private final String availabilitySetName = "availset1";
     private final String availabilitySetName2 = "availset2";
     private final ProximityPlacementGroupType proxGroupType = ProximityPlacementGroupType.STANDARD;
+    private AzureProfile profile;
 
     @Override
     protected void initializeClients(HttpPipeline httpPipeline, AzureProfile profile) {
         rgName = generateRandomResourceName("javacsmrg", 15);
         rgName2 = generateRandomResourceName("javacsmrg2", 15);
+        this.profile = profile;
         super.initializeClients(httpPipeline, profile);
     }
 
@@ -2303,7 +2310,61 @@ public void canCreateDiskWithShares() {
         Assertions.assertEquals(2, disk.virtualMachineIds().size());
     }
 
+    @Test
+    public void canBeginCreateWithContext() {
+        final String vmName = generateRandomResourceName("jvm", 15);
+        final String correlationId = UUID.randomUUID().toString();
+        final String correlationKey = "x-ms-correlation-id";
+        final String publicIpDnsLabel = generateRandomResourceName("pip", 20);
+        final AtomicInteger createCounter = new AtomicInteger(0);
+        HttpPipelinePolicy verificationPolicy = (context, next) -> {
+            if (context.getHttpRequest().getHttpMethod() == HttpMethod.PUT) {
+                // verify that all co-related resource creation requests will have the Context information
+                Object correlationData = context.getContext().getData(correlationKey).get();
+                Assertions.assertEquals(correlationId, correlationData);
+                createCounter.incrementAndGet();
+            }
+            return next.process();
+        };
+        ComputeManager localComputeManager = buildComputeManager(computeManager.httpPipeline(), verificationPolicy);
+
+        Accepted<VirtualMachine> accepted = localComputeManager.virtualMachines()
+            .define(vmName)
+            .withRegion(region)
+            .withNewResourceGroup(rgName)
+            .withNewPrimaryNetwork("10.0.0.0/28")
+            .withPrimaryPrivateIPAddressDynamic()
+            .withNewPrimaryPublicIPAddress(publicIpDnsLabel)
+            .withPopularLinuxImage(KnownLinuxVirtualMachineImage.UBUNTU_SERVER_20_04_LTS)
+            .withRootUsername("Foo12")
+            .withSsh(sshPublicKey())
+            .withSize(VirtualMachineSizeTypes.STANDARD_B1S)
+            .beginCreate(new Context(correlationKey, correlationId));
+        accepted.getFinalResult();
+
+        // resourceGroup + network + neworkInterface + publicIp + VM = 5
+        Assertions.assertEquals(5, createCounter.get());
+    }
+
     // *********************************** helper methods ***********************************
+    private ComputeManager buildComputeManager(HttpPipeline httpPipeline, HttpPipelinePolicy policy) {
+        List<HttpPipelinePolicy> pipelinePolicies = new ArrayList<>();
+        for (int i = 0; i < httpPipeline.getPolicyCount(); i++) {
+            pipelinePolicies.add(httpPipeline.getPolicy(i));
+        }
+
+        pipelinePolicies.add(policy);
+
+        HttpPipeline newPipeline = new HttpPipelineBuilder().httpClient(httpPipeline.getHttpClient())
+            .policies(pipelinePolicies.toArray(new HttpPipelinePolicy[0]))
+            .tracer(httpPipeline.getTracer())
+            .build();
+        ComputeManager manager = ComputeManager.authenticate(newPipeline, profile);
+        ResourceManagerUtils.InternalRuntimeContext internalContext = new ResourceManagerUtils.InternalRuntimeContext();
+        internalContext.setIdentifierFunction(name -> new TestIdentifierProvider(testResourceNamer));
+        setInternalContext(internalContext, manager);
+        return manager;
+    }
 
     private CreatablesInfo prepareCreatableVirtualMachines(Region region, String vmNamePrefix, String networkNamePrefix,
         String publicIpNamePrefix, int vmCount) {

Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`	`"AssetsRepo": "Azure/azure-sdk-assets",`
`3`	`3`	`"AssetsRepoPrefixPath": "java",`
`4`	`4`	`"TagPrefix": "java/resourcemanager/azure-resourcemanager-compute",`
`5`		`- "Tag": "java/resourcemanager/azure-resourcemanager-compute_db05ab8f85"`
	`5`	`+ "Tag": "java/resourcemanager/azure-resourcemanager-compute_6e783cad01"`
`6`	`6`	`}`