diff --git a/context/src/main/java/io/grpc/Context.java b/context/src/main/java/io/grpc/Context.java index abb09427f9..02cf684ff1 100644 --- a/context/src/main/java/io/grpc/Context.java +++ b/context/src/main/java/io/grpc/Context.java @@ -116,42 +116,42 @@ public class Context { */ public static final Context ROOT = new Context(null, EMPTY_ENTRIES); + // Visible For testing + static Storage storage() { + return LazyStorage.storage; + } + // Lazy-loaded storage. Delaying storage initialization until after class initialization makes it // much easier to avoid circular loading since there can still be references to Context as long as // they don't depend on storage, like key() and currentContextExecutor(). It also makes it easier // to handle exceptions. - private static final AtomicReference storage = new AtomicReference<>(); + private static final class LazyStorage { + static final Storage storage; - // For testing - static Storage storage() { - Storage tmp = storage.get(); - if (tmp == null) { - tmp = createStorage(); - } - return tmp; - } - - private static Storage createStorage() { - // Note that this method may be run more than once - try { - Class clazz = Class.forName("io.grpc.override.ContextStorageOverride"); - // The override's constructor is prohibited from triggering any code that can loop back to - // Context - Storage newStorage = (Storage) clazz.getConstructor().newInstance(); - storage.compareAndSet(null, newStorage); - } catch (ClassNotFoundException e) { - Storage newStorage = new ThreadLocalContextStorage(); - // Must set storage before logging, since logging may call Context.current(). - if (storage.compareAndSet(null, newStorage)) { - // Avoid logging if this thread lost the race, to avoid confusion - log.log(Level.FINE, "Storage override doesn't exist. Using default", e); + static { + AtomicReference deferredStorageFailure = new AtomicReference<>(); + storage = createStorage(deferredStorageFailure); + Throwable failure = deferredStorageFailure.get(); + // Logging must happen after storage has been set, as loggers may use Context. + if (failure != null) { + log.log(Level.FINE, "Storage override doesn't exist. Using default", failure); + } + } + + private static Storage createStorage( + AtomicReference deferredStorageFailure) { + try { + Class clazz = Class.forName("io.grpc.override.ContextStorageOverride"); + // The override's constructor is prohibited from triggering any code that can loop back to + // Context + return clazz.asSubclass(Storage.class).getConstructor().newInstance(); + } catch (ClassNotFoundException e) { + deferredStorageFailure.set(e); + return new ThreadLocalContextStorage(); + } catch (Exception e) { + throw new RuntimeException("Storage override failed to initialize", e); } - } catch (Exception e) { - throw new RuntimeException("Storage override failed to initialize", e); } - // Re-retreive from storage since compareAndSet may have failed (returned false) in case of - // race. - return storage.get(); } /** diff --git a/context/src/test/java/io/grpc/ContextTest.java b/context/src/test/java/io/grpc/ContextTest.java index 56ba22b2a7..8feef85c1a 100644 --- a/context/src/test/java/io/grpc/ContextTest.java +++ b/context/src/test/java/io/grpc/ContextTest.java @@ -840,16 +840,19 @@ public class ContextTest { @Test public void storageReturnsNullTest() throws Exception { - Field storage = Context.class.getDeclaredField("storage"); + Class lazyStorageClass = Class.forName("io.grpc.Context$LazyStorage"); + Field storage = lazyStorageClass.getDeclaredField("storage"); assertTrue(Modifier.isFinal(storage.getModifiers())); // use reflection to forcibly change the storage object to a test object storage.setAccessible(true); + Field modifiersField = Field.class.getDeclaredField("modifiers"); + modifiersField.setAccessible(true); + int storageModifiers = modifiersField.getInt(storage); + modifiersField.set(storage, storageModifiers & ~Modifier.FINAL); Object o = storage.get(null); - @SuppressWarnings("unchecked") - AtomicReference storageRef = (AtomicReference) o; - Context.Storage originalStorage = storageRef.get(); + Context.Storage originalStorage = (Context.Storage) o; try { - storageRef.set(new Context.Storage() { + storage.set(null, new Context.Storage() { @Override public Context doAttach(Context toAttach) { return null; @@ -878,8 +881,10 @@ public class ContextTest { assertEquals(Context.ROOT, Context.current()); } finally { // undo the changes - storageRef.set(originalStorage); + storage.set(null, originalStorage); storage.setAccessible(false); + modifiersField.set(storage, storageModifiers | Modifier.FINAL); + modifiersField.setAccessible(false); } }