/**
 * Copyright 2011-2013 the original author or authors.
 *
 * 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 io.informant.local.store;

import io.informant.core.Trace;
import io.informant.core.TraceSink;
import io.informant.util.DaemonExecutors;
import io.informant.util.Singleton;

import java.io.IOException;
import java.util.Collection;
import java.util.Set;
import java.util.concurrent.ExecutorService;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import checkers.lock.quals.GuardedBy;

import com.google.common.base.Stopwatch;
import com.google.common.base.Ticker;
import com.google.common.collect.Sets;

/**
 * Implementation of TraceSink for local storage in embedded H2 database. Some day there may be
 * another implementation for remote storage (e.g. central monitoring system).
 * 
 * @author Trask Stalnaker
 * @since 0.5
 */
@Singleton
public class LocalTraceSink implements TraceSink {

    private static final Logger logger = LoggerFactory.getLogger(LocalTraceSink.class);

    private static final int PENDING_LIMIT = 100;

    private final ExecutorService executorService = DaemonExecutors
            .newSingleThreadExecutor("Informant-TraceSink");

    private final TraceSnapshotService traceSnapshotService;
    private final TraceSnapshotDao traceSnapshotDao;
    private final Ticker ticker;
    private final Set<Trace> pendingCompleteTraces = Sets.newCopyOnWriteArraySet();

    private final Object warningLock = new Object();
    @GuardedBy("warningLock")
    private final Stopwatch lastWarningStopwatch;
    @GuardedBy("warningLock")
    private int countSinceLastWarning;

    LocalTraceSink(TraceSnapshotService traceSnapshotService,
            TraceSnapshotDao traceSnapshotDao, Ticker ticker) {
        this.traceSnapshotService = traceSnapshotService;
        this.traceSnapshotDao = traceSnapshotDao;
        this.ticker = ticker;
        lastWarningStopwatch = new Stopwatch(ticker);
    }

    public void onCompletedTrace(final Trace trace) {
        if (traceSnapshotService.shouldStore(trace)) {
            // promote thread local trace metrics since they will be reset after this method returns
            trace.promoteTraceMetrics();
            if (pendingCompleteTraces.size() >= PENDING_LIMIT) {
                logPendingLimitWarning();
            }
            pendingCompleteTraces.add(trace);
            executorService.execute(new Runnable() {
                public void run() {
                    try {
                        traceSnapshotDao.storeSnapshot(TraceWriter.toTraceSnapshot(trace,
                                Long.MAX_VALUE, false));
                    } catch (IOException e) {
                        logger.error(e.getMessage(), e);
                    }
                    pendingCompleteTraces.remove(trace);
                }
            });
        }
    }

    private void logPendingLimitWarning() {
        // synchronized to prevent two threads from logging the warning at the same time (one should
        // log it and the other should increment the count for the next log)
        synchronized (warningLock) {
            // lastWarningStopwatch is not running the very first time
            if (!lastWarningStopwatch.isRunning() || lastWarningStopwatch.elapsedMillis() > 60000) {
                logger.warn("not storing a trace in the local h2 database because of an excessive"
                        + " backlog of {} traces already waiting to be stored (this warning will"
                        + " appear at most once a minute, there were {} additional traces not"
                        + " stored in the local h2 database since the last warning)",
                        PENDING_LIMIT,
                        countSinceLastWarning);
                lastWarningStopwatch.reset().start();
                countSinceLastWarning = 0;
            } else {
                countSinceLastWarning++;
            }
        }
    }

    // no need to throttle stuck trace storage since throttling is handled upstream by using a
    // single thread executor in StuckTraceCollector
    public void onStuckTrace(Trace trace) {
        try {
            TraceSnapshot snaphsot = TraceWriter.toTraceSnapshot(trace, ticker.read(), false);
            if (!trace.isCompleted()) {
                traceSnapshotDao.storeSnapshot(snaphsot);
            }
        } catch (IOException e) {
            logger.error(e.getMessage(), e);
        }
    }

    public Collection<Trace> getPendingCompleteTraces() {
        return pendingCompleteTraces;
    }

    void close() {
        logger.debug("close()");
        executorService.shutdownNow();
    }
}

@@ -641,24 +641,82 @@
 cal.store;%0A%0A
+import static java.util.concurrent.TimeUnit.MILLISECONDS;%0A
 import io.in
@@ -735,16 +735,16 @@
 .Trace;%0A
-
 import i
@@ -1076,49 +1076,8 @@
 y;%0A%0A
-import com.google.common.base.Stopwatch;%0A
 impo
@@ -1148,16 +1148,70 @@
 ct.Sets;
+%0Aimport com.google.common.util.concurrent.RateLimiter;
 %0A%0A/**%0A *
@@ -2019,121 +2019,69 @@
 nal 
-Object warningLock = new Object();%0A    @GuardedBy(%22warningLock%22)%0A    private final Stopwatch lastWarningStopwatch
+RateLimiter warningRateLimiter = RateLimiter.create(1.0 / 60)
 ;%0A  
@@ -2416,62 +2416,8 @@
 er;%0A
-        lastWarningStopwatch = new Stopwatch(ticker);%0A
     
@@ -2773,24 +2773,48 @@
 tWarning();%0A
+                return;%0A
             
@@ -3407,19 +3407,16 @@
 %0A       
- //
  synchro
@@ -3425,367 +3425,93 @@
 zed 
-to prevent two threads from logging the 
+(
 warning
- at the same time (one should%0A        // log it and the other should increment the count for the next log)%0A        synchronized (warningLock) %7B%0A            // lastWarningStopwatch is not running the very first time%0A            if (!lastWarningStopwatch.isRunning() %7C%7C lastWarningStopwatch.elapsedMillis() %3E 60000
+RateLimiter) %7B%0A            if (warningRateLimiter.tryAcquire(0, MILLISECONDS)
 ) %7B%0A
@@ -3930,108 +3930,30 @@
 MIT,
-%0A                        countSinceLastWarning);%0A                lastWarningStopwatch.reset().start(
+ countSinceLastWarning
 );%0A 

Use guava RateLimiter