<commit_before>/*
 * Copyright © 2015 <NAME>, 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 co.cask.cdap.internal.app.runtime.schedule;

import co.cask.cdap.api.metrics.MetricDataQuery;
import co.cask.cdap.api.metrics.MetricStore;
import co.cask.cdap.api.metrics.MetricTimeSeries;
import co.cask.cdap.api.metrics.MetricType;
import co.cask.cdap.api.metrics.TimeValue;
import co.cask.cdap.api.schedule.SchedulableProgramType;
import co.cask.cdap.api.schedule.Schedule;
import co.cask.cdap.app.runtime.ProgramRuntimeService;
import co.cask.cdap.app.store.Store;
import co.cask.cdap.app.store.StoreFactory;
import co.cask.cdap.common.conf.CConfiguration;
import co.cask.cdap.common.conf.Constants;
import co.cask.cdap.common.exception.NotFoundException;
import co.cask.cdap.common.stream.notification.StreamSizeNotification;
import co.cask.cdap.config.PreferencesStore;
import co.cask.cdap.internal.app.runtime.BasicArguments;
import co.cask.cdap.internal.app.runtime.ProgramOptionConstants;
import co.cask.cdap.internal.app.runtime.schedule.store.DatasetBasedStreamSizeScheduleStore;
import co.cask.cdap.internal.schedule.StreamSizeSchedule;
import co.cask.cdap.notifications.service.NotificationContext;
import co.cask.cdap.notifications.service.NotificationHandler;
import co.cask.cdap.notifications.service.NotificationService;
import co.cask.cdap.proto.Id;
import co.cask.cdap.proto.ProgramType;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.AbstractScheduledService;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import org.apache.twill.common.Cancellable;
import org.apache.twill.common.Threads;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.lang.reflect.Type;
import java.util.Collection;
import java.util.List;
import java.util.NavigableSet;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;

/**
 * {@link Scheduler} that triggers program executions based on data availability in streams.
 */
@Singleton
public class StreamSizeScheduler implements Scheduler {
  private static final Logger LOG = LoggerFactory.getLogger(StreamSizeScheduler.class);
  private static final int STREAM_POLLING_THREAD_POOL_SIZE = 10;
  private static final int POLLING_AFTER_NOTIFICATION_RETRY = 3;

  private final long pollingDelay;
  private final NotificationService notificationService;
  private final MetricStore metricStore;
  private final StoreFactory storeFactory;
  private final ProgramRuntimeService programRuntimeService;
  private final PreferencesStore preferencesStore;
  private final DatasetBasedStreamSizeScheduleStore scheduleStore;
  private final ConcurrentMap<Id.Stream, StreamSubscriber> streamSubscribers;

  // Key is scheduleId
  private final ConcurrentSkipListMap<String, StreamSubscriber> scheduleSubscribers;

  private Store store;
  private Executor sendPollingInfoExecutor;

  // Used to schedule polling of a stream only after a certain time - the time after which the metrics are updated
  private ScheduledExecutorService pollBookingExecutor;

  // Scheduled executor used to poll stream at regular intervals, by querying the metric system
  private ScheduledExecutorService streamPollingExecutor;

  private ListeningExecutorService taskExecutorService;

  @Inject
  public StreamSizeScheduler(CConfiguration cConf, NotificationService notificationService, MetricStore metricStore,
                             StoreFactory storeFactory, ProgramRuntimeService programRuntimeService,
                             PreferencesStore preferencesStore, DatasetBasedStreamSizeScheduleStore scheduleStore) {
    this.pollingDelay = TimeUnit.SECONDS.toMillis(
      cConf.getLong(Constants.Notification.Stream.STREAM_SIZE_SCHEDULE_POLLING_DELAY));
    this.notificationService = notificationService;
    this.metricStore = metricStore;
    this.storeFactory = storeFactory;
    this.programRuntimeService = programRuntimeService;
    this.preferencesStore = preferencesStore;
    this.scheduleStore = scheduleStore;
    this.streamSubscribers = Maps.newConcurrentMap();
    this.scheduleSubscribers = new ConcurrentSkipListMap<String, StreamSubscriber>();
    this.store = null;
  }

  public void start() throws SchedulerException {
    sendPollingInfoExecutor = Executors.newCachedThreadPool(
      Threads.createDaemonThreadFactory("stream-size-scheduler-%d"));
    streamPollingExecutor = Executors.newScheduledThreadPool(STREAM_POLLING_THREAD_POOL_SIZE,
                                                             Threads.createDaemonThreadFactory("stream-polling-%d"));
    pollBookingExecutor = Executors.newSingleThreadScheduledExecutor(
      Threads.createDaemonThreadFactory("polling-booking-executor"));
    taskExecutorService = MoreExecutors.listeningDecorator(
      Executors.newCachedThreadPool(Threads.createDaemonThreadFactory("stream-schedule-task")));

    store = storeFactory.create();

    initializeScheduleStore();
  }

