Skip to content

Instantly share code, notes, and snippets.

@ashb
Last active May 8, 2017 10:26
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save ashb/14f2f4a27f0654cb4be8 to your computer and use it in GitHub Desktop.
Lock using Amazon SimpleDB - not needed since terraform 0.9
#!/usr/bin/env python
from __future__ import print_function
import argparse
import boto
import boto.provider
import boto.sdb
import getpass
import os
import signal
import subprocess
import sys
import time
import uuid
class SimpleDbLock:
def __init__(self, lockDomain, region_name):
self.db = boto.sdb.connect_to_region(region_name)
self.domain = self.db.create_domain(lockDomain)
def acquireLock(self, name, lockDurationSeconds, acquireTimeoutSeconds):
""" Acquires lock and returns a lockId that can be passed to releaseLock()
name - object to lock - can be any string up to 256 chars in length
lockDurationSeconds - Seconds to hold lock for. Once this number of seconds has
elapsed, the lock will expire and other threads will be able to
acquire a lock for the given name
acquireTimeoutSeconds - Seconds to try to acquire lock. If name is already locked, this
method will sleep/retry until acquireTimeoutSeconds is reached
Returns lockId (string) if lock is acquired. If lock cannot be acquired, throws SystemError"""
lockId = "{}-{}".format(uuid.uuid4(), getpass.getuser())
acquireTimeout = time.time() + acquireTimeoutSeconds
while time.time() < acquireTimeout:
try:
# try to create the lock if it doesn't exist
lockTimeout = time.time() + lockDurationSeconds
if self.db.put_attributes(
self.domain,
name,
{'timeout': lockTimeout, 'lockId': lockId, 'lockedBy': getpass.getuser()},
replace=False,
expected_value=['lockId', False]):
return lockId
except boto.exception.SDBResponseError as e:
if e.status != 404 and e.status != 409:
raise e
# couldn't create lock - check for stale lock
attribs = self.db.get_attributes(self.domain, name, consistent_read=True)
if "timeout" in attribs and float(attribs['timeout']) < time.time():
print("lock timed out - releasing with name: {}".format(attribs['lockId']))
self.releaseLock(name, attribs['lockId']) # lock has timed out, so delete it
time.sleep(0.05) # sleep and retry
# couldn't acquire lock - throw error
if attribs and "lockedBy" in attribs:
raise SystemError("Unable to obtain lock for {} after {} seconds -- currently locked by {}".format(name, acquireTimeoutSeconds, attribs["lockedBy"]))
raise SystemError("Unable to obtain lock for {} after {} seconds".format(name, acquireTimeoutSeconds))
def releaseLock(self, name, lockId):
""" Releases previously acquired lock. name - object to lock. lockId - Lock ID returned from acquireLock() """
print("releaseLock({}, {})".format(name, lockId))
try:
return self.db.delete_attributes(self.domain, name, ['timeout', 'lockId', 'lockedBy'], expected_value=['lockId', lockId])
except boto.exception.SDBResponseError as e:
if e.status == 404 or e.status == 409:
return False
else:
raise e
def forceUnlock(self, name):
attribs = self.db.get_attributes(self.domain, name, consistent_read=True)
if attribs is None or "lockId" not in attribs:
return
self.releaseLock(name, attribs['lockId'])
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--environment', required=True)
parser.add_argument('--simple-db', required=True)
parser.add_argument('--lock-duration', type=int, default=1800, help="Lock timeout in seconds")
parser.add_argument('--acquire-timeout', type=int, default=10, help="Time to wait trying to acquire the lock in seconds")
parser.add_argument('--force-unlock', action='store_true', help="Forcibly remove the lock if something went wrong")
parser.add_argument('command', nargs=argparse.REMAINDER)
args = parser.parse_args()
# boto.sdb doesn't support AWS_REGION natively like most of the other classes do.
profile_name = os.environ.get('AWS_PROFILE', 'default')
default_region = os.environ.get('AWS_DEFAULT_REGION', 'eu-west-1')
profile_or_default_region = boto.provider.get_default().shared_credentials.get(profile_name, 'region', default_region)
region = os.environ.get('AWS_REGION', profile_or_default_region)
r = SimpleDbLock(args.simple_db, region)
lock_name = "terraform-lock-{}".format(args.environment)
if args.force_unlock:
r.forceUnlock(lock_name)
sys.exit(0)
if len(args.command) == 0:
print >> sys.stderr, "Commands not provided"
parser.print_help()
sys.exit(1)
lockId = r.acquireLock(lock_name, args.lock_duration, args.acquire_timeout)
returncode = 0
try:
print("Lock acquired, running command")
# Ignore SIGINT so that our child process gets it instead.
signal.signal(signal.SIGINT, signal.SIG_IGN)
returncode = subprocess.call(args.command)
finally:
r.releaseLock(lock_name, lockId)
sys.exit(returncode)
@ashb
Copy link
Author

ashb commented Dec 7, 2015

Call this with ./lock.py --simple-db my-simpledb-domain terraform ... -- the environment is because we have multiple envs from one state tree.

@ashb
Copy link
Author

ashb commented May 8, 2017

Built in to terraform from 0.9 anyway now

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment