/*
 * Copyright 2013 Rackspace
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

package com.rackspacecloud.blueflood.service;

import com.codahale.metrics.Gauge;
import com.codahale.metrics.JmxAttributeGauge;
import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Ticker;
import com.google.common.util.concurrent.Futures;
import com.rackspacecloud.blueflood.concurrent.ThreadPoolBuilder;
import com.rackspacecloud.blueflood.io.Constants;
import com.rackspacecloud.blueflood.utils.Metrics;
import com.rackspacecloud.blueflood.utils.TimeValue;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.framework.state.ConnectionState;
import org.apache.curator.framework.state.ConnectionStateListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.management.MBeanServer;
import javax.management.ObjectName;
import java.lang.management.ManagementFactory;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

class ZKShardLockManager implements ConnectionStateListener, ShardLockManager, ZKShardLockManagerMBean {
    private static final Logger log = LoggerFactory.getLogger(ZKShardLockManager.class);
    private static final AtomicInteger UNIQUE_IDENTIFIER = new AtomicInteger(0);

    private final Random rand = new Random(System.currentTimeMillis());
    /** Unique identifier for this object, used to identify the MBean. */
    private final int id = UNIQUE_IDENTIFIER.getAndIncrement();
    /** Zookeeper client. */
    private final CuratorFramework client;
    private final String ZK_NAMESPACE = "locks/blueflood";
    private final String LOCK_QUALIFIER = "/shards";
    private final long ZK_SESSION_TIMEOUT_MS = new TimeValue(120L, TimeUnit.SECONDS).toMillis();
    private final long ZK_CONN_TIMEOUT_MS = new TimeValue(5L, TimeUnit.SECONDS).toMillis();
    private final long ZK_RETRY_INTERVAL = new TimeValue(50L, TimeUnit.MILLISECONDS).toMillis();
    private final int ZK_MAX_RETRIES = 2;
    private final TimeValue ZK_LOCK_TIMEOUT = new TimeValue(1L, TimeUnit.SECONDS);
    private final ConcurrentHashMap<Integer, Lock> locks; // shard to lock objects
    private final TimeValue shardLockScavengeInterval;
    private final int defaultMaxLocksToAcquirePerCycle;
    private final Ticker ticker = Ticker.systemTicker();

    // modifiable properties.
    private TimeValue minLockHoldTime;
    private TimeValue lockDisinterestedTime;

    /** true if we're connected to zookeeper. */
    private volatile boolean connected = false;
    private volatile long lastScavengedAt = System.currentTimeMillis();

    /** Thread that performs the locking & releasing. */
    private final ThreadPoolExecutor lockWorker;
    /** Number of outstanding acquire & release operations. */
    private final AtomicInteger lockTaskCount = new AtomicInteger(0);

    private final ScheduledThreadPoolExecutor scavengerWorker = new ScheduledThreadPoolExecutor(1, new ThreadFactory() {
        @Override
        public Thread newThread(Runnable r) {
            Thread t = new Thread(r, "ZK Lock Worker " + id);
            t.setDaemon(true);
            return t;
        }
    });

    private final Meter lockAcquisitionFailure = Metrics.meter(ZKShardLockManager.class, "Lock acquisition failures");
    private final com.codahale.metrics.Timer lockAcquisitionTimer = Metrics.timer(ZKShardLockManager.class, "Lock acquisition timer");
    private final Meter lockErrors = Metrics.meter(ZKShardLockManager.class, "Lock errors");

    ZKShardLockManager(String zookeeperCluster, Set<Integer> managedShards) {
        try {
            final MBeanServer mbs = ManagementFactory.getPlatformMBeanServer();
            final String name = String.format("com.rackspacecloud.blueflood.service:type=%s", getClass().getSimpleName() + (id == 0 ? "" : id));
            final ObjectName nameObj = new ObjectName(name);
            mbs.registerMBean(this, nameObj);
            registerMetrics(nameObj, Metrics.getRegistry());
        } catch (Exception exc) {
            log.error("Unable to register mbean for " + getClass().getSimpleName(), exc);
        }

        this.locks = new ConcurrentHashMap<Integer, Lock>();
        RetryPolicy policy = new ExponentialBackoffRetry((int)ZK_RETRY_INTERVAL, ZK_MAX_RETRIES);

        this.client = CuratorFrameworkFactory.
                builder().namespace(ZK_NAMESPACE)
                .connectString(zookeeperCluster)
                .sessionTimeoutMs((int) ZK_SESSION_TIMEOUT_MS)
                .connectionTimeoutMs((int) ZK_CONN_TIMEOUT_MS)
                .retryPolicy(policy).build();
        this.client.getConnectionStateListenable().addListener(this);  // register our listener
        this.client.start();

        Configuration config = Configuration.getInstance();
        for (int shard : managedShards) {
            addShard(shard);
        }
        this.minLockHoldTime = new TimeValue(config.getLongProperty(CoreConfig.SHARD_LOCK_HOLD_PERIOD_MS), TimeUnit.MILLISECONDS);
        this.lockDisinterestedTime = new TimeValue(config.getLongProperty(CoreConfig.SHARD_LOCK_DISINTERESTED_PERIOD_MS), TimeUnit.MILLISECONDS);
        this.shardLockScavengeInterval = new TimeValue(config.getLongProperty(CoreConfig.SHARD_LOCK_SCAVENGE_INTERVAL_MS),
                TimeUnit.MILLISECONDS);
        this.defaultMaxLocksToAcquirePerCycle = config.getIntegerProperty(CoreConfig.MAX_ZK_LOCKS_TO_ACQUIRE_PER_CYCLE);
        this.lockWorker = new ThreadPoolBuilder()
                .withCorePoolSize(1)
                .withMaxPoolSize(1)
                .withKeepAliveTime(new TimeValue(Long.MAX_VALUE, TimeUnit.DAYS))
                .withBoundedQueue(1000)
                .withName("ZkThreadPool")
                .build();
    }

    /**
     * Registers the different ZooKeeper metrics.
     */
    private void registerMetrics(final ObjectName nameObj, MetricRegistry reg) {
        reg.register(MetricRegistry.name(ZKShardLockManager.class, "Lock Disinterested Time Millis"),
                new JmxAttributeGauge(nameObj, "LockDisinterestedTimeMillis"));
        reg.register(MetricRegistry.name(ZKShardLockManager.class, "Min Lock Hold Time Millis"),
                new JmxAttributeGauge(nameObj, "MinLockHoldTimeMillis"));
        reg.register(MetricRegistry.name(ZKShardLockManager.class, "Seconds Since Last Scavenge"),
                new JmxAttributeGauge(nameObj, "SecondsSinceLastScavenge"));

        reg.register(MetricRegistry.name(ZKShardLockManager.class, "Zk Connection Status"),
                new JmxAttributeGauge(nameObj, "ZkConnectionStatus") {
                    @Override
                    public Object getValue() {
                        Object val = super.getValue();
                        if (val.equals("connected")) {
                            return 1;
                        }
                        return 0;
                    }
                });
        reg.register(MetricRegistry.name(ZKShardLockManager.class, "Held Shards"),
                new Gauge<Integer>() {
                    @Override
                    public Integer getValue() {
                        return getHeldShards().size();
                    }
                });

        reg.register(MetricRegistry.name(ZKShardLockManager.class, "Unheld Shards"),
                new Gauge<Integer>() {
                    @Override
                    public Integer getValue() {
                        return getUnheldShards().size();
                    }
                });
        reg.register(MetricRegistry.name(ZKShardLockManager.class, "Error Shards"),
                new Gauge<Integer>() {
                    @Override
                    public Integer getValue() {
                        return getErrorShards().size();
                    }
                });
    }

    /**
     * Initialize the ZKShardLockManager.
     * @param zkWaitTime Time to wait until zookeeper is up.
     */
    public void init(TimeValue zkWaitTime) {
        waitForZKConnections(zkWaitTime.toSeconds());
        prefetchLocks();
        scheduleScavenger();
    }

    /**
     * Waits until the zookeeper connection is available.
     * @param waitTimeSeconds
     */
    @VisibleForTesting boolean waitForZKConnections(long waitTimeSeconds) {
        for (int i = 0; i < waitTimeSeconds; i++) {
            if (connected) {
                return connected;
            }
            log.debug("Waiting for connect");
            try { Thread.sleep(1000); } catch (InterruptedException ex) {}
        }
        return connected;
    }

    /**
     * Only called from {@link #init(TimeValue)}.
     *
     * Try to achieve each lock. This will give us a good initial state where each lock is either held or unheld.
     * The purpose is to avoid a situation at startup where a node assumes it can schedule anything (until a lock
     * is held).
     */
    @VisibleForTesting void prefetchLocks() {

        if (!connected) {
            log.warn("Cannot connect to Zookeeper; will not perform initial lock acquisition");
            for (Lock lock : locks.values()) {
                lock.connectionLost();
            }
        } else {
            log.info("Pre-fetching zookeeper locks for shards");
            boolean isManagingAllShards = locks.size() >= Constants.NUMBER_OF_SHARDS;
            int maxLocksToPrefetch = moreThanHalf();

            if (isManagingAllShards) {
                maxLocksToPrefetch = Constants.NUMBER_OF_SHARDS;
            }

            List<Integer> shards = new ArrayList<Integer>(locks.keySet());
            Collections.shuffle(shards);

            int locksObtained = 0;
            for (int shard : shards) {
                try {
                    log.debug("Initial lock attempt for shard={}", shard);
                    final Lock lock = locks.get(shard);
                    lock.acquire().get();

                    if (lock.isHeld() && ++locksObtained >= maxLocksToPrefetch) {
                        break;
                    }
                } catch (InterruptedException ex) {
                    log.warn("Thread exception while acquiring initial locks: " + ex.getMessage(), ex);
                } catch (ExecutionException ex) {
                    log.error("Problem acquiring lock " + shard + " " + ex.getCause().getMessage(), ex.getCause());
                }
            }
            log.info("Finished pre-fetching zookeeper locks");
        }
    }

    public void scheduleScavenger() {
        scavengerWorker.scheduleAtFixedRate(new Runnable() {
            @Override
            public void run() {
                scavengeLocks();
            }
        }, 30, shardLockScavengeInterval.toSeconds(), TimeUnit.SECONDS);
    }

    /**
     * Determines whether a given shard can be worked on by this blueflood instance.
     */
    public boolean canWork(int shard) {
        return locks.containsKey(shard) && locks.get(shard).canWork();
    }

    private String getLockId(int shard) {
        return  LOCK_QUALIFIER + "/" + shard;
    }

    // This is called when connection to zookeeper is lost
    private void handleZookeeperConnectionFailed() {
        // It is okay for us to proceed with the work we already scheduled (either running or in scheduled queue)
        // for slots. We would duplicate work just for those slots which is fine.
        log.info("Force release all locks as zookeeper connection is lost");
        for (Lock lock : locks.values()) {
            lock.connectionLost();
        }
    }

    private void scavengeLocks() {
        log.debug("Starting scavengeLocks()");
        try {
            int locksAcquiredThisCycle = 0;
            int maxLocksToAcquirePerCycle = defaultMaxLocksToAcquirePerCycle;
            final Integer[] shards = locks.keySet().toArray(new Integer[]{});

            // Linear scan to figure out how many locks are held.
            int locksHeld = 0;
            for (int shard : shards) {
                if (locks.get(shard).isHeld()) {
                    locksHeld++;
                }
            }
            log.debug("Currently holding {} locks.", locksHeld);

            // if the number of locks held is less than 50% of the number of managed shards,
            // be aggressive and acquire more locks
            if (locksHeld <= moreThanHalf()) {
                maxLocksToAcquirePerCycle = moreThanHalf();
            }
            // shouldAttempt

            for (int shard : shards) {
                boolean shouldAttempt = locksAcquiredThisCycle < maxLocksToAcquirePerCycle;
                boolean isAcquired = locks.get(shard).performMaintenance(shouldAttempt);
                if (isAcquired) {
                    locksAcquiredThisCycle++;
                }
            }
            lastScavengedAt = nowMillis();
        } catch (RuntimeException e) {
            log.error("Error while scavengeLocks()", e);
        } finally {
            log.debug("Finishing scavengeLocks().");
        }
    }

    public void stateChanged(CuratorFramework curatorFramework, ConnectionState connectionState) {
        log.info("Connection to Zookeeper toggled to state " + connectionState.toString());
        connected = connectionState == ConnectionState.CONNECTED || connectionState == ConnectionState.RECONNECTED;
        if (connectionState == ConnectionState.LOST) {
            log.error("Connection to Zookeeper toggled to state " + connectionState.toString());
            this.handleZookeeperConnectionFailed();
        } else if (connectionState == ConnectionState.RECONNECTED) {
            log.info("Reconnected to zookeeper, forcing lock scavenge");
            forceLockScavenge();
        } else {
            log.info("Connection to Zookeeper toggled to state " + connectionState.toString());
        }
    }

    public synchronized void addShard(int shard) {
        if (locks.containsKey(shard)) {
            return;
        }
        this.locks.put(shard, new Lock(shard));
    }

    public synchronized void removeShard(int shard) {
        Lock lock = locks.remove(shard);
        if (lock != null) {
            lock.release(false);
        }
    }

    public boolean isConnected() {
        return connected && isCuratorStarted() && client.getZookeeperClient().isConnected();
    }

    /** Unsafe method for testing. */
    @VisibleForTesting
    void waitForQuiesceUnsafe() {
        while (lockTaskCount.get() > 0) {
            log.trace("Waiting for quiesce");
            try { Thread.sleep(100); } catch (InterruptedException ignore) {}
        }
        // this time out needs to be longer than ZK_LOCK_TIMEOUT for valid tests.
        try { Thread.sleep(2000); } catch (InterruptedException ignore) {}
    }

    @VisibleForTesting
    boolean holdsLockUnsafe(int shard) {
        return locks.containsKey(shard) && locks.get(shard).isHeld();
    }

    @VisibleForTesting
    boolean releaseLockUnsafe(int shard) throws Exception {
        return this.release(shard);
    }

    @VisibleForTesting
    void shutdownUnsafe() throws Exception {
        for (Lock lock : locks.values()) {
            this.release(lock.getShard());
        }
        client.close();
    }

    @VisibleForTesting
    Lock getLockUnsafe(int shard) {
        return locks.get(shard);
    }

    //
    // private methods
    //

    private long nowMillis() {
        return ticker.read() / 1000000;
    }

    private int moreThanHalf() {
        return locks.size() / 2 + 1;
    }

    private boolean isCuratorStarted() {
        return client.getState() == CuratorFrameworkState.STARTED;
    }

    //
    // JMX
    //

    public synchronized Collection<Integer> getHeldShards() {
        SortedSet<Integer> held = new TreeSet<Integer>();
        for (Lock lock : locks.values()) {
            if (lock.isHeld()) {
                held.add(lock.getShard());
            }
        }
        return held;
    }

    public synchronized Collection<Integer> getUnheldShards() {
        SortedSet<Integer> unheld = new TreeSet<Integer>();
        for (Lock lock : locks.values()) {
            if (lock.isUnheld()) {
                unheld.add(lock.getShard());
            }
        }
        return unheld;
    }

    public synchronized Collection<Integer> getErrorShards() {
        SortedSet<Integer> errorShards = new TreeSet<Integer>();
        for (Lock lock : locks.values()) {
            if (lock.getLockState() == LockState.ERROR) {
                errorShards.add(lock.getShard());
            }
        }
        return errorShards;
    }

    public synchronized void forceLockScavenge() {
        scavengeLocks();
    }

    public synchronized String getZkConnectionStatus() {
        if (!isCuratorStarted())
            return "not started";
        else if (client.getZookeeperClient().isConnected())
            return "connected";
        else
            return "not connected";
    }

    public synchronized boolean release(int shard) {
        if (locks.containsKey(shard)) {
            try {
                return locks.get(shard).release(false).get();
            } catch (InterruptedException ex) {
                log.error("Thread is interrupted:"+ ex.getMessage(), ex);
                return false;
            } catch (ExecutionException ex) {
                log.error("Release error: " + ex.getCause().getMessage(), ex.getCause());
                return false;
            }
        }
        return false;
    }

    public synchronized boolean acquire(int shard) {
        if (locks.containsKey(shard)) {
            try {
                return locks.get(shard).acquire().get();
            } catch (InterruptedException ex) {
                log.error("Thread is interrupted: "+ ex.getMessage(), ex);
                return false;
            } catch (ExecutionException ex) {
                log.error("Acquire error: " + ex.getCause().getMessage(), ex.getCause());
                return false;
            }
        }
        return false;
    }

    public synchronized long getMinLockHoldTimeMillis() { return minLockHoldTime.toMillis(); }
    public synchronized void setMinLockHoldTimeMillis(long millis) { minLockHoldTime = new TimeValue(millis, TimeUnit.MILLISECONDS); }
    public synchronized long getLockDisinterestedTimeMillis() { return lockDisinterestedTime.toMillis(); }
    public synchronized void setLockDisinterestedTimeMillis(long millis) { lockDisinterestedTime = new TimeValue(millis, TimeUnit.MILLISECONDS); }
    public synchronized long getSecondsSinceLastScavenge() { return ((nowMillis() - lastScavengedAt) / 1000); }

    //
    // Helper classes
    //

    enum LockState {
        /** Locks in UNKNOWN state will be attempt to be acquired. */
        UNKNOWN,
        /** Lock is acquired by the current blueflood instance. No other blueflood server will have this lock. */
        ACQUIRED,
        /** Attempted to acquire, but failed. Lock is already held by another instance. */
        ACQUIRE_FAILED,
        /**
         * Error state. normally means that the connection is lost. {@link Lock#mutex} is set to null.
         */
        ERROR,
        /** The lock was voluntarily released. Will not attempt to acquire this lock during this time. */
        DISINTERESTED
    }

    class Lock {
        private final int shard;
        private LockState state = LockState.UNKNOWN;
        // mutex is non-null unless disconnected from ZooKeeper.
        private InterProcessMutex mutex = null;
        // Timestamp of when the mutex has switched.
        private long stateChanged = nowMillis();

        Lock(int shard) {
            this.shard = shard;
            checkMutex();
        }

        @Override public String toString() {
            return String.format("shard=%d state=%s", shard, state);
        }

        synchronized void checkMutex() {
            if (mutex == null) {
                mutex = new InterProcessMutex(client, getLockId(shard));
            }
        }

        int getShard() { return shard; }
        boolean isHeld() { return connected && mutex != null && state == LockState.ACQUIRED; }
        boolean isUnheld() { return connected && mutex != null && state == LockState.ACQUIRE_FAILED; }

        /**
         * Performs the maintenance of the lock.
         *
         * <ul>
         *     <li>Move the lock state to {@link LockState#UNKNOWN} if necessary.</li>
         *     <li>Attempt to hold the locks in UNKNOWN state.</li>
         *     <li>Attempt to release the locks that were held for too long.</li>
         * </ul>
         *
         * @param shouldAttempt <code>true</code> if the lock should be attempted to be acquired.
         *
         * @return true if the lock was newly acquired during this cycle. otherwise false.
         */
        boolean performMaintenance(boolean shouldAttempt) {
            updateLockState();

            long now = nowMillis();
            if (state == LockState.UNKNOWN && shouldAttempt) {
                acquire();
                return isHeld();
            } else if (state == LockState.ACQUIRED && (now - stateChanged) >= minLockHoldTime.toMillis()) {
                // Lock was held for too long - maybe release to trigger re-balancing.
                final float chance;
                if (now == stateChanged) {
                    chance = 1;
                } else {
                    chance = (float)(now - stateChanged - minLockHoldTime.toMillis()) / (float)(Math.max(1, minLockHoldTime.toMillis()));
                }
                float r = rand.nextFloat();
                if (log.isTraceEnabled()) {
                    log.trace(String.format("Will release %s if, %f < %f", shard, r, chance));
                }
                if (r < chance) {
                    release(false);
                }
            } else if (state == LockState.ERROR && now - stateChanged > minLockHoldTime.toMillis()) {
                log.error("Lock state hasn't toggled from ERROR for " + minLockHoldTime.toString()
                        + "; Client connection status: " + connected);
            }
            return false;
        }

        /**
         * Sees if we should change our lock status to "UNKNOWN". Only locks in UNKNOWN state will periodically
         * be attempted to be acquired.
         */
        synchronized void updateLockState() {
            boolean toUnk = false;
            long now = nowMillis();

            if (state == LockState.DISINTERESTED && now - stateChanged > lockDisinterestedTime.toMillis()) {
                // Lock was voluntarily released long time ago. Attempt to acquire the lock again.
                toUnk = true;
            } else if (state == LockState.ERROR && connected) {
                // Reconnected to Zookeeper.
                toUnk = true;
            } else if (state == LockState.ACQUIRE_FAILED && now - stateChanged > lockDisinterestedTime.toMillis()) {
                // acquisition attempt has failed before long time ago. Attempt to acquire the lock again.
                toUnk = true;
            }

            if (toUnk) {
                setState(LockState.UNKNOWN);
            }
        }

        LockState getLockState() {
            return state;
        }

        /**
         * Only work on the acquired shards. Technically, working on the {@link LockState#ERROR} shard is ok,
         * but it is not recommended as it's going to cause a load problem.
         */
        boolean canWork() {
            return state == LockState.ACQUIRED;
        }

        /**
         * Triggered during the connection loss. Release the lease, and clear the mutex.
         */
        void connectionLost() {
            release(true);
        }

        synchronized void setState(LockState newState) {
            state = newState;
            if (newState == LockState.ERROR) {
                if (mutex.isAcquiredInThisProcess()) {
                    log.error("Assertion failure. You're still holding the lock in an error state.");
                }
                // clear the mutex if we reach the error state.
                mutex = null;
            }
            stateChanged = nowMillis();
        }

        Future<Boolean> acquire() {
            try {
                log.debug("Acquiring lock for shard={}.", shard);
                lockTaskCount.incrementAndGet();
                return lockWorker.submit(acquirer());
            } catch (RejectedExecutionException ex) {
                log.warn(String.format("Rejected lock execution: active:%d queue:%d shard:%d", lockWorker.getActiveCount(), lockWorker.getQueue().size(), shard));
                return Futures.immediateFuture(false);
            }
        }

        Future<Boolean> release(boolean disconnected) {
            try {
                log.debug("Releasing lock for shard={}.", shard);
                lockTaskCount.incrementAndGet();
                return lockWorker.submit(releaser(disconnected));
            } catch (RejectedExecutionException ex) {
                log.warn(String.format("Rejected lock execution: active:%d queue:%d shard:%d", lockWorker.getActiveCount(), lockWorker.getQueue().size(), shard));
                return Futures.immediateFuture(false);
            }
        }

        /**
         * Attempt to acquire a lock.
         * @return <code>true</code> if the lock is in acquired state (already acquired lock will return true).
         */
        Callable<Boolean> acquirer() {
            return new Callable<Boolean>() {
                public Boolean call() throws Exception {
                    com.codahale.metrics.Timer.Context ctx = lockAcquisitionTimer.time();
                    synchronized (Lock.this) {
                        try {
                            checkMutex();
                            if (state == LockState.ACQUIRED) {
                                // Already acquired.
                                return true;
                            } else if (state != LockState.UNKNOWN) {
                                // Prior acquisition have failed. Do not attempt.
                                return false;
                            } else {
                                log.debug("Trying ZK lock for shard={}", shard);
                                if (mutex.isAcquiredInThisProcess()) {
                                    // this should not occur but is here as a safeguard.
                                    log.error("Lock already acquired for shard={}", shard);
                                    setState(LockState.ACQUIRED);
                                    return true;
                                } else {
                                    try {
                                        boolean acquired = mutex.acquire(ZK_LOCK_TIMEOUT.getValue(), ZK_LOCK_TIMEOUT.getUnit());
                                        if (acquired) {
                                            setState(LockState.ACQUIRED);
                                            log.debug("Acquired ZK lock for shard={}", shard);
                                        } else {
                                            setState(LockState.ACQUIRE_FAILED);
                                            lockAcquisitionFailure.mark();
                                            log.debug("Acquire ZK failed for shard={}", shard);
                                        }
                                        return acquired;
                                    } catch (Exception e) {
                                        log.error("Exception on ZK acquire for shard={}", shard);
                                        lockErrors.mark();
                                        setState(LockState.ERROR);
                                        return false;
                                    }
                                }
                            }
                        } finally {
                            lockTaskCount.decrementAndGet();
                            ctx.stop();
                        }
                    }
                }
            };
        }

        /**
         * Attempt to release the lock.
         * @param disconnected If we're releasing the lock due to a ZooKeeper connection loss.
         */
        Callable<Boolean> releaser(final boolean disconnected) {
            return new Callable<Boolean>() {
                public Boolean call() throws Exception {
                    synchronized (Lock.this) {
                        try {
                            checkMutex();
                            if (mutex.isAcquiredInThisProcess()) {
                                if (state != LockState.ACQUIRED) {
                                    log.error("Held lock not held by this process? shard={}.", shard);
                                }
                                mutex.release();
                                return true;
                            }
                            return false;
                        } finally {
                            lockTaskCount.decrementAndGet();
                            if (disconnected) {
                                setState(LockState.ERROR);
                            } else {
                                setState(LockState.DISINTERESTED);
                            }
                        }
                    }
                }
            };
        }
    }
}


