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/lambda_check_send_parameters.tf

@@ -149,10 +149,25 @@ data "aws_iam_policy_document" "check_send_parameters" {
     ]
   }
 
+  statement {
+    sid    = "DynamoDBAccess"
+    effect = "Allow"
+
+    actions = [
+      "dynamodb:GetItem",
+      "dynamodb:PutItem",
+      "dynamodb:DeleteItem",
+      "dynamodb:UpdateItem", 
+    ]
+
+    resources = [
+      "arn:aws:dynamodb:eu-west-2:${var.account_id}:table/${local.locktable_name}"
+    ]
+  }
+
   dynamic "statement" {
     for_each = local.vpc_enabled ? [true] : []
     content {
-
       sid    = "EC2Interfaces"
       effect = "Allow"
 
@@ -216,4 +231,6 @@ data "aws_iam_policy_document" "check_send_parameters_check_sfn" {
       "${replace(aws_sfn_state_machine.send_message.arn, "stateMachine", "execution")}*"
     ]
   }
+
+
 }

module/lambda_fetch_message_chunk.tf

@@ -171,6 +171,22 @@ data "aws_iam_policy_document" "fetch_message_chunk" {
     ]
   }
 
+  statement {
+    sid    = "DynamoDBAccess"
+    effect = "Allow"
+
+    actions = [
+      "dynamodb:GetItem",
+      "dynamodb:PutItem",
+      "dynamodb:DeleteItem",
+      "dynamodb:UpdateItem", 
+    ]
+
+    resources = [
+      "arn:aws:dynamodb:eu-west-2:${var.account_id}:table/${local.locktable_name}"
+    ]
+  }
+
   dynamic "statement" {
     for_each = var.use_secrets_manager ? [true] : []
     content {

module/lambda_lock_manager.tf

@@ -0,0 +1,147 @@
+locals {
+  lock_manager_name = "${local.name}-lock-manager"
+}
+
+resource "aws_lambda_function" "lock_manager" {
+  function_name    = local.lock_manager_name
+  filename         = data.archive_file.app.output_path
+  handler          = "mesh_lock_manager_application.lambda_handler"
+  runtime          = local.python_runtime
+  timeout          = local.lambda_timeout
+  source_code_hash = data.archive_file.app.output_base64sha256
+  role             = aws_iam_role.lock_manager.arn
+  layers           = [aws_lambda_layer_version.mesh_aws_client_dependencies.arn]
+
+  publish = true
+
+  environment {
+    variables = local.common_env_vars
+  }
+
+  dynamic "vpc_config" {
+    for_each = local.vpc_enabled ? [local.vpc_enabled] : []
+    content {
+      subnet_ids         = var.subnet_ids
+      security_group_ids = [aws_security_group.lock_manager[0].id]
+    }
+  }
+
+  depends_on = [
+    aws_cloudwatch_log_group.lock_manager,
+  ]
+}
+
+resource "aws_cloudwatch_log_group" "lock_manager" {
+  name              = "/aws/lambda/${local.lock_manager_name}"
+  retention_in_days = var.cloudwatch_retention_in_days
+  kms_key_id        = aws_kms_key.mesh.arn
+  lifecycle {
+    ignore_changes = [
+      log_group_class, # localstack not currently returning this
+    ]
+  }
+}
+
+resource "aws_iam_role" "lock_manager" {
+  name               = "${local.lock_manager_name}-role"
+  description        = "${local.lock_manager_name}-role"
+  assume_role_policy = data.aws_iam_policy_document.lock_manager_assume.json
+}
+
+data "aws_iam_policy_document" "lock_manager_assume" {
+  statement {
+    actions = ["sts:AssumeRole"]
+
+    principals {
+      type = "Service"
+
+      identifiers = [
+        "lambda.amazonaws.com",
+      ]
+    }
+  }
+}
+
+resource "aws_iam_role_policy_attachment" "lock_manager" {
+  role       = aws_iam_role.lock_manager.name
+  policy_arn = aws_iam_policy.lock_manager.arn
+}
+
+resource "aws_iam_policy" "lock_manager" {
+  name        = "${local.lock_manager_name}-policy"
+  description = "${local.lock_manager_name}-policy"
+  policy      = data.aws_iam_policy_document.lock_manager.json
+}
+
+data "aws_iam_policy_document" "lock_manager" {
+  statement {
+    sid    = "CloudWatchAllow"
+    effect = "Allow"
+
+    actions = [
+      "logs:CreateLogGroup",
+      "logs:CreateLogStream",
+      "logs:PutLogEvents"
+    ]
+
+    resources = [
+      "${aws_cloudwatch_log_group.lock_manager.arn}*"
+    ]
+  }
+
+  statement {
+    sid    = "SSMDescribe"
+    effect = "Allow"
+
+    actions = [
+      "ssm:DescribeParameters"
+    ]
+
+    resources = [
+      "arn:aws:ssm:eu-west-2:${var.account_id}:parameter/${local.name}/*"
+    ]
+  }
+
+  statement {
+    sid    = "SSMGet"
+    effect = "Allow"
+
+    actions = [
+      "ssm:GetParameter",
+      "ssm:GetParameters",
+      "ssm:GetParametersByPath"
+    ]
+
+    resources = [
+      "arn:aws:ssm:eu-west-2:${var.account_id}:parameter/${local.name}/*",
+      "arn:aws:ssm:eu-west-2:${var.account_id}:parameter/${local.name}"
+    ]
+  }
+
+  statement {
+    sid    = "KMSDecrypt"
+    effect = "Allow"
+
+    actions = [
+      "kms:Decrypt"
+    ]
+
+    resources = concat(
+      [aws_kms_alias.mesh.target_key_arn],
+      var.use_secrets_manager ? local.secrets_kms_key_arns : []
+    )
+  }
+
+  statement {
+    sid    = "DynamoDBDelete"
+    effect = "Allow"
+
+    actions = [
+      "dynamodb:DeleteItem"
+    ]
+
+    resources = [
+      "arn:aws:dynamodb:eu-west-2:${var.account_id}:table/${local.locktable_name}"
+    ]
+  }
+}

module/lambda_poll_mailbox.tf

@@ -135,6 +135,22 @@ data "aws_iam_policy_document" "poll_mailbox" {
       var.use_secrets_manager ? local.secrets_kms_key_arns : []
     )
   }
+
+  statement {
+    sid    = "DynamoDBAccess"
+    effect = "Allow"
+
+    actions = [
+      "dynamodb:GetItem",
+      "dynamodb:PutItem",
+      "dynamodb:DeleteItem",
+      "dynamodb:UpdateItem", 
+    ]
+
+    resources = [
+      "arn:aws:dynamodb:eu-west-2:${var.account_id}:table/${local.locktable_name}"
+    ]
+  }
 
   dynamic "statement" {
     for_each = var.use_secrets_manager ? [true] : []
@@ -170,6 +186,7 @@ data "aws_iam_policy_document" "poll_mailbox" {
     }
   }
 
+
 }
 
 resource "aws_iam_role_policy_attachment" "poll_mailbox_lambda_insights" {

module/lambda_send_message_chunk.tf

@@ -154,6 +154,21 @@ data "aws_iam_policy_document" "send_message_chunk" {
       "${aws_s3_bucket.mesh.arn}/*"
     ]
   }
