onStart/onStop events fix jooby-project#289

jknack · jknack · commit b36e046578f3 · 2016-02-26T11:53:35.000-03:00
diff --git a/jooby/src/main/java/org/jooby/Jooby.java b/jooby/src/main/java/org/jooby/Jooby.java
@@ -489,11 +489,6 @@ public EnvDep(final Predicate<String> predicate, final Consumer<Config> callback
    */
   private Set<Object> bag = new LinkedHashSet<>();
 
-  /**
-   * Keep track of modules.
-   */
-  private final Set<Jooby.Module> modules = new LinkedHashSet<>();
-
   /**
    * The override config. Optional.
    */
@@ -511,6 +506,12 @@ public EnvDep(final Predicate<String> predicate, final Consumer<Config> callback
   /** Route's prefix. */
   private String prefix;
 
+  /** startup callback .*/
+  private List<Runnable> onStart = new ArrayList<>();
+
+  /** stop callback .*/
+  private List<Runnable> onStop = new ArrayList<>();
+
   public Jooby() {
     this(null);
   }
@@ -626,6 +627,30 @@ public Jooby env(final Env.Builder env) {
     return this;
   }
 
+  /**
+   * Run code at application startup time.
+   *
+   * @param callback A callback to run.
+   * @return This instance.
+   */
+  public Jooby onStart(final Runnable callback) {
+    requireNonNull(callback, "Callback is required.");
+    onStart.add(callback);
+    return this;
+  }
+
+  /**
+   * Run code at application shutdown time.
+   *
+   * @param callback A callback to run.
+   * @return This instance.
+   */
+  public Jooby onStop(final Runnable callback) {
+    requireNonNull(callback, "Callback is required.");
+    onStop.add(callback);
+    return this;
+  }
+
   /**
    * Run the given callback if and only if, app runs in the given enviroment.
    *
@@ -3016,7 +3041,6 @@ private Route.Definition appendDefinition(final Route.Definition route) {
    */
   public Jooby use(final Jooby.Module module) {
     requireNonNull(module, "A module is required.");
-    modules.add(module);
     bag.add(module);
     return this;
   }
@@ -3201,6 +3225,8 @@ public void start(final String[] args, final Consumer<List<Route.Definition>> ro
         end - start,
         injector.getInstance(AppPrinter.class));
 
+    onStart.forEach(Runnable::run);
+
     boolean join = config.hasPath("server.join") ? config.getBoolean("server.join") : true;
     if (join) {
       server.join();
@@ -3551,7 +3577,7 @@ private static List<Config> modconf(final Collection<Object> bag) {
    */
   public void stop() {
     if (injector != null) {
-      stopManaged();
+      stopManaged(injector, onStop);
 
       try {
         Server server = injector.getInstance(Server.class);
@@ -3565,10 +3591,10 @@ public void stop() {
     }
   }
 
-  private void stopManaged() {
+  private static void stopManaged(final Injector injector, final List<Runnable> onStop) {
     // stop modules
     injector.getInstance(LifecycleProcessor.class).destroy();
-    modules.clear();
+    onStop.forEach(Runnable::run);
   }
 
   /**
diff --git a/jooby/src/test/java/org/jooby/JoobyTest.java b/jooby/src/test/java/org/jooby/JoobyTest.java
@@ -649,6 +649,47 @@ public void defaults() throws Exception {
         } , boot);
   }
 
+  @Test
+  public void onStartStopCallback() throws Exception {
+
+    new MockUnit(Binder.class, Runnable.class)
+        .expect(guice)
+        .expect(shutdown)
+        .expect(config)
+        .expect(env)
+        .expect(classInfo)
+        .expect(ssl)
+        .expect(charset)
+        .expect(locale)
+        .expect(zoneId)
+        .expect(timeZone)
+        .expect(dateTimeFormatter)
+        .expect(numberFormat)
+        .expect(decimalFormat)
+        .expect(renderers)
+        .expect(session)
+        .expect(routes)
+        .expect(routeHandler)
+        .expect(params)
+        .expect(requestScope)
+        .expect(webSockets)
+        .expect(tmpdir)
+        .expect(err)
+        .expect(unit -> {
+          unit.get(Runnable.class).run();
+          unit.get(Runnable.class).run();
+        })
+        .run(unit -> {
+
+          Jooby app = new Jooby()
+              .onStart(unit.get(Runnable.class))
+              .onStop(unit.get(Runnable.class));
+          app.start();
+          app.stop();
+
+        } , boot);
+  }
+
   @Test
   public void defaultsWithCallback() throws Exception {
 
@@ -698,7 +739,8 @@ public void customEnv() throws Exception {
           expect(env.name()).andReturn("dev").times(2);
 
           Env.Builder builder = unit.get(Env.Builder.class);
-          expect(builder.build(isA(Config.class), isA(Jooby.class), isA(Locale.class))).andReturn(env);
+          expect(builder.build(isA(Config.class), isA(Jooby.class), isA(Locale.class)))
+              .andReturn(env);
 
           Binder binder = unit.get(Binder.class);
 
@@ -2017,7 +2059,6 @@ public void assets() throws Exception {
           assertNotNull(route);
           assertTrue(route.isPresent());
 
-
           ((RouteImpl) route.get()).handle(unit.get(Request.class), unit.get(Response.class),
               unit.get(Route.Chain.class));
 
