001 /**
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements. See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership. The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License. You may obtain a copy of the License at
009 *
010 * http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018 package org.apache.hadoop.crypto.key.kms;
019
020 import java.io.IOException;
021 import java.util.HashSet;
022 import java.util.LinkedList;
023 import java.util.List;
024 import java.util.Queue;
025 import java.util.concurrent.ExecutionException;
026 import java.util.concurrent.LinkedBlockingQueue;
027 import java.util.concurrent.ThreadPoolExecutor;
028 import java.util.concurrent.TimeUnit;
029
030 import com.google.common.base.Preconditions;
031 import com.google.common.cache.CacheBuilder;
032 import com.google.common.cache.CacheLoader;
033 import com.google.common.cache.LoadingCache;
034 import com.google.common.util.concurrent.ThreadFactoryBuilder;
035 import org.apache.hadoop.classification.InterfaceAudience;
036
037 /**
038 * A Utility class that maintains a Queue of entries for a given key. It tries
039 * to ensure that there is are always at-least <code>numValues</code> entries
040 * available for the client to consume for a particular key.
041 * It also uses an underlying Cache to evict queues for keys that have not been
042 * accessed for a configurable period of time.
043 * Implementing classes are required to implement the
044 * <code>QueueRefiller</code> interface that exposes a method to refill the
045 * queue, when empty
046 */
047 @InterfaceAudience.Private
048 public class ValueQueue <E> {
049
050 /**
051 * QueueRefiller interface a client must implement to use this class
052 */
053 public interface QueueRefiller <E> {
054 /**
055 * Method that has to be implemented by implementing classes to fill the
056 * Queue.
057 * @param keyName Key name
058 * @param keyQueue Queue that needs to be filled
059 * @param numValues number of Values to be added to the queue.
060 * @throws IOException
061 */
062 public void fillQueueForKey(String keyName,
063 Queue<E> keyQueue, int numValues) throws IOException;
064 }
065
066 private static final String REFILL_THREAD =
067 ValueQueue.class.getName() + "_thread";
068
069 private final LoadingCache<String, LinkedBlockingQueue<E>> keyQueues;
070 private final ThreadPoolExecutor executor;
071 private final UniqueKeyBlockingQueue queue = new UniqueKeyBlockingQueue();
072 private final QueueRefiller<E> refiller;
073 private final SyncGenerationPolicy policy;
074
075 private final int numValues;
076 private final float lowWatermark;
077
078 private volatile boolean executorThreadsStarted = false;
079
080 /**
081 * A <code>Runnable</code> which takes a string name.
082 */
083 private abstract static class NamedRunnable implements Runnable {
084 final String name;
085 private NamedRunnable(String keyName) {
086 this.name = keyName;
087 }
088 }
089
090 /**
091 * This backing blocking queue used in conjunction with the
092 * <code>ThreadPoolExecutor</code> used by the <code>ValueQueue</code>. This
093 * Queue accepts a task only if the task is not currently in the process
094 * of being run by a thread which is implied by the presence of the key
095 * in the <code>keysInProgress</code> set.
096 *
097 * NOTE: Only methods that ware explicitly called by the
098 * <code>ThreadPoolExecutor</code> need to be over-ridden.
099 */
100 private static class UniqueKeyBlockingQueue extends
101 LinkedBlockingQueue<Runnable> {
102
103 private static final long serialVersionUID = -2152747693695890371L;
104 private HashSet<String> keysInProgress = new HashSet<String>();
105
106 @Override
107 public synchronized void put(Runnable e) throws InterruptedException {
108 if (keysInProgress.add(((NamedRunnable)e).name)) {
109 super.put(e);
110 }
111 }
112
113 @Override
114 public Runnable take() throws InterruptedException {
115 Runnable k = super.take();
116 if (k != null) {
117 keysInProgress.remove(((NamedRunnable)k).name);
118 }
119 return k;
120 }
121
122 @Override
123 public Runnable poll(long timeout, TimeUnit unit)
124 throws InterruptedException {
125 Runnable k = super.poll(timeout, unit);
126 if (k != null) {
127 keysInProgress.remove(((NamedRunnable)k).name);
128 }
129 return k;
130 }
131
132 }
133
134 /**
135 * Policy to decide how many values to return to client when client asks for
136 * "n" values and Queue is empty.
137 * This decides how many values to return when client calls "getAtMost"
138 */
139 public static enum SyncGenerationPolicy {
140 ATLEAST_ONE, // Return atleast 1 value
141 LOW_WATERMARK, // Return min(n, lowWatermark * numValues) values
142 ALL // Return n values
143 }
144
145 /**
146 * Constructor takes the following tunable configuration parameters
147 * @param numValues The number of values cached in the Queue for a
148 * particular key.
149 * @param lowWatermark The ratio of (number of current entries/numValues)
150 * below which the <code>fillQueueForKey()</code> funciton will be
151 * invoked to fill the Queue.
152 * @param expiry Expiry time after which the Key and associated Queue are
153 * evicted from the cache.
154 * @param numFillerThreads Number of threads to use for the filler thread
155 * @param policy The SyncGenerationPolicy to use when client
156 * calls "getAtMost"
157 * @param refiller implementation of the QueueRefiller
158 */
159 public ValueQueue(final int numValues, final float lowWatermark,
160 long expiry, int numFillerThreads, SyncGenerationPolicy policy,
161 final QueueRefiller<E> refiller) {
162 Preconditions.checkArgument(numValues > 0, "\"numValues\" must be > 0");
163 Preconditions.checkArgument(((lowWatermark > 0)&&(lowWatermark <= 1)),
164 "\"lowWatermark\" must be > 0 and <= 1");
165 Preconditions.checkArgument(expiry > 0, "\"expiry\" must be > 0");
166 Preconditions.checkArgument(numFillerThreads > 0,
167 "\"numFillerThreads\" must be > 0");
168 Preconditions.checkNotNull(policy, "\"policy\" must not be null");
169 this.refiller = refiller;
170 this.policy = policy;
171 this.numValues = numValues;
172 this.lowWatermark = lowWatermark;
173 keyQueues = CacheBuilder.newBuilder()
174 .expireAfterAccess(expiry, TimeUnit.MILLISECONDS)
175 .build(new CacheLoader<String, LinkedBlockingQueue<E>>() {
176 @Override
177 public LinkedBlockingQueue<E> load(String keyName)
178 throws Exception {
179 LinkedBlockingQueue<E> keyQueue =
180 new LinkedBlockingQueue<E>();
181 refiller.fillQueueForKey(keyName, keyQueue,
182 (int)(lowWatermark * numValues));
183 return keyQueue;
184 }
185 });
186
187 executor =
188 new ThreadPoolExecutor(numFillerThreads, numFillerThreads, 0L,
189 TimeUnit.MILLISECONDS, queue, new ThreadFactoryBuilder()
190 .setDaemon(true)
191 .setNameFormat(REFILL_THREAD).build());
192 }
193
194 public ValueQueue(final int numValues, final float lowWaterMark, long expiry,
195 int numFillerThreads, QueueRefiller<E> fetcher) {
196 this(numValues, lowWaterMark, expiry, numFillerThreads,
197 SyncGenerationPolicy.ALL, fetcher);
198 }
199
200 /**
201 * Initializes the Value Queues for the provided keys by calling the
202 * fill Method with "numInitValues" values
203 * @param keyNames Array of key Names
204 * @throws ExecutionException
205 */
206 public void initializeQueuesForKeys(String... keyNames)
207 throws ExecutionException {
208 for (String keyName : keyNames) {
209 keyQueues.get(keyName);
210 }
211 }
212
213 /**
214 * This removes the value currently at the head of the Queue for the
215 * provided key. Will immediately fire the Queue filler function if key
216 * does not exist.
217 * If Queue exists but all values are drained, It will ask the generator
218 * function to add 1 value to Queue and then drain it.
219 * @param keyName String key name
220 * @return E the next value in the Queue
221 * @throws IOException
222 * @throws ExecutionException
223 */
224 public E getNext(String keyName)
225 throws IOException, ExecutionException {
226 return getAtMost(keyName, 1).get(0);
227 }
228
229 /**
230 * Drains the Queue for the provided key.
231 *
232 * @param keyName the key to drain the Queue for
233 */
234 public void drain(String keyName ) {
235 try {
236 keyQueues.get(keyName).clear();
237 } catch (ExecutionException ex) {
238 //NOP
239 }
240 }
241
242 /**
243 * This removes the "num" values currently at the head of the Queue for the
244 * provided key. Will immediately fire the Queue filler function if key
245 * does not exist
246 * How many values are actually returned is governed by the
247 * <code>SyncGenerationPolicy</code> specified by the user.
248 * @param keyName String key name
249 * @param num Minimum number of values to return.
250 * @return List<E> values returned
251 * @throws IOException
252 * @throws ExecutionException
253 */
254 public List<E> getAtMost(String keyName, int num) throws IOException,
255 ExecutionException {
256 LinkedBlockingQueue<E> keyQueue = keyQueues.get(keyName);
257 // Using poll to avoid race condition..
258 LinkedList<E> ekvs = new LinkedList<E>();
259 try {
260 for (int i = 0; i < num; i++) {
261 E val = keyQueue.poll();
262 // If queue is empty now, Based on the provided SyncGenerationPolicy,
263 // figure out how many new values need to be generated synchronously
264 if (val == null) {
265 // Synchronous call to get remaining values
266 int numToFill = 0;
267 switch (policy) {
268 case ATLEAST_ONE:
269 numToFill = (ekvs.size() < 1) ? 1 : 0;
270 break;
271 case LOW_WATERMARK:
272 numToFill =
273 Math.min(num, (int) (lowWatermark * numValues)) - ekvs.size();
274 break;
275 case ALL:
276 numToFill = num - ekvs.size();
277 break;
278 }
279 // Synchronous fill if not enough values found
280 if (numToFill > 0) {
281 refiller.fillQueueForKey(keyName, ekvs, numToFill);
282 }
283 // Asynch task to fill > lowWatermark
284 if (i <= (int) (lowWatermark * numValues)) {
285 submitRefillTask(keyName, keyQueue);
286 }
287 return ekvs;
288 }
289 ekvs.add(val);
290 }
291 } catch (Exception e) {
292 throw new IOException("Exeption while contacting value generator ", e);
293 }
294 return ekvs;
295 }
296
297 private void submitRefillTask(final String keyName,
298 final Queue<E> keyQueue) throws InterruptedException {
299 if (!executorThreadsStarted) {
300 synchronized (this) {
301 // To ensure all requests are first queued, make coreThreads =
302 // maxThreads
303 // and pre-start all the Core Threads.
304 executor.prestartAllCoreThreads();
305 executorThreadsStarted = true;
306 }
307 }
308 // The submit/execute method of the ThreadPoolExecutor is bypassed and
309 // the Runnable is directly put in the backing BlockingQueue so that we
310 // can control exactly how the runnable is inserted into the queue.
311 queue.put(
312 new NamedRunnable(keyName) {
313 @Override
314 public void run() {
315 int cacheSize = numValues;
316 int threshold = (int) (lowWatermark * (float) cacheSize);
317 // Need to ensure that only one refill task per key is executed
318 try {
319 if (keyQueue.size() < threshold) {
320 refiller.fillQueueForKey(name, keyQueue,
321 cacheSize - keyQueue.size());
322 }
323 } catch (final Exception e) {
324 throw new RuntimeException(e);
325 }
326 }
327 }
328 );
329 }
330
331 /**
332 * Cleanly shutdown
333 */
334 public void shutdown() {
335 executor.shutdownNow();
336 }
337
338 }