Merge pull request #7 from chiradeep/implement_ttl

Use TTL feature of Dynamodb to automatically clear rows after last us…

Author: chiradeep

.gitignore

@@ -2,3 +2,4 @@
 /dist/
 /*.egg-info
 /*.egg
+/*.vscode

README.md

@@ -68,7 +68,7 @@ Although the code is general-purpose and can be used outside of AWS Lambda, note
 * Does not detect/prevent deadlocks. There is no spin lock, but the mutex user could create one by spinning until an acquire succeeds. Within a Lambda function, one should avoid taking more than 1 lock.
 * Not re-entrant. If a thread (e.g.,a lambda function) tries to re-acquire a lock it already holds, it will block
 * Not designed for speed. The DynamoDb table backing the locks is generally provisioned as low throughput (2 ops/sec)
-* No cleanup. Each named mutex gets a row in the Dynamodb table. These rows are never cleaned up even if nobody is using the mutex. It should be trivial to write a lambda to vaccuum these rows, however, given the small size of the row and the infinite capacity of DynamoDb, it is not likely to be a problem, if used for its intended purpose.
+* Cleanup - the mutex table is created with the [TTL](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/TTL.html) feature. By default, the TTL is 2 days; after 2 days of disuse, the row created to hold the mutex should get automatically deleted. You can use set the `ttl_minutes` in the constructor to choose a different TTL
 
 
 # TODO

dyndbmutex/dyndbmutex.py

@@ -20,7 +20,7 @@ def setup_logging():
 
 DEFAULT_MUTEX_TABLE_NAME = 'Mutex'
 NO_HOLDER = '__empty__'
-
+TWO_DAYS_IN_MINUTES = 2*24*60
 
 class AcquireLockFailedError(Exception):
         pass
@@ -33,11 +33,12 @@ def timestamp_millis():
 
 class MutexTable:
 
-    def __init__(self, region_name='us-west-2'):
+    def __init__(self, region_name='us-west-2', ttl_minutes=TWO_DAYS_IN_MINUTES):
         self.dbresource = boto3.resource('dynamodb', region_name=region_name)
         self.dbclient = boto3.client('dynamodb', region_name=region_name)
         self.table_name = os.environ.get('DD_MUTEX_TABLE_NAME', DEFAULT_MUTEX_TABLE_NAME)
         logger.info("Mutex table name is " + self.table_name)
+        self.ttl_minutes = ttl_minutes
         self.get_table()
 
     def get_table(self):
@@ -53,7 +54,10 @@ def get_table(self):
 
     def delete_table(self):
         self.dbclient.delete_table(TableName=self.table_name)
-        logger.debug("Deleted table")
+        logger.info("Deleted table " + self.table_name)
+
+    def get_lock(self, lockname):
+        return self.get_table().get_item(Key={'lockname': lockname})
 
     def create_table(self):
         try:
@@ -84,26 +88,37 @@ def create_table(self):
         else:
             logger.debug("Called create_table")
             table.wait_until_exists()
-            logger.info("Created table")
+            logger.info("Created table " + self.table_name)
+            try: 
+                self.dbclient.update_time_to_live(
+                TableName=self.table_name,
+                TimeToLiveSpecification={
+                 'Enabled': True,
+                 'AttributeName': 'ttl'
+                })
+            except botocore.exceptions.ClientError as e:
+                logger.error("Error setting TTL on table", exc_info=e)
             return table
 
     def write_lock_item(self, lockname, caller, waitms):
         expire_ts = timestamp_millis() + waitms
+        ttl = expire_ts + self.ttl_minutes*60
         logger.debug("Write_item: lockname=" + lockname + ", caller=" +
                      caller + ", Expire time is " + str(expire_ts))
         try:
             self.get_table().put_item(
                 Item={
                     'lockname': lockname,
                     'expire_ts': expire_ts,
-                    'holder': caller
+                    'holder': caller,
+                    'ttl': ttl
                 },
                 # TODO: adding Attr("holder").eq(caller) should make it re-entrant
                 ConditionExpression=Attr("holder").eq(NO_HOLDER) | Attr('lockname').not_exists()
             )
         except botocore.exceptions.ClientError as e:
             if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
-                logger.warn("Write_item: lockname=" + lockname +
+                logger.warning("Write_item: lockname=" + lockname +
                              ", caller=" + caller + ", lock is being held")
                 return False
         logger.debug("Write_item: lockname=" + lockname +
@@ -122,7 +137,7 @@ def clear_lock_item(self, lockname, caller):
             )
         except botocore.exceptions.ClientError as e:
             if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
-                logger.warn("clear_lock_item: lockname=" + lockname + ", caller=" + caller +
+                logger.warning("clear_lock_item: lockname=" + lockname + ", caller=" + caller +
                              " release failed")
                 return False
         logger.debug("clear_lock_item: lockname=" + lockname + ", caller=" + caller + " release succeeded")
@@ -143,7 +158,7 @@ def prune_expired(self, lockname, caller):
             )
         except botocore.exceptions.ClientError as e:
             if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
-                logger.warn("Prune: lockname=" + lockname + ", caller=" + caller +
+                logger.warning("Prune: lockname=" + lockname + ", caller=" + caller +
                              " Prune failed")
                 return False
         logger.debug("Prune: lockname=" + lockname + ", caller=" + caller + " Prune succeeded")
@@ -153,13 +168,13 @@ def prune_expired(self, lockname, caller):
 class DynamoDbMutex:
 
     def __init__(self, name, holder=None,
-                 timeoutms=30 * 1000, region_name='us-west-2'):
+                 timeoutms=30 * 1000, region_name='us-west-2', ttl_minutes=TWO_DAYS_IN_MINUTES):
         if holder is None:
             holder = str(uuid.uuid4())
         self.lockname = name
         self.holder = holder
         self.timeoutms = timeoutms
-        self.table = MutexTable(region_name=region_name)
+        self.table = MutexTable(region_name=region_name, ttl_minutes=ttl_minutes)
         self.locked = False
 
     def lock(self):
@@ -184,6 +199,9 @@ def __exit__(self, exc_type, exc_value, traceback):
 
     def is_locked(self):
         return self.locked
+    
+    def get_raw_lock(self):
+        return self.table.get_lock(self.lockname)
 
     @staticmethod
     def delete_table(region_name='us-west-2'):

scripts/create-table.sh

@@ -9,3 +9,4 @@ aws dynamodb create-table \
         AttributeName=lockname,AttributeType=S \
     --key-schema AttributeName=lockname,KeyType=HASH \
     --provisioned-throughput ReadCapacityUnits=2,WriteCapacityUnits=2
+aws dynamodb update-time-to-live --table-name  $MUTEX_TABLE_NAME --time-to-live-specification "Enabled=true, AttributeName=ttl"

test/test_mutex.py

@@ -2,6 +2,7 @@
 import unittest
 import base64
 import sys
+import warnings
 import os
 from os import path
 
@@ -22,6 +23,10 @@ def setUpClass(cls):
         setup_logging()
         super(TestDynamoDbMutex, cls).setUpClass()
 
+    def setUp(self):
+        #python 3.6 emits warnings due to requests which can be ignoredz    
+        warnings.filterwarnings("ignore", category=ResourceWarning, message="unclosed.*<ssl.SSLSocket.*>")
+
     def test_create(self):
         m = DynamoDbMutex(random_name(), "myself", 3 * 1000)
         assert(m.lock())
@@ -38,13 +43,13 @@ def test_timeout(self):
         m = DynamoDbMutex(random_name(), "myself", 3 * 1000)
         m.lock()
         time.sleep(5)
-        assert(m.lock())
+        self.assertTrue(m.lock())
         m.release()
 
     def test_mutual_exclusion(self):
         m = DynamoDbMutex(random_name(), holder=random_name())
         m.lock()
-        assert(m.lock() == False)
+        self.assertFalse(m.lock())
         m.release()
 
     def test_with(self):
@@ -55,7 +60,7 @@ def test_with(self):
                 raise
         except:
             print("In exception handler")
-            assert(m.is_locked() == False)
+            self.assertFalse(m.is_locked())
 
     def test_with_fail(self):
         name = random_name()
@@ -68,9 +73,9 @@ def test_with_fail(self):
                 time.sleep(3)
         except AcquireLockFailedError:
             print("In exception handler")
-            assert(m2.is_locked() == False)
+            self.assertFalse(m2.is_locked())
             exceptionHappened = True
-        assert(exceptionHappened)
+        self.assertTrue(exceptionHappened)
 
     def test_release_expired(self):
         name = random_name()
@@ -80,11 +85,24 @@ def test_release_expired(self):
         time.sleep(3)
         caller = "caller2"
         m2 = DynamoDbMutex(name=name, holder=caller, timeoutms=2 * 1000)
-        assert(m2.lock())
+        self.assertTrue(m2.lock())
         m1.release()
-        assert(m2.is_locked())
+        self.assertTrue(m2.is_locked())
         m2.release()
 
+    def test_ttl(self):
+        os.environ['DD_MUTEX_TABLE_NAME'] = random_name()
+        name = random_name()
+        caller = "myself"
+        m1 = DynamoDbMutex(name=name, holder=caller, timeoutms=2 * 1000, ttl_minutes=2)
+        m1.lock()
+        item = m1.get_raw_lock()
+        diff = item['Item']['ttl'] - item['Item']['expire_ts']
+        print (diff)
+        self.assertTrue(diff > 0 and diff <= 2*60)
+        DynamoDbMutex.delete_table()
+        del os.environ['DD_MUTEX_TABLE_NAME']
+
     def test_table_creation_race(self):
         old_table_name = os.environ.get('DD_MUTEX_TABLE_NAME')
         try:
@@ -95,7 +113,7 @@ def test_table_creation_race(self):
             def create_mutex(i):
                 mutex_name = random_name()
                 mutex = cls(name=mutex_name, holder='caller%i' % i)
-                self.assertEquals(mutex.table.table_name, table_name)
+                self.assertEqual(mutex.table.table_name, table_name)
                 return mutex
 
             num_threads = 2