11/*
2- * Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+ * Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
33 *
44 * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except
55 * in compliance with the License. A copy of the License is located at
3030public class MostRecentProvider implements EncryptionMaterialsProvider {
3131 private static final long MILLI_TO_NANO = 1000000L ;
3232 private static final long TTL_GRACE_IN_NANO = 500 * MILLI_TO_NANO ;
33- private final ReentrantLock lock = new ReentrantLock (true );
3433 private final ProviderStore keystore ;
35- private final String materialName ;
34+ protected final String defaultMaterialName ;
3635 private final long ttlInNanos ;
3736 private final LRUCache <EncryptionMaterialsProvider > cache ;
38- private final AtomicReference < State > state = new AtomicReference <>( new State ()) ;
37+ private final LRUCache < LockedState > currentVersions ;
3938
4039 /**
4140 * Creates a new {@link MostRecentProvider}.
@@ -45,22 +44,26 @@ public class MostRecentProvider implements EncryptionMaterialsProvider {
4544 */
4645 public MostRecentProvider (final ProviderStore keystore , final String materialName , final long ttlInMillis ) {
4746 this .keystore = checkNotNull (keystore , "keystore must not be null" );
48- this .materialName = checkNotNull ( materialName , "materialName must not be null" ) ;
47+ this .defaultMaterialName = materialName ;
4948 this .ttlInNanos = ttlInMillis * MILLI_TO_NANO ;
5049 this .cache = new LRUCache <EncryptionMaterialsProvider >(1000 );
50+ this .currentVersions = new LRUCache <>(1000 );
5151 }
5252
5353 @ Override
5454 public EncryptionMaterials getEncryptionMaterials (EncryptionContext context ) {
55- State s = state .get ();
55+ final String materialName = getMaterialName (context );
56+ final LockedState ls = getCurrentVersion (materialName );
57+
58+ final State s = ls .getState ();
5659 if (s .provider != null && System .nanoTime () - s .lastUpdated <= ttlInNanos ) {
5760 return s .provider .getEncryptionMaterials (context );
5861 }
5962 if (s .provider == null || System .nanoTime () - s .lastUpdated > ttlInNanos + TTL_GRACE_IN_NANO ) {
6063 // Either we don't have a provider at all, or we're more than 500 milliseconds past
6164 // our update time. Either way, grab the lock and force an update.
62- lock .lock ();
63- } else if (!lock .tryLock ()) {
65+ ls .lock ();
66+ } else if (!ls .tryLock ()) {
6467 // If we can't get the lock immediately, just use the current provider
6568 return s .provider .getEncryptionMaterials (context );
6669 }
@@ -73,36 +76,37 @@ public EncryptionMaterials getEncryptionMaterials(EncryptionContext context) {
7376 // First version of the material, so we want to allow creation
7477 currentVersion = 0 ;
7578 currentProvider = keystore .getOrCreate (materialName , currentVersion );
76- cache .add (Long . toString ( currentVersion ), currentProvider );
79+ cache .add (buildCacheKey ( materialName , currentVersion ), currentProvider );
7780 } else if (newVersion != s .currentVersion ) {
7881 // We're retrieving an existing version, so we avoid the creation
7982 // flow as it is slower
8083 currentVersion = newVersion ;
8184 currentProvider = keystore .getProvider (materialName , currentVersion );
82- cache .add (Long . toString ( currentVersion ), currentProvider );
85+ cache .add (buildCacheKey ( materialName , currentVersion ), currentProvider );
8386 } else {
8487 // Our version hasn't changed, so we'll just re-use the existing
8588 // provider to avoid the overhead of retrieving and building a new one
8689 currentVersion = newVersion ;
8790 currentProvider = s .provider ;
8891 // There is no need to add this to the cache as it's already there
8992 }
90- s = new State (currentProvider , currentVersion );
91- state .set (s );
9293
93- return s .provider .getEncryptionMaterials (context );
94+ ls .update (currentProvider , currentVersion );
95+
96+ return ls .getState ().provider .getEncryptionMaterials (context );
9497 } finally {
95- lock .unlock ();
98+ ls .unlock ();
9699 }
97100 }
98101
99102 public DecryptionMaterials getDecryptionMaterials (EncryptionContext context ) {
103+ final String materialName = getMaterialName (context );
100104 final long version = keystore .getVersionFromMaterialDescription (
101105 context .getMaterialDescription ());
102- EncryptionMaterialsProvider provider = cache .get (Long . toString ( version ));
106+ EncryptionMaterialsProvider provider = cache .get (buildCacheKey ( materialName , version ));
103107 if (provider == null ) {
104108 provider = keystore .getProvider (materialName , version );
105- cache .add (Long . toString ( version ), provider );
109+ cache .add (buildCacheKey ( materialName , version ), provider );
106110 }
107111 return provider .getDecryptionMaterials (context );
108112 }
@@ -112,12 +116,12 @@ public DecryptionMaterials getDecryptionMaterials(EncryptionContext context) {
112116 */
113117 @ Override
114118 public void refresh () {
115- state . set ( new State () );
119+ currentVersions . clear ( );
116120 cache .clear ();
117121 }
118122
119123 public String getMaterialName () {
120- return materialName ;
124+ return defaultMaterialName ;
121125 }
122126
123127 public long getTtlInMills () {
@@ -129,15 +133,36 @@ public long getTtlInMills() {
129133 * currently have a current version.
130134 */
131135 public long getCurrentVersion () {
132- return state . get ().currentVersion ;
136+ return getCurrentVersion ( getMaterialName ()). getState ().currentVersion ;
133137 }
134138
135139 /**
136140 * The last time the current version was updated. Returns 0 if we do not currently have a
137141 * current version.
138142 */
139143 public long getLastUpdated () {
140- return state .get ().lastUpdated / MILLI_TO_NANO ;
144+ return getCurrentVersion (getMaterialName ()).getState ().lastUpdated / MILLI_TO_NANO ;
145+ }
146+
147+ protected String getMaterialName (final EncryptionContext context ) {
148+ return defaultMaterialName ;
149+ }
150+
151+ private LockedState getCurrentVersion (final String materialName ) {
152+ final LockedState result = currentVersions .get (materialName );
153+ if (result == null ) {
154+ currentVersions .add (materialName , new LockedState ());
155+ return currentVersions .get (materialName );
156+ } else {
157+ return result ;
158+ }
159+ }
160+
161+ private static String buildCacheKey (final String materialName , final long version ) {
162+ StringBuilder result = new StringBuilder (materialName );
163+ result .append ('#' );
164+ result .append (version );
165+ return result .toString ();
141166 }
142167
143168 private static <V > V checkNotNull (final V ref , final String errMsg ) {
@@ -148,6 +173,34 @@ private static <V> V checkNotNull(final V ref, final String errMsg) {
148173 }
149174 }
150175
176+ private static class LockedState {
177+ private final ReentrantLock lock = new ReentrantLock (true );
178+ private volatile AtomicReference <State > state = new AtomicReference <>(new State ());
179+
180+ public State getState () {
181+ return state .get ();
182+ }
183+
184+ public void unlock () {
185+ lock .unlock ();
186+ }
187+
188+ public boolean tryLock () {
189+ return lock .tryLock ();
190+ }
191+
192+ public void lock () {
193+ lock .lock ();
194+ }
195+
196+ public void update (EncryptionMaterialsProvider provider , long currentVersion ) {
197+ if (!lock .isHeldByCurrentThread ()) {
198+ throw new IllegalStateException ("Lock not held by current thread" );
199+ }
200+ state .set (new State (provider , currentVersion ));
201+ }
202+ }
203+
151204 private static class State {
152205 public final EncryptionMaterialsProvider provider ;
153206 public final long currentVersion ;
0 commit comments