103 lines
3.2 KiB
Python
Executable File
103 lines
3.2 KiB
Python
Executable File
#!/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)
|