/* * Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved. * * Licensed under the Apache License, Version 2.0 (the "License"). You may not * use this file except in compliance with the License. A copy of the License * is located at * * http://aws.amazon.com/apache2.0/ * * or in the "license" file accompanying this file. This file is distributed * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either * express or implied. See the License for the specific language governing * permissions and limitations under the License. */ package software.amazon.swage.threadcontext; import java.util.concurrent.Callable; import java.util.concurrent.CompletionService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; /** * A CompletionService that propagates the current ThreadContext to tasks it executes. * * @param <V> type of result produced */ public class ContextAwareCompletionService<V> implements CompletionService<V> { private final CompletionService<V> delegate; /** * Constructor providing the CompletionService to use to execute the tasks. * * @param delegate an CompletionService to delegate to */ public ContextAwareCompletionService(CompletionService<V> delegate) { this.delegate = delegate; } @Override public Future<V> submit(Callable<V> task) { return delegate.submit(ThreadContext.current().wrap(task)); } @Override public Future<V> submit(Runnable task, V result) { return delegate.submit(ThreadContext.current().wrap(task), result); } @Override public Future<V> take() throws InterruptedException { return delegate.take(); } @Override public Future<V> poll() { return delegate.poll(); } @Override public Future<V> poll(long timeout, TimeUnit unit) throws InterruptedException { return delegate.poll(timeout, unit); } }