Coordinating singleton token between AWS Lambda instances

Personally, I would try to de-couple read and write operations of the token. There should only be one "process" that writes tokens and multiple "processes" that read the token from your storage.

One solution could be for example, that you use AWS EventBridge to periodically trigger a Lambda (say every 30 minutes). This Lambda calls the /auth endpoint, gets a fresh token and then stores it in DynamoDB.

The Lambdas with the business logic then should use consistent/strong reads to avoid race conditions (writing a new token while another Lambda reads a now stale token).

The business logic in your message Lambda just has to read the token from the database now, allowing you to "scale" to as many Lambdas as you want without them getting in each others way.

While there is no native locking implemented in DynamoDB, here is a 2017 AWS blog post showcasing DynamoDB locking lib for Java SDK. https://aws.amazon.com/blogs/database/building-distributed-locks-with-the-dynamodb-lock-client/

This excerpt shines some light on the mechanism:

The DynamoDB Lock Client uses the DynamoDB UpdateItem API to heartbeat and extend locks each host owns. Additionally, the lock client uses client-side TTL to expire locks. The lock client uses a recent version of the AWS SDK for Java. The locking protocol employs conditional updates extensively and throughout the lifecycle of a lock (creation, renewal by heartbeat, and expiration by deletion).

There are some forks and adaptations for other languages like Python.

I'd like to offer a slightly different take on the @Jens' answer:

1. Use DynamoDB to store the current auth token: A DDB table with a single entry can be used as an authoritative source for the auth token. More on how to update the tokens at #3 below.

2. Local caching for the auth token: While Lambda is supposed to be stateless, it does actually have the ability to do some local caching, so each Lambda instance should follow some logic like this:

   • if it has a NULL value for the auth token in the local cache, read it from DDB and update the local cache
   • if it has a non-NULL value for the auth token, then attempt to use the auth token to make requests
   • if a request fails due to invalidated token, attempt to do a token refresh (see #3 below)

3. Take advantage of conditional writes to refresh tokens: For this, any lambda instance that finds a bad token should attempt to "designate itself" to refresh the token by doing a conditional write on the token record in the DDB table, attempting to claim itself as the self-designated updater

   • if the conditional write succeeds, the lambda continues to make the POST /auth request, get a new token and subsequently update the DDB record and removes itself from the lock;
   • if the conditional write fails, the lambda should back-off for some time (maybe 200 ms or so) and then attempt to re-read the record to see if another lambda successfully refreshed the token in the mean time
   • if upon reading the record the lambda determines that the token had been refreshed (ie. token value is different than what it had), it can proceed to using the new token and all is good; if it find that the record is still locked, it can wait some more and retry
   • you should probably come up with a safe back-off period of time (say 1-2 seconds) after which a lambda should "take over" the responsibility of updating the token, under the assumption that the other instance that had previously locked it has been terminated for whatever reason; this was you don't end up in a deadlock where all instances wait forever

This solution has the advantage of reducing the need to read from DDB for each request and not requiring any other services (other the Lambda and DDB).

When performing the conditional write, lambda can generate a UUID for the operation and store this UUID and a timestamp on the items to indicate the time when the item was "locked". Subsequent instances will read the UUID and timestamp and determine that some other instance must be working on the item. And if too much time has passed since the item was locked, another instance can "take over".

I ended up implementing this with a distributed-lock algorithm, hinging on DynamoDB's conditional writes, as largely outlined in Mike's and Oleksii's answers. I'll provide some of the details below (implemented in Python's boto3 API, omitting a number of implementation-specific details):

1. Read the token from DynamoDB and store it locally.

2. If a request made using a token is being rejected, mark the token as "invalid" internally.

3. If no token exists or it is known to be invalid, follow this procedure:

   if old_token:
       condition = Attr('AccessToken').eq(old_token)
   else:
       condition = Attr('AccessToken').not_exists()
   
   release_code = uuid4().hex
   now = int(time())
   lock = {'release_code': release_code, 'expires': now + 15}
   condition &= Attr('AccessTokenUpdateLock').not_exists() | Attr('AccessTokenUpdateLock.expires').lte(now)
   
   table.update_item(
       Key=key,
       UpdateExpression='SET AccessTokenUpdateLock = :l',
       ExpressionAttributeValues={':l': lock},
       ConditionExpression=condition,
       ReturnValues='ALL_NEW'
   )
   

   In a nutshell, this adds an AccessTokenUpdateLock property to the specific bot's record in the table, if no such lock already exists or it has already expired. The lock as such is basically a random string. If this succeeds, it is safe to request a new token from the API. Otherwise this will fail with a DynamoDB.Client.exceptions.ConditionalCheckFailedException.

   In the latter case, the client goes into a exponentially backing-off read loop:

   @backoff.on_predicate(backoff.expo, jitter=backoff.full_jitter, max_value=10, max_time=60)
   def await_updated_access_token():
       item = table.get_item(Key=key, ConsistentRead=True)
   
       if item['Item'].get('AccessToken') != old_token:
           # store new token
           return # new token
   
       try:
           if item['Item']['AccessTokenUpdateLock']['expires'] <= time():
               raise AccessTokenUpdateHasFailed
       except KeyError:
           return None
   

   The backoff library will keep repeating this until the function returns a token or raises an exception. If AccessTokenUpdateHasFailed is raised, it starts again and tries to acquire a lock itself.

   If the client did successfully acquire a lock, it then fetches a new API token, and writes it to DynamoDB:

   if old_token:
       condition = Attr('AccessToken').eq(old_token)
   else:
       condition = Attr('AccessToken').not_exists()
   
   condition &= Attr('AccessTokenUpdateLock.release_code').eq(release_code)
   
   response = table.update_item(
       Key=key,
       UpdateExpression='SET AccessToken = :t REMOVE AccessTokenUpdateLock',
       ExpressionAttributeValues={':t': new_token},
       ConditionExpression=condition,
       ReturnValues='ALL_NEW'
   )
   

This appears to be reasonably robust with minimal retries to the API itself.