/*

 * Copyright 2008 Sun Microsystems, Inc.  All Rights Reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code 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 General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Sun Microsystems, Inc., 4150 Network Circle, Santa Clara,

 * CA 95054 USA or visit www.sun.com if you need additional information or

 * have any questions.

 */

/*

 * @test VirtualMBeanNotifTest.java

 * @bug 5108776

 * @build VirtualMBeanNotifTest Wombat WombatMBean

 * @summary Test that Virtual MBeans can be implemented and emit notifs.

 * @author  Daniel Fuchs

 */

import java.lang.management.ManagementFactory;

import java.lang.reflect.InvocationHandler;

import java.lang.reflect.InvocationTargetException;

import java.lang.reflect.Method;

import java.lang.reflect.Proxy;

import java.util.ArrayList;

import java.util.Arrays;

import java.util.HashMap;

import java.util.HashSet;

import java.util.List;

import java.util.Map;

import java.util.Set;

import java.util.concurrent.ArrayBlockingQueue;

import java.util.concurrent.BlockingQueue;

import java.util.concurrent.CopyOnWriteArrayList;

import java.util.concurrent.TimeUnit;

import javax.management.Attribute;

import javax.management.DynamicMBean;

import javax.management.InstanceNotFoundException;

import javax.management.JMException;

import javax.management.MBeanAttributeInfo;

import javax.management.MBeanInfo;

import javax.management.MBeanNotificationInfo;

import javax.management.MBeanServer;

import javax.management.Notification;

import javax.management.NotificationBroadcaster;

import javax.management.NotificationEmitter;

import javax.management.NotificationListener;

import javax.management.ObjectName;

import javax.management.RuntimeOperationsException;

import javax.management.StandardMBean;

import javax.management.event.EventSubscriber;

import javax.management.namespace.VirtualEventManager;

import javax.management.namespace.MBeanServerSupport;

public class VirtualMBeanNotifTest {

    /**

     * An invocation handler that can implement DynamicMBean,

     * NotificationBroadcaster, NotificationEmitter.

     * The invocation handler works by forwarding all calls received from

     * the implemented interfaces through a wrapped MBeanServer object.

     */

    public static class DynamicWrapper

