twill-zookeeper/src/main/java/org/apache/twill/internal/zookeeper/ReentrantDistributedLock.java

/*
 * 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.twill.internal.zookeeper;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import com.google.common.util.concurrent.Uninterruptibles;
import org.apache.twill.common.Cancellable;
import org.apache.twill.common.Threads;
import org.apache.twill.zookeeper.NodeChildren;
import org.apache.twill.zookeeper.NodeData;
import org.apache.twill.zookeeper.OperationFuture;
import org.apache.twill.zookeeper.ZKClient;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.UUID;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import javax.annotation.Nullable;

/**
 * A reentrant distributed lock implementation that uses ZooKeeper. It uses the recipe described in
 *
 * http://zookeeper.apache.org/doc/trunk/recipes.html#sc_recipes_Locks
 */
public class ReentrantDistributedLock implements Lock {
  private static final Logger LOG = LoggerFactory.getLogger(ReentrantDistributedLock.class);

  private final ZKClient zkClient;
  private final String path;
  private final ThreadLocal<String> localLockNode;
  private final ReentrantLock lock;

  /**
   * Creates a distributed lock instance.
   *
   * @param zkClient the {@link ZKClient} to interact with the ZooKeeper used for the lock coordination
   * @param path the path in ZooKeeper where the lock coordination happens
   */
  public ReentrantDistributedLock(ZKClient zkClient, String path) {
    this.zkClient = zkClient;
    this.path = path.startsWith("/") ? path : "/" + path;
    this.localLockNode = new ThreadLocal<String>();
    this.lock = new ReentrantLock();
  }

  @Override
  public void lock() {
    lock.lock();
    try {
      acquire(false, true);
    } catch (Exception e) {
      lock.unlock();
      throw Throwables.propagate(e);
    }
  }

  @Override
  public void lockInterruptibly() throws InterruptedException {
    lock.lockInterruptibly();
    try {
      acquire(true, true);
    } catch (Exception e) {
      lock.unlock();
      Throwables.propagateIfInstanceOf(e, InterruptedException.class);
      throw Throwables.propagate(e);
    }
  }

  @Override
  public boolean tryLock() {
    if (!lock.tryLock()) {
      return false;
    }
    try {
      if (acquire(false, false)) {
        return true;
      }
      lock.unlock();
      return false;
    } catch (Exception e) {
      lock.unlock();
      throw Throwables.propagate(e);
    }
  }

