package com.metamx.druid.realtime;

import com.google.common.base.Throwables;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.metamx.common.concurrent.ScheduledExecutors;
import com.metamx.common.logger.Logger;
import com.metamx.druid.index.v1.IndexGranularity;
import com.metamx.druid.input.InputRow;
import com.metamx.druid.realtime.plumber.IntervalRejectionPolicyFactory;
import com.metamx.druid.realtime.plumber.RejectionPolicy;
import org.joda.time.DateTime;
import org.joda.time.Duration;
import org.joda.time.Interval;
import org.joda.time.Period;

import java.io.IOException;
import java.util.concurrent.Callable;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

/**
 */
public class GracefulShutdownFirehose implements Firehose
{
  private static final Logger log = new Logger(GracefulShutdownFirehose.class);

  private final Firehose firehose;
  private final IndexGranularity segmentGranularity;
  private final long windowMillis;
  private final ScheduledExecutorService scheduledExecutor;
  private final RejectionPolicy rejectionPolicy;

  private volatile boolean shutdown = false;

  public GracefulShutdownFirehose(
      Firehose firehose,
      IndexGranularity segmentGranularity,
      Period windowPeriod
  )
  {
    this.firehose = firehose;
    this.segmentGranularity = segmentGranularity;
    this.windowMillis = windowPeriod.toStandardDuration().getMillis() * 2;
    this.scheduledExecutor = Executors.newScheduledThreadPool(
        1,
        new ThreadFactoryBuilder()
            .setDaemon(true)
            .setNameFormat("firehose_scheduled_%d")
            .build()
    );

    final long truncatedNow = segmentGranularity.truncate(new DateTime()).getMillis();
    final long end = segmentGranularity.increment(truncatedNow);

    this.rejectionPolicy = new IntervalRejectionPolicyFactory(new Interval(truncatedNow, end)).create(windowPeriod);
  }

  public void shutdown() throws IOException
  {
    final long truncatedNow = segmentGranularity.truncate(new DateTime()).getMillis();
    final long end = segmentGranularity.increment(truncatedNow) + windowMillis;
    final Duration timeUntilShutdown = new Duration(System.currentTimeMillis(), end);

    log.info("Shutting down in %s", timeUntilShutdown);

    ScheduledExecutors.scheduleWithFixedDelay(
        scheduledExecutor,
        timeUntilShutdown,
        new Callable<ScheduledExecutors.Signal>()
        {
          @Override
          public ScheduledExecutors.Signal call() throws Exception
          {
            try {
              firehose.close();
            }
            catch (Exception e) {
              throw Throwables.propagate(e);
            }

            return ScheduledExecutors.Signal.STOP;
          }
        }
    );

    shutdown = true;
  }

  @Override
  public boolean hasMore()
  {
    return firehose.hasMore();
  }

  @Override
  public InputRow nextRow()
  {
    InputRow next = firehose.nextRow();

    if (!shutdown || rejectionPolicy.accept(next.getTimestampFromEpoch())) {
      return next;
    }

    return null;
  }

  @Override
  public Runnable commit()
  {
    return firehose.commit();
  }

  @Override
  public void close() throws IOException
  {
    firehose.close();
  }
}

@@ -452,16 +452,71 @@
 Policy;%0A
+import org.apache.commons.lang.mutable.MutableBoolean;%0A
 import o
@@ -1172,16 +1172,83 @@
 olicy;%0A%0A
+  private final MutableBoolean hasMore = new MutableBoolean(true);%0A
   privat
@@ -2718,39 +2718,46 @@
             
-firehose.clo
+hasMore.setValue(fal
 se
-(
 );%0A         
@@ -3007,32 +3007,58 @@
 )%0A  %7B%0A    return
+ hasMore.booleanValue() &&
  firehose.hasMor

fix shutdown problem