ZkConnectionManagerImpl.java

/*
 * Copyright (C) 2012 Facebook, 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 com.facebook.zookeeper.connection;

import com.facebook.concurrency.ErrorLoggingRunnable;
import com.facebook.concurrency.NamedThreadFactory;
import com.facebook.zookeeper.ZooKeeperFactory;
import com.facebook.zookeeper.ZooKeeperIface;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

public class ZkConnectionManagerImpl implements ZkConnectionManager {
  private static final Logger LOG =
    LoggerFactory.getLogger(ZkConnectionManagerImpl.class);

  private final ZooKeeperFactory zooKeeperFactory;
  private final List<Watcher> registeredWatchers =
    new CopyOnWriteArrayList<Watcher>();
  private final ConnectionWatcher connectionWatcher = new ConnectionWatcher();
  private final ConnectionRenewer connectionRenewer =
    new ConnectionRenewer();
  private final int connectTimeoutMillis;
  private final int retryIntervalMillis;
  private volatile ZooKeeperIface zk = null;
  private volatile boolean isStarted = false;
  private volatile boolean isShutDown = false;

  public ZkConnectionManagerImpl(
    ZooKeeperFactory zooKeeperFactory,
    int connectTimeoutMillis,
    int retryIntervalMillis
  ) {
    this.zooKeeperFactory = zooKeeperFactory;
    this.connectTimeoutMillis = connectTimeoutMillis;
    this.retryIntervalMillis = retryIntervalMillis;
  }

  // Instance must be started before it becomes valid for use
  public synchronized void start() {
    if (isStarted) {
      throw new IllegalStateException("Should only be started once");
    }
    // Connect to the ZooKeeper cluster
    try {
      // Call connect first to guarantee zk is not null when this method exits
      connect();
    } catch (IOException e) {
      connectionRenewer.activate();
    }
    isStarted = true;
  }

  private void verifyOperational() {
    if (!isStarted) {
      throw new IllegalStateException("Not yet started");
    }
    if (isShutDown) {
      throw new IllegalStateException("Already closed");
    }
  }

  private boolean isAlive() {
    return zk != null && zk.getState().isAlive();
  }

  private synchronized void connect() throws IOException {
    if (!isShutDown && !isAlive()) {
      LOG.info("Initializing ZooKeeper connection");
      connectionWatcher.reset();
      zk = zooKeeperFactory.create(connectionWatcher);
    }
  }

  @Override
  public ZooKeeperIface getClient() throws InterruptedException {
    verifyOperational();
    // Wait if we are in the process of connecting
    if (!connectionWatcher.waitForConnect(connectTimeoutMillis, TimeUnit.MILLISECONDS)) {
      LOG.error("Exceeded " + connectTimeoutMillis + " ms waiting for " +
        "connection to be established! Using disconnected client...");
    }
    return zk;
  }

  @Override
  public ZooKeeper.States registerWatcher(Watcher watcher) {
    // This operation may legitimately happen at any time without problems

    // Note: the setting of the watch MUST precede the reading of the state
    // in order to ensure that the caller doesn't miss any watches after
    // getting the state.
    registeredWatchers.add(watcher);
    return (zk == null) ? ZooKeeper.States.CLOSED : zk.getState();
  }

  public synchronized void shutdown() throws InterruptedException {
    LOG.info("Closing ZooKeeper connection");
    verifyOperational();
    try {
      connectionRenewer.shutdown();
      if (isAlive()) {
        zk.close();
      }
    } finally {
      isShutDown = true;
    }
  }

  /**
   * When activated, will try to connect to ZooKeeper until it succeeds.
   */
  private class ConnectionRenewer {
    private final ScheduledExecutorService retryExecutor =
      Executors.newSingleThreadScheduledExecutor(
        new NamedThreadFactory("ZkConnectionManager-renewer")
      );
    private final AtomicBoolean isScheduled = new AtomicBoolean(false);

    public void activate() {
      if (isScheduled.compareAndSet(false, true)) {
        retryExecutor.execute(new ErrorLoggingRunnable(new Runnable() {
          @Override
          public void run() {
            try {
              connect();
              return; // Success, don't reschedule
            } catch (IOException e) {
              // Try again after sleeping...
              LOG.error("Failed to connect to ZooKeeper", e);
            } finally {
              isScheduled.set(false);
            }
            if (isScheduled.compareAndSet(false, true)) {
              retryExecutor.schedule(
                this, retryIntervalMillis, TimeUnit.MILLISECONDS
              );
            }
          }
        }));
      }
    }

    public void shutdown() {
      retryExecutor.shutdown();
    }
  }

  private class ConnectionWatcher implements Watcher {
    private volatile CountDownLatch connectedSignal = new CountDownLatch(1);

    public synchronized void reset() {
      connectedSignal.countDown(); // Unblock any prior threads
      connectedSignal = new CountDownLatch(1);
    }

    public boolean waitForConnect(int timeout, TimeUnit timeUnit)
      throws InterruptedException {
      return connectedSignal.await(timeout, timeUnit);
    }

    @Override
    public void process(WatchedEvent event) {
      // Handle the connection event signals. No thread needed because all
      // operations are non-blocking
      switch (event.getState()) {
        case SyncConnected:
          LOG.info("ZooKeeper connected");
          connectedSignal.countDown();
          break;
        case Disconnected:
          LOG.warn("ZooKeeper disconnected!");
          break;
        case Expired:
          LOG.warn("ZooKeeper session expired!");
          // Reset connection to the ZooKeeper cluster
          connectionRenewer.activate();
          break;
      }

      // Signal the registered watchers with the connection event
      for (Watcher watcher : registeredWatchers) {
        try {
          // All watchers should have non-blocking implementations
          watcher.process(event);
        } catch (Throwable t) {
          LOG.error("Registered watcher failed handling connection event", t);
        }
      }
    }
  }
}