Refactor caching logic, add test coverage

Author: RanVaknin

.changes/next-release/bugfix-foo-85b64de.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Amazon DynamoDB Enhanced Client",
+    "contributor": "",
+    "description": "Improved performance by caching partition and sort key name lookups in StaticTableMetadata."
+}

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/StaticTableMetadata.java

@@ -87,7 +87,13 @@ public <T> Optional<T> customMetadataObject(String key, Class<? extends T> objec
 
     @Override
     public List<String> indexPartitionKeys(String indexName) {
-        return partitionKeyCache.computeIfAbsent(indexName, this::computePartitionKeys);
+        List<String> cached = partitionKeyCache.get(indexName);
+        if (cached != null) {
+            return cached;
+        }
+        List<String> computedKeys = computePartitionKeys(indexName);
+        List<String> existingKeys = partitionKeyCache.putIfAbsent(indexName, computedKeys);
+        return existingKeys == null ? computedKeys : existingKeys;
     }
 
     private List<String> computePartitionKeys(String indexName) {

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/mapper/StaticTableMetadataTest.java

@@ -397,6 +397,17 @@ public void indexSortKeys_shouldReturnUnmodifiableList() {
         assertThatThrownBy(() -> result.add("foo")).isInstanceOf(UnsupportedOperationException.class);
     }
 
+    @Test
+    public void indexPartitionKeys_shouldReturnUnmodifiableList() {
+        StaticTableMetadata metadata = StaticTableMetadata.builder()
+                                                          .addIndexPartitionKey(primaryIndexName(),
+                                                                           ATTRIBUTE_NAME,
+                                                                           AttributeValueType.S)
+                                                          .build();
+        List<String> result = metadata.indexPartitionKeys(primaryIndexName());
+        assertThatThrownBy(() -> result.add("foo")).isInstanceOf(UnsupportedOperationException.class);
+    }
+
     @Test
     public void getIndexKeys_partitionAndSort() {
         TableMetadata tableMetadata = StaticTableMetadata.builder()