1 /*
2  * JBoss, Home of Professional Open Source.
3  * Copyright 2016 Red Hat, Inc., and individual contributors
4  * as indicated by the @author tags.
5  *
6  * Licensed under the Apache License, Version 2.0 (the "License");
7  * you may not use this file except in compliance with the License.
8  * You may obtain a copy of the License at
9  *
10  *     http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing, software
13  * distributed under the License is distributed on an "AS IS" BASIS,
14  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15  * See the License for the specific language governing permissions and
16  * limitations under the License.
17  */

18
19 package org.wildfly.common.context;
20
21 import static java.security.AccessController.doPrivileged;
22
23 import java.security.PrivilegedAction;
24 import java.util.concurrent.ConcurrentHashMap;
25 import java.util.concurrent.atomic.AtomicReference;
26 import java.util.function.Supplier;
27
28 import org.wildfly.common.Assert;
29
30 /**
31  * A context manager for a {@link Contextual} type.
32  *
33  * @param <C> the public type of the contextual object
34  * @author <a href="mailto:david.lloyd@redhat.com">David M. Lloyd</a>
35  */

36 public final class ContextManager<C extends Contextual<C>> implements Supplier<C> {
37     private final AtomicReference<Supplier<C>> globalDefaultSupplierRef = new AtomicReference<>();
38     private final ConcurrentHashMap<ClassLoader, Supplier<C>> perClassLoaderDefault = new ConcurrentHashMap<>();
39     private final Class<C> type;
40     private final String name;
41     private final ThreadLocal<State<C>> stateRef = ThreadLocal.withInitial(State::new);
42     private final ContextPermission getPermission;
43
44     /**
45      * Construct a new instance, with a name matching the class name of the given {@code type}.
46      *
47      * @param type the type class of the context object (must not be {@code null})
48      */

49     public ContextManager(final Class<C> type) {
50         this(type, type.getName());
51     }
52
53     /**
54      * Construct a new instance.
55      *
56      * @param type the type class of the context object (must not be {@code null})
57      * @param name the name to use for permission checks (must not be {@code null} or empty)
58      */

59     public ContextManager(final Class<C> type, final String name) {
60         Assert.checkNotNullParam("type", type);
61         Assert.checkNotNullParam("name", name);
62         Assert.checkNotEmptyParam("name", name);
63         this.type = type;
64         this.name = name;
65         // construct commonly-used permission object
66         getPermission = new ContextPermission(name, ContextPermission.STR_GET);
67     }
68
69     /**
70      * Get the global default context instance.  Note that the global default is determined by way of a {@link Supplier} so
71      * the returned value may vary from call to call, depending on the policy of that {@code Supplier}.
72      *
73      * @return the global default, or {@code nullif none is installed or available
74      */

75     public C getGlobalDefault() {
76         final SecurityManager sm = System.getSecurityManager();
77         if (sm != null) {
78             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_GET_GLOBAL_DEF));
79         }
80         final Supplier<C> globalDefault = globalDefaultSupplierRef.get();
81         return globalDefault == null ? null : globalDefault.get();
82     }
83
84     /**
85      * Set the global default instance supplier.  The supplier, if one is given, should have a reasonable policy such
86      * that callers of {@link #getGlobalDefault()} will obtain results consistent with a general expectation of stability.
87      *
88      * @param supplier the supplier, or {@code null} to remove the global default
89      */

90     public void setGlobalDefaultSupplier(final Supplier<C> supplier) {
91         final SecurityManager sm = System.getSecurityManager();
92         if (sm != null) {
93             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_SET_GLOBAL_DEF_SUP));
94         }
95         globalDefaultSupplierRef.set(supplier);
96     }
97
98     /**
99      * Set the global default instance supplier, but only if it was not already set.  If no supplier is set, the given
100      * supplier supplier is queried to get the new value to set.
101      *
102      * @param supplierSupplier the supplier supplier (must not be {@code null})
103      * @return {@code trueif the supplier was set, {@code falseif it was already set to something else
104      * @see #setGlobalDefaultSupplier(Supplier)
105      */

