Add first draft

Signed-off-by: Marc Schöchlin <marc.schoechlin@uhurutec.com>

Author: Marc Schöchlin

src/openstack_workload_generator/__main__.py

@@ -20,6 +20,7 @@
     iso_timestamp,
     deep_merge_dict,
 )
+from .entities.loadbalancer import WorkloadGeneratorLoadBalancer
 
 
 def establish_connection():
@@ -180,6 +181,20 @@ def generated_clouds_yaml():
     help="A list of vms to be deleted in the created projects",
 )
 
+exclusive_group_load_balancers = parser.add_mutually_exclusive_group(required=False)
+exclusive_group_load_balancers.add_argument(
+    "--create_load_balancers",
+    action="store_true",
+    help="Create a load balancer per project",
+)
+
+exclusive_group_load_balancers.add_argument(
+    "--delete_load_balancers",
+    action="store_true",
+    help="Delete load balancers per project",
+)
+
+
 args = parser.parse_args()
 
 if args.os_cloud == "":
@@ -217,6 +232,15 @@ def generated_clouds_yaml():
                     workload_project.get_and_create_machines(
                         args.create_machines, args.wait_for_machines
                     )
+
+                    if args.create_load_balancers:
+                        lb = WorkloadGeneratorLoadBalancer(
+                            workload_project
+                        ).get_and_create_load_balancer()
+                        lb.add_members()
+                    if args.delete_load_balancers:
+                        lb = WorkloadGeneratorLoadBalancer(workload_project)
+
                     if args.ansible_inventory:
                         workload_project.dump_inventory_hosts(args.ansible_inventory)
 

src/openstack_workload_generator/entities/helpers.py

@@ -32,6 +32,11 @@ class Config:
         "verify_ssl_certificate": "false",
         "cloud_init_extra_script": """#!/bin/bash\necho "HELLO WORLD"; date > READY; whoami >> READY""",
         "wait_for_server_timeout": "300",
+        "lb_name": "workload-generator-lb",
+        "lb_listener_name": "workload-generator-listener",
+        "lb_pool_name": "workload-generator-pool",
+        "lb_member_tcp_port": "22",
+        "lb_tcp_port": "22",
     }
 
     _file: str | None = None
@@ -142,8 +147,10 @@ def get_number_of_floating_ips_per_project() -> int:
     @staticmethod
     def get_admin_vm_password() -> str:
         if Config.get("admin_vm_password").upper() == "ASK_PASSWORD":
-            Config._config["admin_vm_password"] = getpass.getpass("Enter the wanted admin_vm_password: ")
-        return Config.get("admin_vm_password",  regex=r".{5,}")
+            Config._config["admin_vm_password"] = getpass.getpass(
+                "Enter the wanted admin_vm_password: "
+            )
+        return Config.get("admin_vm_password", regex=r".{5,}")
 
     @staticmethod
     def get_vm_flavor() -> str:
@@ -180,7 +187,9 @@ def get_admin_vm_ssh_key() -> str:
     @staticmethod
     def get_admin_domain_password() -> str:
         if Config.get("admin_domain_password").upper() == "ASK_PASSWORD":
-            Config._config["admin_domain_password"] = getpass.getpass("Enter the wanted admin_domain_password: ")
+            Config._config["admin_domain_password"] = getpass.getpass(
+                "Enter the wanted admin_domain_password: "
+            )
         return Config.get("admin_domain_password", regex=r".{5,}")
 
     @staticmethod
@@ -216,9 +225,29 @@ def quota(quota_name: str, quota_category: str, default_value: int) -> int:
             return default_value
 
     @staticmethod
-    def get_network_mtu():
+    def get_network_mtu() -> int:
         return int(Config.get("network_mtu", regex=r"\d+"))
 
+    @staticmethod
+    def get_lb_name() -> str:
+        return Config.get("lb_name", regex=r"\s+")
+
+    @staticmethod
+    def get_lb_listener_name() -> str:
+        return Config.get("lb_listener_name", regex=r"\s+")
+
+    @staticmethod
+    def get_lb_pool_name() -> str:
+        return Config.get("lb_pool_name", regex=r"\s+")
+
+    @staticmethod
+    def get_lb_member_tcp_port() -> int:
+        return int(Config.get("lb_member_tcp_port", regex=r"\d+"))
+
+    @staticmethod
+    def get_lb_tcp_port() -> int:
+        return int(Config.get("lb_tcp_port", regex=r"\d+"))
+
 
 class DomainCache:
     _domains: dict[str, str] = dict()