/**
 *    Copyright 2014 Square, Inc.
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

package com.rackspacecloud.blueflood.rollup;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.rackspacecloud.blueflood.exceptions.GranularityException;
import com.rackspacecloud.blueflood.io.Constants;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

/**
 * Immutable data structure representing the current slot being checked.
 * 3-tuple of (shard, slot, granularity).
 *
 * @author Jeeyoung Kim
 */
public final class SlotKey {
    private final int shard;
    private final int slot;
    private final Granularity granularity;

    private SlotKey(Granularity granularity, int slot, int shard) {
        Preconditions.checkNotNull(granularity);
        Preconditions.checkArgument(shard >= 0, "shard");
        Preconditions.checkArgument(shard < Constants.NUMBER_OF_SHARDS, "shard");
        Preconditions.checkArgument(slot >= 0, "slot");
        Preconditions.checkArgument(slot < granularity.numSlots(), "slot");

        this.shard = shard;
        this.slot = slot;
        this.granularity = granularity;
    }

    public int getShard() {
        return shard;
    }

    public int getSlot() {
        return slot;
    }

    public Granularity getGranularity() {
        return granularity;
    }

    public static SlotKey of(Granularity granularity, int slot, int shard) {
        return new SlotKey(granularity, slot, shard);
    }

    /**
     * Given the encoded slot key, returns the java object of it.
     * For the valid slot keys, this function is inverse of {@link #toString()}.
     * @return decoded {@link SlotKey}, null if it's an invalid slotkey.
     *
     * This function is mostly used in the tests.
     */
    public static SlotKey parse(String string) {
        String[] tokens = string.split(",");
        if (tokens.length != 3) {
            return null;
        }
        Granularity granularity = Granularity.fromString(tokens[0]);
        if (granularity == null) {
            return null;
        }
        try {
            int slot = Integer.parseInt(tokens[1]);
            int shard = Integer.parseInt(tokens[2]);
            return of(granularity, slot, shard);
        } catch (IllegalArgumentException e) {
            return null;
        }
    }

