content: /**
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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
   * <p>
   * http://www.apache.org/licenses/LICENSE-2.0
   * <p>
   * 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 org.apache.atlas.repository.graph;
  
  import com.google.common.annotations.VisibleForTesting;
  import org.apache.atlas.ApplicationProperties;
  import org.apache.atlas.AtlasConfiguration;
  import org.apache.atlas.AtlasException;
  import org.apache.atlas.ha.HAConfiguration;
  import org.apache.atlas.listener.ActiveStateChangeHandler;
  import org.apache.atlas.repository.graphdb.AtlasGraph;
  import org.apache.atlas.repository.graphdb.AtlasGraphManagement;
  import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
  import org.apache.atlas.repository.graphdb.AtlasVertex;
  import org.apache.atlas.service.Service;
  import org.apache.commons.configuration.Configuration;
  import org.apache.commons.lang.StringUtils;
  import org.apache.commons.lang.math.NumberUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.springframework.core.annotation.Order;
  import org.springframework.stereotype.Component;
  
  import javax.inject.Inject;
  
  import java.text.SimpleDateFormat;
  import java.time.Instant;
  import java.time.temporal.ChronoUnit;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.Map;
  import java.util.TimeZone;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import static org.apache.atlas.ApplicationProperties.DEFAULT_INDEX_RECOVERY;
  import static org.apache.atlas.repository.Constants.PROPERTY_KEY_INDEX_RECOVERY_CUSTOM_TIME;
  import static org.apache.atlas.repository.Constants.PROPERTY_KEY_INDEX_RECOVERY_NAME;
  import static org.apache.atlas.repository.Constants.PROPERTY_KEY_INDEX_RECOVERY_PREV_TIME;
  import static org.apache.atlas.repository.Constants.PROPERTY_KEY_INDEX_RECOVERY_START_TIME;
  import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.setEncodedProperty;
  
  @Component
  @Order(8)
  public class IndexRecoveryService implements Service, ActiveStateChangeHandler {
      private static final Logger LOG = LoggerFactory.getLogger(IndexRecoveryService.class);
  
      private static final String DATE_FORMAT                               = "yyyy-MM-dd'T'HH:mm:ss.SSS'Z'";
      private static final String INDEX_HEALTH_MONITOR_THREAD_NAME          = "index-health-monitor";
      private static final String SOLR_STATUS_CHECK_RETRY_INTERVAL          = "atlas.graph.index.status.check.frequency";
      private static final String SOLR_INDEX_RECOVERY_CONFIGURED_START_TIME = "atlas.index.recovery.start.time";
      private static final long   SOLR_STATUS_RETRY_DEFAULT_MS              = 30000; // 30 secs default
  
      public final  RecoveryInfoManagement recoveryInfoManagement;
      public        RecoveryThread         recoveryThread;
  
      private final Thread        indexHealthMonitor;
      private final Configuration configuration;
      private final boolean       isIndexRecoveryEnabled;
  
      @Inject
      public IndexRecoveryService(Configuration config, AtlasGraph graph) {
          this.configuration          = config;
          this.isIndexRecoveryEnabled = config.getBoolean(ApplicationProperties.INDEX_RECOVERY_CONF, DEFAULT_INDEX_RECOVERY);
  
          long recoveryStartTimeFromConfig = getRecoveryStartTimeFromConfig(config);
          long healthCheckFrequencyMillis  = config.getLong(SOLR_STATUS_CHECK_RETRY_INTERVAL, SOLR_STATUS_RETRY_DEFAULT_MS);
  
          this.recoveryInfoManagement = new RecoveryInfoManagement(graph);
          this.recoveryThread         = new RecoveryThread(recoveryInfoManagement, graph, recoveryStartTimeFromConfig, healthCheckFrequencyMillis);
          this.indexHealthMonitor     = new Thread(recoveryThread, INDEX_HEALTH_MONITOR_THREAD_NAME);
      }
  
      @Override
      public void start() throws AtlasException {
          if (configuration == null || !HAConfiguration.isHAEnabled(configuration)) {
              LOG.info("==> IndexRecoveryService.start()");
  
              startTxLogMonitoring();
  
              LOG.info("<== IndexRecoveryService.start()");
          }
      }
  
      @Override
      public void stop() throws AtlasException {
          try {
              recoveryThread.shutdown();
  
              indexHealthMonitor.join();
          } catch (InterruptedException e) {
              LOG.error("indexHealthMonitor: Interrupted", e);
          }
      }
  
      @Override
      public void instanceIsActive() throws AtlasException {
          LOG.info("==> IndexRecoveryService.instanceIsActive()");
  
          startTxLogMonitoring();
  
          LOG.info("<== IndexRecoveryService.instanceIsActive()");
      }
  
      @Override
      public void instanceIsPassive() throws AtlasException {
          LOG.info("==> IndexRecoveryService.instanceIsPassive()");
  
          stop();
  
          LOG.info("<== IndexRecoveryService.instanceIsPassive()");
      }
  
      @Override
      public int getHandlerOrder() {
          return ActiveStateChangeHandler.HandlerOrder.INDEX_RECOVERY.getOrder();
      }
  
      private long getRecoveryStartTimeFromConfig(Configuration config) {
          long ret = 0L;
  
          try {
              String time = config.getString(SOLR_INDEX_RECOVERY_CONFIGURED_START_TIME);
              if (StringUtils.isNotBlank(time)) {
                  SimpleDateFormat dateFormat = new SimpleDateFormat(DATE_FORMAT);
                  dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
  
                  ret = dateFormat.parse(time).toInstant().toEpochMilli();
              }
          } catch (Exception e) {
              LOG.error("Error fetching: {}", SOLR_INDEX_RECOVERY_CONFIGURED_START_TIME, e);
          }
  
          return ret;
      }
  
      private void startTxLogMonitoring() {
          if (!isIndexRecoveryEnabled) {
              LOG.warn("IndexRecoveryService: Recovery should be enabled.");
  
              return;
          }
  
          try {
              if (indexHealthMonitor.getState() == Thread.State.NEW) {
                  indexHealthMonitor.start();
              }
          } catch (Exception ex) {
              LOG.error("Error while starting Index Health Monitor", ex);
          }
      }
  
      public static class RecoveryThread implements Runnable {
          private final AtlasGraph             graph;
          private final RecoveryInfoManagement recoveryInfoManagement;
          private final AtomicBoolean          shouldRun = new AtomicBoolean(false);
          private final long                   indexStatusCheckRetryMillis;
          private       Object                 txRecoveryObject;
  
          private RecoveryThread(RecoveryInfoManagement recoveryInfoManagement, AtlasGraph graph, long startTimeFromConfig, long healthCheckFrequencyMillis) {
              this.graph                       = graph;
              this.recoveryInfoManagement      = recoveryInfoManagement;
              this.indexStatusCheckRetryMillis = healthCheckFrequencyMillis;
  
              if (startTimeFromConfig > 0) {
                  this.recoveryInfoManagement.updateStartTime(startTimeFromConfig);
              }
          }
  
          public void run() {
              shouldRun.set(true);
  
              LOG.info("Index Health Monitor: Starting...");
  
              while (shouldRun.get()) {
                  try {
                      boolean isIdxHealthy = waitAndCheckIfIndexBackendHealthy();
  
                      if (this.txRecoveryObject == null && isIdxHealthy) {
                          startMonitoring();
                      }
  
                      if (this.txRecoveryObject != null && !isIdxHealthy) {
                          stopMonitoring();
                      }
                  } catch (Exception e) {
                      LOG.error("Error: Index recovery monitoring!", e);
                  }
              }
          }
  
          public void shutdown() {
              try {
                  LOG.info("Index Health Monitor: Shutdown: Starting...");
  
                  // handle the case where thread was not started at all
                  // and shutdown called
                  if (!shouldRun.get()) {
                      return;
                  }
  
                  shouldRun.set(false);
              } finally {
                  LOG.info("Index Health Monitor: Shutdown: Done!");
              }
          }
  
          public boolean isIndexBackendHealthy() throws AtlasException {
              return this.graph.getGraphIndexClient().isHealthy();
          }
  
          public void startMonitoringByUserRequest(Long startTime) {
              startMonitoring(startTime);
          }
  
          public void stopMonitoringByUserRequest() {
              stopIndexRecovery();
  
              LOG.info("Index Recovery: Stopped!");
          }
  
          private boolean waitAndCheckIfIndexBackendHealthy() throws AtlasException, InterruptedException {
              Thread.sleep(indexStatusCheckRetryMillis);
  
              return isIndexBackendHealthy();
          }
  
          private void startMonitoring() {
              startMonitoring(recoveryInfoManagement.getStartTime());
          }
  
          private void startMonitoring(Long startTime) {
              if (startTime == null || startTime == 0L) {
                  LOG.error("Index Recovery requested without start time");
  
                  return;
              }
  
              try (AtlasGraphManagement management = this.graph.getManagementSystem()) {
                  txRecoveryObject = management.startIndexRecovery(startTime);
  
                  printIndexRecoveryStats();
  
                  management.setIsSuccess(true);
  
                  LOG.info("Index Recovery: Started! Recovery time: {}", Instant.ofEpochMilli(startTime));
              } catch (Exception e) {
                  LOG.error("Index Recovery with recovery time: {} failed", Instant.ofEpochMilli(startTime), e);
              }
          }
  
          private void stopMonitoring() {
              stopIndexRecoveryAndUpdateStartTime();
          }
  
          private void stopIndexRecoveryAndUpdateStartTime() {
              Instant newStartTime = Instant.now().minusMillis(2 * indexStatusCheckRetryMillis);
  
              stopIndexRecovery();
  
              recoveryInfoManagement.updateStartTime(newStartTime.toEpochMilli());
  
              LOG.info("Index Recovery: Stopped! Recovery time: {}", newStartTime);
          }
  
          private void stopIndexRecovery() {
              try (AtlasGraphManagement management = this.graph.getManagementSystem()) {
                  management.stopIndexRecovery(txRecoveryObject);
  
                  printIndexRecoveryStats();
  
                  management.setIsSuccess(true);
              } catch (Exception e) {
                  LOG.info("Index Recovery: Stopped! Error!", e);
              } finally {
                  this.txRecoveryObject = null;
              }
          }
  
          private void printIndexRecoveryStats() {
              try (AtlasGraphManagement management = this.graph.getManagementSystem()) {
                  management.printIndexRecoveryStats(txRecoveryObject);
                  management.setIsSuccess(true);
              } catch (Exception e) {
                  LOG.error("Index Recovery: printIndexRecoveryStats() failed!", e);
              }
          }
      }
  
      @VisibleForTesting
      public static class RecoveryInfoManagement {
          private static final String INDEX_RECOVERY_TYPE_NAME = "__solrIndexRecoveryInfo";
  
          private final AtlasGraph graph;
  
          public RecoveryInfoManagement(AtlasGraph graph) {
              this.graph = graph;
          }
  
          public void updateStartTime(long time) {
              updateIndexRecoveryTime(PROPERTY_KEY_INDEX_RECOVERY_START_TIME, time);
          }
  
          public void updateCustomStartTime(long time) {
              updateIndexRecoveryTime(PROPERTY_KEY_INDEX_RECOVERY_CUSTOM_TIME, time);
          }
  
          public void updateIndexRecoveryTime(String timePropertyKey, long time) {
              Map<String, String> indexRecoveryData = new HashMap<>();
  
              indexRecoveryData.put(timePropertyKey, String.valueOf(time));
  
              updateIndexRecoveryData(indexRecoveryData);
          }
  
          public void updateIndexRecoveryData(Map<String, String> indexRecoveryData) {
              try {
                  Long        startTime          = NumberUtils.createLong(indexRecoveryData.get(PROPERTY_KEY_INDEX_RECOVERY_START_TIME));
                  Long        prevStartTime      = NumberUtils.createLong(indexRecoveryData.get(PROPERTY_KEY_INDEX_RECOVERY_PREV_TIME));
                  Long        customStartTime    = NumberUtils.createLong(indexRecoveryData.get(PROPERTY_KEY_INDEX_RECOVERY_CUSTOM_TIME));
                  boolean     isStartTimeUpdated = startTime != null;
                  AtlasVertex vertex             = findVertex();
  
                  if (vertex == null) {
                      vertex = graph.addVertex();
  
                      setEncodedProperty(vertex, PROPERTY_KEY_INDEX_RECOVERY_NAME, INDEX_RECOVERY_TYPE_NAME);
                  } else {
                      prevStartTime = isStartTimeUpdated ? getStartTime(vertex) : prevStartTime;
                  }
  
                  if (startTime != null) {
                      setEncodedProperty(vertex, PROPERTY_KEY_INDEX_RECOVERY_START_TIME, startTime);
                  }
  
                  if (prevStartTime != null) {
                      setEncodedProperty(vertex, PROPERTY_KEY_INDEX_RECOVERY_PREV_TIME, prevStartTime);
                  }
  
                  if (customStartTime != null) {
                      setEncodedProperty(vertex, PROPERTY_KEY_INDEX_RECOVERY_CUSTOM_TIME, customStartTime);
                  }
              } catch (Exception ex) {
                  LOG.error("Error updating index recovery data", ex);
              } finally {
                  graph.commit();
              }
          }
  
          public Long getStartTime() {
              AtlasVertex vertex = findVertex();
  
              return getStartTime(vertex);
          }
  
          public AtlasVertex findVertex() {
              AtlasGraphQuery       query   = graph.query().has(PROPERTY_KEY_INDEX_RECOVERY_NAME, INDEX_RECOVERY_TYPE_NAME);
              Iterator<AtlasVertex> results = query.vertices().iterator();
  
              return results.hasNext() ? results.next() : null;
          }
  
          private Long getStartTime(AtlasVertex vertex) {
              Long defaultStartTime = getStartTimeByTxLogTTL();
  
              if (vertex == null) {
                  LOG.warn("Vertex passed is NULL: Returned is startTime by TTL {}", Instant.ofEpochMilli(defaultStartTime));
  
                  return defaultStartTime;
              }
  
              Long startTime = 0L;
  
              try {
                  startTime = vertex.getProperty(PROPERTY_KEY_INDEX_RECOVERY_START_TIME, Long.class);
              } catch (Exception e) {
                  LOG.error("Error retrieving startTime", e);
              }
  
              return startTime == null || startTime == 0L ? defaultStartTime : startTime;
          }
  
          private Long getStartTimeByTxLogTTL() {
              long ttl = AtlasConfiguration.SOLR_INDEX_TX_LOG_TTL_CONF.getLong();
              return Instant.now().minus(ttl, ChronoUnit.HOURS).toEpochMilli();
          }
      }
  }
  