/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "License").
 * You may not use this file except in compliance with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */

/*
 * Copyright 2008 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

import org.opensolaris.os.dtrace.*;
import java.util.*;

/**
 * Assert getAggregate() can explicitly specify the anonymous aggregation.
 */
public class TestGetAggregate {
    static final String programString =
            "profile:::tick-50ms" +
            "{" +
            "        @ = count();" +
            "        @a = count();" +
            "}";

    static final String ANONYMOUS_AGGREGATION = "";
    static final int TICK = 50;
    static final int EXPECTED_TICKS = 3;
    static final int INTERVALS = 4;

    static void
    testIncluded(Consumer consumer, String ... aggregationNames)
            throws DTraceException, InterruptedException
    {
        Aggregate aggregate;
        Set <String> included = new HashSet <String> ();
        int n = 1;

        for (String name : aggregationNames) {
            included.add(name);
        }

        // Wait up to a full second to obtain aggregate data. Without a
        // time limit, we'll loop forever if no aggregation was
        // successfully included.
        do {
            Thread.sleep(TICK);
            aggregate = consumer.getAggregate(included, null);
        } while (aggregate.asMap().isEmpty() && n++ < (1000 / TICK));

        for (String name : included) {
            if (aggregate.getAggregation(name) == null) {
                throw new IllegalStateException("@" + name +
                        " was explicitly included but did not appear " +
                        "in the aggregate");
            }
        }
        for (Aggregation a : aggregate.getAggregations()) {
            if (!included.contains(a.getName())) {
                throw new IllegalStateException("@" + a.getName() +
                        " was not explicitly included but appeared " +
                        "in the aggregate anyway");
            }
        }

        if (!consumer.isRunning()) {
            throw new IllegalStateException("consumer exited");
        }
    }

    static void
    testCleared(Consumer consumer, String ... aggregationNames)
            throws DTraceException, InterruptedException
    {
        Aggregate aggregate;
        AggregationRecord rec;
        long value;
        Long firstValue;
        int n = 1;
        Map <String, Long> firstValues = new HashMap <String, Long> ();
        Set <String> cleared = new HashSet <String> ();

        for (String name : aggregationNames) {
            cleared.add(name);
        }

        do {
            Thread.sleep(TICK);
            aggregate = consumer.getAggregate(null, cleared);
        } while (aggregate.asMap().isEmpty() && n++ < (1000 / TICK));
        n = 1;

        do {
            Thread.sleep(TICK * EXPECTED_TICKS);
            aggregate = consumer.getAggregate(null, cleared);

            for (Aggregation a : aggregate.getAggregations()) {
                if (!firstValues.containsKey(a.getName())) {
                    rec = a.getRecord(Tuple.EMPTY);
                    value = rec.getValue().getValue().longValue();
                    firstValues.put(a.getName(), value);
                }
            }
        } while (consumer.isRunning() && n++ < INTERVALS);

        for (Aggregation a : aggregate.getAggregations()) {
            rec = a.getRecord(Tuple.EMPTY);
            value = rec.getValue().getValue().longValue();
            firstValue = firstValues.get(a.getName());

            if (cleared.contains(a.getName())) {
                // last value should be about the same as first value
                if (value > (firstValue * 2)) {
                    throw new IllegalStateException(
                            "@" + a.getName() + " should have " +
                            "been cleared but instead grew from " +
                            firstValue + " to " + value);
                }
            } else {
                // last value should be about (INTERVALS * firstValue)
                if (value < (firstValue * 2)) {
                    throw new IllegalStateException(
                            "@" + a.getName() + " should have " +
                            "accumulated a running total but " +
                            "instead went from " +
                            firstValue + " to " + value);
                }
            }
        }

        if (!consumer.isRunning()) {
            throw new IllegalStateException("consumer exited");
        }
    }

    static Integer includedStatus;
    static Integer clearedStatus;

    static void
    startIncludedTest()
    {
        final Consumer consumer = new LocalConsumer();
        consumer.addConsumerListener(new ConsumerAdapter() {
            public void consumerStarted(ConsumerEvent e) {
                new Thread(new Runnable() {
                    public void run() {
                        try {
                            testIncluded(consumer, ANONYMOUS_AGGREGATION);
                            includedStatus = 0;
                        } catch (Exception e) {
                            includedStatus = 1;
                            e.printStackTrace();
                        } finally {
                            consumer.abort();
                        }
                    }
                }).start();
            }
        });

        try {
            consumer.open();
            consumer.setOption(Option.aggrate, Option.millis(TICK));
            consumer.compile(programString);
            consumer.enable();
            consumer.go();
        } catch (Exception e) {
            includedStatus = 1;
            e.printStackTrace();
        }
    }

    static void
    startClearedTest()
    {
        final Consumer consumer = new LocalConsumer();
        consumer.addConsumerListener(new ConsumerAdapter() {
            public void consumerStarted(ConsumerEvent e) {
                new Thread(new Runnable() {
                    public void run() {
                        try {
                            testCleared(consumer, ANONYMOUS_AGGREGATION);
                            clearedStatus = 0;
                        } catch (Exception e) {
                            clearedStatus = 1;
                            e.printStackTrace();
                        } finally {
                            consumer.abort();
                        }
                    }
                }).start();
            }
        });

        try {
            consumer.open();
            consumer.setOption(Option.aggrate, Option.millis(TICK));
            consumer.compile(programString);
            consumer.enable();
            consumer.go();
        } catch (Exception e) {
            clearedStatus = 1;
            e.printStackTrace();
        }
    }

    public static void
    main(String[] args)
    {
        startIncludedTest();

        do {
            try {
                Thread.sleep(TICK);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        } while (includedStatus == null);

        startClearedTest();

        do {
            try {
                Thread.sleep(TICK);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        } while (clearedStatus == null);

        if (includedStatus != 0 || clearedStatus != 0) {
            System.out.println("Failure");
            System.exit(1);
        }

        System.out.println("Success");
    }
}