    /**
     * This method creates a collection of slot keys that are within the same
     * timespan of the current slot key but of a finer granularity. For example,
     * a slot key of granularity {@link Granularity#MIN_20 MIN_20} will have
     * four child keys, each of granularity {@link Granularity#MIN_5 MIN_5}.
     *

     *
     * This method is recursive, so that it also includes the current slot
     * key's children and their children, and so on, down to the finest
     * granularity. For example, a slot key of granularity
     * {@link Granularity#MIN_1440 MIN_1440} will have 6 immediate children of
     * granularity {@link Granularity#MIN_240 MIN_240}, 6*4=24 descendant slot
     * key of granularity {@link Granularity#MIN_60 MIN_60}, 6*4*3=72
     * descendant slot key of granularity {@link Granularity#MIN_20 MIN_20},
     * 6*4*3*4=288 descendant slot key of granularity
     * {@link Granularity#MIN_5 MIN_5},and 6*4*3*4*1=288 descendant slot key of
     * granularity {@link Granularity#FULL FULL}; in that case,the top-level
     * call of this method will return a collection 6+24+72+288+288=678
     * descendants.
     *
     * @return a Collection of all descendant {@link SlotKey}s
     */
    public Collection<SlotKey> getChildrenKeys() {
        if (granularity == Granularity.FULL) {
            return ImmutableList.of();
        }

        List<SlotKey> result = new ArrayList<SlotKey>();
        Granularity finer;
        try {
            finer = granularity.finer();
        } catch (GranularityException e) {
            throw new AssertionError("Should not occur.");
        }

        int factor = finer.numSlots() / granularity.numSlots();
        for (int i = 0; i < factor; i++) {
            int childSlot = slot * factor + i;
            SlotKey child = SlotKey.of(finer, childSlot, shard);
            result.add(child);
            result.addAll(child.getChildrenKeys());
        }
        return result;
    }

    /**
     * Same as the method {@link #getChildrenKeys()} except this method returns only
     * the children corresponding to the destination granularity.
     *
     * For example, a slot key of granularity
     * {@link Granularity#MIN_1440 MIN_1440}, for destination granularity of
     * {@link Granularity#MIN_60 MIN_60}, will have 6*4=24 children.
     *
     * @param destGranularity
     * @return
     */
    public Collection<SlotKey> getChildrenKeys(Granularity destGranularity) {

        if (!getGranularity().isCoarser(destGranularity)) {
            throw new IllegalArgumentException(String.format("Current granularity [%s] must be coarser than the destination granularity [%s]", getGranularity(), destGranularity));
        }

        List<SlotKey> result = new ArrayList<SlotKey>();
        for(SlotKey slotKey: this.getChildrenKeys()) {
            if (slotKey.getGranularity().equals(destGranularity)) {
                result.add(slotKey);
            }
        }

        return Collections.unmodifiableList(result);
    }

    /**
     * Returns the string representation used to store in the database.
     */
    @Override public String toString() {
        return String.format("%s,%d,%d", granularity.name(), slot, shard).intern();
    }

    @Override public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;

        SlotKey slotKey = (SlotKey) o;

        if (shard != slotKey.shard) return false;
        if (slot != slotKey.slot) return false;
        if (!granularity.equals(slotKey.granularity)) return false;

        return true;
    }

    @Override public int hashCode() {
        int result = shard;
        result = 31 * result + slot;
        result = 31 * result + granularity.hashCode();
        return result;
    }

    /**
     * This method would extrapolate a given slotkey to a corresponding parent slotkey of a given(higher) granularity.
     *
     * For example: For a given slotkey 'metrics_5m,1,30' which corresponds to granularity=5m, slot=1, shard=30, if we
     *              extrapolate this to a destination granularity of 20m, this would return 'metrics_20m,0,30',  which
     *              corresponds to granularity=20m, slot=0, shard=30 which is a parent key of the given slotkey
     *
     * @param destGranularity
     * @return
     */
    public SlotKey extrapolate(Granularity destGranularity) {

        if (destGranularity.equals(this.getGranularity())) {
            return this;
        }
        if (!destGranularity.isCoarser(getGranularity())) {
            throw new IllegalArgumentException("Destination granularity must be coarser than the current granularity");
        }

        int factor = getGranularity().numSlots() / destGranularity.numSlots();
        int parentSlot = getSlot() / factor;

        return SlotKey.of(destGranularity, parentSlot, getShard());
    }
}