  /**
   * Initialize the stream size schedule store, and restart all the schedule tasks that were persisted in it.
   *
   * @throws SchedulerException if the persistent store could not be initialized, or if we couldn't list the
   *                            persisted tasks
   */
  private void initializeScheduleStore() throws SchedulerException {
    try {
      scheduleStore.initialize();
    } catch (Throwable t) {
      throw new SchedulerException("Error when initializing store for stream size schedules", t);
    }

    // Reschedule the persisted schedules
    List<StreamSizeScheduleState> scheduleStates;
    try {
      scheduleStates = scheduleStore.list();
    } catch (Throwable t) {
      throw new SchedulerException("Error when listing stream size schedules from store", t);
    }

    for (StreamSizeScheduleState scheduleState : scheduleStates) {
      try {
        restoreScheduleFromStore(scheduleState.getProgramId(), scheduleState.getProgramType(),
                                 scheduleState.getStreamSizeSchedule(), scheduleState.isRunning(),
                                 scheduleState.getBaseRunSize(), scheduleState.getBaseRunTs(),
                                 scheduleState.getLastRunSize(), scheduleState.getLastRunTs());
      } catch (SchedulerException e) {
        // We should never enter this, but if we do, we still want the other schedule tasks to be handled
        LOG.error("Could not schedule task '{}' from persistent store", scheduleState, e);
      }
    }

    // Poll all the Streams for active tasks
    for (StreamSubscriber streamSubscriber : streamSubscribers.values()) {
      if (streamSubscriber.getActiveTasksCount() <= 0) {
        continue;
      }

      try {
        StreamSize streamSize = streamSubscriber.pollOnce();
        streamSubscriber.sendPollingInfoToActiveTasks(streamSize);
      } catch (IOException e) {
        // Failing to poll should not make this init fail
        LOG.warn("Could not poll size for stream '{}'", streamSubscriber.getStreamId(), e);
      }
    }
  }

  public void stop() {
    for (StreamSubscriber subscriber : streamSubscribers.values()) {
      subscriber.stopAndWait();
    }
    pollBookingExecutor.shutdownNow();
    streamPollingExecutor.shutdownNow();
  }