106     public boolean setGlobalDefaultSupplierIfNotSet(final Supplier<Supplier<C>> supplierSupplier) {
107         final SecurityManager sm = System.getSecurityManager();
108         if (sm != null) {
109             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_SET_GLOBAL_DEF_SUP));
110         }
111         final AtomicReference<Supplier<C>> ref = this.globalDefaultSupplierRef;
112         // try not to compute the value if not needed
113         return ref.get() == null && ref.compareAndSet(null, supplierSupplier.get());
114     }
115
116     /**
117      * Set the global default instance.  This instance will be returned from all subsequent calls to {@link #getGlobalDefault()},
118      * replacing any previous instance or {@linkplain #setGlobalDefaultSupplier(Supplier) supplier} that was set.
119      *
120      * @param globalDefault the global default value, or {@code null} to remove the global default
121      */

122     public void setGlobalDefault(final C globalDefault) {
123         final SecurityManager sm = System.getSecurityManager();
124         if (sm != null) {
125             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_GET_GLOBAL_DEF));
126         }
127         globalDefaultSupplierRef.set(globalDefault == null ? null : () -> globalDefault);
128     }
129
130     /**
131      * Get the class loader default instance.  Note that the class loader default is determined by way of a {@link Supplier} so
132      * the returned value may vary from call to call, depending on the policy of that {@code Supplier}.
133      *
134      * @param classLoader the class loader
135      * @return the global default, or {@code nullif none is installed or available
136      */

137     public C getClassLoaderDefault(final ClassLoader classLoader) {
138         final SecurityManager sm = System.getSecurityManager();
139         if (sm != null) {
140             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_GET_CLASSLOADER_DEF));
141         }
142         final Supplier<C> supplier;
143         if (classLoader == null) {
144             return null;
145         }
146         supplier = perClassLoaderDefault.get(classLoader);
147         return supplier == null ? null : supplier.get();
148     }
149
150     /**
151      * Set the per-class loader default instance supplier.  The supplier, if one is given, should have a reasonable policy such
152      * that callers of {@link #getClassLoaderDefault(ClassLoader)} will obtain results consistent with a general expectation of stability.
153      *
154      * @param classLoader the class loader (must not be {@code null})
155      * @param supplier the supplier, or {@code null} to remove the default for this class loader
156      */

157     public void setClassLoaderDefaultSupplier(final ClassLoader classLoader, final Supplier<C> supplier) {
158         Assert.checkNotNullParam("classLoader", classLoader);
159         final SecurityManager sm = System.getSecurityManager();
160         if (sm != null) {
161             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_SET_CLASSLOADER_DEF_SUP));
162         }
163         if (supplier == null) {
164             perClassLoaderDefault.remove(classLoader);
165         } else {
166             perClassLoaderDefault.put(classLoader, supplier);
167         }
168     }
169
170     /**
171      * Set the per-class loader default instance supplier.  The supplier, if one is given, should have a reasonable policy such
172      * that callers of {@link #getClassLoaderDefault(ClassLoader)} will obtain results consistent with a general expectation of stability.
173      *
174      * @param classLoader the class loader (must not be {@code null})
175      * @param classLoaderDefault the class loader default value, or {@code null} to remove the default
176      */

177     public void setClassLoaderDefault(final ClassLoader classLoader, final C classLoaderDefault) {
178         Assert.checkNotNullParam("classLoader", classLoader);
179         final SecurityManager sm = System.getSecurityManager();
180         if (sm != null) {
181             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_SET_CLASSLOADER_DEF));
182         }
183         if (classLoaderDefault == null) {
184             perClassLoaderDefault.remove(classLoader);
185         } else {
186             perClassLoaderDefault.put(classLoader, () -> classLoaderDefault);
187         }
188     }
189
190     /**
191      * Get the per-thread default context instance.  Note that the per-thread default is determined by way of a {@link Supplier} so
192      * the returned value may vary from call to call, depending on the policy of that {@code Supplier}.
193      *
194      * @return the per-thread default, or {@code nullif none is installed or available
195      */