            implements InvocationHandler {

        /**

         * Inserts an additional class at the head of a signature array.

         * @param first The first class in the signature

         * @param rest  The other classes in the signature

         * @return A signature array, of length rest.length+1.

         */

        static Class[] concat(Class first, Class... rest) {

            if (rest == null || rest.length == 0) {

                return new Class[] { first };

            }

            final Class[] sig = new Class[rest.length+1];

            sig[0] = first;

            System.arraycopy(rest, 0, sig, 1, rest.length);

            return sig;

        }

        /**

         * Inserts an additional object at the head of a parameters array.

         * @param first The first object in the parameter array.

         * @param rest  The other objects in the parameter array.

         * @return A parameter array, of length rest.length+1.

         */

        static Object[] concat(Object first, Object... rest) {

            if (rest == null || rest.length == 0) {

                return new Object[] { first };

            }

            final Object[] params = new Object[rest.length+1];

            params[0] = first;

            System.arraycopy(rest, 0, params, 1, rest.length);

            return params;

        }

        /**

         * These two sets are used to check that all methods from

         * implemented interfaces are mapped.

         * unmapped is the set of methods that couldn't be mapped.

         * mapped is the set of methods that could be mapped.

         */

        final static Set<Method> unmapped = new HashSet<Method>();

        final static Set<Method> mapped = new HashSet<Method>();

        /**

         * For each method define in one of the interfaces intf, tries

         * to find a corresponding method in the reference class ref, where

         * the method in ref has the same name, and takes an additional

         * ObjectName as first parameter.

         *

         * So for instance, if ref is MBeanServer and intf is {DynamicMBean}

         * the result map is:

         *     DynamicMBean.getAttribute -> MBeanServer.getAttribute

         *     DynamicMBean.setAttribute -> MBeanServer.setAttribute

         *     etc...

         * If a method was mapped, it is additionally added to 'mapped'

         * If a method couldn't be mapped, it is added to 'unmmapped'.

         * In our example above, DynamicMBean.getNotificationInfo will end

         * up in 'unmapped'.

         *

         * @param ref   The reference class - to which calls will be forwarded

         *              with an additional ObjectName parameter inserted.

         * @param intf  The proxy interface classes - for which we must find an

         *              equivalent in 'ref'

         * @return A map mapping the methods from intfs to the method of ref.

         */

        static Map<Method,Method> makeMapFor(Class<?> ref, Class<?>... intf) {

            final Map<Method,Method> map = new HashMap<Method,Method>();

            for (Class<?> clazz : intf) {

                for (Method m : clazz.getMethods()) {

                    try {

                        final Method m2 =

                            ref.getMethod(m.getName(),

                            concat(ObjectName.class,m.getParameterTypes()));

                        map.put(m,m2);

                        mapped.add(m);

                    } catch (Exception x) {

                        unmapped.add(m);

                    }

                }

            }

            return map;

        }

        /**

         * Tries to map all methods from DynamicMBean.class and

         * NotificationEmitter.class to their equivalent in MBeanServer.

         * This should be all the methods except

         * DynamicMBean.getNotificationInfo.

         */

        static final Map<Method,Method> mbeanmap =

                makeMapFor(MBeanServer.class,DynamicMBean.class,

                NotificationEmitter.class);

        /**

         * Tries to map all methods from DynamicMBean.class and

         * NotificationEmitter.class to an equivalent in DynamicWrapper.

         * This time only DynamicMBean.getNotificationInfo will be mapped.

         */

        static final Map<Method,Method> selfmap =

                makeMapFor(DynamicWrapper.class,DynamicMBean.class,

                NotificationEmitter.class);

        /**

         * Now check that we have mapped all methods.

         */

        static {

            unmapped.removeAll(mapped);

            if (unmapped.size() > 0)

                throw new ExceptionInInitializerError("Couldn't map "+ unmapped);

        }

        /**

         * The wrapped MBeanServer to which everything is delegated.

         */

        private final MBeanServer server;

        /**

         * The name of the MBean we're proxying.

         */

        private final ObjectName name;

        DynamicWrapper(MBeanServer server, ObjectName name) {

            this.server=server;

            this.name=name;

        }

        /**

         * Creates a new proxy for the given MBean. Implements

         * NotificationEmitter/NotificationBroadcaster if the proxied

         * MBean also does.

         * @param name    the name of the proxied MBean

         * @param server  the wrapped server

         * @return a DynamicMBean proxy

         * @throws javax.management.InstanceNotFoundException

         */

        public static DynamicMBean newProxy(ObjectName name, MBeanServer server)

            throws InstanceNotFoundException {

            if (server.isInstanceOf(name,

                    NotificationEmitter.class.getName())) {

                // implements NotificationEmitter

                return (DynamicMBean)

                        Proxy.newProxyInstance(

                        DynamicWrapper.class.getClassLoader(),

                        new Class[] {NotificationEmitter.class,

                        DynamicMBean.class},

                        new DynamicWrapper(server, name));

            }

            if (server.isInstanceOf(name,

                    NotificationBroadcaster.class.getName())) {

                // implements NotificationBroadcaster

                return (DynamicMBean)

                        Proxy.newProxyInstance(

                        DynamicWrapper.class.getClassLoader(),

                        new Class[] {NotificationBroadcaster.class,

                        DynamicMBean.class},

                        new DynamicWrapper(server, name));

            }

            // Only implements DynamicMBean.

            return (DynamicMBean)

                        Proxy.newProxyInstance(

                        DynamicWrapper.class.getClassLoader(),

                        new Class[] {DynamicMBean.class},

                        new DynamicWrapper(server, name));

        }

        public Object invoke(Object proxy, Method method, Object[] args)

                throws Throwable {

            // Look for a method on this class (takes precedence)

            final Method self = selfmap.get(method);

            if (self != null)

                return call(this,self,concat(name,args));

            // no method found on this class, look for the same method

            // on the wrapped MBeanServer

            final Method mbean = mbeanmap.get(method);

            if (mbean != null)

                return call(server,mbean,concat(name,args));

            // This isn't a method that can be forwarded to MBeanServer.

            // If it's a method from Object, call it on this.

            if (method.getDeclaringClass().equals(Object.class))

                return call(this,method,args);

            throw new NoSuchMethodException(method.getName());

        }

        // Call a method using reflection, unwraps invocation target exceptions

        public Object call(Object handle, Method m, Object[] args)

                throws Throwable {

            try {

                return m.invoke(handle, args);

            } catch (InvocationTargetException x) {

               throw x.getCause();

            }

        }

        // this method is called when DynamicMBean.getNotificationInfo() is

        // called. This is the method that should be mapped in

        // 'selfmap'

        public MBeanNotificationInfo[] getNotificationInfo(ObjectName name)

            throws JMException {

            return server.getMBeanInfo(name).getNotifications();

        }

    }