  @Override
  public void schedule(Id.Program program, SchedulableProgramType programType, Schedule schedule)
    throws SchedulerException {<commit_msg>Stop task executor service<commit_after>/*
 * Copyright © 2015 <NAME>, 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 co.cask.cdap.internal.app.runtime.schedule;

import co.cask.cdap.api.metrics.MetricDataQuery;
import co.cask.cdap.api.metrics.MetricStore;
import co.cask.cdap.api.metrics.MetricTimeSeries;
import co.cask.cdap.api.metrics.MetricType;
import co.cask.cdap.api.metrics.TimeValue;
import co.cask.cdap.api.schedule.SchedulableProgramType;
import co.cask.cdap.api.schedule.Schedule;
import co.cask.cdap.app.runtime.ProgramRuntimeService;
import co.cask.cdap.app.store.Store;
import co.cask.cdap.app.store.StoreFactory;
import co.cask.cdap.common.conf.CConfiguration;
import co.cask.cdap.common.conf.Constants;
import co.cask.cdap.common.exception.NotFoundException;
import co.cask.cdap.common.stream.notification.StreamSizeNotification;
import co.cask.cdap.config.PreferencesStore;
import co.cask.cdap.internal.app.runtime.BasicArguments;
import co.cask.cdap.internal.app.runtime.ProgramOptionConstants;
import co.cask.cdap.internal.app.runtime.schedule.store.DatasetBasedStreamSizeScheduleStore;
import co.cask.cdap.internal.schedule.StreamSizeSchedule;
import co.cask.cdap.notifications.service.NotificationContext;
import co.cask.cdap.notifications.service.NotificationHandler;
import co.cask.cdap.notifications.service.NotificationService;
import co.cask.cdap.proto.Id;
import co.cask.cdap.proto.ProgramType;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.AbstractScheduledService;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import org.apache.twill.common.Cancellable;
import org.apache.twill.common.Threads;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.lang.reflect.Type;
import java.util.Collection;
import java.util.List;
import java.util.NavigableSet;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;

/**
 * {@link Scheduler} that triggers program executions based on data availability in streams.
 */
@Singleton
public class StreamSizeScheduler implements Scheduler {
  private static final Logger LOG = LoggerFactory.getLogger(StreamSizeScheduler.class);
  private static final int STREAM_POLLING_THREAD_POOL_SIZE = 10;
  private static final int POLLING_AFTER_NOTIFICATION_RETRY = 3;

  private final long pollingDelay;
  private final NotificationService notificationService;
  private final MetricStore metricStore;
  private final StoreFactory storeFactory;
  private final ProgramRuntimeService programRuntimeService;
  private final PreferencesStore preferencesStore;
  private final DatasetBasedStreamSizeScheduleStore scheduleStore;
  private final ConcurrentMap<Id.Stream, StreamSubscriber> streamSubscribers;

  // Key is scheduleId
  private final ConcurrentSkipListMap<String, StreamSubscriber> scheduleSubscribers;

  private Store store;
  private Executor sendPollingInfoExecutor;

  // Used to schedule polling of a stream only after a certain time - the time after which the metrics are updated
  private ScheduledExecutorService pollBookingExecutor;

  // Scheduled executor used to poll stream at regular intervals, by querying the metric system
  private ScheduledExecutorService streamPollingExecutor;

  private ListeningExecutorService taskExecutorService;

  @Inject
  public StreamSizeScheduler(CConfiguration cConf, NotificationService notificationService, MetricStore metricStore,
                             StoreFactory storeFactory, ProgramRuntimeService programRuntimeService,
                             PreferencesStore preferencesStore, DatasetBasedStreamSizeScheduleStore scheduleStore) {
    this.pollingDelay = TimeUnit.SECONDS.toMillis(
      cConf.getLong(Constants.Notification.Stream.STREAM_SIZE_SCHEDULE_POLLING_DELAY));
    this.notificationService = notificationService;
    this.metricStore = metricStore;
    this.storeFactory = storeFactory;
    this.programRuntimeService = programRuntimeService;
    this.preferencesStore = preferencesStore;
    this.scheduleStore = scheduleStore;
    this.streamSubscribers = Maps.newConcurrentMap();
    this.scheduleSubscribers = new ConcurrentSkipListMap<String, StreamSubscriber>();
    this.store = null;
  }

  public void start() throws SchedulerException {
    sendPollingInfoExecutor = Executors.newCachedThreadPool(
      Threads.createDaemonThreadFactory("stream-size-scheduler-%d"));
    streamPollingExecutor = Executors.newScheduledThreadPool(STREAM_POLLING_THREAD_POOL_SIZE,
                                                             Threads.createDaemonThreadFactory("stream-polling-%d"));
    pollBookingExecutor = Executors.newSingleThreadScheduledExecutor(
      Threads.createDaemonThreadFactory("polling-booking-executor"));
    taskExecutorService = MoreExecutors.listeningDecorator(
      Executors.newCachedThreadPool(Threads.createDaemonThreadFactory("stream-schedule-task")));

    store = storeFactory.create();

    initializeScheduleStore();
  }

  /**
   * Initialize the stream size schedule store, and restart all the schedule tasks that were persisted in it.
   *
   * @throws SchedulerException if the persistent store could not be initialized, or if we couldn't list the
   *                            persisted tasks
   */
  private void initializeScheduleStore() throws SchedulerException {
    try {
      scheduleStore.initialize();
    } catch (Throwable t) {
      throw new SchedulerException("Error when initializing store for stream size schedules", t);
    }

    // Reschedule the persisted schedules
    List<StreamSizeScheduleState> scheduleStates;
    try {
      scheduleStates = scheduleStore.list();
    } catch (Throwable t) {
      throw new SchedulerException("Error when listing stream size schedules from store", t);
    }

    for (StreamSizeScheduleState scheduleState : scheduleStates) {
      try {
        restoreScheduleFromStore(scheduleState.getProgramId(), scheduleState.getProgramType(),
                                 scheduleState.getStreamSizeSchedule(), scheduleState.isRunning(),
                                 scheduleState.getBaseRunSize(), scheduleState.getBaseRunTs(),
                                 scheduleState.getLastRunSize(), scheduleState.getLastRunTs());
      } catch (SchedulerException e) {
        // We should never enter this, but if we do, we still want the other schedule tasks to be handled
        LOG.error("Could not schedule task '{}' from persistent store", scheduleState, e);
      }
    }

    // Poll all the Streams for active tasks
    for (StreamSubscriber streamSubscriber : streamSubscribers.values()) {
      if (streamSubscriber.getActiveTasksCount() <= 0) {
        continue;
      }

      try {
        StreamSize streamSize = streamSubscriber.pollOnce();
        streamSubscriber.sendPollingInfoToActiveTasks(streamSize);
      } catch (IOException e) {
        // Failing to poll should not make this init fail
        LOG.warn("Could not poll size for stream '{}'", streamSubscriber.getStreamId(), e);
      }
    }
  }

  public void stop() {
    for (StreamSubscriber subscriber : streamSubscribers.values()) {
      subscriber.stopAndWait();
    }
    if (pollBookingExecutor != null) {
      pollBookingExecutor.shutdownNow();
    }
    if (streamPollingExecutor != null) {
      streamPollingExecutor.shutdownNow();
    }
    if (taskExecutorService != null) {
      taskExecutorService.shutdownNow();
    }
  }

  @Override
  public void schedule(Id.Program program, SchedulableProgramType programType, Schedule schedule)
    throws SchedulerException {
{'fraction_non_alphanumeric': 0.0819230107280133, 'fraction_numerical': 0.0016063335436865354, 'mean_word_length': 4.241130487071557, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4945533', 'n_tokens_mistral': 4690, 'n_tokens_neox': 4604, 'n_words': 1135}