node-external-ip-controller/node_external_ip_controller_async.py

import asyncio
from kubernetes_asyncio import client, config, watch

# Configuration
SERVICE_NAME_PATTERN = (
    "traefik"  # Define the service name pattern or label to identify services
)
NAMESPACE = "kube-system"
ANNOTATION_KEY = "kube-vip.io/loadbalancerIPs"
ZERO_GATEWAY_IP = "172.28.10.1"
NODE_LABEL = "svccontroller.k3s.cattle.io/enablelb=true"


async def update_service_annotation(v1, service_name, external_ip):
    try:
        # Get the current service object
        service = await v1.read_namespaced_service(service_name, NAMESPACE)

        # Check if the annotation needs to be updated
        current_annotation = service.metadata.annotations.get(ANNOTATION_KEY)
        target_annotation = ",".join({external_ip, ZERO_GATEWAY_IP})
        if current_annotation != target_annotation:
            # Update the annotation
            body = {"metadata": {"annotations": {ANNOTATION_KEY: target_annotation}}}
            await v1.patch_namespaced_service(service_name, NAMESPACE, body)
            print(
                f"Updated service {service_name} with new external IP: {target_annotation}"
            )

    except client.exceptions.ApiException as e:
        print(f"API Exception in update_service_annotation: {e}")


async def watch_nodes():
    await config.load_incluster_config()
    v1 = client.CoreV1Api()
    w = watch.Watch()

    while True:
        try:
            async for event in w.stream(
                v1.list_node, label_selector=NODE_LABEL, _request_timeout=300
            ):
                node = event["object"]
                node_name = node.metadata.name

                # Check for external IP
                external_ip = None
                for address in node.status.addresses:
                    if address.type == "ExternalIP":
                        external_ip = address.address
                        break

                if external_ip:
                    print(f"Detected external IP {external_ip} for node {node_name}")

                    # Get all services that need to be updated with this external IP
                    services = await v1.list_namespaced_service(
                        NAMESPACE, label_selector=SERVICE_NAME_PATTERN
                    )
                    for service in services.items:
                        await update_service_annotation(
                            v1, service.metadata.name, external_ip
                        )

        except client.exceptions.ApiException as e:
            print(f"API Exception in watch_nodes: {e}")
            await asyncio.sleep(5)

        except asyncio.CancelledError:
            print("Watch task was cancelled.")
            break

        except Exception as e:
            print(f"Unexpected error in watch_nodes: {e}")
            await asyncio.sleep(5)


async def watch_services():
    await config.load_incluster_config()
    v1 = client.CoreV1Api()
    w = watch.Watch()

    while True:
        try:
            async for event in w.stream(
                v1.list_namespaced_service,
                NAMESPACE,
                label_selector=SERVICE_NAME_PATTERN,
                _request_timeout=300,
            ):
                if event["type"] == "ADDED":
                    service = event["object"]
                    service_name = service.metadata.name
                    print(f"New service detected: {service_name}")

                    # Fetch current external IP for nodes that match the label
                    nodes = await v1.list_node(label_selector=NODE_LABEL)
                    for node in nodes.items:
                        for address in node.status.addresses:
                            if address.type == "ExternalIP":
                                external_ip = address.address
                                await update_service_annotation(
                                    v1, service_name, external_ip
                                )
                                break

        except client.exceptions.ApiException as e:
            print(f"API Exception in watch_services: {e}")
            await asyncio.sleep(5)

        except asyncio.CancelledError:
            print("Watch task was cancelled.")
            break

        except Exception as e:
            print(f"Unexpected error in watch_services: {e}")
            await asyncio.sleep(5)


async def main():
    await config.load_incluster_config()

    # Run watch_nodes and watch_services concurrently
    await asyncio.gather(watch_nodes(), watch_services())


if __name__ == "__main__":
    asyncio.run(main())