/*
 * Copyright 2013 Rackspace
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

package com.rackspacecloud.blueflood.concurrent;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.rackspacecloud.blueflood.utils.TimeValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class ThreadPoolBuilder {
    private static final Logger log = LoggerFactory.getLogger(ThreadPoolBuilder.class);
    /** Used to ensure that the thread pools have unique name. */
    private static final ConcurrentHashMap<String, AtomicInteger> nameMap = new ConcurrentHashMap<String, AtomicInteger>();
    private static final String DEFAULT_NAME = "Threadpool";
    private int corePoolSize = 10;
    private int maxPoolSize = 10;
    private int queueSize = 0;
    private TimeValue keepAliveTime = new TimeValue(30, TimeUnit.SECONDS);

    private RejectedExecutionHandler rejectedHandler = new ThreadPoolExecutor.CallerRunsPolicy();
    private Thread.UncaughtExceptionHandler exceptionHandler = new Thread.UncaughtExceptionHandler() {
        public void uncaughtException(Thread t, Throwable e) {
            log.error(e.getMessage(), e);
        }
    };

    private String threadNameFormat = null;
    private String poolName = null;

    public ThreadPoolBuilder() {
        withName(DEFAULT_NAME);
    }

    public ThreadPoolBuilder withCorePoolSize(int size) {
        this.corePoolSize = size;
        return this;
    }

    public ThreadPoolBuilder withMaxPoolSize(int size) {
        this.maxPoolSize = size;
        return this;
    }

    public ThreadPoolBuilder withSynchronousQueue() {
        this.queueSize = 0;
        return this;
    }

    public ThreadPoolBuilder withUnboundedQueue() {
        this.queueSize = -1;
        return this;
    }

    public ThreadPoolBuilder withBoundedQueue(int size) {
        this.queueSize = size;
        return this;
    }

    public ThreadPoolBuilder withKeepAliveTime(TimeValue time) {
        this.keepAliveTime = time;
        return this;
    }

    public ThreadPoolBuilder withRejectedHandler(RejectedExecutionHandler rejectedHandler) {
        this.rejectedHandler = rejectedHandler;
        return this;
    }

    /**
     * Set the threadpool name. Used to generate metric names and thread names.
     */
    public ThreadPoolBuilder withName(String name) {
        // ensure we've got a spot to put the thread id.
        if (!name.contains("%d")) {
            name = name + "-%d";
        }
        nameMap.putIfAbsent(name, new AtomicInteger(0));
        int id = nameMap.get(name).incrementAndGet();
        this.poolName = String.format(name, id);
        if (id > 1) {
            this.threadNameFormat = name.replace("%d", id + "-%d");
        } else {
            this.threadNameFormat = name;
        }
        return this;
    }

    public ThreadPoolBuilder withExceptionHandler(Thread.UncaughtExceptionHandler exceptionHandler) {
        this.exceptionHandler = exceptionHandler;
        return this;
    }

    public ThreadPoolExecutor build() {
        BlockingQueue<Runnable> workQueue;
        switch (this.queueSize) {
            case 0: workQueue = new SynchronousQueue<Runnable>();
                break;
            case -1: workQueue = new LinkedBlockingQueue<Runnable>();
                break;
            default: workQueue = new ArrayBlockingQueue<Runnable>(queueSize);
                break;
        };

        ThreadPoolExecutor executor = new ThreadPoolExecutor(
                corePoolSize, maxPoolSize,
                keepAliveTime.getValue(), keepAliveTime.getUnit(),
                workQueue,
                new ThreadFactoryBuilder().setNameFormat(threadNameFormat).setPriority(Thread.NORM_PRIORITY).setUncaughtExceptionHandler(exceptionHandler).build(),
                rejectedHandler);
        InstrumentedThreadPoolExecutor.instrument(executor, poolName);
        return executor;
    }
}


/*
 * Copyright 2013-2015 Rackspace
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

package com.rackspacecloud.blueflood.inputs.handlers;

import com.codahale.metrics.Counter;
import com.codahale.metrics.Timer;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.gson.*;
import com.rackspacecloud.blueflood.exceptions.InvalidDataException;
import com.rackspacecloud.blueflood.http.DefaultHandler;
import com.rackspacecloud.blueflood.http.HttpRequestHandler;
import com.rackspacecloud.blueflood.inputs.formats.AggregatedPayload;
import com.rackspacecloud.blueflood.io.Constants;
import com.rackspacecloud.blueflood.io.Instrumentation;
import com.rackspacecloud.blueflood.outputs.formats.ErrorResponse;
import com.rackspacecloud.blueflood.tracker.Tracker;
import com.rackspacecloud.blueflood.types.MetricsCollection;
import com.rackspacecloud.blueflood.utils.Clock;
import com.rackspacecloud.blueflood.utils.DefaultClockImpl;
import com.rackspacecloud.blueflood.utils.Metrics;
import com.rackspacecloud.blueflood.utils.TimeValue;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http.FullHttpRequest;
import io.netty.handler.codec.http.HttpResponseStatus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeoutException;

public class HttpAggregatedMultiIngestionHandler implements HttpRequestHandler {

    private static final Logger log = LoggerFactory.getLogger(HttpAggregatedMultiIngestionHandler.class);

    private static final Timer handlerTimer = Metrics.timer(HttpAggregatedMultiIngestionHandler.class, "HTTP aggregated multi metrics ingestion timer");
    private static final Counter requestCount = Metrics.counter(HttpAggregatedMultiIngestionHandler.class, "HTTP aggregated multi Request Count");

    private final HttpMetricsIngestionServer.Processor processor;
    private final TimeValue timeout;
    private final boolean enablePerTenantMetrics;
    private final Clock clock = new DefaultClockImpl();

    public HttpAggregatedMultiIngestionHandler(HttpMetricsIngestionServer.Processor processor, TimeValue timeout) {
        this(processor, timeout, false);
    }

    public HttpAggregatedMultiIngestionHandler(HttpMetricsIngestionServer.Processor processor, TimeValue timeout, boolean enablePerTenantMetrics) {
        this.processor = processor;
        this.timeout = timeout;
        this.enablePerTenantMetrics = enablePerTenantMetrics;
    }

    // our own stuff.
    @Override
    public void handle(ChannelHandlerContext ctx, FullHttpRequest request) {

        Tracker.getInstance().track(request);

        requestCount.inc();

        final Timer.Context timerContext = handlerTimer.time();
        long ingestTime = clock.now().getMillis();

        // this is all JSON.
        String body = null;
        try {
            String submitterTenantId = request.headers().get(HttpMetricsIngestionServer.TENANT_ID_HEADER);

            body = request.content().toString(Constants.DEFAULT_CHARSET);
            List<AggregatedPayload> bundleList = createBundleList(body);

            if (bundleList.size() > 0) {
                // has aggregated metric bundle in body
                // convert and add metric bundle to MetricsCollection if valid
                MetricsCollection collection = new MetricsCollection();
                List<ErrorResponse.ErrorData> errors = new ArrayList<ErrorResponse.ErrorData>();

                // for each metric bundle
                int delayedMetricsCount = 0;
                int metricsCount = 0;
                for (AggregatedPayload bundle : bundleList) {
                    // validate, convert, and add to collection
                    List<ErrorResponse.ErrorData> bundleValidationErrors = bundle.getValidationErrors();
                    if (bundleValidationErrors.isEmpty()) {
                        // no validation error, add to collection
                        collection.add(PreaggregateConversions.buildMetricsCollection(bundle));
                    } else {
                        // failed validation, add to error
                        errors.addAll(bundleValidationErrors);
                    }

                    if (bundle.hasDelayedMetrics(ingestTime)) {
                        Tracker.getInstance().trackDelayedAggregatedMetricsTenant(bundle.getTenantId(),
                                bundle.getTimestamp(),
                                bundle.getDelayTime(ingestTime),
                                bundle.getAllMetricNames());
                        bundle.markDelayMetricsReceived(ingestTime);
                        delayedMetricsCount += bundle.getAllMetricNames().size();
                    } else {
                        metricsCount += bundle.getAllMetricNames().size();
                    }
                }

                // if has validation errors and no valid metrics
                if (!errors.isEmpty() && collection.size() == 0) {
                    // return BAD_REQUEST and error
                    DefaultHandler.sendErrorResponse(ctx, request,
                            errors,
                            HttpResponseStatus.BAD_REQUEST);
                    return;
                }

                // process valid metrics in collection
                ListenableFuture<List<Boolean>> futures = processor.apply(collection);
                List<Boolean> persisteds = futures.get(timeout.getValue(), timeout.getUnit());
                for (Boolean persisted : persisteds) {
                    if (!persisted) {
                        DefaultHandler.sendErrorResponse(ctx, request, "Internal error persisting data", HttpResponseStatus.INTERNAL_SERVER_ERROR);
                        return;
                    }
                }

                recordPerTenantMetrics(submitterTenantId, metricsCount, delayedMetricsCount);

                // return OK or MULTI_STATUS response depending if there were validation errors
                if (errors.isEmpty()) {
                    // no validation error, response OK
                    DefaultHandler.sendResponse(ctx, request, null, HttpResponseStatus.OK);
                    return;
                } else {
                    // has some validation errors, response MULTI_STATUS
                    DefaultHandler.sendErrorResponse(ctx, request, errors, HttpResponseStatus.MULTI_STATUS);
                    return;
                }

            } else {
                // no aggregated metric bundles in body, response OK
                DefaultHandler.sendResponse(ctx, request, "No valid metrics", HttpResponseStatus.BAD_REQUEST);
                return;
            }
        } catch (JsonParseException ex) {
            log.debug(String.format("BAD JSON: %s", body));
            DefaultHandler.sendErrorResponse(ctx, request, ex.getMessage(), HttpResponseStatus.BAD_REQUEST);
        } catch (InvalidDataException ex) {
            log.debug(String.format("Invalid request body: %s", body));
            DefaultHandler.sendErrorResponse(ctx, request, ex.getMessage(), HttpResponseStatus.BAD_REQUEST);
        } catch (TimeoutException ex) {
            DefaultHandler.sendErrorResponse(ctx, request, "Timed out persisting metrics", HttpResponseStatus.ACCEPTED);
        } catch (Exception ex) {
            log.debug(String.format("Exception processing: %s", body));
            log.error("Other exception while trying to parse content", ex);
            DefaultHandler.sendErrorResponse(ctx, request, "Internal error saving data", HttpResponseStatus.INTERNAL_SERVER_ERROR);
        } finally {
            timerContext.stop();
            requestCount.dec();
        }

    }

    public static List<AggregatedPayload> createBundleList(String json) {
        Gson gson = new Gson();
        JsonParser parser = new JsonParser();
        JsonElement element = parser.parse(json);

        if (!element.isJsonArray()) {
            throw new InvalidDataException("Invalid request body");
        }

        JsonArray jArray = element.getAsJsonArray();

        ArrayList<AggregatedPayload> bundleList = new ArrayList<AggregatedPayload>();

        for(JsonElement obj : jArray )
        {
            AggregatedPayload bundle = AggregatedPayload.create(obj);
            bundleList.add(bundle);
        }

        return bundleList;
    }

    @VisibleForTesting
    protected void recordPerTenantMetrics(String submitterTenantId, int metricsCount, int delayedMetricsCount) {
        if ( enablePerTenantMetrics ) {
            if ( metricsCount > 0 ) {
                Instrumentation.getIngestedMetricsMeter(submitterTenantId).mark(metricsCount);
            }
            if ( delayedMetricsCount > 0 ) {
                Instrumentation.getIngestedDelayedMetricsMeter(submitterTenantId).mark(delayedMetricsCount);
            }

        }
    }
}


package com.rackspacecloud.blueflood.io;

import java.util.*;

import static com.rackspacecloud.blueflood.types.Locator.METRIC_TOKEN_SEPARATOR_REGEX;

/**
 * When we ingest these below metrics, ES generates the following indexes
 *
 *     metric           indexes
 *   -----------        -------
 *   foo.bar.baz.aux -> [foo, foo.bar, foo.bar.baz, foo.bar.baz.aux, bar, baz, aux]
 *   foo.bar.baz     -> [foo, foo.bar, foo.bar.baz, bar, baz]
 *
 * If we request ES to return us upto the next level of the metric indexes for foo.bar.*,
 * we get this below response from ES where key is the metric index and doc_count is
 * the number of documents(think metric names) it points to.
 *
 * "buckets" : [ {
 *     "key" : "foo.bar.baz",
 *     "doc_count" : 2
 *   }, {
 *     "key" : "foo.bar.baz.aux",
 *     "doc_count" : 1
 * } ]
 *
 * If we feed each metric index to this class by calling the add(metricIndex, docCount) method,
 * this class analyzes the data and provides methods to retrieve different aspects of this data
 * with respect to baseLevel(if set to 3 with above data, uses foo.bar.baz as reference).
 *
 */
