Fix issue with weld shutdown hook

quarkusio · Jul 26, 2018 · dc63226 · dc63226
1 parent 5a7d800
commit dc63226
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 4 deletions.
diff --git a/.../deployment/src/main/java/org/jboss/shamrock/deployment/codegen/BytecodeRecorderImpl.java b/.../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/.../deployment/src/main/java/org/jboss/shamrock/weld/deployment/WeldAnnotationProcessor.java b/.../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
@@ -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,16 +36,16 @@ public void connect() throws IOException {
         Class<?> clazz = Class.forName("sun.net.www.protocol.jar.JarFileFactory");
         Field field = clazz.getDeclaredField("fileCache");
         field.setAccessible(true);
-        Map<String, JarFile > fileCache = (Map<String, JarFile>) field.get(null);
-        for(Map.Entry<String, JarFile> e : new HashSet<>(fileCache.entrySet())) {
+        Map<String, JarFile> fileCache = (Map<String, JarFile>) field.get(null);
+        for (Map.Entry<String, JarFile> e : new HashSet<>(fileCache.entrySet())) {
             e.getValue().close();
         }
         fileCache.clear();
 
         field = clazz.getDeclaredField("urlCache");
         field.setAccessible(true);
         Map<JarFile, URL> urlCache = (Map<JarFile, URL>) field.get(null);
-        for(Map.Entry<JarFile, URL> e : new HashSet<>(urlCache.entrySet())) {
+        for (Map.Entry<JarFile, URL> 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<Bean<?>> 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