Cache result of endpoint partitioning

smithy-lang · Nov 28, 2023 · efd5032 · efd5032
1 parent baa4186
commit efd5032
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/...main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/rulesgen/StdLib.kt b/...main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/rulesgen/StdLib.kt
@@ -8,6 +8,8 @@ package software.amazon.smithy.rust.codegen.client.smithy.endpoint.rulesgen
 import software.amazon.smithy.model.node.Node
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointsLib
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.CustomRuntimeFunction
+import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.EndpointStdLib
+import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
@@ -47,13 +49,25 @@ class AwsPartitionResolver(runtimeConfig: RuntimeConfig, private val partitionsD
     override val id: String = "aws.partition"
     private val codegenScope = arrayOf(
         "PartitionResolver" to EndpointsLib.PartitionResolver(runtimeConfig),
+        "Lazy" to CargoDependency.OnceCell.toType().resolve("sync::Lazy"),
     )
 
     override fun structFieldInit() = writable {
         val json = Node.printJson(partitionsDotJson).dq()
         rustTemplate(
-            """partition_resolver: #{PartitionResolver}::new_from_json(b$json).expect("valid JSON")""",
+            """partition_resolver: #{DEFAULT_PARTITION_RESOLVER}.clone()""",
             *codegenScope,
+            "DEFAULT_PARTITION_RESOLVER" to RuntimeType.forInlineFun("DEFAULT_PARTITION_RESOLVER", EndpointStdLib) {
+                rustTemplate(
+                    """
+                    // Loading the partition JSON is expensive since it involves many regex compilations,
+                    // so cache the result so that it only need to be paid for the first constructed client.
+                    pub(crate) static DEFAULT_PARTITION_RESOLVER: #{Lazy}<#{PartitionResolver}> =
+                        #{Lazy}::new(|| #{PartitionResolver}::new_from_json(b$json).expect("valid JSON"));
+                    """,
+                    *codegenScope,
+                )
+            },
         )
     }
 

diff --git a/rust-runtime/inlineable/src/endpoint_lib/partition.rs b/rust-runtime/inlineable/src/endpoint_lib/partition.rs
@@ -17,7 +17,7 @@ use std::borrow::Cow;
 use std::collections::HashMap;
 
 /// Determine the AWS partition metadata for a given region
-#[derive(Debug, Default)]
+#[derive(Clone, Debug, Default)]
 pub(crate) struct PartitionResolver {
     partitions: Vec<PartitionMetadata>,
 }
@@ -151,7 +151,7 @@ impl PartitionResolver {
 
 type Str = Cow<'static, str>;
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) struct PartitionMetadata {
     id: Str,
     region_regex: Regex,
@@ -204,7 +204,7 @@ impl PartitionMetadata {
     }
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) struct PartitionOutput {
     name: Str,
     dns_suffix: Str,
@@ -213,7 +213,7 @@ pub(crate) struct PartitionOutput {
     supports_dual_stack: bool,
 }
 
-#[derive(Debug, Default)]
+#[derive(Clone, Debug, Default)]
 pub(crate) struct PartitionOutputOverride {
     name: Option<Str>,
     dns_suffix: Option<Str>,