public class MetricIndexData {

    private static class MetricIndexDocCount {

        private long actualDocCount;
        private long childrenTotalDocCount;

        public MetricIndexDocCount(long actualDocCount, long childrenTotalDocCount) {
            this.actualDocCount = actualDocCount;
            this.childrenTotalDocCount = childrenTotalDocCount;
        }

        public void setActualDocCount(long actualDocCount) {
            this.actualDocCount = actualDocCount;
        }

        public void addChildrenTotalDocCount(long childrenDocCount) {
            this.childrenTotalDocCount += childrenDocCount;
        }
    }

    private final int baseLevel;

    //for metric names which have next level. (compared to baseLevel)
    private final Set<String> metricNamesWithNextLevelSet = new LinkedHashSet<String>();

    //contains all metric indexes which are of the same length as baseLevel. This is to determine if we
    //have any complete metric names at base level itself.
    final Map<String, MetricIndexDocCount> metricNameBaseLevelMap = new HashMap<>();

    public MetricIndexData(int baseLevel) {
        this.baseLevel = baseLevel;
    }

    /**
     * For a given metricIndex and docCount, classifies the data with respect to baseLevel
     * and stores it accordingly.
     *
     * @param metricIndex
     * @param docCount is document(metric name) count
     */
    public void add(String metricIndex, long docCount) {

        final String[] tokens = metricIndex.split(METRIC_TOKEN_SEPARATOR_REGEX);

        /**
         *
         * For the ES response shown in class description, for a baseLevel of 2,
         * the data is classified as follows
         *
         * metricNamesWithNextLevelSet    -> {foo.bar.baz}   (Metric Names which are at base + 1 level and also have a subsequent level.)
         *
         * For count map, data is in the form {metricIndex, (actualDocCount, childrenTotalDocCount)}
         *
         * metricNameBaseLevelMap   -> {foo.bar.baz -> (2, 1)}     (all indexes which are of same length as baseLevel)
         *
         */

        switch (tokens.length - baseLevel) {
            case 1:
                if (baseLevel > 0) {
                    metricNamesWithNextLevelSet.add(metricIndex.substring(0, metricIndex.lastIndexOf(".")));
                } else {
                    metricNamesWithNextLevelSet.add(metricIndex.substring(0, metricIndex.indexOf(".")));
                }

                //For foo.bar.baz, baseLevel=3 we update children doc count of foo.bar.baz
                addChildrenDocCount(metricNameBaseLevelMap,
                                    metricIndex.substring(0, metricIndex.lastIndexOf(".")),
                                    docCount);
                break;

            case 0:

                setActualDocCount(metricNameBaseLevelMap, metricIndex, docCount);
                break;

            default:
                break;
        }
    }

    /**
     * Metric Names which have next level compared to baseLevel.
     *
     * Ex: For metrics foo.bar.baz.qux, foo.bar.baz
     *      if baseLevel = 3, returns foo.bar.baz as there it has next level token 'qux'
     *
     * @return
     */
    public Set<String> getMetricNamesWithNextLevel() {
        return Collections.unmodifiableSet(metricNamesWithNextLevelSet);
    }

    /**
     * Returns complete metric names which are of same length as baseLevel
     *
     * Ex: For metrics foo.bar.baz.qux, foo.bar.baz
     *      if baseLevel = 3, returns foo.bar.baz
     *
     * @return
     */
    public Set<String> getCompleteMetricNamesAtBaseLevel() {
        return getCompleteMetricNames(metricNameBaseLevelMap);
    }

    /**
     *  Compares actualDocCount and total docCount of its immediate children of an index
     *  to determine if the metric index is a complete metric name or not.
     *
     *  For the ES response shown in class description, for a baseLevel of 2,
     *  foo.bar.baz has actualDocCount of 2, but total doc count of all its children,
     *  which in this case is only foo.bar.baz, is 1. So foo.bar.baz must be metric by itself
     *
     * @param metricIndexMap
     * @return
     */
    private Set<String> getCompleteMetricNames(Map<String, MetricIndexDocCount> metricIndexMap) {

        Set<String> completeMetricNames = new HashSet<String>();
        for (Map.Entry<String, MetricIndexDocCount> entry : metricIndexMap.entrySet()) {

            MetricIndexDocCount metricIndexDocCount = entry.getValue();
            if (metricIndexDocCount != null) {

                //if total doc count is greater than its children docs, its a complete metric name
                if (metricIndexDocCount.actualDocCount > 0 &&
                        metricIndexDocCount.actualDocCount > metricIndexDocCount.childrenTotalDocCount) {
                    completeMetricNames.add(entry.getKey());
                }
            }
        }

        return Collections.unmodifiableSet(completeMetricNames);
    }

    private void setActualDocCount(Map<String, MetricIndexDocCount> metricIndexMap, String key, final long count) {
        MetricIndexDocCount metricIndexDocCount = metricIndexMap.containsKey(key) ? metricIndexMap.get(key) : new MetricIndexDocCount(0, 0);
        metricIndexDocCount.setActualDocCount(count);

        metricIndexMap.put(key, metricIndexDocCount);
    }

    private void addChildrenDocCount(Map<String, MetricIndexDocCount> metricIndexMap, String key, final long childrenDocCount) {
        MetricIndexDocCount metricIndexDocCount = metricIndexMap.containsKey(key) ? metricIndexMap.get(key) : new MetricIndexDocCount(0, 0);
        metricIndexDocCount.addChildrenTotalDocCount(childrenDocCount);

        metricIndexMap.put(key, metricIndexDocCount);
    }
}


