EM-933 Throttling errors withing send-message state machine

module/dynamodb_table_locktable.tf

@@ -0,0 +1,44 @@
+locals {
+    locktable_name = "${local.name}-lock-table"
+}
+
+#tfsec:ignore:aws-dynamodb-enable-at-rest-encryption
+resource "aws_dynamodb_table" "lock_table" {
+    name            = local.locktable_name
+    billing_mode    = "PROVISIONED"
+    read_capacity   = 20
+    write_capacity  = 20
+    hash_key        = "LockName"
+    stream_enabled = false
+    point_in_time_recovery {
+      enabled = true
+    }
+
+    server_side_encryption {
+      enabled = true
+    }
+
+    attribute {
+      name = "LockName"
+      type = "S"
+    }
+
+    attribute {
+        name = "LockType"
+        type = "S"
+    }
+
+    attribute {
+        name = "LockOwner"
+        type = "S"
+    }
+
+    global_secondary_index {
+      name = "LockTypeOwnerTableIndex"
+      hash_key = "LockType"
+      range_key = "LockOwner"
+      write_capacity = 5
+      read_capacity = 5
+      projection_type = "KEYS_ONLY"
+    }
+}

module/locals.tf

@@ -21,6 +21,8 @@ locals {
     MESH_URL    = local.mesh_url[var.mesh_env]
     MESH_BUCKET = aws_s3_bucket.mesh.bucket
 
+    DDB_LOCK_TABLE_NAME = aws_dynamodb_table.lock_table.name
+
     CHUNK_SIZE         = var.chunk_size
     CRUMB_SIZE         = var.crumb_size == null ? var.chunk_size : var.crumb_size
     NEVER_COMPRESS     = var.never_compress

module/stepfunctions_get_messages.tf

@@ -99,7 +99,10 @@ resource "aws_sfn_state_machine" "get_messages" {
         OutputPath = "$.Payload"
         Parameters = {
           FunctionName = "${aws_lambda_function.poll_mailbox.arn}:${aws_lambda_function.poll_mailbox.version}"
-          "Payload.$"  = "$"
+          Payload = {
+            "EventDetail.$" = "$"
+            "ExecutionId.$" = "$$.Execution.Id"
+          }
         }
         Resource = "arn:aws:states:::lambda:invoke"
         Retry = [

module/stepfunctions_send_message.tf

@@ -29,7 +29,10 @@ resource "aws_sfn_state_machine" "send_message" {
         OutputPath = "$.Payload"
         Parameters = {
           FunctionName = "${aws_lambda_function.check_send_parameters.arn}:${aws_lambda_function.check_send_parameters.version}"
-          "Payload.$"  = "$"
+          Payload = {
+            "EventDetail.$" = "$"
+            "ExecutionId.$" = "$$.Execution.Id"
+          }
         }
         Resource = "arn:aws:states:::lambda:invoke"
         Retry = [

pyproject.toml

@@ -27,12 +27,13 @@ mypy = "^1.4.0"
 coverage = "^7.2.7"
 pytest = "^8.1.1"
 pytest-asyncio = "^0.23.6"
-moto = {extras = ["s3", "ssm", "stepfunctions", "secretsmanager"], version = "^5.0.5"}
+moto = {extras = ["s3", "ssm", "stepfunctions", "secretsmanager", "dynamodb"], version = "^5.0.14"}
 boto3-stubs = {extras = ["s3", "ssm", "secretsmanager", "dynamodb", "stepfunctions", "sqs", "lambda", "logs",], version = "^1.34.32"}
 ruff = "^0"
 petname = "^2.6"
 black = "^24"
 tox = "^4.18"
+freezegun = "^1.5.1"
 
 [tool.poetry.group.local.dependencies]
 
@@ -169,6 +170,7 @@ legacy_tox_ini = """
         moto
         boto3
         petname
+        freezegun
     commands =
         pytest
 """

src/cloudlogbase.cfg

@@ -106,6 +106,10 @@ Log Text = Will send file='{file}' filesize="{file_size}" from bucket='{bucket}'
 Log Level = INFO
 Log Text = For send for mapping destMailbox='{dest_mailbox}' from srcMailbox='{src_mailbox}' with workflowId='{workflow_id}'
 
+[MESHSEND0004b]
+Log Level = INFO
+Log Text = Debug info: debug_str='{debug_str}'
+
 [MESHSEND0005]
 Log Level = INFO
 Log Text = Send chunk started for file='{file}' from bucket='{bucket}' chunk='{chunk_num}' of max_chunk='{max_chunk}'
@@ -122,6 +126,22 @@ Log Text = Sent chunk='{chunk_num}' of max_chunk='{max_chunk}' for file='{file}'
 Log Level = INFO
 Log Text = chunk='{chunk_num}' is the final chunk out of max_chunk='{max_chunk}' for file='{file}' from bucket='{bucket}' and has been sent
 
+[MESHSEND0009]
+Log Level = INFO
+Log Text = Acquiring lock_name='{lock_name}' with owner_id='{owner_id}'
+
+[MESHSEND0010]
+Log Level = INFO
+Log Text = Released lock_name='{lock_name}' with owner_id='{owner_id}'
+
+[MESHSEND0011]
+Log Level = ERROR
+Log Text = Failed to release lock_name='{lock_name}' as owner_id='{owner_id}' because it was owned by lock_owner='{lock_owner}'
+
+[MESHSEND0012]
+Log Level = INFO
+Log Text = Skipping lock release due to missing lock_name='{lock_name}' or owner_id='{owner_id}'
+
 [MESHPOLL0001]
 Log Level = INFO
 Log Text = mailbox='{mailbox}' has polled message_count='{message_count}' many messages
@@ -132,7 +152,11 @@ Log Text = msg='{error}' raised when polling mailbox='{mailbox}'
 
 [MESHPOLL0002a]
 Log Level = WARN
-Log Text = Performing sending singleton check for mailbox='{mailbox) sf_mailbox='{sf_mailbox}' missing a mailbox
+Log Text = Performing sending singleton check for mailbox='{mailbox)' sf_mailbox='{sf_mailbox}' missing a mailbox
+
+[MESHPOLL0003]
+Log Level = INFO
+Log Text = Acquiring lock_name='{lock_name}' for owner_id='{owner_id}'
 
 [MESHFETCH0001]
 Log Level = INFO
@@ -220,4 +244,12 @@ Log Text = Message with message_id='{message_id}' acknowledged and successfully
 
 [MESHFETCH0013]
 Log Level = INFO
-Log Text = File with multiple chunks received message_id='{message_id}'
+Log Text = File with multiple chunks received message_id='{message_id}'
+
+[MESHFETCH0014]
+Log Level = INFO
+Log Text = Released lock_name='{lock_name}' with owner_id='{owner_id}'
+
+[MESHSEND0015]
+Log Level = INFO
+Log Text = Skipping lock release due to missing lock_name='{lock_name}' or owner_id='{owner_id}'

src/mesh_check_send_parameters_application.py

@@ -1,10 +1,13 @@
 from dataclasses import asdict
-from functools import partial
 from http import HTTPStatus
 from typing import Any
 
 from shared.application import MESHLambdaApplication
-from shared.common import SingletonCheckFailure, return_failure, singleton_check
+from shared.common import (
+    SingletonCheckFailure,
+    acquire_lock,
+    return_failure,
+)
 from shared.send_parameters import get_send_parameters
 from spine_aws_common.utilities import human_readable_bytes
 
@@ -22,11 +25,23 @@ def _get_internal_id(self):
         return self._create_new_internal_id()
 
     def start(self):
+
+        self.log_object.write_log(
+            "MESHSEND0004b",
+            None,
+            {"debug_str": str(self.event.raw_event)},
+        )
+
+        event_details = self.event["EventDetail"]["detail"]
+        owner_id = (
+            self.event.get("ExecutionId") or f"intenalID_{self._get_internal_id()}"
+        )
+
         # in case of crash, set to internal server error so next stage fails
         self.response = {"statusCode": int(HTTPStatus.INTERNAL_SERVER_ERROR)}
 
-        bucket = self.event["detail"]["requestParameters"]["bucketName"]
-        key = self.event["detail"]["requestParameters"]["key"]
+        bucket = event_details["requestParameters"]["bucketName"]
+        key = event_details["requestParameters"]["key"]
 
         self.log_object.write_log("MESHSEND0001", None, {"bucket": bucket, "file": key})
 
@@ -54,11 +69,19 @@ def start(self):
             },
         )
         try:
-            check = partial(self.is_send_for_same_file, send_params=send_params)
-            singleton_check(
-                self.config.send_message_step_function_arn,
-                check,
-                self.sfn,
+
+            lock_name = f"SendLock_{send_params.s3_bucket}_{send_params.s3_key}"
+
+            self.log_object.write_log(
+                "MESHSEND0009",
+                None,
+                {"lock_name": lock_name, "owner_id": owner_id},
+            )
+
+            acquire_lock(
+                self.ddb_client,
+                lock_name,
+                owner_id,
             )
 
         except SingletonCheckFailure as e:
@@ -104,6 +127,8 @@ def start(self):
                 "message_id": None,
                 "current_byte_position": 0,
                 "send_params": asdict(send_params),
+                "lock_name": lock_name,
+                "owner_id": owner_id,
             },
         }
 

src/mesh_fetch_message_chunk_application.py

@@ -9,7 +9,7 @@
 from requests import Response
 from requests.structures import CaseInsensitiveDict
 from shared.application import INBOUND_BUCKET, INBOUND_FOLDER, MESHLambdaApplication
-from shared.common import nullsafe_quote
+from shared.common import nullsafe_quote, release_lock
 from shared.config import MiB
 
 _METADATA_HEADERS = {
@@ -87,6 +87,8 @@ def initialise(self):
         self.log_object.internal_id = self.internal_id
         self.s3_bucket = self.input.get("s3_bucket", "")  # will be empty on first chunk
         self.s3_key = self.input.get("s3_key", "")  # will be empty on first chunk
+        self.lock_name = self.response.get("lock_name")
+        self.owner_id = self.response.get("owner_id")
 
     @property
     def http_response(self) -> Response:
@@ -117,9 +119,21 @@ def start(self):
             self._ensure_s3_bucket_and_key(is_report)
             if is_report or self.number_of_chunks < 2:
                 self._handle_un_chunked_message(is_report)
-                return
-
-            self._handle_multiple_chunk_message()
+            else:
+                self._handle_multiple_chunk_message()
+
+            if self.owner_id and self.lock_name:
+                release_lock(
+                    self.ddb_client,
+                    self.lock_name,
+                    self.owner_id,
+                )
+            else:
+                self.log_object.write_log(
+                    "MESHSEND0015",
+                    None,
+                    {"lock_name": self.lock_name, "owner_id": self.owner_id},
+                )
 
     def _retrieve_current_chunk(self):
         self._http_response = self.get_chunk(

src/mesh_poll_mailbox_application.py

@@ -4,7 +4,11 @@
 from aws_lambda_powertools.shared.functions import strtobool
 from requests import HTTPError
 from shared.application import MESHLambdaApplication
-from shared.common import SingletonCheckFailure, return_failure, singleton_check
+from shared.common import (
+    SingletonCheckFailure,
+    acquire_lock,
+    return_failure,
+)
 
 
 class HandshakeFailure(Exception):
@@ -28,13 +32,25 @@ def __init__(self, additional_log_config=None, load_ssm_params=False):
 
         self.handshake: bool = False
         self.response: dict[str, Any] = {}
+        self.execution_id = None
 
     def initialise(self):
         # initialise
         self.mailbox_id = self.event["mailbox"]
         self.handshake = bool(strtobool(self.event.get("handshake", "false")))
         self.response = {}
 
+        print("POLL_EVENT", self.event.raw_event)
+        print("EXECUTION_ID", self.execution_id)
+
+    def process_event(self, event):
+        event_detail = event.get("EventDetail", {})
+        if event_detail:
+            # Enhanced event detail format from the step function
+            self.execution_id = event.get("ExecutionId")
+
+        return self.EVENT_TYPE(event_detail or event)
+
     def start(self):
         # in case of crash
         self.response = {"statusCode": int(HTTPStatus.INTERNAL_SERVER_ERROR)}
@@ -48,10 +64,20 @@ def start(self):
                 return
 
         try:
-            singleton_check(
-                self.config.get_messages_step_function_arn,
-                self.is_same_mailbox_check,
-                self.sfn,
+            lock_name = f"FetchLock_{self.mailbox_id}"
+
+            owner_id = self.execution_id or f"internalID_{self._get_internal_id()}"
+
+            self.log_object.write_log(
+                "MESHPOLL0003",
+                None,
+                {"lock_name": lock_name, "owner_id": owner_id},
+            )
+
+            acquire_lock(
+                self.ddb_client,
+                lock_name,
+                owner_id,
             )
 
         except SingletonCheckFailure as e:
@@ -104,6 +130,8 @@ def start(self):
                 "internal_id": self.log_object.internal_id,
                 "message_count": message_count,
                 "message_list": output_list,
+                "lock_name": lock_name,
+                "execution_id": self.execution_id,
             },
             # Parameters for a follow-up iteration through the messages in this execution
             "mailbox": self.mailbox_id,