Fix performance bug in MostRecentProvider

Author: SalusaSecondus

src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/providers/MostRecentProvider.java

@@ -70,17 +70,25 @@ public EncryptionMaterials getEncryptionMaterials(EncryptionContext context) {
             final long currentVersion;
             final EncryptionMaterialsProvider currentProvider;
             if (newVersion < 0) {
+                // First version of the material, so we want to allow creation
                 currentVersion = 0;
                 currentProvider = keystore.getOrCreate(materialName, currentVersion);
-                s = new State(currentProvider, currentVersion);
-                state.set(s);
+                cache.add(Long.toString(currentVersion), currentProvider);
             } else if (newVersion != s.currentVersion) {
+                // We're retrieving an existing version, so we avoid the creation
+                // flow as it is slower
                 currentVersion = newVersion;
                 currentProvider = keystore.getProvider(materialName, currentVersion);
                 cache.add(Long.toString(currentVersion), currentProvider);
-                s = new State(currentProvider, currentVersion);
-                state.set(s);
+            } else {
+                // Our version hasn't changed, so we'll just re-use the existing
+                // provider to avoid the overhead of retrieving and building a new one
+                currentVersion = newVersion;
+                currentProvider = s.provider;
+                // There is no need to add this to the cache as it's already there
             }
+            s = new State(currentProvider, currentVersion);
+            state.set(s);
 
             return s.provider.getEncryptionMaterials(context);
         } finally {

src/test/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/providers/MostRecentProviderTests.java

@@ -120,7 +120,7 @@ public void singleMaterial() throws InterruptedException {
 
     @Test
     public void singleMaterialWithRefresh() throws InterruptedException {
-        final MostRecentProvider prov = new MostRecentProvider(store, MATERIAL_NAME, 100);
+        final MostRecentProvider prov = new MostRecentProvider(store, MATERIAL_NAME, 500);
         assertNull(methodCalls.get("putItem"));
         final EncryptionMaterials eMat1 = prov.getEncryptionMaterials(ctx);
         // It's a new provider, so we see a single putItem
@@ -136,10 +136,25 @@ public void singleMaterialWithRefresh() throws InterruptedException {
         assertEquals(1, (int) methodCalls.get("query")); // To find current version
         assertEquals(1, (int) methodCalls.get("getItem"));
         assertEquals(0, store.getVersionFromMaterialDescription(eMat3.getMaterialDescription()));
+        prov.refresh();
 
         assertEquals(eMat1.getSigningKey(), eMat2.getSigningKey());
         assertEquals(eMat1.getSigningKey(), eMat3.getSigningKey());
 
+        // Ensure that after cache refresh  we only get one more hit as opposed to multiple
+        prov.getEncryptionMaterials(ctx);
+        Thread.sleep(700);
+        // Force refresh
+        prov.getEncryptionMaterials(ctx);
+        methodCalls.clear();
+        // Check to ensure no more hits
+        assertEquals(eMat1.getSigningKey(), prov.getEncryptionMaterials(ctx).getSigningKey());
+        assertEquals(eMat1.getSigningKey(), prov.getEncryptionMaterials(ctx).getSigningKey());
+        assertEquals(eMat1.getSigningKey(), prov.getEncryptionMaterials(ctx).getSigningKey());
+        assertEquals(eMat1.getSigningKey(), prov.getEncryptionMaterials(ctx).getSigningKey());
+        assertEquals(eMat1.getSigningKey(), prov.getEncryptionMaterials(ctx).getSigningKey());
+        assertTrue(methodCalls.isEmpty());
+
         // Ensure we can decrypt all of them without hitting ddb more than the minimum
         final MostRecentProvider prov2 = new MostRecentProvider(store, MATERIAL_NAME, 500);
         final DecryptionMaterials dMat1 = prov2.getDecryptionMaterials(ctx(eMat1));