diff --git a/core/deployment/src/main/java/org/jboss/shamrock/deployment/codegen/BytecodeRecorderImpl.java b/core/deployment/src/main/java/org/jboss/shamrock/deployment/codegen/BytecodeRecorderImpl.java index 8b203cabaeb89..443083777ebf0 100644 --- a/core/deployment/src/main/java/org/jboss/shamrock/deployment/codegen/BytecodeRecorderImpl.java +++ b/core/deployment/src/main/java/org/jboss/shamrock/deployment/codegen/BytecodeRecorderImpl.java @@ -307,11 +307,19 @@ public void close() throws IOException { if (call.method.getReturnType() != void.class) { ContextObject annotation = call.method.getAnnotation(ContextObject.class); if (annotation != null) { + ca.dup(); ca.aload(1); ca.swap(); ca.ldc(annotation.value()); ca.swap(); ca.invokevirtual(StartupContext.class.getName(), "putValue", "(Ljava/lang/String;Ljava/lang/Object;)V"); + if(call.returnedProxy != null) { + Integer pos = returnValuePositions.get(call.returnedProxy); + if (pos == null) { + returnValuePositions.put(call.returnedProxy, pos = localVarCounter++); + } + ca.astore(pos); + } } else if (call.returnedProxy != null) { Integer pos = returnValuePositions.get(call.returnedProxy); if (pos == null) { diff --git a/weld/deployment/src/main/java/org/jboss/shamrock/weld/deployment/WeldAnnotationProcessor.java b/weld/deployment/src/main/java/org/jboss/shamrock/weld/deployment/WeldAnnotationProcessor.java index 1620f1d8664b6..1e50f9e7b07b4 100644 --- a/weld/deployment/src/main/java/org/jboss/shamrock/weld/deployment/WeldAnnotationProcessor.java +++ b/weld/deployment/src/main/java/org/jboss/shamrock/weld/deployment/WeldAnnotationProcessor.java @@ -43,6 +43,10 @@ public void process(ArchiveContext archiveContext, ProcessorContext processorCon SeContainer weld = template.doBoot(init); template.setupInjection(weld); } + try(BytecodeRecorder recorder = processorContext.addDeploymentTask(RuntimePriority.WELD_DEPLOYMENT)) { + WeldDeploymentTemplate template = recorder.getRecordingProxy(WeldDeploymentTemplate.class); + template.registerShutdownHook(null); + } } diff --git a/weld/runtime/src/main/java/org/jboss/shamrock/weld/runtime/WeldDeploymentTemplate.java b/weld/runtime/src/main/java/org/jboss/shamrock/weld/runtime/WeldDeploymentTemplate.java index cecac8cdd61c0..e68b21bb2b0db 100644 --- a/weld/runtime/src/main/java/org/jboss/shamrock/weld/runtime/WeldDeploymentTemplate.java +++ b/weld/runtime/src/main/java/org/jboss/shamrock/weld/runtime/WeldDeploymentTemplate.java @@ -14,6 +14,7 @@ import javax.enterprise.inject.se.SeContainerInitializer; import javax.enterprise.inject.spi.Bean; +import org.jboss.shamrock.runtime.ContextObject; import org.jboss.shamrock.runtime.InjectionFactory; import org.jboss.shamrock.runtime.InjectionInstance; import org.jboss.shamrock.runtime.RuntimeInjector; @@ -35,8 +36,8 @@ public void connect() throws IOException { Class clazz = Class.forName("sun.net.www.protocol.jar.JarFileFactory"); Field field = clazz.getDeclaredField("fileCache"); field.setAccessible(true); - Map fileCache = (Map) field.get(null); - for(Map.Entry e : new HashSet<>(fileCache.entrySet())) { + Map fileCache = (Map) field.get(null); + for (Map.Entry e : new HashSet<>(fileCache.entrySet())) { e.getValue().close(); } fileCache.clear(); @@ -44,7 +45,7 @@ public void connect() throws IOException { field = clazz.getDeclaredField("urlCache"); field.setAccessible(true); Map urlCache = (Map) field.get(null); - for(Map.Entry e : new HashSet<>(urlCache.entrySet())) { + for (Map.Entry e : new HashSet<>(urlCache.entrySet())) { e.getKey().close(); } urlCache.clear(); @@ -62,7 +63,8 @@ public void addClass(SeContainerInitializer initializer, Class clazz) { initializer.addBeanClasses(clazz); } - public SeContainer doBoot(SeContainerInitializer initializer) throws Exception { + @ContextObject("weld.container") + public SeContainer doBoot( SeContainerInitializer initializer) throws Exception { SeContainer container = initializer.initialize(); // Force client proxy init to run Set> instance = container.getBeanManager().getBeans(Object.class); @@ -74,6 +76,15 @@ public SeContainer doBoot(SeContainerInitializer initializer) throws Exception { return container; } + public void registerShutdownHook(@ContextObject("weld.container") SeContainer container) { + Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() { + @Override + public void run() { + container.close(); + } + }, "Weld Shutdown Hook Thread")); + } + public void setupInjection(SeContainer container) { RuntimeInjector.setFactory(new InjectionFactory() { @Override