Refactor internals move shutdown hook into DBeanScope

Author: rbygrave

inject/src/main/java/io/avaje/inject/DBeanScopeBuilder.java

@@ -116,12 +116,7 @@ public BeanScope build() {
     for (BeanContextFactory factory : factoryOrder.factories()) {
       factory.build(builder);
     }
-
-    BeanScope beanScope = builder.build();
-    if (shutdownHook) {
-      return new ShutdownAwareBeanScope(beanScope);
-    }
-    return beanScope;
+    return builder.build(shutdownHook);
   }
 
   /**
@@ -137,106 +132,6 @@ private Class<?> suppliedType(Class<?> suppliedClass) {
     }
   }
 
-  /**
-   * Internal shutdown hook.
-   */
-  private static class Hook extends Thread {
-
-    private final ShutdownAwareBeanScope context;
-
-    Hook(ShutdownAwareBeanScope context) {
-      this.context = context;
-    }
-
-    @Override
-    public void run() {
-      context.shutdown();
-    }
-  }
-
-  /**
-   * Proxy that handles shutdown hook registration and de-registration.
-   */
-  private static class ShutdownAwareBeanScope implements BeanScope {
-
-    private final ReentrantLock lock = new ReentrantLock();
-    private final BeanScope beanScope;
-    private final Hook hook;
-    private boolean shutdown;
-
-    ShutdownAwareBeanScope(BeanScope beanScope) {
-      this.beanScope = beanScope;
-      this.hook = new Hook(this);
-      Runtime.getRuntime().addShutdownHook(hook);
-    }
-
-    @Override
-    public RequestScopeBuilder newRequestScope() {
-      return beanScope.newRequestScope();
-    }
-
-    @Override
-    public <T> RequestScopeMatch<T> requestProvider(Class<T> type, String name) {
-      return beanScope.requestProvider(type, name);
-    }
-
-    @Override
-    public <T> T get(Class<T> beanClass) {
-      return beanScope.get(beanClass);
-    }
-
-    @Override
-    public <T> T get(Class<T> beanClass, String name) {
-      return beanScope.get(beanClass, name);
-    }
-
-    @Override
-    public List<Object> listByAnnotation(Class<?> annotation) {
-      return beanScope.listByAnnotation(annotation);
-    }
-
-    @Override
-    public <T> List<T> list(Class<T> interfaceType) {
-      return beanScope.list(interfaceType);
-    }
-
-    @Override
-    public <T> List<T> listByPriority(Class<T> interfaceType) {
-      return beanScope.listByPriority(interfaceType);
-    }
-
-    @Override
-    public <T> List<T> listByPriority(Class<T> interfaceType, Class<? extends Annotation> priority) {
-      return beanScope.listByPriority(interfaceType, priority);
-    }
-
-    @Override
-    public void close() {
-      lock.lock();
-      try {
-        if (!shutdown) {
-          Runtime.getRuntime().removeShutdownHook(hook);
-        }
-        beanScope.close();
-      } finally {
-        lock.unlock();
-      }
-    }
-
-    /**
-     * Close via shutdown hook.
-     */
-    void shutdown() {
-      lock.lock();
-      try {
-        shutdown = true;
-        close();
-      } finally {
-        lock.unlock();
-      }
-    }
-  }
-
   /**
    * Helper to order the BeanContextFactory based on dependsOn.
    */

inject/src/main/java/io/avaje/inject/spi/Builder.java

@@ -158,5 +158,5 @@ static Builder newBuilder(List<SuppliedBean> suppliedBeans, List<EnrichBean> enr
   /**
    * Build and return the bean context.
    */
-  BeanScope build();
+  BeanScope build(boolean withShutdownHook);
 }

inject/src/main/java/io/avaje/inject/spi/DBeanScope.java

@@ -21,13 +21,20 @@ class DBeanScope implements BeanScope {
   private final List<BeanLifecycle> lifecycleList;
   private final DBeanMap beans;
   private final Map<String, RequestScopeMatch<?>> reqScopeProviders;
-
+  private final ShutdownHook shutdownHook;
+  private boolean shutdown;
   private boolean closed;
 
-  DBeanScope(List<BeanLifecycle> lifecycleList, DBeanMap beans, Map<String, RequestScopeMatch<?>> reqScopeProviders) {
+  DBeanScope(boolean withShutdownHook, List<BeanLifecycle> lifecycleList, DBeanMap beans, Map<String, RequestScopeMatch<?>> reqScopeProviders) {
     this.lifecycleList = lifecycleList;
     this.beans = beans;
     this.reqScopeProviders = reqScopeProviders;
+    if (withShutdownHook) {
+      this.shutdownHook = new ShutdownHook(this);
+      Runtime.getRuntime().addShutdownHook(shutdownHook);
+    } else {
+      this.shutdownHook = null;
+    }
   }
 
   @Override
@@ -68,7 +75,7 @@ public <T> List<T> listByPriority(Class<T> interfaceType, Class<? extends Annota
     return list.size() > 1 ? sortByPriority(list, priorityAnnotation) : list;
   }
 
-  private  <T> List<T> sortByPriority(List<T> list, final Class<? extends Annotation> priorityAnnotation) {
+  private <T> List<T> sortByPriority(List<T> list, final Class<? extends Annotation> priorityAnnotation) {
     boolean priorityUsed = false;
     List<SortBean<T>> tempList = new ArrayList<>(list.size());
     for (T bean : list) {
@@ -113,6 +120,9 @@ DBeanScope start() {
   public void close() {
     lock.lock();
     try {
+      if (shutdownHook != null && !shutdown) {
+        Runtime.getRuntime().removeShutdownHook(shutdownHook);
+      }
       if (!closed) {
         // we only allow one call to preDestroy
         closed = true;
@@ -126,6 +136,30 @@ public void close() {
     }
   }
 
+  private void shutdown() {
+    lock.lock();
+    try {
+      shutdown = true;
+      close();
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  private static class ShutdownHook extends Thread {
+    private final DBeanScope scope;
+
+    ShutdownHook(DBeanScope scope) {
+      this.scope = scope;
+    }
+
+    @Override
+    public void run() {
+      scope.shutdown();
+    }
+  }
+
+
   private static class SortBean<T> implements Comparable<SortBean<T>> {
 
     private final T bean;

inject/src/main/java/io/avaje/inject/spi/DBuilder.java

@@ -195,8 +195,8 @@ private void runInjectors() {
     }
   }
 
-  public BeanScope build() {
+  public BeanScope build(boolean withShutdownHook) {
     runInjectors();
-    return new DBeanScope(lifecycleList, beanMap, reqScopeProviders).start();
+    return new DBeanScope(withShutdownHook, lifecycleList, beanMap, reqScopeProviders).start();
   }
 }