<gh_stars>0
# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
# Licensed under the Apache License, Version 2.0 https://aws.amazon.com/apache-2-0/

import time
import os
import logging

from utils.performance_tracker import EventsCounter, performance_tracker_initializer
from utils.state_table_common import *
from api.queue_manager import queue_manager

logging.basicConfig(format="%(asctime)s - %(levelname)s - %(filename)s - %(funcName)s  - %(lineno)d - %(message)s",
                    datefmt='%H:%M:%S', level=logging.INFO)

region = os.environ.get('REGION', None)

perf_tracker = performance_tracker_initializer(
    os.environ["METRICS_ARE_ENABLED"],
    os.environ["METRICS_TTL_CHECKER_LAMBDA_CONNECTION_STRING"],
    os.environ["METRICS_GRAFANA_PRIVATE_IP"])

from api.state_table_manager import state_table_manager

state_table = state_table_manager(grid_state_table_service=os.environ.get('TASKS_STATUS_TABLE_SERVICE', None),
                                  grid_state_table_config=os.environ.get('TASKS_STATUS_TABLE_CONFIG', None),
                                  tasks_state_table_name=os.environ.get('TASKS_STATUS_TABLE_NAME', None),
                                  endpoint_url=os.environ.get('DB_ENDPOINT_URL', None),
                                  region=region)

queue = queue_manager(
    grid_queue_service=os.environ['GRID_QUEUE_SERVICE'],
    grid_queue_config=os.environ['GRID_QUEUE_CONFIG'],
    endpoint_url=os.environ["QUEUE_ENDPOINT_URL"],
    queue_name=os.environ['TASKS_QUEUE_NAME'],
    region=region)

dlq = queue_manager(
    grid_queue_service=os.environ['GRID_QUEUE_SERVICE'],  # TODO extend parameters to configure this queue.
    grid_queue_config=os.environ['GRID_QUEUE_CONFIG'],
    endpoint_url=os.environ["QUEUE_ENDPOINT_URL"],
    queue_name=os.environ['TASKS_QUEUE_DLQ_NAME'],
    region=region)

MAX_RETRIES = 5
RETRIEVE_EXPIRED_TASKS_LIMIT = 200


# TODO: implement archival after 10 days in S3

def lambda_handler(event, context):
    """Handler called by AWS Lambda runtime

    Args:
      event(dict): a CloudWatch Event generated every minute
      context:

    Returns:

    """
    stats_obj = {'01_invocation_tstmp': {"label": "None", "tstmp": int(round(time.time() * 1000))}}
    event_counter = EventsCounter(
        ["counter_expired_tasks", "counter_failed_to_acquire",
         "counter_failed_tasks", "counter_released_tasks", "counter_inconsistent_state", "counter_tasks_queue_size"])

    for expired_tasks in state_table.query_expired_tasks():

        event_counter.increment("counter_expired_tasks", len(expired_tasks))
        event_counter.increment("counter_tasks_queue_size", queue.get_queue_length())

        for item in expired_tasks:
            print("Processing expired task: {}".format(item))
            task_id = item.get('task_id')
            owner_id = item.get('task_owner')
            current_heartbeat_timestamp = item.get('heartbeat_expiration_timestamp')
            try:
                is_acquired = state_table.acquire_task_for_ttl_lambda(task_id, owner_id, current_heartbeat_timestamp)

                if not is_acquired:
                    # task has been updated at the very last second...
                    event_counter.increment("counter_failed_to_acquire")
                    continue

                # retreive current number of retries and SQS_handler
                retries, sqs_handler_id, task_priority = retreive_retries_and_sqs_handler_and_priority(task_id)
                print("Number of retires for task[{}]: {} Priority: {}".format(task_id, retries, task_priority))
                print("Last owner for task [{}]: {}".format(task_id, owner_id))

                # TODO: MAX_RETRIES should be extracted from task definition... Store in DDB?
                if retries == MAX_RETRIES:
                    print("Failing task {} after {} retries".format(task_id, retries))
                    event_counter.increment("counter_failed_tasks")
                    fail_task(task_id, sqs_handler_id, task_priority)
                    continue

                event_counter.increment("counter_released_tasks")
                # else
                state_table.retry_task(task_id, retries + 1)

                try:
                    # Task can be acquired by an agent from this point
                    reset_sqs_vto(sqs_handler_id, task_priority)
                    print("SUCCESS FIX for {}".format(task_id))

                except Exception:
                    try:
                        logging.error('Failed to reset VTO trying to delete: {} '.format(task_id))
                        delete_message_from_queue(sqs_handler_id)
                    except Exception:
                        logging.error('Inconsistent task: {} sending do DLQ'.format(task_id))

                        event_counter.increment("counter_inconsistent_state")
                        set_task_inconsistent(task_id)
                        send_to_dlq(item)

            except Exception as e:
                print("Cannot process task {} : {}".format(task_id, e))
                print("Sending task {} to DLQ...".format(task_id))
                logging.error('Lambda ttl error: {}'.format(e))
                send_to_dlq(item)

    stats_obj['02_completion_tstmp'] = {"label": "ttl_execution_time", "tstmp": int(round(time.time() * 1000))}
    perf_tracker.add_metric_sample(stats_obj, event_counter=event_counter, from_event="01_invocation_tstmp",
                                   to_event="02_completion_tstmp")
    perf_tracker.submit_measurements()


