Fix callback handling; wrong equals.

pojo/src/main/java/org/jboss/as/pojo/descriptor/CallbackConfig.java

@@ -23,13 +23,8 @@
 package org.jboss.as.pojo.descriptor;
 
 import org.jboss.as.pojo.BeanState;
-import org.jboss.msc.service.ServiceName;
-import org.jboss.msc.value.InjectedValue;
 
 import java.io.Serializable;
-import java.lang.reflect.Method;
-import java.util.Collections;
-import java.util.Set;
 
 /**
  * Callback meta data.
@@ -45,24 +40,6 @@ public class CallbackConfig extends AbstractConfigVisitorNode implements Seriali
     private BeanState state = BeanState.INSTALLED;
     private String signature;
 
-    private final InjectedValue<Set<Object>> beans = new InjectedValue<Set<Object>>();
-
-    @Override
-    public void visit(ConfigVisitor visitor) {
-        if (visitor.getState().next() == whenRequired) {
-            Method m = visitor.getBeanInfo().findMethod(methodName, signature);
-            if (m.getParameterTypes().length != 1)
-                throw new IllegalArgumentException("Illegal method parameter length: " + m);
-            ServiceName dependency = BeanMetaDataConfig.toInstancesName(m.getParameterTypes()[0], state);
-            visitor.addOptionalDependency(dependency, beans);
-        }
-    }
-
-    public Set<Object> getBeans() {
-        Set<Object> set = beans.getOptionalValue();
-        return (set != null) ? set : Collections.emptySet();
-    }
-
     public String getMethodName() {
         return methodName;
     }

pojo/src/main/java/org/jboss/as/pojo/service/Callback.java

@@ -65,7 +65,7 @@ public BeanState getState() {
     }
 
     public void dispatch() throws Throwable {
-        for (Object bean : config.getBeans())
+        for (Object bean : InstancesService.getBeans(getType(), getState()))
             dispatch(bean);
     }
 
@@ -80,4 +80,18 @@ protected Object getClassValue(Class<?> type) {
         joinpoint.setParameters(new ValueConfig[]{param});
         joinpoint.dispatch();
     }
+
+    @Override
+    public int hashCode() {
+        return config.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj instanceof Callback == false)
+            return false;
+
+        Callback callback = (Callback) obj;
+        return config.equals(callback.config);
+    }
 }

pojo/src/main/java/org/jboss/as/pojo/service/InstancesService.java

@@ -56,8 +56,9 @@ final class InstancesService implements Service<Set<Object>> {
     private final Class<?> type;
     private final Set<Object> instances = new HashSet<Object>(); // we do own locking, per type (service is also per type)
 
-    private static Map<Class<?>, Map<BeanState, List<Callback>>> incallbacks = new HashMap<Class<?>, Map<BeanState, List<Callback>>>();
-    private static Map<Class<?>, Map<BeanState, List<Callback>>> uncallbacks = new HashMap<Class<?>, Map<BeanState, List<Callback>>>();
+    private static Map<TypeBeanStateKey, Set<Object>> beans = new HashMap<TypeBeanStateKey, Set<Object>>();
+    private static Map<TypeBeanStateKey, List<Callback>> incallbacks = new HashMap<TypeBeanStateKey, List<Callback>>();
+    private static Map<TypeBeanStateKey, List<Callback>> uncallbacks = new HashMap<TypeBeanStateKey, List<Callback>>();
 
     private InstancesService(Class<?> type) {
         this.type = type;
@@ -67,9 +68,9 @@ private InstancesService(Class<?> type) {
      * Add bean instance.
      *
      * @param registry the service registry
-     * @param target the service target
-     * @param state the bean state
-     * @param bean the bean
+     * @param target   the service target
+     * @param state    the bean state
+     * @param bean     the bean
      * @throws StartException for any error
      */
     static void addInstance(ServiceRegistry registry, ServiceTarget target, BeanState state, Object bean) throws StartException {
@@ -80,8 +81,8 @@ static void addInstance(ServiceRegistry registry, ServiceTarget target, BeanStat
      * Remove bean instance.
      *
      * @param registry the service registry
-     * @param state the bean state
-     * @param bean the bean
+     * @param state    the bean state
+     * @param bean     the bean
      */
     static void removeInstance(ServiceRegistry registry, BeanState state, Object bean) {
         removeInstance(registry, state, bean.getClass(), bean);
@@ -123,7 +124,7 @@ static void removeUncallback(Callback callback) {
         removeCallback(uncallbacks, callback);
     }
 
-    private static void addCallback(Map<Class<?>, Map<BeanState, List<Callback>>> map, Callback callback) {
+    private static void addCallback(Map<TypeBeanStateKey, List<Callback>> map, Callback callback) {
         final Class<?> type = callback.getType();
         synchronized (type) {
             if (map == incallbacks) {
@@ -134,33 +135,25 @@ private static void addCallback(Map<Class<?>, Map<BeanState, List<Callback>>> ma
                 }
             }
 
-            Map<BeanState, List<Callback>> states = map.get(type);
-            if (states == null) {
-                states = new HashMap<BeanState, List<Callback>>();
-                map.put(type, states);
-            }
-            List<Callback> callbacks = states.get(callback.getState());
+            TypeBeanStateKey key = new TypeBeanStateKey(type, callback.getState());
+            List<Callback> callbacks = map.get(key);
             if (callbacks == null) {
                 callbacks = new ArrayList<Callback>();
-                states.put(callback.getState(), callbacks);
+                map.put(key, callbacks);
             }
             callbacks.add(callback);
         }
     }
 
-    private static void removeCallback(Map<Class<?>, Map<BeanState, List<Callback>>> map, Callback callback) {
+    private static void removeCallback(Map<TypeBeanStateKey, List<Callback>> map, Callback callback) {
         final Class<?> type = callback.getType();
         synchronized (type) {
-            Map<BeanState, List<Callback>> states = map.get(type);
-            if (states != null) {
-                List<Callback> callbacks = states.get(callback.getState());
-                if (callbacks != null) {
-                    callbacks.remove(callback);
-                    if (callbacks.isEmpty())
-                        states.remove(callback.getState());
-                }
-                if (states.isEmpty())
-                    map.remove(type);
+            TypeBeanStateKey key = new TypeBeanStateKey(type, callback.getState());
+            List<Callback> callbacks = map.get(key);
+            if (callbacks != null) {
+                callbacks.remove(callback);
+                if (callbacks.isEmpty())
+                    map.remove(key);
             }
 
             if (map == uncallbacks) {
@@ -173,18 +166,18 @@ private static void removeCallback(Map<Class<?>, Map<BeanState, List<Callback>>>
         }
     }
 
-    private static void invokeCallbacks(Map<Class<?>, Map<BeanState, List<Callback>>> map, BeanState state, final Class<?> clazz, Object bean) {
+    private static void invokeCallbacks(Map<TypeBeanStateKey, List<Callback>> map, BeanState state, final Class<?> clazz, Object bean) {
         synchronized (clazz) {
-            Map<BeanState, List<Callback>> states = map.get(clazz);
-            if (states != null) {
-                List<Callback> callbacks = states.get(state);
-                if (callbacks != null)
-                    for (Callback c : callbacks)
-                        try {
-                            c.dispatch(bean);
-                        } catch (Throwable t) {
-                            log.warn("Error invoking callback: " + c, t);
-                        }
+            TypeBeanStateKey key = new TypeBeanStateKey(clazz, state);
+            List<Callback> callbacks = map.get(key);
+            if (callbacks != null) {
+                for (Callback c : callbacks) {
+                    try {
+                        c.dispatch(bean);
+                    } catch (Throwable t) {
+                        log.warn("Error invoking callback: " + c, t);
+                    }
+                }
             }
         }
     }
@@ -198,6 +191,11 @@ private static void addInstance(ServiceRegistry registry, ServiceTarget target,
         InstancesService service = putIfAbsent(registry, name, builder);
         synchronized (clazz) {
             service.instances.add(bean);
+
+            TypeBeanStateKey key = new TypeBeanStateKey(clazz, state);
+            if (beans.containsKey(key) == false)
+                beans.put(key, service.instances);
+
             invokeCallbacks(incallbacks, state, clazz, bean);
         }
 
@@ -208,7 +206,7 @@ private static void addInstance(ServiceRegistry registry, ServiceTarget target,
     }
 
     private static InstancesService putIfAbsent(ServiceRegistry registry, ServiceName name, ServiceBuilder builder) throws StartException {
-        for (;;) {
+        for (; ; ) {
             try {
                 ServiceController sc = registry.getService(name);
                 if (sc == null) {
@@ -232,8 +230,11 @@ private static void removeInstance(ServiceRegistry registry, BeanState state, fi
             synchronized (clazz) {
                 service.instances.remove(bean);
                 invokeCallbacks(uncallbacks, state, clazz, bean);
-                if (service.instances.isEmpty())
+                if (service.instances.isEmpty()) {
+                    beans.remove(new TypeBeanStateKey(clazz, state));
+
                     controller.setMode(ServiceController.Mode.REMOVE);
+                }
             }
         }
 
@@ -257,4 +258,10 @@ public Set<Object> getValue() throws IllegalStateException, IllegalArgumentExcep
             return Collections.unmodifiableSet(instances);
         }
     }
+
+    static Set<Object> getBeans(Class<?> type, BeanState state) {
+        TypeBeanStateKey key = new TypeBeanStateKey(type, state);
+        Set<Object> objects = beans.get(key);
+        return (objects != null) ? objects : Collections.emptySet();
+    }
 }

pojo/src/main/java/org/jboss/as/pojo/service/TypeBeanStateKey.java

@@ -0,0 +1,52 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2011, Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags. See the copyright.txt file in the
+ * distribution for a full listing of individual contributors.
+ *
+ * This is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as
+ * published by the Free Software Foundation; either version 2.1 of
+ * the License, or (at your option) any later version.
+ *
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this software; if not, write to the Free
+ * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
+ * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
+ */
+
+package org.jboss.as.pojo.service;
+
+import org.jboss.as.pojo.BeanState;
+
+/**
+ * @author <a href="mailto:ales.justin@jboss.org">Ales Justin</a>
+ */
+class TypeBeanStateKey {
+    private final Class<?> type;
+    private final BeanState state;
+
+    public TypeBeanStateKey(Class<?> type, BeanState state) {
+        this.type = type;
+        this.state = state;
+    }
+
+    @Override
+    public int hashCode() {
+        return type.hashCode() + 7 * state.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj instanceof TypeBeanStateKey == false)
+            return false;
+
+        TypeBeanStateKey tbsk = (TypeBeanStateKey) obj;
+        return type.equals(tbsk.type) && state == tbsk.state;
+    }
+}