  @Override
  public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
    long startTime = System.nanoTime();
    if (!lock.tryLock(time, unit)) {
      return false;
    }
    long timeoutNano = unit.toNanos(time) - (System.nanoTime() - startTime);
    try {
      if (acquire(true, true, timeoutNano, TimeUnit.NANOSECONDS)) {
        return true;
      }
      lock.unlock();
      return false;
    } catch (ExecutionException e) {
      lock.unlock();
      throw Throwables.propagate(e.getCause());
    } catch (TimeoutException e) {
      lock.unlock();
      return false;
    }
  }

  @Override
  public void unlock() {
    if (!lock.isHeldByCurrentThread()) {
      throw new IllegalStateException("Cannot unlock without holding a lock by thread " + Thread.currentThread());
    }

    try {
      if (lock.getHoldCount() == 1) {
        // If it is the last lock entry for this thread, remove the zk node as well.
        try {
          Uninterruptibles.getUninterruptibly(zkClient.delete(localLockNode.get()));
        } catch (ExecutionException e) {
          throw Throwables.propagate(e.getCause());
        } finally {
          localLockNode.remove();
        }
      }
    } finally {
      lock.unlock();
    }
  }

  /**
   * Currently not supported and will always throw {@link UnsupportedOperationException}.
   */
  @Override
  public Condition newCondition() {
    // TODO: Add the support of it when needed
    throw new UnsupportedOperationException("Condition not supported.");
  }

  /**
   * Acquires a distributed lock through ZooKeeper. It's the same as calling
   * {@link #acquire(boolean, boolean, long, TimeUnit)} with {@link Long#MAX_VALUE} as timeout.
   */
  private boolean acquire(boolean interruptible, boolean waitForLock) throws InterruptedException, ExecutionException {
    try {
      return acquire(interruptible, waitForLock, Long.MAX_VALUE, TimeUnit.SECONDS);
    } catch (TimeoutException e) {
      // Should never happen
      throw Throwables.propagate(e);
    }
  }

  /**
   * Acquires a distributed lock through ZooKeeper.
   *
   * @param interruptible true if acquisition of lock can be interrupted
   * @param waitForLock true if wants to wait for the lock when not able to acquire it
   * @param timeout time to wait for the lock before giving up
   * @param unit unit for the timeout
   * @throws InterruptedException if {@code interruptible} is set to {@code true} and the current thread is interrupted
   *                              while acquiring the lock
   * @throws ExecutionException if there is failure while trying to acquire the lock
   */
  private boolean acquire(boolean interruptible,
                          final boolean waitForLock,
                          long timeout,
                          TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
    Preconditions.checkState(lock.isHeldByCurrentThread(), "Not owner of local lock.");
    if (lock.getHoldCount() > 1) {
      // Already owner of the lock, simply return.
      return true;
    }

    // Use a Future to help deal with different variants of locking
    // (lock, lockInterruptibly, tryLock, tryLock with timeout)
    // When the completion future is completed successfully, it means the lock is acquired and the future contains
    // the ZK node path to the ephemeral node that is representing this lock.
    // If it is failed, it means there is exception while trying to acquire the lock
    // If it is cancelled, it means to abort the acquisition logic (due to timeout / interrupt).
    final SettableFuture<String> completion = SettableFuture.create();

    // If the connection expired, fail the locking process if it is still in progress
    final Cancellable watcherCancellable = zkClient.addConnectionWatcher(new Watcher() {
      @Override
      public void process(WatchedEvent event) {
        if (event.getState() == Event.KeeperState.Expired) {
          completion.setException(new IllegalStateException("ZK session expired"));
        }
      }
    });
    // Always remove the watcher on completion
    completion.addListener(new Runnable() {
      @Override
      public void run() {
        watcherCancellable.cancel();
      }
    }, Threads.SAME_THREAD_EXECUTOR);


    // Step 1. Create a ephemeral sequential node
    final String guid = UUID.randomUUID().toString();
    final String lockPath = String.format("%s/%s-", path, guid);
    OperationFuture<String> future = zkClient.create(lockPath, null, CreateMode.EPHEMERAL_SEQUENTIAL, true);

    Futures.addCallback(future, new FutureCallback<String>() {
      @Override
      public void onSuccess(final String lockNode) {
        // If lock failed due to whatever reason, delete the lock node.
        deleteNodeOnFailure(completion, lockNode);

        // If the lock is completed (mainly due to cancellation), simply abort the lock acquisition logic.
        if (completion.isDone()) {
          return;
        }

        // Step 2-5. Try to determine who is the lock owner and watch for ZK node changes if itself is not the owner.
        doAcquire(completion, waitForLock, guid, lockNode);
      }

      @Override
      public void onFailure(Throwable t) {
        if (t instanceof KeeperException.ConnectionLossException) {
          // Ignore connection exception in create. Going to handle it in next step.
          // See the ZK receipt for details about the possible failure situation that can cause this.
          doAcquire(completion, waitForLock, guid, null);
        } else {
          LOG.error("Exception raised when creating lock node at {}", lockPath, t);
          completion.setException(t);
        }
      }
    });

    // Gets the result from the completion
    try {
      if (interruptible) {
        localLockNode.set(completion.get(timeout, unit));
      } else {
        localLockNode.set(Uninterruptibles.getUninterruptibly(completion, timeout, unit));
      }
      return true;
    } catch (InterruptedException e) {
      completion.cancel(true);
      throw e;
    } catch (TimeoutException e) {
      completion.cancel(true);
      throw e;
    } catch (CancellationException e) {
      // If the completion get cancelled, meaning the lock acquisition is aborted.
      return false;
    }
  }

  /**
   * Executes the lock acquisition process. This corresponds to step 2-5 of the distributed lock receipt.
   */
  private void doAcquire(final SettableFuture<String> completion, final boolean waitForLock,
                         final String guid, @Nullable final String lockPath) {
    // Step 2. Get all children under the lock parent.
    Futures.addCallback(zkClient.getChildren(path), new FutureCallback<NodeChildren>() {

      @Override
      public void onSuccess(NodeChildren children) {
        // Find the lock node in case the creation step failed by matching the guid
        // See "Recoverable Errors and the GUID" in the ZooKeeper guide
        final String lockNode = lockPath == null ? findLockNode(children.getChildren(), guid) : lockPath;
        if (lockNode == null) {
          // If not able to find the lock node, fail the locking procedure.
          completion.setException(new IllegalStateException("Failed to acquire lock").fillInStackTrace());
          return;
        }

        if (lockPath == null) {
          // If lock node was not determined in step 1 due to connection loss exception, need to add the
          // node deletion handler in here after the actual lockNode is determined.
          deleteNodeOnFailure(completion, lockNode);
        }

        // Find the node to watch, which is the one with the largest id that is smaller than currentId
        // If the current id is the smallest one, nodeToWatch will be null
        final String nodeToWatch = findNodeToWatch(children, lockNode, guid);

        // Step 3a. lockNode is the lowest, hence this become lock owner.
        if (nodeToWatch == null) {
          // Acquired lock
          completion.set(lockNode);
        } else if (!waitForLock) {
          // This is for the case of tryLock() without timeout.
          completion.cancel(true);
        }
        // If the lock acquisition is completed, due to whatever reason, we don't need to watch for any other nodes
        if (completion.isDone()) {
          return;
        }

        // Step 3b and 4. See if the the next lowest sequence ID exists. If it does, leave a watch
        // Use getData() instead of exists() to avoid leaking Watcher resources (if the node is gone, there will
        // be a watch left on the ZK server if exists() is used).
        OperationFuture<NodeData> getDataFuture = zkClient.getData(nodeToWatch, new Watcher() {
          @Override
          public void process(WatchedEvent event) {
            if (!completion.isDone()) {
              // If the watching node changed, go to step 2.
              doAcquire(completion, waitForLock, guid, lockNode);
            }
          }
        });

        // Step 5. Depends on the exists call result, either go to step 2 if the nodeToWatch is gone or just
        // let the watcher to trigger step 2 when there is change to the nodeToWatch.
        Futures.addCallback(getDataFuture, new FutureCallback<NodeData>() {
          @Override
          public void onSuccess(NodeData nodeData) {
            // No-op
          }

          @Override
          public void onFailure(Throwable t) {
            // See if the failure is due to node not exists. If that's the case, go to step 2.
            if (t instanceof KeeperException.NoNodeException && !completion.isDone()) {
              doAcquire(completion, waitForLock, guid, lockNode);
            } else {
              // If failed due to something else, fail the lock acquisition.
              completion.setException(t);
            }
          }
        });
      }

      @Override
      public void onFailure(Throwable t) {
        if (lockPath != null) {
          completion.setException(t);
        } else {
          doAcquire(completion, waitForLock, guid, null);
        }
      }
    });
  }

  /**
   * Deletes the given node if the given future failed.
   */
  private void deleteNodeOnFailure(final ListenableFuture<?> future, final String node) {
    future.addListener(new Runnable() {
      @Override
      public void run() {
        try {
          future.get();
        } catch (Exception e) {
          zkClient.delete(node);
        }
      }
    }, Threads.SAME_THREAD_EXECUTOR);
  }

  /**
   * Find the maximum node that is smaller than node of the given lockPath.
   *
   * @return the node found or {@code null} if no such node exist
   */
  private String findNodeToWatch(NodeChildren children, String lockPath, String guid) {
    // Lock path is "path/guid-id"
    int guidLen = guid.length();
    int id = Integer.parseInt(lockPath.substring(path.length() + guidLen + 2));

    String nodeToWatch = null;
    int maxOfMins = Integer.MIN_VALUE;
    for (String node : children.getChildren()) {
      int nodeId = Integer.parseInt(node.substring(guidLen + 1));
      if (nodeId < id && nodeId > maxOfMins) {
        maxOfMins = nodeId;
        nodeToWatch = path + "/" + node;
      }
    }
    return nodeToWatch;
  }

  /**
   * Finds the node path that matches the given prefix.
   */
  private String findLockNode(Iterable<String> nodes, String prefix) {
    for (String child : nodes) {
      if (child.startsWith(prefix)) {
        return path + "/" + child;
      }
    }
    return null;
  }
}