196     public C getThreadDefault() {
197         final SecurityManager sm = System.getSecurityManager();
198         if (sm != null) {
199             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_GET_THREAD_DEF));
200         }
201         final Supplier<C> defaultSupplier = stateRef.get().defaultSupplier;
202         return defaultSupplier == null ? null : defaultSupplier.get();
203     }
204
205     /**
206      * Set the per-thread default instance supplier.  The supplier, if one is given, should have a reasonable policy such
207      * that callers of {@link #getThreadDefault()} will obtain results consistent with a general expectation of stability.
208      *
209      * @param supplier the supplier, or {@code null} to remove the per-thread default
210      */

211     public void setThreadDefaultSupplier(final Supplier<C> supplier) {
212         final SecurityManager sm = System.getSecurityManager();
213         if (sm != null) {
214             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_SET_THREAD_DEF_SUP));
215         }
216         stateRef.get().defaultSupplier = supplier;
217     }
218
219     /**
220      * Set the per-thread default instance.  This instance will be returned from all subsequent calls to {@link #getThreadDefault()},
221      * replacing any previous instance or {@linkplain #setThreadDefaultSupplier(Supplier) supplier} that was set.
222      *
223      * @param threadDefault the per-thread default value, or {@code null} to remove the per-thread default
224      */

225     public void setThreadDefault(final C threadDefault) {
226         final SecurityManager sm = System.getSecurityManager();
227         if (sm != null) {
228             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_SET_THREAD_DEF));
229         }
230         stateRef.get().defaultSupplier = threadDefault == null ? null : () -> threadDefault;
231     }
232
233     /**
234      * Get the currently active context, possibly examining the per-thread or global defaults.
235      *
236      * @return the current context, or {@code nullif none is active
237      */

238     public C get() {
239         final SecurityManager sm = System.getSecurityManager();
240         if (sm != null) {
241             sm.checkPermission(getPermission);
242         }
243         return getPrivileged();
244     }
245
246     /**
247      * Get a privileged supplier for this context manager which returns the currently active context without a permission
248      * check.
249      *
250      * @return the privileged supplier
251      */

252     public Supplier<C> getPrivilegedSupplier() {
253         final SecurityManager sm = System.getSecurityManager();
254         if (sm != null) {
255             sm.checkPermission(new ContextPermission(name, ContextPermission.STR_GET_PRIV_SUP));
256         }
257         return this::getPrivileged;
258     }
259
260     private C getPrivileged() {
261         final State<C> state = stateRef.get();
262         C c = state.current;
263         if (c != nullreturn c;
264         final Thread currentThread = Thread.currentThread();
265         final SecurityManager sm = System.getSecurityManager();
266         ClassLoader classLoader;
267         if (sm != null) {
268             classLoader = doPrivileged((PrivilegedAction<ClassLoader>) currentThread::getContextClassLoader);
269         } else {
270             classLoader = currentThread.getContextClassLoader();
271         }
272         Supplier<C> supplier;
273         if (classLoader != null) {
274             supplier = perClassLoaderDefault.get(classLoader);
275             if (supplier != null) {
276                 c = supplier.get();
277                 if (c != nullreturn c;
278             }
279         }
280         supplier = state.defaultSupplier;
281         if (supplier != null) {
282             c = supplier.get();
283             if (c != nullreturn c;
284         }
285         supplier = globalDefaultSupplierRef.get();
286         return supplier != null ? supplier.get() : null;
287     }
288
289     C getAndSetCurrent(Contextual<C> newVal) {
290         final C cast = type.cast(newVal);
291         final State<C> state = stateRef.get();
292         try {
293             return state.current;
294         } finally {
295             state.current = cast;
296         }
297     }
298
299     void restoreCurrent(C oldVal) {
300         stateRef.get().current = oldVal;
301     }
302
303     static class State<T> {
304         T current;
305         Supplier<T> defaultSupplier;
306
307         State() {
308         }
309     }
310 }
311