    /**

     * Just so that we can call the same test twice but with two

     * different implementations of VirtualMBeanServerSupport.

     */

    public static interface MBeanServerWrapperFactory {

        public MBeanServer wrapMBeanServer(MBeanServer wrapped);

    }

    /**

     * A VirtualMBeanServerSupport that wrapps an MBeanServer and does not

     * use VirtualEventManager.

     */

    public static class VirtualMBeanServerTest

            extends MBeanServerSupport {

        final MBeanServer wrapped;

        public VirtualMBeanServerTest(MBeanServer wrapped) {

            this.wrapped=wrapped;

        }

        @Override

        public DynamicMBean getDynamicMBeanFor(final ObjectName name)

                throws InstanceNotFoundException {

            if (wrapped.isRegistered(name))

                return DynamicWrapper.newProxy(name,wrapped);

            throw new InstanceNotFoundException(String.valueOf(name));

        }

        @Override

        protected Set<ObjectName> getNames() {

            return wrapped.queryNames(null, null);

        }

        public final static MBeanServerWrapperFactory factory =

                new MBeanServerWrapperFactory() {

            public MBeanServer wrapMBeanServer(MBeanServer wrapped) {

                return new VirtualMBeanServerTest(wrapped);

            }

            @Override

            public String toString() {

                return VirtualMBeanServerTest.class.getName();

            }

        };

    }

     /**

     * A VirtualMBeanServerSupport that wrapps an MBeanServer and

     * uses a VirtualEventManager.

     */

    public static class VirtualMBeanServerTest2

            extends VirtualMBeanServerTest {

        final EventSubscriber sub;

        final NotificationListener nl;

        final VirtualEventManager  mgr;

        /**

         * We use an EventSubscriber to subscribe for all notifications from

         * the wrapped MBeanServer, and publish them through a

         * VirtualEventManager. Not a very efficient way of doing things.

         * @param wrapped

         */

        public VirtualMBeanServerTest2(MBeanServer wrapped) {

            super(wrapped);

            this.sub = EventSubscriber.getEventSubscriber(wrapped);

            this.mgr = new VirtualEventManager();

            this.nl = new NotificationListener() {

                public void handleNotification(Notification notification, Object handback) {

                    mgr.publish((ObjectName)notification.getSource(), notification);

                }

            };

            try {

                sub.subscribe(ObjectName.WILDCARD, nl, null, null);

            } catch (RuntimeException x) {

                throw x;

            } catch (Exception x) {

                throw new IllegalStateException("can't subscribe for notifications!");

            }

        }

        @Override

        public NotificationEmitter

                getNotificationEmitterFor(ObjectName name)

                throws InstanceNotFoundException {

            final DynamicMBean mbean = getDynamicMBeanFor(name);

            if (mbean instanceof NotificationEmitter)

                return mgr.getNotificationEmitterFor(name);

            return null;

        }

        public final static MBeanServerWrapperFactory factory =

                new MBeanServerWrapperFactory() {

            public MBeanServer wrapMBeanServer(MBeanServer wrapped) {

                return new VirtualMBeanServerTest2(wrapped);

            }

            @Override

            public String toString() {

                return VirtualMBeanServerTest2.class.getName();

            }

        };

    }

