Adding support for loadbalancer policy update in hns. (#2085)

princepereira · web-flow · commit 0f7d8de9948c · 2024-04-15T11:26:12.000-04:00
Signed-off-by: Prince Pereira &lt;ppereira@microsoft.com&gt;
diff --git a/hcn/hcnloadbalancer.go b/hcn/hcnloadbalancer.go
@@ -163,6 +163,49 @@ func createLoadBalancer(settings string) (*HostComputeLoadBalancer, error) {
 	return &outputLoadBalancer, nil
 }
 
+func updateLoadBalancer(loadbalancerId string, settings string) (*HostComputeLoadBalancer, error) {
+	loadBalancerGuid, err := guid.FromString(loadbalancerId)
+	if err != nil {
+		return nil, errInvalidLoadBalancerID
+	}
+	// Update loadBalancer.
+	var (
+		loadBalancerHandle hcnLoadBalancer
+		resultBuffer       *uint16
+		propertiesBuffer   *uint16
+	)
+	hr := hcnOpenLoadBalancer(&loadBalancerGuid, &loadBalancerHandle, &resultBuffer)
+	if err := checkForErrors("hcnOpenLoadBalancer", hr, resultBuffer); err != nil {
+		return nil, err
+	}
+	hr = hcnModifyLoadBalancer(loadBalancerHandle, settings, &resultBuffer)
+	if err := checkForErrors("hcnModifyLoadBalancer", hr, resultBuffer); err != nil {
+		return nil, err
+	}
+	// Query loadBalancer.
+	hcnQuery := defaultQuery()
+	query, err := json.Marshal(hcnQuery)
+	if err != nil {
+		return nil, err
+	}
+	hr = hcnQueryLoadBalancerProperties(loadBalancerHandle, string(query), &propertiesBuffer, &resultBuffer)
+	if err := checkForErrors("hcnQueryLoadBalancerProperties", hr, resultBuffer); err != nil {
+		return nil, err
+	}
+	properties := interop.ConvertAndFreeCoTaskMemString(propertiesBuffer)
+	// Close loadBalancer.
+	hr = hcnCloseLoadBalancer(loadBalancerHandle)
+	if err := checkForErrors("hcnCloseLoadBalancer", hr, nil); err != nil {
+		return nil, err
+	}
+	// Convert output to HostComputeLoadBalancer
+	var outputLoadBalancer HostComputeLoadBalancer
+	if err := json.Unmarshal([]byte(properties), &outputLoadBalancer); err != nil {
+		return nil, err
+	}
+	return &outputLoadBalancer, nil
+}
+
 func deleteLoadBalancer(loadBalancerID string) error {
 	loadBalancerGUID, err := guid.FromString(loadBalancerID)
 	if err != nil {
@@ -237,6 +280,23 @@ func (loadBalancer *HostComputeLoadBalancer) Create() (*HostComputeLoadBalancer,
 	return loadBalancer, nil
 }
 
+// Update Loadbalancer.
+func (loadBalancer *HostComputeLoadBalancer) Update(hnsLoadbalancerID string) (*HostComputeLoadBalancer, error) {
+	logrus.Debugf("hcn::HostComputeLoadBalancer::Create id=%s", hnsLoadbalancerID)
+
+	jsonString, err := json.Marshal(loadBalancer)
+	if err != nil {
+		return nil, err
+	}
+
+	logrus.Debugf("hcn::HostComputeLoadBalancer::Update JSON: %s", jsonString)
+	loadBalancer, hcnErr := updateLoadBalancer(hnsLoadbalancerID, string(jsonString))
+	if hcnErr != nil {
+		return nil, hcnErr
+	}
+	return loadBalancer, nil
+}
+
 // Delete LoadBalancer.
 func (loadBalancer *HostComputeLoadBalancer) Delete() error {
 	logrus.Debugf("hcn::HostComputeLoadBalancer::Delete id=%s", loadBalancer.Id)
diff --git a/hcn/hcnloadbalancer_test.go b/hcn/hcnloadbalancer_test.go
@@ -42,6 +42,70 @@ func TestCreateDeleteLoadBalancer(t *testing.T) {
 	}
 }
 
+func TestCreateUpdateDeleteLoadBalancer(t *testing.T) {
+	network, err := CreateTestOverlayNetwork()
+	if err != nil {
+		t.Fatal(err)
+	}
+	endpoint, err := HcnCreateTestEndpoint(network)
+	if err != nil {
+		t.Fatal(err)
+	}
+	loadBalancer, err := HcnCreateTestLoadBalancer(endpoint)
+	if err != nil {
+		t.Fatal(err)
+	}
+	jsonString, err := json.Marshal(loadBalancer)
+	if err != nil {
+		t.Fatal(err)
+	}
+	fmt.Printf("LoadBalancer JSON:\n%s \n", jsonString)
+
+	secondEndpoint, err := HcnCreateTestEndpoint(network)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	HcnLoadBalancerTestAddBackend(loadBalancer, secondEndpoint.Id)
+
+	loadBalancer, err = loadBalancer.Update(loadBalancer.Id)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(loadBalancer.HostComputeEndpoints) != 2 {
+		t.Fatalf("Update loadBalancer with backend add failed")
+	}
+
+	HcnLoadBalancerTestRemoveBackend(loadBalancer, secondEndpoint.Id)
+
+	loadBalancer, err = loadBalancer.Update(loadBalancer.Id)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(loadBalancer.HostComputeEndpoints) != 1 {
+		t.Fatalf("Update loadBalancer with backend remove failed")
+	}
+
+	err = loadBalancer.Delete()
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = secondEndpoint.Delete()
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = endpoint.Delete()
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = network.Delete()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestGetLoadBalancerById(t *testing.T) {
 	network, err := CreateTestOverlayNetwork()
 	if err != nil {
diff --git a/hcn/hcnutils_test.go b/hcn/hcnutils_test.go
@@ -317,6 +317,23 @@ func HcnCreateTestLoadBalancer(endpoint *HostComputeEndpoint) (*HostComputeLoadB
 	return loadBalancer.Create()
 }
 
+func HcnLoadBalancerTestAddBackend(loadBalancer *HostComputeLoadBalancer, endpointId string) {
+	endpointIds := loadBalancer.HostComputeEndpoints
+	endpointIds = append(endpointIds, endpointId)
+	loadBalancer.HostComputeEndpoints = endpointIds
+}
+
+func HcnLoadBalancerTestRemoveBackend(loadBalancer *HostComputeLoadBalancer, endpointId string) {
+	endpointIds := loadBalancer.HostComputeEndpoints
+	for i, v := range endpointIds {
+		if v == endpointId {
+			endpointIds = append(endpointIds[:i], endpointIds[i+1:]...)
+			break
+		}
+	}
+	loadBalancer.HostComputeEndpoints = endpointIds
+}
+
 func HcnCreateTestRemoteSubnetRoute() (*PolicyNetworkRequest, error) {
 	rsr := RemoteSubnetRoutePolicySetting{
 		DestinationPrefix:           "192.168.2.0/24",