def fail_task(task_id, sqs_handler_id, task_priority):
    """This function set the task_status of task to fail

    Args:
      task_id(str): the id of the task to update
      sqs_handler_id(str): the sqs handler associated to this task
      task_priority(int): the priority of the task.

    Returns:
      Nothing

    Raises:
      Exception: if DynamoDB table cannot be updated

    """
    try:
        delete_message_from_queue(sqs_handler_id, task_priority)

        state_table.update_task_status_to_failed(task_id)


    except Exception as e:
        logging.error("Cannot fail task {} : {}".format(task_id, e))
        raise e


def set_task_inconsistent(task_id):
    """This function set the task_status of task to inconsistent

    Args:
      task_id(str): the id of the task to update

    Returns:
      Nothing

    Raises:
      Exception: if DynamoDB table cannot be updated

    """
    try:
        state_table.update_task_status_to_inconsistent(task_id)

    except Exception as e:
        logging.error("Cannot set task to inconsystent {} : {}".format(task_id, e))
        raise e


def delete_message_from_queue(sqs_handler_id, task_priority):
    """This function delete a message from a SQS queue

    Args:
      sqs_handler_id(str): the sqs handler associated of the message to be deleted
      task_priority(int): priority of the task

    Returns:
      Nothing

    Raises:
      Exception: if SQS queue cannot be updated

    """

    try:
        queue.delete_message(sqs_handler_id, task_priority)

    except Exception as e:
        logging.error("Cannot delete message {} : {}".format(sqs_handler_id, e))
        raise e


def retreive_retries_and_sqs_handler_and_priority(task_id):
    """This function retrieve (i) the number of retries,
    (ii) the SQS handler associated to an expired task
    and (iii) and the priority under which this task was executed.

    Args:
      task_id(str): the id of the expired task

    Returns:
      rtype: 3 variables

    Raises:
      Exception: if DynamoDB query failed

    """
    try:

        resp_task = state_table.get_task_by_id(task_id)
        # CHeck if 1 and only 1
        return resp_task.get('retries'), \
               resp_task.get('sqs_handler_id'), \
               resp_task.get('task_priority')

    except Exception as e:
        logging.error("Cannot retreive retries and handler for task {} : {}".format(task_id, e))
        raise e


def reset_sqs_vto(handler_id, task_priority):
    """

    Args:
      handler_id:

    Returns:

    """
    try:
        visibility_timeout_sec = 0
        queue.change_visibility(handler_id, visibility_timeout_sec, task_priority)

    except Exception as e:
        logging.error("Cannot reset VTO for message {} : {}".format(handler_id, e))
        raise e


def send_to_dlq(task):
    """

    Args:
      task:

    Returns:

    """
    print("Sending task [{}] to DLQ".format(task))
    dlq.send_message(message_bodies=[str(task)])

# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
# Licensed under the Apache License, Version 2.0 https://aws.amazon.com/apache-2-0/

import time
import os
import logging

from utils.performance_tracker import EventsCounter, performance_tracker_initializer
from utils.state_table_common import *
from api.queue_manager import queue_manager

logging.basicConfig(format="%(asctime)s - %(levelname)s - %(filename)s - %(funcName)s  - %(lineno)d - %(message)s",
                    datefmt='%H:%M:%S', level=logging.INFO)

region = os.environ.get('REGION', None)

perf_tracker = performance_tracker_initializer(
    os.environ["METRICS_ARE_ENABLED"],
    os.environ["METRICS_TTL_CHECKER_LAMBDA_CONNECTION_STRING"],
    os.environ["METRICS_GRAFANA_PRIVATE_IP"])

