From 1b340927878947fe28adccdbdda91b2bbc738c30 Mon Sep 17 00:00:00 2001
From: fossand <fossand@amazon.com>
Date: Tue, 22 Aug 2023 16:17:14 -0700
Subject: [PATCH] Add retry to runtime extension

---
 .changeset/chilled-lies-care.md               |  6 +++++
 .../defaultExtensionConfiguration.ts          |  6 ++++-
 .../smithy-client/src/extensions/retry.ts     | 27 +++++++++++++++++++
 .../defaultExtensionConfiguration.ts          |  5 ++--
 packages/types/src/extensions/index.ts        |  1 +
 packages/types/src/extensions/retry.ts        | 10 +++++++
 6 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 .changeset/chilled-lies-care.md
 create mode 100644 packages/smithy-client/src/extensions/retry.ts
 create mode 100644 packages/types/src/extensions/retry.ts

diff --git a/.changeset/chilled-lies-care.md b/.changeset/chilled-lies-care.md
new file mode 100644
index 00000000000..58d4a5cdb03
--- /dev/null
+++ b/.changeset/chilled-lies-care.md
@@ -0,0 +1,6 @@
+---
+"@smithy/smithy-client": minor
+"@smithy/types": minor
+---
+
+Add retry to runtime extension
diff --git a/packages/smithy-client/src/extensions/defaultExtensionConfiguration.ts b/packages/smithy-client/src/extensions/defaultExtensionConfiguration.ts
index 42a3d4243b5..33a521de049 100644
--- a/packages/smithy-client/src/extensions/defaultExtensionConfiguration.ts
+++ b/packages/smithy-client/src/extensions/defaultExtensionConfiguration.ts
@@ -1,11 +1,13 @@
 import type { DefaultExtensionConfiguration } from "@smithy/types";
 
 import { getChecksumConfiguration, resolveChecksumRuntimeConfig } from "./checksum";
+import { getRetryConfiguration, resolveRetryRuntimeConfig } from "./retry";
 
 /**
  * @internal
  */
-export type DefaultExtensionConfigType = Parameters<typeof getChecksumConfiguration>[0];
+export type DefaultExtensionConfigType = Parameters<typeof getChecksumConfiguration>[0] &
+  Parameters<typeof getRetryConfiguration>[0];
 
 /**
  * @internal
@@ -15,6 +17,7 @@ export type DefaultExtensionConfigType = Parameters<typeof getChecksumConfigurat
 export const getDefaultExtensionConfiguration = (runtimeConfig: DefaultExtensionConfigType) => {
   return {
     ...getChecksumConfiguration(runtimeConfig),
+    ...getRetryConfiguration(runtimeConfig),
   };
 };
 
@@ -34,5 +37,6 @@ export const getDefaultClientConfiguration = getDefaultExtensionConfiguration;
 export const resolveDefaultRuntimeConfig = (config: DefaultExtensionConfiguration) => {
   return {
     ...resolveChecksumRuntimeConfig(config),
+    ...resolveRetryRuntimeConfig(config),
   };
 };
diff --git a/packages/smithy-client/src/extensions/retry.ts b/packages/smithy-client/src/extensions/retry.ts
new file mode 100644
index 00000000000..b00edb4386e
--- /dev/null
+++ b/packages/smithy-client/src/extensions/retry.ts
@@ -0,0 +1,27 @@
+import { Provider, RetryStrategy, RetryStrategyConfiguration, RetryStrategyV2 } from "@smithy/types";
+
+/**
+ * @internal
+ */
+export const getRetryConfiguration = (
+  runtimeConfig: Partial<{ retryStrategy: Provider<RetryStrategyV2 | RetryStrategy> }>
+) => {
+  return {
+    _retryStrategy: runtimeConfig.retryStrategy!,
+    setRetryStrategy(retryStrategy: Provider<RetryStrategyV2 | RetryStrategy>): void {
+      this._retryStrategy = retryStrategy;
+    },
+    retryStrategy(): Provider<RetryStrategyV2 | RetryStrategy> {
+      return this._retryStrategy;
+    },
+  };
+};
+
+/**
+ * @internal
+ */
+export const resolveRetryRuntimeConfig = (retryStrategyConfiguration: RetryStrategyConfiguration) => {
+  const runtimeConfig: Partial<Record<string, Provider<RetryStrategyV2 | RetryStrategy>>> = {};
+  runtimeConfig.retryStrategy = retryStrategyConfiguration.retryStrategy();
+  return runtimeConfig;
+};
diff --git a/packages/types/src/extensions/defaultExtensionConfiguration.ts b/packages/types/src/extensions/defaultExtensionConfiguration.ts
index 942cb9bee4c..3aefdbf6e84 100644
--- a/packages/types/src/extensions/defaultExtensionConfiguration.ts
+++ b/packages/types/src/extensions/defaultExtensionConfiguration.ts
@@ -1,10 +1,9 @@
 import { ChecksumConfiguration } from "./checksum";
+import { RetryStrategyConfiguration } from "./retry";
 
 /**
  * @internal
  *
  * Default extension configuration consisting various configurations for modifying a service client
  */
-export interface DefaultExtensionConfiguration extends ChecksumConfiguration {}
-
-type GetDefaultConfigurationType = (runtimeConfig: any) => DefaultExtensionConfiguration;
+export interface DefaultExtensionConfiguration extends ChecksumConfiguration, RetryStrategyConfiguration {}
diff --git a/packages/types/src/extensions/index.ts b/packages/types/src/extensions/index.ts
index 38e7a4dc2cb..cce65a1ce2a 100644
--- a/packages/types/src/extensions/index.ts
+++ b/packages/types/src/extensions/index.ts
@@ -1,3 +1,4 @@
 export * from "./defaultClientConfiguration";
 export * from "./defaultExtensionConfiguration";
 export { AlgorithmId, ChecksumAlgorithm, ChecksumConfiguration } from "./checksum";
+export { RetryStrategyConfiguration } from "./retry";
diff --git a/packages/types/src/extensions/retry.ts b/packages/types/src/extensions/retry.ts
new file mode 100644
index 00000000000..3f0a985e4a4
--- /dev/null
+++ b/packages/types/src/extensions/retry.ts
@@ -0,0 +1,10 @@
+import { RetryStrategyV2 } from "../retry";
+import { Provider, RetryStrategy } from "../util";
+
+/**
+ * @internal
+ */
+export interface RetryStrategyConfiguration {
+  setRetryStrategy(retryStrategy: Provider<RetryStrategyV2 | RetryStrategy>): void;
+  retryStrategy(): Provider<RetryStrategyV2 | RetryStrategy>;
+}