    public static void test(MBeanServerWrapperFactory factory) throws Exception {

        final MBeanServer server = ManagementFactory.getPlatformMBeanServer();

        // names[] are NotificationEmitters

        final ObjectName[] emitters = new ObjectName[2];

        // shields[] have been shielded by wrapping them in a StandardMBean,

        // so although the resource is an MBean that implements

        // NotificationEmitter, the registered MBean (the wrapper) doesn't.

        final ObjectName[] shielded = new ObjectName[2];

        final List<ObjectName> registered = new ArrayList<ObjectName>(4);

        try {

            // register two MBeans before wrapping

            server.registerMBean(new Wombat(),

                    emitters[0] = new ObjectName("bush:type=Wombat,name=wom"));

            registered.add(emitters[0]);

            // we shield the second MBean in a StandardMBean so that it does

            // not appear as a NotificationEmitter.

            server.registerMBean(

                    new StandardMBean(new Wombat(), WombatMBean.class),

                    shielded[0] = new ObjectName("bush:type=Wombat,name=womshield"));

            registered.add(shielded[0]);

            final MBeanServer vserver = factory.wrapMBeanServer(server);

            // register two other MBeans after wrapping

            server.registerMBean(new Wombat(),

                    emitters[1] = new ObjectName("bush:type=Wombat,name=bat"));

            registered.add(emitters[1]);

            // we shield the second MBean in a StandardMBean so that it does

            // not appear as a NotificationEmitter.

            server.registerMBean(

                    new StandardMBean(new Wombat(), WombatMBean.class),

                    shielded[1] = new ObjectName("bush:type=Wombat,name=batshield"));

            registered.add(shielded[1]);

            // Call test with this config - we have two wombats who broadcast

            // notifs (emitters) and two wombats who don't (shielded).

            test(vserver, emitters, shielded);

            System.out.println("*** Test passed for: " + factory);

        } finally {

            // Clean up the platform mbean server for the next test...

            for (ObjectName n : registered) {

                try {

                    server.unregisterMBean(n);

                } catch (Exception x) {

                    x.printStackTrace();

                }

            }

        }

    }

    /**

     * Perform the actual test.

     * @param vserver    A virtual MBeanServerSupport implementation

     * @param emitters   Names of NotificationBroadcaster MBeans

     * @param shielded   Names of non NotificationBroadcaster MBeans

     * @throws java.lang.Exception

     */

    public static void test(MBeanServer vserver, ObjectName[] emitters,

            ObjectName[] shielded) throws Exception {

        // To catch exception in NotificationListener

        final List<Exception> fail = new CopyOnWriteArrayList<Exception>();

        // A queue of received notifications

        final BlockingQueue<Notification> notifs =

                new ArrayBlockingQueue<Notification>(50);

        // A notification listener that puts the notification it receives

        // in the queue.

        final NotificationListener handler = new NotificationListener() {

            public void handleNotification(Notification notification,

                    Object handback) {

                try {

                    notifs.put(notification);

                } catch (Exception x) {

                    fail.add(x);

                }

            }

        };

        // A list of attribute names for which we might receive an

        // exception. If an exception is received when getting these

        // attributes - the test will not fail.

        final List<String> exceptions = Arrays.asList( new String[] {

           "UsageThresholdCount","UsageThreshold","UsageThresholdExceeded",

           "CollectionUsageThresholdCount","CollectionUsageThreshold",

           "CollectionUsageThresholdExceeded"

        });

        // This is just a sanity check. Get all attributes of all MBeans.

        for (ObjectName n : vserver.queryNames(null, null)) {

            final MBeanInfo m = vserver.getMBeanInfo(n);

            for (MBeanAttributeInfo mba : m.getAttributes()) {

                // System.out.println(n+":");

                Object val;

                try {

                    val = vserver.getAttribute(n, mba.getName());

                } catch (Exception x) {

                    // only accept exception for those attributes that

                    // have a valid reason to fail...

                    if (exceptions.contains(mba.getName())) val = x;

                    else throw new Exception("Failed to get " +

                            mba.getName() + " from " + n,x);

                }

                // System.out.println("\t "+mba.getName()+": "+ val);

            }

        }

        // The actual tests. Register for notifications with notif emitters

        for (ObjectName n : emitters) {

            vserver.addNotificationListener(n, handler, null, n);

        }

        // Trigger the emission of notifications, check that we received them.

        for (ObjectName n : emitters) {

            vserver.setAttribute(n,

                    new Attribute("Caption","I am a new wombat!"));

            final Notification notif = notifs.poll(4, TimeUnit.SECONDS);

            if (!notif.getSource().equals(n))

                throw new Exception("Bad source for "+ notif);

            if (fail.size() > 0)

                throw new Exception("Failed to handle notif",fail.remove(0));

        }

        // Check that we didn't get more notifs than expected

        if (notifs.size() > 0)

            throw new Exception("Extra notifications in queue: "+notifs);

        // Check that if the MBean doesn't exist, we get InstanceNotFound.

        try {

            vserver.addNotificationListener(new ObjectName("toto:toto=toto"),

                    handler, null, null);

            throw new Exception("toto:toto=toto doesn't throw INFE");

        } catch (InstanceNotFoundException x) {

            System.out.println("Received "+x+" as expected.");

        }

        // For those MBeans that shouldn't be NotificationEmitters, check that

        // we get IllegalArgumentException