alpine-overlay/kubezero/zdt-base/uniq_hostname.py

#!/usr/bin/env python3

# use pyminify: pyminifier --obfuscate-variables $0 > minified_$0

import sys
import boto3
import time

my_tag = 'CustomHostname'

ec2_client = boto3.client('ec2')


def get_tag(resource, tag_name):
    try:
        for tag in resource['Tags']:
            if tag['Key'] == tag_name:
                return tag['Value']

    except KeyError:
        pass

    return None


def enumerate_instances(launchtemplate_id, hostname_format):
    """  Search for all names in use as well as all instances in flight """
    in_use = []
    in_flight = []

    try:
        # look for other instances sharing the same launch template which works across multiple ASGs eg. kube-control
        for r in ec2_client.describe_instances(Filters=[{'Name': 'instance-state-name', 'Values': ['pending', 'running', 'stopping', 'stopped']}, {'Name': 'tag:aws:ec2launchtemplate:id', 'Values': [launchtemplate_id]}])['Reservations']:
            for instance in r["Instances"]:
                # If an instance already has the tag mark index as used
                hostname = get_tag(instance, my_tag)
                if hostname:
                    in_use.append(hostname)
                else:
                    in_flight.append(instance['InstanceId'])

    except (KeyError, IndexError):
        pass

    in_flight = sorted(in_flight)

    return in_use, in_flight


# Test format string
try:
    launchtemplate_id = sys.argv[1]
    my_instance_id = sys.argv[2]
    hostname_format = sys.argv[3]

    hostname_format.format(1)
    print("Using {0} as Hostname format string".format(
        hostname_format), file=sys.stderr)
except (IndexError):
    print("Invalid or missing format string, aborting", file=sys.stderr)
    sys.exit(0)

timeout = 0
while True:
    in_use, in_flight = enumerate_instances(launchtemplate_id, hostname_format)

    # Ideally in_flight has only one element, my_instance_id
    # otherwise we have a race condition
    # if my_instance_id is the first element in the sorted list we proceed to break the race
    # otherwise wait for 5 seconds and try again
    if my_instance_id not in in_flight:
        print("Seems like instance {0} is already in_use, aborting".format(
            my_instance_id), file=sys.stderr)
        break

    if in_flight[0] == my_instance_id:
        # Now lets find the lowest free index
        index = 0
        my_hostname = hostname_format.format(index)
        for name in sorted(in_use):
            if name != my_hostname:
                break
            index += 1
            my_hostname = hostname_format.format(index)

        print("Assigning {0} to {1}".format(
            my_hostname, my_instance_id), file=sys.stderr)

        ec2_client.create_tags(Resources=[str(my_instance_id)], Tags=[
                               {'Key': str(my_tag), 'Value': str(my_hostname)}])
        print(my_hostname)
        break

    print("{0} are in flight and we are not the first, {1} are currently in use, retrying...".format(
        in_flight, in_use), file=sys.stderr)
    timeout += 2

    if timeout > 300:
        print("Parallel booting instances did not settle within time limit (300s). Giving up!", file=sys.stderr)
        break

    time.sleep(2)