From 68ff606838196967032aa2a40eb011269ad171db Mon Sep 17 00:00:00 2001
From: Pooyan Razian <prazian@gmail.com>
Date: Sat, 15 Jun 2024 16:12:04 +0200
Subject: [PATCH] feat: Add service timeout to custom resources

---
 .../aws-cdk-lib/core/lib/custom-resource.ts   | 31 +++++++++++++++++++
 .../core/test/custom-resource.test.ts         | 26 +++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/core/lib/custom-resource.ts b/packages/aws-cdk-lib/core/lib/custom-resource.ts
index 5e77f7f77e220..f2dfa68a2d504 100644
--- a/packages/aws-cdk-lib/core/lib/custom-resource.ts
+++ b/packages/aws-cdk-lib/core/lib/custom-resource.ts
@@ -1,6 +1,7 @@
 import { Construct } from 'constructs';
 import { CfnResource } from './cfn-resource';
 import { RemovalPolicy } from './removal-policy';
+import { Duration } from './duration';
 import { Resource } from './resource';
 import { Token } from './token';
 
@@ -91,6 +92,16 @@ export interface CustomResourceProps {
    */
   readonly removalPolicy?: RemovalPolicy;
 
+  /**
+   * The ServiceTimeout property from Cloudformation
+   *
+   * The value must be a duration between 1 and 3600 seconds.
+   * The default value is 3600 seconds (1 hour).
+   *
+   * @default Duration.minutes(60)
+   */
+  readonly serviceTimeout?: Duration;
+
   /**
    * Convert all property keys to pascal case.
    *
@@ -131,6 +142,7 @@ export class CustomResource extends Resource {
     const type = renderResourceType(props.resourceType);
     const pascalCaseProperties = props.pascalCaseProperties ?? false;
     const properties = pascalCaseProperties ? uppercaseProperties(props.properties || {}) : (props.properties || {});
+    const serviceTimeout = renderServiceTimeout(props.serviceTimeout);
 
     this.resource = new CfnResource(this, 'Default', {
       type,
@@ -143,6 +155,10 @@ export class CustomResource extends Resource {
     this.resource.applyRemovalPolicy(props.removalPolicy, {
       default: RemovalPolicy.DESTROY,
     });
+
+    if (serviceTimeout) {
+      this.resource.addPropertyOverride('ServiceTimeout', serviceTimeout);
+    }
   }
 
   /**
@@ -214,3 +230,18 @@ function renderResourceType(resourceType?: string) {
 
   return resourceType;
 }
+
+function renderServiceTimeout(serviceTimeout?: Duration): number|undefined {
+
+  if (!serviceTimeout) {
+    return undefined;
+  }
+
+  let timeoutSeconds = serviceTimeout.toSeconds();
+
+  if (timeoutSeconds < 1 || timeoutSeconds > 3600) {
+        throw new Error('ServiceTimeout must be an integer from 1 to 3600');
+  }
+
+  return timeoutSeconds;
+}
diff --git a/packages/aws-cdk-lib/core/test/custom-resource.test.ts b/packages/aws-cdk-lib/core/test/custom-resource.test.ts
index c792e30834cc8..997745e4b02c2 100644
--- a/packages/aws-cdk-lib/core/test/custom-resource.test.ts
+++ b/packages/aws-cdk-lib/core/test/custom-resource.test.ts
@@ -1,5 +1,5 @@
 import { toCloudFormation } from './util';
-import { CustomResource, RemovalPolicy, Stack } from '../lib';
+import {CustomResource, Duration, RemovalPolicy, Stack} from '../lib';
 
 describe('custom resource', () => {
   test('simple case provider identified by service token', () => {
@@ -82,6 +82,30 @@ describe('custom resource', () => {
     });
   });
 
+  test('custom timeout', () => {
+    // GIVEN
+    const stack = new Stack();
+
+    // WHEN
+    new CustomResource(stack, 'MyCustomResource', {
+      serviceToken: 'MyServiceToken',
+      serviceTimeout: Duration.minutes(5),
+    });
+
+    // THEN
+    expect(toCloudFormation(stack)).toEqual({
+      Resources: {
+        MyCustomResource: {
+          Type: 'AWS::CloudFormation::CustomResource',
+          Properties: {
+            ServiceToken: 'MyServiceToken',
+          },
+          ServiceTimeout: 300,
+        },
+      },
+    });
+  });
+
   test('resource type must begin with "Custom::"', () => {
     // GIVEN
     const stack = new Stack();