core/org.eclipse.smila.zookeeper/code/src/org/eclipse/smila/zookeeper/ZkLock.java

/*******************************************************************************
 * Copyright (c) 2008, 2011 Attensity Europe GmbH and brox IT Solutions GmbH. All rights reserved. This program and the
 * accompanying materials are made available under the terms of the Eclipse Public License v1.0 which accompanies this
 * distribution, and is available at http://www.eclipse.org/legal/epl-v10.html
 * 
 * Contributors: Andreas Weber (Attensity Europe GmbH) - initial implementation
 **********************************************************************************************************************/

package org.eclipse.smila.zookeeper;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.data.Stat;

/**
 * Zookeeper Lock implementation. Assumes that this is called during a running zookeeper session.
 * 
 * @author aweber
 */
public class ZkLock {

  /** local logger. */
  private final Log _log = LogFactory.getLog(getClass());

  /** Connection to ZooKeeper server, includes Zookeeper client. */
  private final ZkConnection _zk;

  /** Path of the lock node. */
  private final String _lockPath;

  /** helps to find out if a zookeeper session where we got a lock is still valid. */
  private long _currentZkClientSessionId;

  /**
   * @param zk
   *          zookeeper connection
   * @param lockRootNodePath
   *          zookeeper node where to set the lock
   */
  public ZkLock(final ZkConnection zk, final String lockRootNodePath) {
    if (lockRootNodePath == null) {
      throw new NullPointerException("parameter lockRootNodePath is <null>");
    }
    _zk = zk;
    _lockPath = lockRootNodePath + "/" + "lock";
    // ensure root node for lock exists
    Stat stat = null;
    try {
      stat = _zk.exists(lockRootNodePath);
    } catch (final Exception e) {
      throw new RuntimeException(e);
    }
    if (stat == null) {
      throw new IllegalArgumentException("root zookeeper node '" + lockRootNodePath + "' doesn't exist!");
    }
  }

  /**
   * Acquires the lock only if it is free at the time of invocation.
   * 
   * Assumes that this is called during a running zookeeper session.
   * 
   * @return 'true' if the lock could be successfully acquired. If the lock is not available then this method will
   *         return immediately with the value 'false'.
   */
  public boolean tryLock() {
    if (_log.isTraceEnabled()) {
      _log.trace("tryLock()");
    }
    try {
      if (_zk.exists(_lockPath) != null) {
        return false;
      }
      try {
        // use ephemeral node to set lock -> if zookeeper session disconnects, node will be removed by zookeeper      
        _zk.createNode(_lockPath, ZkConnection.NO_DATA, CreateMode.EPHEMERAL);
        _currentZkClientSessionId = _zk.getSessionId();
        return true;
      } catch (final KeeperException.NodeExistsException e) {
        // lock is already set
        return false;
      }
    } catch (final Exception e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * Releases the lock.
   * 
   * Assumes that this is called during a running zookeeper session.
   */
  public void unlock() {
    if (_log.isTraceEnabled()) {
      _log.trace("unlock()");
    }
    try {
      _zk.deleteNode(_lockPath);
    } catch (final KeeperException.NoNodeException e) {
      if (_zk.getSessionId() == _currentZkClientSessionId) {
        _log.warn("Could not remove lock because lock wasn't set before.");
      } else {
        _log.warn("Could not remove lock because session expired in between");
      }
    } catch (final Exception e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * Can be called by a client to find out if the acquired lock is still valid. Could be called at the last possible
   * moment where client could do a rollback of its operations done after acquiring the lock.
   */
  public boolean isLockStillValid() {
    // if zookeeper session expired in between, lock isn't valid anymore
    return (_zk.getSessionId() == _currentZkClientSessionId);
  }
}