/*
 * Copyright 2013 Rackspace
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

package com.rackspacecloud.blueflood.service;

import com.codahale.metrics.Histogram;
import com.codahale.metrics.Meter;
import com.google.common.base.Ticker;
import com.rackspacecloud.blueflood.exceptions.GranularityException;
import com.rackspacecloud.blueflood.rollup.Granularity;
import com.rackspacecloud.blueflood.rollup.SlotKey;
import com.rackspacecloud.blueflood.utils.Clock;
import com.rackspacecloud.blueflood.utils.DefaultClockImpl;
import com.rackspacecloud.blueflood.utils.Metrics;
import com.rackspacecloud.blueflood.utils.Util;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public class ShardStateManager {
    private static final Logger log = LoggerFactory.getLogger(ShardStateManager.class);
    private static final Set<Integer> ALL_SHARDS = new HashSet<Integer>(Util.parseShards("ALL"));
    final Set<Integer> shards; // Managed shards
    final Map<Integer, ShardToGranularityMap> shardToGranularityStates = new HashMap<Integer, ShardToGranularityMap>();
    private final Ticker serverTimeMillisecondTicker;
    private static final long millisInADay = 24 * 60 * 60 * 1000;

    private static final Histogram timeSinceUpdate = Metrics.histogram(RollupService.class, "Shard Slot Time Elapsed scheduleEligibleSlots");
    // todo: CM_SPECIFIC verify changing metric class name doesn't break things.
    private static final Meter updateStampMeter = Metrics.meter(ShardStateManager.class, "Shard Slot Update Meter");
    private final Meter parentBeforeChild = Metrics.meter(RollupService.class, "Parent slot executed before child");
    private static final Map<Granularity, Meter> granToReRollMeters = new HashMap<Granularity, Meter>();
    private static final Map<Granularity, Meter> reRollForShortDelayMetricsMeters = new HashMap<Granularity, Meter>();
    private static final Map<Granularity, Meter> reRollForLongDelayMetricsMeters = new HashMap<Granularity, Meter>();
    private static final Map<Granularity, Meter> granToDelayedMetricsMeter = new HashMap<Granularity, Meter>();

    // If there are no delayed metrics, a slot should only get rolled up again after 14 days since its last rollup. Since we only allow
    // delayed data that is BEFORE_CURRENT_COLLECTIONTIME_MS (typically 3 days) old, we are assuming that if rollup happens again within
    // that time, its a re-roll because of delayed data.
    public static final long REROLL_TIME_SPAN_ASSUMED_VALUE = Configuration.getInstance().getLongProperty( CoreConfig.BEFORE_CURRENT_COLLECTIONTIME_MS );

    private final Clock clock;

    static {
        for (Granularity rollupGranularity : Granularity.rollupGranularities()) {
            granToReRollMeters.put(rollupGranularity, Metrics.meter(RollupService.class, String.format("%s Re-rolling up because of delayed metrics", rollupGranularity.shortName())));
            granToDelayedMetricsMeter.put(rollupGranularity, Metrics.meter(RollupService.class, String.format("Delayed metric that has a danger of TTLing", rollupGranularity.shortName())));

            reRollForShortDelayMetricsMeters.put(rollupGranularity, Metrics.meter(RollupService.class, String.format("%s Slots waiting to be re-rolled because of short delay metrics", rollupGranularity.shortName())));
            reRollForLongDelayMetricsMeters.put(rollupGranularity, Metrics.meter(RollupService.class, String.format("%s Slots waiting to be re-rolled because of long delay metrics", rollupGranularity.shortName())));
        }
    }

    protected ShardStateManager(Collection<Integer> shards, Ticker ticker) {
        this(shards, ticker, new DefaultClockImpl());
    }

    protected ShardStateManager(Collection<Integer> shards, Ticker ticker, Clock clock) {
        this.shards = new HashSet<Integer>(shards);
        for (Integer shard : ALL_SHARDS) { // Why not just do this for managed shards?
            shardToGranularityStates.put(shard, new ShardToGranularityMap(shard));
        }
        this.serverTimeMillisecondTicker = ticker;
        this.clock = clock;
    }

    protected Collection<Integer> getManagedShards() {
        return Collections.unmodifiableCollection(this.shards);
    }

    protected Boolean contains(int shard) {
        return shards.size() != 0 && shards.contains(shard);
    }

    protected void add(int shard) {
        if (contains(shard))
            return;
        shards.add(shard);
    }

    protected void remove(int shard) {
        if (!contains(shard))
            return;
        this.shards.remove(shard);
    }

    public SlotStateManager getSlotStateManager(int shard, Granularity granularity) {
        return shardToGranularityStates.get(shard).granularityToSlots.get(granularity);
    }

    protected UpdateStamp getUpdateStamp(SlotKey slotKey) {
        SlotStateManager slotStateManager = this.getSlotStateManager(slotKey.getShard(), slotKey.getGranularity());
        UpdateStamp stamp = slotStateManager.slotToUpdateStampMap.get(slotKey.getSlot());
        return stamp;

    }

    // Side effect: mark dirty slots as clean
    protected Map<Granularity, Map<Integer, UpdateStamp>> getDirtySlotsToPersist(int shard) {
        Map<Granularity, Map<Integer, UpdateStamp>> slotTimes = new HashMap<Granularity, Map<Integer, UpdateStamp>>();
        int numUpdates = 0;
        for (Granularity gran : Granularity.rollupGranularities()) {
            Map<Integer, UpdateStamp> dirty = getSlotStateManager(shard, gran).getDirtySlotStampsAndMarkClean();
            slotTimes.put(gran, dirty);

            if (dirty.size() > 0) {
                numUpdates += dirty.size();
            }
        }
        if (numUpdates > 0) {
            // for updates that come by way of scribe, you'll typically see 5 as the number of updates (one for
            // each granularity).  On rollup slaves the situation is a bit different. You'll see only the slot
            // of the granularity just written to marked dirty (so 1).
            log.debug("Found {} dirty slots for shard {}", new Object[]{numUpdates, shard});
            return slotTimes;
        }
        return null;
    }

    public void updateSlotOnRead(int shard, SlotState slotState) {
        getSlotStateManager(shard, slotState.getGranularity()).updateSlotOnRead(slotState);
    }

    public void setAllCoarserSlotsDirtyForSlot(SlotKey slotKey) {
        boolean done = false;
        Granularity coarserGran = slotKey.getGranularity();
        int coarserSlot = slotKey.getSlot();

        while (!done) {
            try {
                coarserGran = coarserGran.coarser();
                coarserSlot = coarserGran.slotFromFinerSlot(coarserSlot);
                ConcurrentMap<Integer, UpdateStamp> updateStampsBySlotMap = getSlotStateManager(slotKey.getShard(), coarserGran).slotToUpdateStampMap;
                UpdateStamp coarseSlotStamp = updateStampsBySlotMap.get(coarserSlot);

                if (coarseSlotStamp == null) {
                    log.debug("No stamp for coarser slot: {}; supplied slot: {}",
                            SlotKey.of(coarserGran, coarserSlot, slotKey.getShard()),
                            slotKey);
                    updateStampsBySlotMap.putIfAbsent(coarserSlot,
                            new UpdateStamp(serverTimeMillisecondTicker.read(), UpdateStamp.State.Active, true));
                    continue;
                }

                UpdateStamp.State coarseSlotState = coarseSlotStamp.getState();
                if (coarseSlotState != UpdateStamp.State.Active) {
                    parentBeforeChild.mark();
                    log.debug("Coarser slot not in active state when finer slot {} just got rolled up. Marking coarser slot {} dirty.",
                            slotKey, SlotKey.of(coarserGran, coarserSlot, slotKey.getShard()));
                    coarseSlotStamp.setState(UpdateStamp.State.Active);
                    coarseSlotStamp.setDirty(true);
                    coarseSlotStamp.setTimestamp(serverTimeMillisecondTicker.read());
                }
            } catch (GranularityException ex) {
                done = true;
            }
        }
    }

    private class ShardToGranularityMap {
        final int shard;
        final Map<Granularity, SlotStateManager> granularityToSlots = new HashMap<Granularity, SlotStateManager>();

        protected ShardToGranularityMap(int shard) {
            this.shard = shard;

            for (Granularity granularity : Granularity.rollupGranularities()) {
                granularityToSlots.put(granularity, new SlotStateManager(shard, granularity));
            }
        }
    }

    protected class SlotStateManager {
        private final int shard;
        final Granularity granularity;
        final ConcurrentMap<Integer, UpdateStamp> slotToUpdateStampMap;

        protected SlotStateManager(int shard, Granularity granularity) {
            this.shard = shard;
            this.granularity = granularity;
            slotToUpdateStampMap = new ConcurrentHashMap<Integer, UpdateStamp>(granularity.numSlots());
        }

        /**
         *
         * Imagine metrics are flowing in from multiple ingestor nodes. The
         * ingestion path updates schedule context while writing metrics to
         * cassandra.(See
         * {@link com.rackspacecloud.blueflood.inputs.processors.BatchWriter BatchWriter}).
         * We cannot make any ordering guarantees on the metrics. So every
         * metric that comes in updates the slot state to its collection time.
         *

         *
         * This state gets pushed in cassandra by {@link ShardStatePusher} and read on
         * the rollup slave. Rollup slave is going to update its state to
         * {@link com.rackspacecloud.blueflood.service.UpdateStamp.State#Active ACTIVE}
         * as long as the timestamp does not match. Rollup slave shard
         * map can be in 3 states:
         * 1) {@link com.rackspacecloud.blueflood.service.UpdateStamp.State#Active Active}
         * 2) {@link com.rackspacecloud.blueflood.service.UpdateStamp.State#Rolled Rolled}
         * 3) {@link com.rackspacecloud.blueflood.service.UpdateStamp.State#Running Running}.
         * Every {@code ACTIVE} update is taken for Rolled and Running states,
         * but if the shard map is already in an {@code ACTIVE} state, then the
         * update happens only if the timestamp of update coming in is greater
         * than what we have. On Rollup slave it means eventually when it rolls
         * up data for the {@code ACTIVE} slot, it will be marked with the
         * collection time belonging to a metric which was generated later.
         *

         *
         * For a case of multiple ingestors, it means eventually higher
         * timestamp will win, and will be updated even if that ingestor did
         * not receive metric with that timestamp and will stop triggering the
         * state to {@code ACTIVE} on rollup host. After this convergence is
         * reached the last rollup time match with the last active times on all
         * ingestor nodes.
         *
         * LastUpdateTime of an active slot is used as last ingest time and lastUpdateTime
         * of a rolled slot is used as last rollup time. UpdateStamp which is in memory
         * for each slot has the last ingest time and last rollup time of that slot.
         */
        protected void updateSlotOnRead(SlotState slotState) {
            final int slot = slotState.getSlot();
            final long timestamp = slotState.getTimestamp();
            UpdateStamp.State state = slotState.getState();

            //For slots in state "A", this would be last ingest time
            //For slots in state "X", this would be last rollup time
            final long lastUpdateTimestamp = slotState.getLastUpdatedTimestamp();

            UpdateStamp stampInMemory = slotToUpdateStampMap.get(slot);
            if (stampInMemory == null) {
                // haven't seen this slot before, take the update. This happens when a blueflood service is just started.
                slotToUpdateStampMap.put(slot, new UpdateStamp(timestamp, state, false, 0, lastUpdateTimestamp));
            } else if (stampInMemory.getTimestamp() != timestamp && state.equals(UpdateStamp.State.Active)) {
                // 1) new update coming in. We can be in 3 states 1) Active 2) Rolled 3) Running. Apply the update in all cases except when we are already active and
                //    the triggering timestamp we have is greater or the stampInMemory is yet to be persisted i.e still dirty

                // This "if" is equivalent to:
                //  if (current is not active) || (current is older && clean)
                if (!(stampInMemory.getState().equals(UpdateStamp.State.Active) && (stampInMemory.getTimestamp() > timestamp || stampInMemory.isDirty()))) {
                    slotToUpdateStampMap.put(slot, new UpdateStamp(timestamp, state, false, stampInMemory.getLastRollupTimestamp(), lastUpdateTimestamp));
                } else {
                    // keep rewriting the newer timestamp, in case it has been overwritten:
                    stampInMemory.setDirty(true); // This is crucial for convergence, we need to superimpose a higher timestamp which can be done only if we set it to dirty
                }
            } else if (stampInMemory.getTimestamp() == timestamp && state.equals(UpdateStamp.State.Rolled)) {
                // 2) if current value is same but value being applied is a remove, remove wins.
                stampInMemory.setState(UpdateStamp.State.Rolled);

                //For incoming update(from metrics_state) of "Rolled" status, we use its last updated time as the last rollup time.
                if (lastUpdateTimestamp > stampInMemory.getLastRollupTimestamp())
                    stampInMemory.setLastRollupTimestamp(lastUpdateTimestamp);
            } else if (state.equals(UpdateStamp.State.Rolled)) {

                //For incoming update(from metrics_state) of "Rolled" status, we use its last updated time as the last rollup time.
                if (lastUpdateTimestamp > stampInMemory.getLastRollupTimestamp())
                    stampInMemory.setLastRollupTimestamp(lastUpdateTimestamp);
            }
        }

        protected void createOrUpdateForSlotAndMillisecond(int slot, long millis) {
            long nowMillis = clock.now().getMillis();
            if (slotToUpdateStampMap.containsKey(slot)) {
                UpdateStamp stamp = slotToUpdateStampMap.get(slot);
                stamp.setTimestamp(millis);

                // Temporarily setting last ingest time to current time until we get more accurate value from db.
                // This will not be persisted.
                stamp.setLastIngestTimestamp(nowMillis);

                stamp.setState(UpdateStamp.State.Active);
                stamp.setDirty(true);
            } else {
                slotToUpdateStampMap.put(slot, new UpdateStamp(millis, UpdateStamp.State.Active, true, 0, nowMillis));
            }
            updateStampMeter.mark();
        }

        protected Map<Integer, UpdateStamp> getDirtySlotStampsAndMarkClean() {
            HashMap<Integer, UpdateStamp> dirtySlots = new HashMap<Integer, UpdateStamp>();
            for (Map.Entry<Integer, UpdateStamp> entry : slotToUpdateStampMap.entrySet()) {
                if (entry.getValue().isDirty()) {
                    dirtySlots.put(entry.getKey(), entry.getValue());
                    entry.getValue().setDirty(false);
                }
            }
            return dirtySlots;
        }

        protected UpdateStamp getAndSetState(int slot, UpdateStamp.State state) {
            UpdateStamp stamp = slotToUpdateStampMap.get(slot);
            stamp.setState(state);
            return stamp;
        }

        // gets a snapshot of the last updates
        protected Map<Integer, UpdateStamp> getSlotStamps() {
            // essentially a copy on read map.
            return Collections.unmodifiableMap(slotToUpdateStampMap);
        }

        /**
         * Determines if a slot is being re-rolled or not.
         *
         * Since we only allow delayed metrics upto 3 days(BEFORE_CURRENT_COLLECTIONTIME_MS), a slot can be
         * identified as being re-rolled, if the last rollup is within those last 3 days.
         *
         * @param slot
         * @param now
         * @return
         */
        protected boolean isReroll(int slot, long now) {
            final UpdateStamp updateStamp = slotToUpdateStampMap.get(slot);
            final long timeElapsedSinceLastRollup = now - updateStamp.getLastRollupTimestamp();

            if (updateStamp.getLastRollupTimestamp() > 0 &&
                    timeElapsedSinceLastRollup < REROLL_TIME_SPAN_ASSUMED_VALUE) {
                return true;
            }

            return false;
        }

        /**
         * A slot will become eligible for rollup/re-roll based on the below three configs.
         *
         * 1) 1st rollup  -> Eligible after ROLLUP_DELAY_MILLIS from collection time.
         * 2) 1st re-roll -> This happens for metrics with short delay (within SHORT_DELAY_METRICS_ROLLUP_DELAY_MILLIS).
         *                   Eligible after SHORT_DELAY_METRICS_ROLLUP_DELAY_MILLIS from collection time.
         * 3) nth re-roll -> This happens for metrics with long delay(more than SHORT_DELAY_METRICS_ROLLUP_DELAY_MILLIS).
         *                   Eligible after LONG_DELAY_METRICS_ROLLUP_WAIT_MILLIS from last ingest time.
         *                   This re-roll repeats as we keep getting delayed metrics.
         *
         *
         *     |<-- SHORT_DELAY_METRICS_ROLLUP_DELAY--->|
         *     |                                        |
         *     |<-- ROLLUP_DELAY--->|                   |     |<--ROLLUP_WAIT-->|                |<--ROLLUP_WAIT-->|
         *     |                    |                   |     |                 |                |                 |
         * -------------------------------------------------------------------------------------------------------------
         * |slot|                   ^          ^        ^     ^                 ^    ^   ^   ^   ^                 ^
         * | X  |                   |          |        |     |                 |    |   |   |   |                 |
         *                          |   delayed metric  | delayed metric        |   delayed metrics                |
         *                          |                   |                       |                                  |
         *                      1st rollup         1st re-roll              2nd re-roll                      3rd re-roll
         *
         *
         * @param now is current time
         * @param maxAgeMillis is ROLLUP_DELAY_MILLIS
         * @param rollupDelayForMetricsWithShortDelay is SHORT_DELAY_METRICS_ROLLUP_DELAY_MILLIS
         * @param rollupWaitForMetricsWithLongDelay is LONG_DELAY_METRICS_ROLLUP_WAIT_MILLIS
         * @return list of slots that are eligible for rollup
         */
        protected List<Integer> getSlotsEligibleForRollup(long now,
                                                          long maxAgeMillis,
                                                          long rollupDelayForMetricsWithShortDelay,
                                                          long rollupWaitForMetricsWithLongDelay) {
            List<Integer> outputKeys = new ArrayList<Integer>();
            long nowMillis = clock.now().getMillis();
            for (Map.Entry<Integer, UpdateStamp> entry : slotToUpdateStampMap.entrySet()) {
                final int slot = entry.getKey();
                final UpdateStamp update = entry.getValue();
                final long timeElapsed = now - update.getTimestamp();
                timeSinceUpdate.update(timeElapsed);

                if (update.getState() == UpdateStamp.State.Rolled) {
                    continue;
                }
                if (timeElapsed <= maxAgeMillis) {
                    continue;
                }

                //Handling re-rolls:
                if (isReroll(slot, now)) {
                    SlotKey slotKey = SlotKey.of(granularity, entry.getKey(), shard);

                    //short delay
                    if (timeElapsed <= rollupDelayForMetricsWithShortDelay) {

                        reRollForShortDelayMetricsMeters.get(granularity).mark();
                        log.debug(String.format("Short delay: Delaying re-roll of slotKey [%s] as [%d] millis " +
                                "haven't elapsed since collection time:[%d] now: [%d] time elapsed: [%d] last " +
                                "rollup time: [%d]", slotKey, rollupDelayForMetricsWithShortDelay,
                                update.getTimestamp(), now, timeElapsed, update.getLastRollupTimestamp()));
                        continue;
                    }

                    if (update.getLastIngestTimestamp() > 0 ) {
                        long delayOfLastIngestedMetric = update.getLastIngestTimestamp() - update.getTimestamp();
                        final long timeElapsedSinceLastIngest = now - update.getLastIngestTimestamp();

                        //long delay
                        if (delayOfLastIngestedMetric > rollupDelayForMetricsWithShortDelay &&
                                timeElapsedSinceLastIngest <= rollupWaitForMetricsWithLongDelay) {

                            reRollForLongDelayMetricsMeters.get(granularity).mark();
                            log.debug(String.format("Long delay: Delaying re-roll of slotKey [%s] as we received " +
                                            "delayed metrics within the last [%d] millis with rollup_wait of [%d] millis. last " +
                                            "ingest time: [%d]", slotKey, timeElapsedSinceLastIngest,
                                    rollupWaitForMetricsWithLongDelay, update.getLastIngestTimestamp()));
                            continue;
                        }
                    }

                    granToReRollMeters.get(granularity).mark();
                    if (nowMillis - update.getTimestamp() >= millisInADay) {
                        granToDelayedMetricsMeter.get(granularity).mark();
                    }
                }
                outputKeys.add(entry.getKey());
            }

            return outputKeys;
        }
    }
}


