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

Signed-off-by: Prince Pereira <ppereira@microsoft.com>
microsoft · Jun 14, 2024 · 507cff5 · 507cff5
1 parent 7b891f7
commit 507cff5
Show file tree

Hide file tree

Showing 3 changed files with 143 additions and 2 deletions.
diff --git a/hcn/hcnloadbalancer.go b/hcn/hcnloadbalancer.go
@@ -28,14 +28,14 @@ type HostComputeLoadBalancer struct {
 	Flags                LoadBalancerFlags         `json:",omitempty"` // 0: None, 1: EnableDirectServerReturn
 }
 
-//LoadBalancerFlags modify settings for a loadbalancer.
+// LoadBalancerFlags modify settings for a loadbalancer.
 type LoadBalancerFlags uint32
 
 var (
 	// LoadBalancerFlagsNone is the default.
 	LoadBalancerFlagsNone LoadBalancerFlags = 0
 	// LoadBalancerFlagsDSR enables Direct Server Return (DSR)
-	LoadBalancerFlagsDSR LoadBalancerFlags = 1
+	LoadBalancerFlagsDSR  LoadBalancerFlags = 1
 	LoadBalancerFlagsIPv6 LoadBalancerFlags = 2
 )
 
@@ -161,6 +161,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 {
@@ -235,6 +278,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
@@ -41,6 +41,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
@@ -319,6 +319,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",