+  statement {
+    sid    = "DynamoDBAccess"
+    effect = "Allow"
+
+    actions = [
+      "dynamodb:GetItem",
+      "dynamodb:PutItem",
+      "dynamodb:DeleteItem",
+      "dynamodb:UpdateItem", 
+    ]
+
+    resources = [
+      "arn:aws:dynamodb:eu-west-2:${var.account_id}:table/${local.locktable_name}"
+    ]
+  }
 
   dynamic "statement" {
     for_each = var.use_secrets_manager ? [true] : []

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

@@ -30,7 +30,7 @@ resource "aws_sfn_state_machine" "get_messages" {
       "Failed?" = {
         Choices = [
           {
-            Next          = "Poll complete"
+            Next          = "Poll complete release lock"
             NumericEquals = 204
             Variable      = "$.statusCode"
           },
@@ -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 = [
@@ -132,9 +135,42 @@ resource "aws_sfn_state_machine" "get_messages" {
             Variable      = "$.body.message_count"
           },
         ]
-        Default = "Poll complete"
+        Default = "Poll complete release lock"
         Type    = "Choice"
       }
+      "Poll complete release lock" = {
+        Next       = "Poll complete"
+        OutputPath = "$.Payload"
+        Parameters = {
+          FunctionName = "${aws_lambda_function.lock_manager.arn}:${aws_lambda_function.lock_manager.version}"
+          Payload = {
+            "EventDetail.$" = "$"
+            "Operation"     = "release"
+          }
+        }
+        Resource = "arn:aws:states:::lambda:invoke"
+        Retry = [
+          {
+            BackoffRate = 2
+            ErrorEquals = [
+              "Lambda.ServiceException",
+              "Lambda.AWSLambdaException",
+              "Lambda.SdkClientException",
+            ]
+            IntervalSeconds = 2
+            MaxAttempts     = 3
+          },
+          {
+            ErrorEquals = [
+              "States.TaskFailed"
+            ],
+            BackoffRate     = 1,
+            IntervalSeconds = 300,
+            MaxAttempts     = 2
+          },
+        ]
+        Type = "Task"
+      }
     }
   })
 }
@@ -214,7 +250,9 @@ data "aws_iam_policy_document" "get_messages" {
       aws_lambda_function.fetch_message_chunk.arn,
       "${aws_lambda_function.fetch_message_chunk.arn}:*",
       aws_lambda_function.poll_mailbox.arn,
-      "${aws_lambda_function.poll_mailbox.arn}:*"
+      "${aws_lambda_function.poll_mailbox.arn}:*",
+      aws_lambda_function.lock_manager.arn,
+      "${aws_lambda_function.lock_manager.arn}:*"
     ]
 
     actions = ["lambda:InvokeFunction"]