from api.state_table_manager import state_table_manager

state_table = state_table_manager(grid_state_table_service=os.environ.get('TASKS_STATUS_TABLE_SERVICE', None),
                                  grid_state_table_config=os.environ.get('TASKS_STATUS_TABLE_CONFIG', None),
                                  tasks_state_table_name=os.environ.get('TASKS_STATUS_TABLE_NAME', None),
                                  endpoint_url=os.environ.get('DB_ENDPOINT_URL', None),
                                  region=region)

queue = queue_manager(
    grid_queue_service=os.environ['GRID_QUEUE_SERVICE'],
    grid_queue_config=os.environ['GRID_QUEUE_CONFIG'],
    endpoint_url=os.environ["QUEUE_ENDPOINT_URL"],
    queue_name=os.environ['TASKS_QUEUE_NAME'],
    region=region)

dlq = queue_manager(
    grid_queue_service=os.environ['GRID_QUEUE_SERVICE'],  # TODO extend parameters to configure this queue.
    grid_queue_config=os.environ['GRID_QUEUE_CONFIG'],
    endpoint_url=os.environ["QUEUE_ENDPOINT_URL"],
    queue_name=os.environ['TASKS_QUEUE_DLQ_NAME'],
    region=region)

MAX_RETRIES = 5
RETRIEVE_EXPIRED_TASKS_LIMIT = 200


# TODO: implement archival after 10 days in S3

def lambda_handler(event, context):
    """Handler called by AWS Lambda runtime

    Args:
      event(dict): a CloudWatch Event generated every minute
      context:

    Returns:

    """
    stats_obj = {'01_invocation_tstmp': {"label": "None", "tstmp": int(round(time.time() * 1000))}}
    event_counter = EventsCounter(
        ["counter_expired_tasks", "counter_failed_to_acquire",
         "counter_failed_tasks", "counter_released_tasks", "counter_inconsistent_state", "counter_tasks_queue_size"])

    for expired_tasks in state_table.query_expired_tasks():

        event_counter.increment("counter_expired_tasks", len(expired_tasks))
        event_counter.increment("counter_tasks_queue_size", queue.get_queue_length())

        for item in expired_tasks:
            print("Processing expired task: {}".format(item))
            task_id = item.get('task_id')
            owner_id = item.get('task_owner')
            current_heartbeat_timestamp = item.get('heartbeat_expiration_timestamp')
            try:
                is_acquired = state_table.acquire_task_for_ttl_lambda(task_id, owner_id, current_heartbeat_timestamp)

                if not is_acquired:
                    # task has been updated at the very last second...
                    event_counter.increment("counter_failed_to_acquire")
                    continue

                # retreive current number of retries and SQS_handler
                retries, sqs_handler_id, task_priority = retreive_retries_and_sqs_handler_and_priority(task_id)
                print("Number of retires for task[{}]: {} Priority: {}".format(task_id, retries, task_priority))
                print("Last owner for task [{}]: {}".format(task_id, owner_id))

                # TODO: MAX_RETRIES should be extracted from task definition... Store in DDB?
                if retries == MAX_RETRIES:
                    print("Failing task {} after {} retries".format(task_id, retries))
                    event_counter.increment("counter_failed_tasks")
                    fail_task(task_id, sqs_handler_id, task_priority)
                    continue

                event_counter.increment("counter_released_tasks")
                # else
                state_table.retry_task(task_id, retries + 1)

                try:
                    # Task can be acquired by an agent from this point
                    reset_sqs_vto(sqs_handler_id, task_priority)
                    print("SUCCESS FIX for {}".format(task_id))

                except Exception:
                    try:
                        logging.error('Failed to reset VTO trying to delete: {} '.format(task_id))
                        delete_message_from_queue(sqs_handler_id)
                    except Exception:
                        logging.error('Inconsistent task: {} sending do DLQ'.format(task_id))

                        event_counter.increment("counter_inconsistent_state")
                        set_task_inconsistent(task_id)
                        send_to_dlq(item)

            except Exception as e:
                print("Cannot process task {} : {}".format(task_id, e))
                print("Sending task {} to DLQ...".format(task_id))
                logging.error('Lambda ttl error: {}'.format(e))
                send_to_dlq(item)

    stats_obj['02_completion_tstmp'] = {"label": "ttl_execution_time", "tstmp": int(round(time.time() * 1000))}
    perf_tracker.add_metric_sample(stats_obj, event_counter=event_counter, from_event="01_invocation_tstmp",
                                   to_event="02_completion_tstmp")
    perf_tracker.submit_measurements()


