//设置package包名称以及导入依赖的类
package com.groupon.dse.testutils

import java.io.File
import java.net.InetSocketAddress

import kafka.utils.Utils
import org.apache.zookeeper.server.{NIOServerCnxnFactory, ZooKeeperServer}


class EmbeddedZookeeper(zkPort: Int, zkSnapshotDir: File, zkLogDir: File, autoStart: Boolean) {

  val connectString = s"127.0.0.1:${zkPort}"
  val tickTime = 500
  val factory = new NIOServerCnxnFactory()
  val maxZkConnections = 100
  factory.configure(new InetSocketAddress("127.0.0.1", zkPort), maxZkConnections)

  var zookeeper: ZooKeeperServer = null

  def this(port: Int, autoStart: Boolean = true) = this(port, TestUtils.tempDir, TestUtils.tempDir, autoStart)

  if (autoStart) {
    start()
  }

  def start(): Unit = {
    // With Zookeeper 3.4, the startup logic of the ZookeeperServer has changed where a sequence of:
    // zookeeper.start() -> zookeeper.shutdown() -> zookeeper.start()
    // will fail to restart the ZookeeperServer. Because of this, a new ZookeeperServer needs to be instantiated if
    // we want to simulate Zookeeper unavailability during tests
    zookeeper = new ZooKeeperServer(zkSnapshotDir, zkLogDir, tickTime)
    factory.startup(zookeeper)
  }

  def stop(): Unit = {
    zookeeper.shutdown()
    zookeeper = null
  }

  def snapshotDir: File = zkSnapshotDir

  def logDir: File = zkLogDir

  def cleanShutdown(): Unit = {
    shutdown()
    Utils.rm(zkLogDir)
    Utils.rm(zkSnapshotDir)
  }

  def shutdown(): Unit = {
    Utils.swallow(zookeeper.shutdown())
    Utils.swallow(factory.shutdown())
    zookeeper = null
  }
}