src/openstack_workload_generator/entities/loadbalancer.py

@@ -0,0 +1,120 @@
+import logging
+
+from openstack.connection import Connection
+from . import WorkloadGeneratorProject
+from .helpers import Config
+
+LOGGER = logging.getLogger()
+
+
+class WorkloadGeneratorLoadBalancer:
+    def __init__(
+        self,
+        project: WorkloadGeneratorProject,
+    ):
+        self.conn: Connection = project.project_conn
+        self.project = project
+        self.lb = None
+        self.listener = None
+        self.pool = None
+        self.health_monitor = None
+
+    def get_load_balancer_by_name(self, name: str):
+        for lb in self.conn.load_balancer.load_balancers(name=name):
+            if lb.name == name:
+                return lb
+        return None
+
+    def get_listener_by_name(self, name: str):
+        for listener in self.conn.load_balancer.listeners(name=name):
+            if listener.name == name:
+                return listener
+        return None
+
+    def get_pool_by_name(self, name: str):
+        for pool in self.conn.load_balancer.pools(name=name):
+            if pool.name == name:
+                return pool
+        return None
+
+    def get_health_monitor_for_pool(self, pool_id: str):
+        for hm in self.conn.load_balancer.health_monitors():
+            if hm.pool_id == pool_id:
+                return hm
+        return None
+
+    def get_and_create_load_balancer(
+        self,
+    ):
+        self.lb = self.get_load_balancer_by_name(Config.get_lb_name())
+        if self.lb is not None:
+            LOGGER.info(
+                f"Load balancer with name {Config.get_lb_name()} already exists"
+            )
+        else:
+            LOGGER.info(f"Create load balancer with name {Config.get_lb_name()}")
+            self.lb = self.conn.load_balancer.create_load_balancer(
+                name=Config.get_lb_name(),
+                vip_subnet_id=self.project.vip_subnet_id,
+            )
+            self.conn.load_balancer.wait_for_load_balancer(self.lb.id)
+
+        self.listener = self.get_listener_by_name(Config.get_lb_listener_name())
+        if self.listener is not None:
+            LOGGER.info(
+                f"Load balancer listener with name {Config.get_lb_listener_name()} already exists"
+            )
+        else:
+            LOGGER.info(
+                f"Create load balancer listener with name {Config.get_lb_listener_name()} on tcp port {Config.get_lb_listener_port()}"
+            )
+            self.listener = self.conn.load_balancer.create_listener(
+                name=Config.get_lb_listener_name(),
+                loadbalancer_id=self.lb.id,
+                protocol="TCP",
+                protocol_port=Config.get_lb_member_tcp_port(),
+            )
+            self.conn.load_balancer.wait_for_load_balancer(self.lb.id)
+
+        self.pool = self.get_pool_by_name(Config.get_lb_pool_name())
+        if self.pool is not None:
+            LOGGER.info(
+                f"Load balancer pool with name {Config.get_lb_pool_name()} already exists"
+            )
+        else:
+            LOGGER.info(
+                f"Create load balancer pool with name {Config.get_lb_pool_name()}"
+            )
+            self.pool = self.conn.load_balancer.create_pool(
+                name=Config.get_lb_pool_name(),
+                listener_id=self.listener.id,
+                protocol="TCP",
+                lb_algorithm="ROUND_ROBIN",
+            )
+            self.conn.load_balancer.wait_for_load_balancer(self.lb.id)
+
+        self.health_monitor = self.get_health_monitor_for_pool(self.pool.id)
+        if self.health_monitor is not None:
+            LOGGER.info(
+                f"Health monitor for pool with name {Config.get_lb_pool_name()} already exists"
+            )
+        else:
+            self.conn.load_balancer.create_health_monitor(
+                self.pool.id, type="TCP", delay=5, timeout=3, max_retries=3
+            )
+
+    def add_member_to_load_balancer(self, ip: str):
+        if self.lb is None:
+            raise RuntimeError("Loadbalancer has not been created of gathered")
+
+        self.conn.load_balancer.create_member(
+            self.pool.id,
+            address=ip,
+            protocol_port=Config.get_lb_member_tcp_port(),
+            subnet_id=self.project.vip_subnet_id,
+        )
+
+    def add_members_to_load_balancer(self):
+        for machine in self.project.workload_machines:
+            LOGGER.info(f"Adding member to load balancer with ip {machine.internal_ip}")
+            self.add_member_to_load_balancer(machine.internal_ip)