package com.rackspacecloud.blueflood.types;

import org.apache.commons.lang.StringUtils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;
import static java.util.stream.Collectors.toSet;

/**
 *
 * For a given tenantId = '111111', metricName = 'a.b.c.d',
 *
 * level 2 token will have the following information.
 *
 *      token = 'c'
 *      parent = 'a.b'
 *      isLeaf = false
 *      id = '111111:a.b.c'
 *      tenantId = '111111'
 *
 * level 3 token which is a leaf node will have the following information
 *
 *      token = 'd'
 *      parent = 'a.b.c'
 *      isLeaf = true
 *      id = '111111:a.b.c.d:$'
 *      tenantId = '111111'
 *
 */
public class Token {

    public static final String LEAF_NODE_SUFFIX = ":$";
    public static final String SEPARATOR = ":";

    private final String token;
    private final String parent;
    private final boolean isLeaf;

    private final String id;
    private final Locator locator;

    /**
     * Constructor used to create {@link Token} instances.
     *
     * @param locator is a metric locator
     * @param tokens, array of tokens. For metric name a.b.c.d, tokens would be the array ["a", "b", "c", "d"]
     * @param level level of token this object represents. Ex: level 1 of metric a.b.c.d, would be "b"
     */
    public Token(Locator locator, String[] tokens, int level) {

        if (tokens == null || tokens.length <= 0)
            throw new IllegalArgumentException("Invalid tokens. Must be an array of size " +
                                                       "greater than 0, representing tokens of a metric name.");

        if (level < 0 || level > tokens.length - 1)
            throw new IllegalArgumentException("Invalid level for the given tokens");

        this.locator = locator;
        this.token = tokens[level];
        this.isLeaf = level == tokens.length - 1;

        String prefix = locator.getTenantId() + Token.SEPARATOR;
        this.parent = joinTokens("", "", tokens, level);

        String suffix = isLeaf ? LEAF_NODE_SUFFIX : "";
        this.id = joinTokens(prefix, suffix, tokens, level + 1);
    }

    private String joinTokens(String prefix, String suffix, String[] tokens, int level) {
        return Arrays.stream(tokens)
                     .limit(level)
                     .collect(joining(Locator.METRIC_TOKEN_SEPARATOR, prefix, suffix));
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;

        if (o == null || getClass() != o.getClass()) return false;

        Token token = (Token) o;

        return id.equalsIgnoreCase(token.getId());
    }

    @Override
    public int hashCode() {
        return id == null ? 0 : id.hashCode();
    }

    public String getId() {
        return id;
    }

    public Locator getLocator() {
        return locator;
    }

    public String getToken() {
        return token;
    }

    public String getParent() {
        return parent;
    }

    public boolean isLeaf() {
        return isLeaf;
    }

