support for LB backend CRUD and enhanced VM assignment logic

azure-mgmt-network/src/main/java/com/microsoft/azure/management/network/LoadBalancer.java

@@ -106,6 +106,12 @@ interface WithInternetFrontendOrBackend extends WithInternetFrontends, WithBacke
          * The stage of the definition allowing to add a backend.
          */
         interface WithBackend extends WithVirtualMachine {
+            /**
+             * Adds a new empty backend to the load balancer.
+             * @param name the name to assign to the backend
+             * @return the next stage of the update
+             */
+            WithBackendOrCreate withBackend(String name);
         }
 
         /**
@@ -137,10 +143,34 @@ interface WithVirtualMachine {
              * <p>
              * If the virtual machines are not in the same availability set, the load balancer will still
              * be created, but the virtual machines will not associated with its back end.
+             * <p>
+             * Only those virtual machines will be associated with the load balancer that already have an existing
+             * network interface. Virtual machines without a network interface will be skipped.
              * @param vms existing virtual machines
              * @return the next stage of the update
              */
             WithBackendOrCreate withExistingVirtualMachines(SupportsNetworkInterfaces...vms);
+
+            /**
+             * Adds the specified set of virtual machines, assuming they are from the same
+             * availability set, to the specified backend of this load balancer.
+             * <p>
+             * If an existing backend with the provided name does not exist on this load balancer, it will be created.
+             * If the name is null, a new name for the backend will be generated automatically.
+             * <p>
+             * Reference to the primary IP configurations of the primary network interfaces of each of the provided set of
+             * virtual machines will be added to the backend.
+             * <p>
+             * If the virtual machines are not in the same availability set, the load balancer will still
+             * be created, but the virtual machines will not be associated it.
+             * <p>
+             * Only those virtual machines will be associated with the load balancer that already have an existing
+             * network interface. Virtual machines without a network interface will be skipped.
+             * @param backendName the name of the backend to associate the virtual machines with
+             * @param vms existing virtual machines
+             * @return the next stage of the update
+             */
+            WithBackendOrCreate withExistingVirtualMachines(String backendName, SupportsNetworkInterfaces...vms);
         }
 
         /**
@@ -299,6 +329,25 @@ interface WithCreateAndRule extends
      * Grouping of load balancer update stages.
      */
     interface UpdateStages {
+        /**
+         * The stage of the load balancer update allowing to add or remove backends.
+         */
+        interface WithBackend {
+            /**
+             * Adds a new empty backend to the load balancer.
+             * @param name the name to assign to the backend
+             * @return the next stage of the update
+             */
+            Update withBackend(String name);
+
+            /**
+             * Removes the specified backend from the load balancer.
+             * @param name the name of the backend to remove
+             * @return the next stage of the update
+             */
+            Update withoutBackend(String name);
+        }
+
         /**
          * The stage of the load balancer update allowing to add, remove or modify probes.
          */
@@ -396,6 +445,7 @@ interface WithProbe {
     interface Update extends
         Appliable<LoadBalancer>,
         Resource.UpdateWithTags<Update>,
-        UpdateStages.WithProbe {
+        UpdateStages.WithProbe,
+        UpdateStages.WithBackend {
     }
 }

azure-mgmt-network/src/main/java/com/microsoft/azure/management/network/implementation/LoadBalancerImpl.java

@@ -7,8 +7,10 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.TreeMap;
 
 import com.microsoft.azure.SubResource;
@@ -46,7 +48,7 @@ class LoadBalancerImpl
         LoadBalancer.Update {
 
     private final LoadBalancersInner innerCollection;
-    private final List<SupportsNetworkInterfaces> vms = new ArrayList<>();
+    private final HashMap<String, String> nicsInBackends = new HashMap<>();
     private List<String> creatablePIPKeys = new ArrayList<>();
     private final TreeMap<String, Backend> backends = new TreeMap<>();
     private final TreeMap<String, TcpProbe> tcpProbes = new TreeMap<>();
@@ -181,22 +183,14 @@ public void success(ServiceResponse<LoadBalancerInner> response) {
     }
 
     private void beforeCreating()  {
-        // Ensure backend pools list
-        List<BackendAddressPoolInner> backendPools = this.inner().backendAddressPools();
-        if (backendPools == null) {
-            backendPools = new ArrayList<>();
-            this.inner().withBackendAddressPools(backendPools);
-        }
-
-        // Ensure first backend pool
-        BackendAddressPoolInner backendPool;
-        if (backendPools.size() == 0) {
-            backendPool = new BackendAddressPoolInner()
-                    .withName("backendpool" + (backendPools.size() + 1));
-            backendPools.add(backendPool);
+        // Ensure existence of backends referred to by VMs
+        for (String backendName : this.nicsInBackends.values()) {
+            if (!this.backends().containsKey(backendName)) {
+                this.withBackend(backendName);
+            }
         }
 
-        // Account for the newly created public IPs
+        // Account for the newly created public IPs5
         for (String pipKey : this.creatablePIPKeys) {
             PublicIpAddress pip = (PublicIpAddress) this.createdResource(pipKey);
             if (pip != null) {
@@ -237,18 +231,20 @@ private void beforeCreating()  {
 
     private void afterCreating() throws Exception {
         // Update the NICs to point to the backend pool
-        for (SupportsNetworkInterfaces vm : this.vms) {
-            NetworkInterface primaryNIC = vm.primaryNetworkInterface();
-            NicIpConfiguration nicIp = primaryNIC.primaryIpConfiguration();
-            primaryNIC.update()
+        for (Entry<String, String> nicInBackend : this.nicsInBackends.entrySet()) {
+            String nicId = nicInBackend.getKey();
+            String backendName = nicInBackend.getValue();
+            NetworkInterface nic = this.myManager().networkInterfaces().getById(nicId);
+            NicIpConfiguration nicIp = nic.primaryIpConfiguration();
+            nic.update()
                 .updateIpConfiguration(nicIp.name())
                     .withExistingLoadBalancer(this)
-                    .withBackendAddressPool(this.inner().backendAddressPools().get(0).name())
+                    .withBackendAddressPool(backendName)
                     .parent()
                 .apply();
         }
 
-        this.vms.clear();
+        this.nicsInBackends.clear();
         this.refresh();
     }
 
@@ -348,15 +344,27 @@ public LoadBalancerImpl withExistingPublicIpAddresses(PublicIpAddress... publicI
         return this;
     }
 
-    private LoadBalancerImpl withExistingVirtualMachine(SupportsNetworkInterfaces vm) {
-      this.vms.add(vm);
-      return this;
+    private LoadBalancerImpl withExistingVirtualMachine(SupportsNetworkInterfaces vm, String backendName) {
+        if (vm.primaryNetworkInterfaceId() != null) {
+            this.nicsInBackends.put(vm.primaryNetworkInterfaceId(), backendName.toLowerCase());
+        }
+        return this;
     }
 
     @Override
     public LoadBalancerImpl withExistingVirtualMachines(SupportsNetworkInterfaces... vms) {
-        for (SupportsNetworkInterfaces vm : vms) {
-            withExistingVirtualMachine(vm);
+        return this.withExistingVirtualMachines(null, vms);
+    }
+
+    @Override public LoadBalancerImpl withExistingVirtualMachines(String backendName, SupportsNetworkInterfaces... vms) {
+        if (backendName == null) {
+            backendName = "backend" + (this.backends().size() + 1);
+        }
+
+        if (vms != null) {
+            for (SupportsNetworkInterfaces vm : vms) {
+                withExistingVirtualMachine(vm, backendName);
+            }
         }
         return this;
     }
@@ -457,6 +465,32 @@ public LoadBalancerImpl withoutProbe(Probe probe) {
         return this.withoutProbe(probe.name());
     }
 
+    @Override
+    public LoadBalancerImpl withBackend(String name) {
+        BackendAddressPoolInner inner = new BackendAddressPoolInner()
+                .withName(name);
+        BackendImpl backend = new BackendImpl(inner.name(), inner, this);
+        ensureInnerBackends().add(inner);
+        this.backends.put(inner.name(), backend);
+        return this;
+    }
+
+    @Override
+    public LoadBalancerImpl withoutBackend(String name) {
+        this.backends.remove(name);
+        List<BackendAddressPoolInner> inners = this.inner().backendAddressPools();
+        if (inners != null) {
+            for (int i = 0; i < inners.size(); i++) {
+                if (inners.get(i).name().equalsIgnoreCase(name)) {
+                    inners.remove(i);
+                    break;
+                }
+            }
+        }
+
+        return this;
+    }
+
     // Getters
 
     @Override

azure/src/test/java/com/microsoft/azure/TestLoadBalancer.java

@@ -126,6 +126,7 @@ public LoadBalancer createResource(LoadBalancers resources) throws Exception {
                 .withExistingResourceGroup(groupName)
                 .withExistingPublicIpAddresses(pip1)
                 .withExistingVirtualMachines(existingVMs)
+                .withBackend("backend2")
                 .withTcpProbe(80, "tcp1")
                 .withLoadBalancedPort(80, TransportProtocol.TCP)
                 .defineTcpProbe("tcp2")
@@ -140,7 +141,8 @@ public LoadBalancer createResource(LoadBalancers resources) throws Exception {
                     .attach()
                 .create();
 
-        Assert.assertTrue(lb.backends().size() == 1);
+        Assert.assertTrue(lb.backends().size() == 2);
+        Assert.assertTrue(lb.backends().containsKey("backend2"));
         return lb;
     }
 
@@ -164,13 +166,17 @@ public LoadBalancer updateResource(LoadBalancer resource) throws Exception {
                     .withIntervalInSeconds(14)
                     .withNumberOfProbes(5)
                     .parent()
+                .withBackend("backend3")
+                .withoutBackend("backend2")
                 .withTag("tag1", "value1")
                 .withTag("tag2", "value2")
                 .apply();
         Assert.assertTrue(resource.tags().containsKey("tag1"));
         Assert.assertTrue(resource.httpProbes().containsKey("http2"));
         Assert.assertTrue(resource.tcpProbes().containsKey("tcp3"));
         Assert.assertTrue(!resource.tcpProbes().containsKey("tcp2"));
+        Assert.assertTrue(resource.backends().containsKey("backend3"));
+        Assert.assertTrue(!resource.backends().containsKey("backend2"));
 
         return resource;
     }