/*
    * 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
    *
   *     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 org.apache.giraph.zk;
  
  import org.apache.giraph.conf.DefaultImmutableClassesGiraphConfigurable;
  import org.apache.giraph.conf.GiraphConstants;
  import org.apache.giraph.utils.ThreadUtils;
  import org.apache.log4j.Logger;
  import org.apache.zookeeper.jmx.ManagedUtil;
  import org.apache.zookeeper.server.DatadirCleanupManager;
  import org.apache.zookeeper.server.ServerCnxnFactory;
  import org.apache.zookeeper.server.ZooKeeperServer;
  import org.apache.zookeeper.server.persistence.FileTxnSnapLog;
  import org.apache.zookeeper.server.quorum.QuorumPeerMain;
  
  import javax.management.JMException;
  import java.io.File;
  import java.io.IOException;
  
  /**
   * Zookeeper wrapper that starts zookeeper withing master process.
   */
  public class InProcessZooKeeperRunner
      extends DefaultImmutableClassesGiraphConfigurable
      implements ZooKeeperRunner {
  
    /** Class logger */
    private static final Logger LOG =
        Logger.getLogger(InProcessZooKeeperRunner.class);
    /**
     * Wrapper for zookeeper quorum.
     */
    private QuorumRunner quorumRunner = new QuorumRunner();
  
    @Override
    public int start(String zkDir, ZookeeperConfig config) throws IOException {
      return quorumRunner.start(config);
    }
  
    @Override
    public void stop() {
      try {
        quorumRunner.stop();
      } catch (InterruptedException e) {
        LOG.error("Unable to cleanly shutdown zookeeper", e);
      }
    }
  
    @Override
    public void cleanup() {
    }
  
    /**
     * Wrapper around zookeeper quorum. Does not necessarily
     * starts quorum, if there is only one server in config file
     * will only start zookeeper.
     */
    private static class QuorumRunner extends QuorumPeerMain {
  
      /**
       * ZooKeeper server wrapper.
       */
      private ZooKeeperServerRunner serverRunner;
  
      /**
       * Starts quorum and/or zookeeper service.
       * @param config quorum and zookeeper configuration
       * @return zookeeper port
       * @throws IOException if can't start zookeeper
       */
      public int start(ZookeeperConfig config) throws IOException {
        serverRunner = new ZooKeeperServerRunner();
        //Make sure zookeeper starts first and purge manager last
        //This is important because zookeeper creates a folder
        //strucutre on the local disk. Purge manager also tries
        //to create it but from a different thread and can run into
        //race condition. See FileTxnSnapLog source code for details.
        int port = serverRunner.start(config);
        // Start and schedule the the purge task
        DatadirCleanupManager purgeMgr = new DatadirCleanupManager(
            config
                .getDataDir(), config.getDataLogDir(),
            GiraphConstants.ZOOKEEPER_SNAP_RETAIN_COUNT,
            GiraphConstants.ZOOKEEPER_PURGE_INTERVAL);
       purgeMgr.start();
 
       return port;
     }
 
     /**
      * Stop quorum and/or zookeeper.
      * @throws InterruptedException
      */
     public void stop() throws InterruptedException {
       if (quorumPeer != null) {
         quorumPeer.shutdown();
         quorumPeer.join();
       } else if (serverRunner != null) {
         serverRunner.stop();
       } else {
         LOG.warn("Neither quorum nor server is set");
       }
     }
   }
 
   /**
    * Wrapper around zookeeper service.
    */
   public static class ZooKeeperServerRunner  {
     /**
      * Reference to zookeeper factory.
      */
     private ServerCnxnFactory cnxnFactory;
     /**
      * Reference to zookeeper server.
      */
     private ZooKeeperServer zkServer;
 
     /**
      * Start zookeeper service.
      * @param config zookeeper configuration
      * formatted properly
      * @return the port zookeeper has started on.
      * @throws IOException
      */
     public int start(ZookeeperConfig config) throws IOException {
       LOG.warn("Either no config or no quorum defined in config, " +
           "running in process");
       try {
         ManagedUtil.registerLog4jMBeans();
       } catch (JMException e) {
         LOG.warn("Unable to register log4j JMX control", e);
       }
 
       runFromConfig(config);
       ThreadUtils.startThread(new Runnable() {
         @Override
         public void run() {
           try {
             cnxnFactory.join();
             if (zkServer.isRunning()) {
               zkServer.shutdown();
             }
           } catch (InterruptedException e) {
             LOG.error(e.getMessage(), e);
           }
 
         }
       }, "zk-thread");
       return zkServer.getClientPort();
     }
 
 
     /**
      * Run from a ServerConfig.
      * @param config ServerConfig to use.
      * @throws IOException
      */
     public void runFromConfig(ZookeeperConfig config) throws IOException {
       LOG.info("Starting server");
       try {
         // Note that this thread isn't going to be doing anything else,
         // so rather than spawning another thread, we will just call
         // run() in this thread.
         // create a file logger url from the command line args
         zkServer = new ZooKeeperServer();
 
         FileTxnSnapLog ftxn = new FileTxnSnapLog(new
             File(config.getDataLogDir()), new File(config.getDataDir()));
         zkServer.setTxnLogFactory(ftxn);
         zkServer.setTickTime(GiraphConstants.DEFAULT_ZOOKEEPER_TICK_TIME);
         zkServer.setMinSessionTimeout(config.getMinSessionTimeout());
         zkServer.setMaxSessionTimeout(config.getMaxSessionTimeout());
         cnxnFactory = ServerCnxnFactory.createFactory();
         cnxnFactory.configure(config.getClientPortAddress(),
             GiraphConstants.DEFAULT_ZOOKEEPER_MAX_CLIENT_CNXNS);
         cnxnFactory.startup(zkServer);
       } catch (InterruptedException e) {
         // warn, but generally this is ok
         LOG.warn("Server interrupted", e);
       }
     }
 
 
     /**
      * Stop zookeeper service.
      */
     public void stop() {
       cnxnFactory.shutdown();
     }
   }
 }