    /**
     * This method return list of tokens with their parents for a current Discovery object.
     *
     * For example: A locator of 1111:a.b.c.d would generate the following tokens
     *
     * Token{token='a', parent='', isLeaf=false, id='111111:a', locator=111111.a.b.c.d}
     * Token{token='b', parent='a', isLeaf=false, id='111111:a.b', locator=111111.a.b.c.d}
     * Token{token='c', parent='a.b', isLeaf=false, id='111111:a.b.c', locator=111111.a.b.c.d}
     * Token{token='d', parent='a.b.c', isLeaf=true, id='111111:a.b.c.d:$', locator=111111.a.b.c.d}
     *
     * @return
     */
    public static List<Token> getTokens(Locator locator) {

        if (StringUtils.isEmpty(locator.getMetricName()) || StringUtils.isEmpty(locator.getTenantId()))
            return new ArrayList<>();

        String[] tokens = locator.getMetricName().split(Locator.METRIC_TOKEN_SEPARATOR_REGEX);

        return IntStream.range(0, tokens.length)
                        .mapToObj(index -> new Token(locator, tokens, index))
                        .collect(toList());
    }

    public static Stream<Token> getUniqueTokens(Stream<Locator> locators) {
        return locators.flatMap(locator -> Token.getTokens(locator).stream())
                       .collect(toSet())
                       .stream();
    }

    @Override
    public String toString() {
        return "Token{" +
                "token='" + token + '\'' +
                ", parent='" + parent + '\'' +
                ", isLeaf=" + isLeaf +
                ", id='" + id + '\'' +
                ", locator=" + locator +
                '}';
    }
}


/*
 * Copyright (c) 2016 Rackspace.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * 	http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.rackspacecloud.blueflood.io.datastax;

import com.datastax.driver.core.*;
import com.google.common.collect.HashBasedTable;
import com.google.common.collect.Table;
import com.rackspacecloud.blueflood.io.CassandraModel;
import com.rackspacecloud.blueflood.io.Instrumentation;
import com.rackspacecloud.blueflood.rollup.Granularity;
import com.rackspacecloud.blueflood.types.*;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.Long;
import java.nio.ByteBuffer;
import java.util.*;

/**
 * This is an abstract class that collects the common behavior of
 * reading/writing preaggregated & basic numeric metrics objects.
 */
public abstract class DAbstractMetricIO {

    private static final Logger LOG = LoggerFactory.getLogger(DAbstractMetricIO.class);

    protected Session session;

    protected final DMetricsCFPreparedStatements metricsCFPreparedStatements;

    protected DAbstractMetricIO() {
        metricsCFPreparedStatements = DMetricsCFPreparedStatements.getInstance();
        session = DatastaxIO.getSession();
    }

    /**
     * Asynchronously insert a rolled up metric to the appropriate column family
     * for a particular granularity
     *
     * @param locator
     * @param collectionTime
     * @param rollup
     * @param granularity
     * @return
     */
    public ResultSetFuture putAsync(Locator locator, long collectionTime, Rollup rollup, Granularity granularity, int ttl) {

        Session session = DatastaxIO.getSession();

        // we use batch statement here in case sub classes
        // override the addRollupToBatch() and provide
        // multiple statements
        BatchStatement batch = new BatchStatement();
        addRollupToBatch(batch, locator, rollup, collectionTime, granularity, ttl);

        Collection<Statement> statements = batch.getStatements();
        if ( statements.size() == 1 ) {
            Statement oneStatement = statements.iterator().next();
            return session.executeAsync(oneStatement);
        } else {
            LOG.debug(String.format("Using BatchStatement for %d statements", statements.size()));
            return session.executeAsync(batch);
        }
    }

    public Statement createStatement(Locator locator, long collectionTime, Rollup rollup, Granularity granularity, int ttl) {
        final PreparedStatement statement;

        if( rollup.getRollupType() == RollupType.BF_BASIC ) {

            // Strings and Booleans don't get rolled up.  I'd like to verify
            // that none are passed in, but that would require a db access

            statement = metricsCFPreparedStatements.basicGranToInsertStatement.get( granularity );
        }
        else {
            statement = metricsCFPreparedStatements.preaggrGranToInsertStatement.get(granularity);
        }

        BoundStatement bound = statement.bind(locator.toString(),
                collectionTime,
                toByteBuffer(rollup),
                ttl);

        return bound;
    }

    /**
     * Fetch rollup objects for a {@link com.rackspacecloud.blueflood.types.Locator}
     * from the specified column family and range.
     *
     * @param locator
     * @param columnFamily
     * @param range
     * @return
     */
    protected <T extends Object> Table<Locator, Long, T> getRollupsForLocator(final Locator locator,
                                                                              String columnFamily,
                                                                              Range range) {
        return getValuesForLocators( new ArrayList<Locator>() {{
            add( locator );
        }}, columnFamily, range );
    }

    /**
     * Fetch values for a list of {@link com.rackspacecloud.blueflood.types.Locator}
     * from the specified column family and range.
     *
     * This is a base behavior for most rollup types. IO subclasses can override
     * this behavior as they see fit.
     *
     * @param locators
     * @param columnFamily
     * @param range
     * @return
     */
    protected <T extends Object> Table<Locator, Long, T> getValuesForLocators( final List<Locator> locators,
                                                                               String columnFamily,
                                                                               Range range ) {

        Table<Locator, Long, T> locatorTimestampRollup = HashBasedTable.create();

        Map<Locator, List<ResultSetFuture>> resultSetFuturesMap = selectForLocatorListAndRange(columnFamily, locators, range);

        for (Map.Entry<Locator, List<ResultSetFuture>> entry : resultSetFuturesMap.entrySet() ) {
            Locator locator = entry.getKey();
            List<ResultSetFuture> futures = entry.getValue();

            Table<Locator, Long, T> result = toLocatorTimestampValue(futures, locator,
                                                    columnFamily,
                                                    range);
            locatorTimestampRollup.putAll(result);
        }
        return locatorTimestampRollup;
    }

    /**
     * Provides a way for the sub class to get a {@link java.nio.ByteBuffer}
     * representation of a certain value.
     *
     * @param value
     * @return
     */
    protected abstract <T extends Object> ByteBuffer toByteBuffer(T value);

    /**
     * Provides a way for the sub class to construct the right Rollup
     * object from a {@link java.nio.ByteBuffer}
     *
     * @param byteBuffer
     * @return
     */
    protected abstract <T extends Object> T fromByteBuffer(ByteBuffer byteBuffer);

    /**
     * Add a {@link com.datastax.driver.core.PreparedStatement} statement to the
     * {@link com.datastax.driver.core.BatchStatement} to insert this Rollup
     * object to metrics_preaggregated_{granularity} column family
     *
     * @param batch
     * @param locator
     * @param collectionTime
     * @param granularity
     * @param ttl
     */
    protected void addRollupToBatch(BatchStatement batch, Locator locator,
                                    Rollup rollup,
                                    long collectionTime,
                                    Granularity granularity,
                                    int ttl) {
        Statement statement = createStatement(locator, collectionTime, rollup, granularity, ttl);
        batch.add(statement);
    }

    /**
     * Asynchronously execute select statements against the specified
     * column family for a specific list of
     * {@link com.rackspacecloud.blueflood.types.Locator} and
     * {@link com.rackspacecloud.blueflood.types.Range}
     *
     * @param locators
     * @param range
     * @return a map of Locator -> a list of ResultSetFuture
     */
    protected Map<Locator, List<ResultSetFuture>> selectForLocatorListAndRange(String columnFamilyName,
                                                                               List<Locator> locators,
                                                                               Range range) {
        Map<Locator, List<ResultSetFuture>> locatorFuturesMap = new HashMap<Locator, List<ResultSetFuture>>();
        for (Locator locator : locators) {
            List<ResultSetFuture> existing = locatorFuturesMap.get(locator);
            if ( existing == null ) {
                existing = new ArrayList<ResultSetFuture>();
                locatorFuturesMap.put(locator, existing);
            }
            existing.addAll(selectForLocatorAndRange(columnFamilyName, locator, range));
        }
        return locatorFuturesMap;
    }

    /**
     * Execute a select statement against the specified column family for a specific
     * {@link com.rackspacecloud.blueflood.types.Locator} and
     * {@link com.rackspacecloud.blueflood.types.Range}
     *
     * @param columnFamily the column family name to do select against
     * @param locator
     * @param range

     * @return
     */
    protected List<ResultSetFuture> selectForLocatorAndRange(String columnFamily, Locator locator, Range range) {
        List<ResultSetFuture> resultsFutures = new ArrayList<ResultSetFuture>();
        PreparedStatement statement = metricsCFPreparedStatements.cfNameToSelectStatement.get(columnFamily);
        resultsFutures.add(
                session.executeAsync(statement.bind(locator.toString(), range.getStart(), range.getStop())));
        return resultsFutures;
    }

    /**
     *  Give a {@link com.datastax.driver.core.ResultSetFuture}, get
     *  the corresponding data from it and return it as a
     *  Table of locator, long and rollup.
     */
    public <T extends Object> Table<Locator, Long, T> toLocatorTimestampValue( List<ResultSetFuture> futures,
                                                                               Locator locator,
                                                                               String columnFamily,
                                                                               Range range) {
        Table<Locator, Long, T> locatorTimestampRollup = HashBasedTable.create();
        for ( ResultSetFuture future : futures ) {
            try {
                List<Row> rows = future.getUninterruptibly().all();

                // we only want to count the number of points we
                // get when we're querying the metrics_full
                if ( StringUtils.isNotEmpty(columnFamily) && columnFamily.equals(CassandraModel.CF_METRICS_FULL_NAME) ) {
                    Instrumentation.getRawPointsIn5MinHistogram().update(rows.size());
                }

                for (Row row : rows) {
                    String key = row.getString(DMetricsCFPreparedStatements.KEY);
                    Locator loc = Locator.createLocatorFromDbKey(key);
                    Long hash = row.getLong(DMetricsCFPreparedStatements.COLUMN1);
                    locatorTimestampRollup.put(loc, hash, (T)fromByteBuffer(row.getBytes(DMetricsCFPreparedStatements.VALUE)));
                }
            } catch (Exception ex) {
                Instrumentation.markReadError();
                LOG.error(String.format("error reading metric for locator %s, column family '%s', range %s",
                        locator, columnFamily, range.toString()), ex);
            }
        }
        return locatorTimestampRollup;
    }

    /**
     * Retrieves the {@link BoundStatement} for a particular metric and granularity.
     * Subclasses will implement this.
     *
     * @param metric
     * @param granularity
     * @return
     */
    protected abstract BoundStatement getBoundStatementForMetric(IMetric metric, Granularity granularity);
}