def fail_task(task_id, sqs_handler_id, task_priority):
    """This function set the task_status of task to fail

    Args:
      task_id(str): the id of the task to update
      sqs_handler_id(str): the sqs handler associated to this task
      task_priority(int): the priority of the task.

    Returns:
      Nothing

    Raises:
      Exception: if DynamoDB table cannot be updated

    """
    try:
        delete_message_from_queue(sqs_handler_id, task_priority)

        state_table.update_task_status_to_failed(task_id)


    except Exception as e:
        logging.error("Cannot fail task {} : {}".format(task_id, e))
        raise e


def set_task_inconsistent(task_id):
    """This function set the task_status of task to inconsistent

    Args:
      task_id(str): the id of the task to update

    Returns:
      Nothing

    Raises:
      Exception: if DynamoDB table cannot be updated

    """
    try:
        state_table.update_task_status_to_inconsistent(task_id)

    except Exception as e:
        logging.error("Cannot set task to inconsystent {} : {}".format(task_id, e))
        raise e


def delete_message_from_queue(sqs_handler_id, task_priority):
    """This function delete a message from a SQS queue

    Args:
      sqs_handler_id(str): the sqs handler associated of the message to be deleted
      task_priority(int): priority of the task

    Returns:
      Nothing

    Raises:
      Exception: if SQS queue cannot be updated

    """

    try:
        queue.delete_message(sqs_handler_id, task_priority)

    except Exception as e:
        logging.error("Cannot delete message {} : {}".format(sqs_handler_id, e))
        raise e


def retreive_retries_and_sqs_handler_and_priority(task_id):
    """This function retrieve (i) the number of retries,
    (ii) the SQS handler associated to an expired task
    and (iii) and the priority under which this task was executed.

    Args:
      task_id(str): the id of the expired task

    Returns:
      rtype: 3 variables

    Raises:
      Exception: if DynamoDB query failed

    """
    try:

        resp_task = state_table.get_task_by_id(task_id)
        # CHeck if 1 and only 1
        return resp_task.get('retries'), \
               resp_task.get('sqs_handler_id'), \
               resp_task.get('task_priority')

    except Exception as e:
        logging.error("Cannot retreive retries and handler for task {} : {}".format(task_id, e))
        raise e


def reset_sqs_vto(handler_id, task_priority):
    """

    Args:
      handler_id:

    Returns:

    """
    try:
        visibility_timeout_sec = 0
        queue.change_visibility(handler_id, visibility_timeout_sec, task_priority)

    except Exception as e:
        logging.error("Cannot reset VTO for message {} : {}".format(handler_id, e))
        raise e


def send_to_dlq(task):
    """

    Args:
      task:

    Returns:

    """
    print("Sending task [{}] to DLQ".format(task))
    dlq.send_message(message_bodies=[str(task)])
# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 # Licensed under the Apache License, Version 2.0 https://aws.amazon.com/apache-2-0/ # TODO extend parameters to configure this queue. # TODO: implement archival after 10 days in S3 Handler called by AWS Lambda runtime      Args:       event(dict): a CloudWatch Event generated every minute       context:      Returns: # task has been updated at the very last second... # retreive current number of retries and SQS_handler # TODO: MAX_RETRIES should be extracted from task definition... Store in DDB? # else # Task can be acquired by an agent from this point This function set the task_status of task to fail      Args:       task_id(str): the id of the task to update       sqs_handler_id(str): the sqs handler associated to this task       task_priority(int): the priority of the task.      Returns:       Nothing      Raises:       Exception: if DynamoDB table cannot be updated This function set the task_status of task to inconsistent      Args:       task_id(str): the id of the task to update      Returns:       Nothing      Raises:       Exception: if DynamoDB table cannot be updated This function delete a message from a SQS queue      Args:       sqs_handler_id(str): the sqs handler associated of the message to be deleted       task_priority(int): priority of the task      Returns:       Nothing      Raises:       Exception: if SQS queue cannot be updated This function retrieve (i) the number of retries,     (ii) the SQS handler associated to an expired task     and (iii) and the priority under which this task was executed.      Args:       task_id(str): the id of the expired task      Returns:       rtype: 3 variables      Raises:       Exception: if DynamoDB query failed # CHeck if 1 and only 1 Args:       handler_id:      Returns: Args:       task:      Returns: