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
   *
   *    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 io.jafka.server;
  
  import com.github.zkclient.IZkStateListener;
  import com.github.zkclient.ZkClient;
  import com.github.zkclient.exception.ZkNodeExistsException;
  import io.jafka.cluster.Broker;
  import io.jafka.log.LogManager;
  import io.jafka.server.TopicTask.TaskType;
  import io.jafka.utils.zookeeper.ZkUtils;
  import org.apache.zookeeper.Watcher.Event.KeeperState;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.Closeable;
  import java.net.InetAddress;
  import java.net.UnknownHostException;
  import java.util.LinkedHashSet;
  import java.util.List;
  import java.util.Set;
  
  /**
   * Handles the server's interaction with zookeeper. The server needs to register the following
   * paths:
   *
   * <pre>
   *   /topics/[topic]/[node_id-partition_num]
   *   /brokers/[0...N] -- host:port
   * </pre>
   *
   * @author adyliu (imxylz@gmail.com)
   * @since 1.0
   */
  public class ServerRegister implements IZkStateListener, Closeable {
  
      private final ServerConfig config;
  
      private final LogManager logManager;
  
      private static final Logger logger = LoggerFactory.getLogger(ServerRegister.class);
  
      private final String brokerIdPath;
  
      private ZkClient zkClient;
  
      private Set<String> topics = new LinkedHashSet<String>();
  
      private final Object lock = new Object();
  
      public ServerRegister(ServerConfig config, LogManager logManager) {
          this.config = config;
          this.logManager = logManager;
          //
          this.brokerIdPath = ZkUtils.BrokerIdsPath + "/" + config.getBrokerId();
      }
  
      public void startup() {
          logger.info("connecting to zookeeper: " + config.getZkConnect());
          zkClient = new ZkClient(config.getZkConnect(), config.getZkSessionTimeoutMs(),
                  config.getZkConnectionTimeoutMs());
          zkClient.subscribeStateChanges(this);
      }
  
      public void processTask(TopicTask task) {
          final String topicPath = ZkUtils.BrokerTopicsPath + "/" + task.topic;
          final String brokerTopicPath = ZkUtils.BrokerTopicsPath + "/" + task.topic + "/" + config.getBrokerId();
          synchronized (lock) {
              switch (task.type) {
                  case DELETE:
                      topics.remove(task.topic);
                      ZkUtils.deletePath(zkClient, brokerTopicPath);
                      List<String> brokers = ZkUtils.getChildrenParentMayNotExist(zkClient, topicPath);
                      if (brokers != null && brokers.size() == 0) {
                          ZkUtils.deletePath(zkClient, topicPath);
                          logger.info("delete topic={} zkpath={}", task.topic, brokerTopicPath);
                      }
                      break;
                  case CREATE:
                      final int partitionsCreated = getPartitions(task.topic);
                      topics.add(task.topic);
                      ZkUtils.createEphemeralPathExpectConflict(zkClient, brokerTopicPath, "" + partitionsCreated);
                      logger.info("create topic={} partitions={} zkpath={}", task.topic, partitionsCreated, brokerTopicPath);
                      break;
                  case ENLARGE:
                      final int partitionsEnlarged = getPartitions(task.topic);
                      ZkUtils.deletePath(zkClient, brokerTopicPath);
                      ZkUtils.createEphemeralPathExpectConflict(zkClient, brokerTopicPath, "" + getPartitions(task.topic));
                      logger.info("enlarge topic={} partitions={} zkpath={}", task.topic, partitionsEnlarged, brokerTopicPath);
                      break;
                  default:
                      logger.error("unknow task: " + task);
                      break;
              }
          }
      }
  
      private int getPartitions(String topic) {
          Integer numParts = logManager.getTopicPartitionsMap().get(topic);
          return numParts == null ? config.getNumPartitions() : numParts.intValue();
      }
  
      /**
       * register broker in the zookeeper
       * <p>
       * path: /brokers/ids/&lt;id&gt; <br>
       * data: creator:host:port
       */
      public void registerBrokerInZk() {
          logger.info("Registering broker " + brokerIdPath);
          String hostname = config.getHostName();
          if (hostname == null) {
              try {
                  hostname = InetAddress.getLocalHost().getHostName();
              } catch (UnknownHostException e) {
                  throw new RuntimeException("cannot get local host, setting 'hostname' in configuration");
              }
          }
          //
          final String creatorId = hostname + "-" + System.currentTimeMillis();
          final Broker broker = new Broker(config.getBrokerId(), creatorId, hostname, config.getPort(),config.isTopicAutoCreated());
          try {
              ZkUtils.createEphemeralPathExpectConflict(zkClient, brokerIdPath, broker.getZKString());
          } catch (ZkNodeExistsException e) {
              String oldServerInfo = ZkUtils.readDataMaybeNull(zkClient, brokerIdPath);
              String message = "A broker (%s) is already registered on the path %s." //
                      + " This probably indicates that you either have configured a brokerid that is already in use, or "//
                      + "else you have shutdown this broker and restarted it faster than the zookeeper " ///
                      + "timeout so it appears to be re-registering.";
              message = String.format(message, oldServerInfo, brokerIdPath);
              throw new RuntimeException(message);
          }
          //
          logger.info("Registering broker " + brokerIdPath + " succeeded with " + broker);
      }
  
      /**
       *
       */
      public void close() {
          if (zkClient != null) {
              logger.info("closing zookeeper client...");
              zkClient.close();
          }
      }
  
      public void handleNewSession() throws Exception {
          logger.info("re-registering broker info in zookeeper for broker " + config.getBrokerId());
          registerBrokerInZk();
          synchronized (lock) {
              logger.info("re-registering broker topics in zookeeper for broker " + config.getBrokerId());
              for (String topic : topics) {
                  processTask(new TopicTask(TaskType.CREATE, topic));
              }
          }
  
      }
  
      public void handleStateChanged(KeeperState state) throws Exception {
          // do nothing, since zkclient will do reconnect for us.
      }
  }
  