lambda recycler

recycler.py

"""
Waiter Lambda

This lambda waits indefinitely for an action, sleeping every n seconds
until it succeeds, fails or times out.
"""
import json
import os
import time
import boto3


TIMEOUT = os.getenv("WAITER_TIMEOUT", 120000)
SLEEP = os.getenv("WAITER_SLEEP", 5)
DEBUG = os.getenv("DEBUG", "False")
LOCAL = os.getenv("LOCAL", False)

if DEBUG.lower() in ["f", "false", "no", "n"]:
    DEBUG = False
else:
    DEBUG = True

def almost_dead(event, context):
    """
    Check if the lambda is less than 10 seconds from dying, if so, exit and recall this lambda
    """
    return context.get_remaining_time_in_millis() < 10000

def waiter(event, context, func, *args, **kwargs):
    """
    Takes a lambda `event` and `context` object and a function as `func`.
    `func` needs to return None if it should continue to be executed or any other value to exit the waiter

    If the waiter function was successful, returns a tuple of (event, context, func output)

    If the waiter function timed out (because the parent lambda timed out) returns None
    In order to have enough time to clean up, this function exits 10 seconds before the lambda is set to end.
    By default your lambdas should be set to the max time and the waiter will properly exit when the function is complete.
    """

    if not "timeout" in event:
        event["timeout"] = TIMEOUT

    if not "time_remaining" in event:
        event["time_remaining"] = TIMEOUT

    while event["time_remaining"] <= TIMEOUT and event["time_remaining"] > 0:
        loop_start = int(time.time()) * 1000

        if almost_dead(event, context):
            if LOCAL:
                return lambda_handler(event, EmptyContext())
            else:
                lambda_client = boto3.client("lambda")
                lambda_client.invoke(
                    FunctionName=context.function_name,
                    InvocationType='Event',
                    Payload=bytes(json.dumps(event), "utf-8"),
                    Qualifier=context.function_version,
                )
                if DEBUG:
                    print("restarted!")
                return

        else:
            output = func(*args, **kwargs)
            if output is not None:
                return event, context, output
            time.sleep(SLEEP)

        if DEBUG:
            print("Timeout remaining is {time_remaining}".format(**event))

        time_elapsed = int(time.time()) * 1000 - loop_start
        event["time_remaining"] = event["time_remaining"] - time_elapsed


        
class EmptyContext(object):

  timeout = 60000
  start = None

  def __init__(self):
      self.start = int(time.time()) * 1000

  def get_remaining_time_in_millis(self):
      return self.timeout - ((int(time.time()) * 1000) - self.start)

        
if __name__ == "__main__":

  
    lambda_handler({}, EmptyContext())