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

103 lines
3.2 KiB
Python
Executable File

#!/usr/bin/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)