/**
 * 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.hadoop.hbase.util;


import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * A utility class to manage a set of locks. Each lock is identified by a String which serves
 * as a key. Typical usage is: <p>
 * class Example{
 * private final static KeyLocker<String> locker = new Locker<String>();
 * <p/>
 * public void foo(String s){
 * Lock lock = locker.acquireLock(s);
 * try {
 * // whatever
 * }finally{
 * lock.unlock();
 * }
 * }
 * }
 * </p>
 */
public class KeyLocker<K extends Comparable<? super K>> {
  private static final Log LOG = LogFactory.getLog(KeyLocker.class);

  // The number of lock we want to easily support. It's not a maximum.
  private static final int NB_CONCURRENT_LOCKS = 1000;

  // We need an atomic counter to manage the number of users using the lock and free it when
  //  it's equal to zero.
  private final Map<K, Pair<KeyLock<K>, AtomicInteger>> locks =
    new HashMap<K, Pair<KeyLock<K>, AtomicInteger>>(NB_CONCURRENT_LOCKS);

  /**
   * Return a lock for the given key. The lock is already locked.
   *
   * @param key
   */
  public ReentrantLock acquireLock(K key) {
    if (key == null) throw new IllegalArgumentException("key must not be null");

    Pair<KeyLock<K>, AtomicInteger> lock;
    synchronized (this) {
      lock = locks.get(key);
      if (lock == null) {
        lock = new Pair<KeyLock<K>, AtomicInteger>(
          new KeyLock<K>(this, key), new AtomicInteger(1));
        locks.put(key, lock);
      } else {
        lock.getSecond().incrementAndGet();
      }
    }
    lock.getFirst().lock();
    return lock.getFirst();
  }

  /**
   * Acquire locks for a set of keys. The keys will be
   * sorted internally to avoid possible deadlock.
   */
  public Map<K, Lock> acquireLocks(final Set<K> keys) {
    Map<K, Lock> locks = new HashMap<K, Lock>(keys.size());
    SortedSet<K> sortedKeys = new TreeSet<K>(keys);
    for (K key : sortedKeys) {
      locks.put(key, acquireLock(key));
    }
    return locks;
  }

  /**
   * Free the lock for the given key.
   */
  private synchronized void releaseLock(K key) {
    Pair<KeyLock<K>, AtomicInteger> lock = locks.get(key);
    if (lock != null) {
      if (lock.getSecond().decrementAndGet() == 0) {
        locks.remove(key);
      }
    } else {
      String message = "Can't release the lock for " + key+", this key is not in the key list." +
        " known keys are: "+ locks.keySet();
      LOG.error(message);
      throw new RuntimeException(message);
    }
  }

  static class KeyLock<K extends Comparable<? super K>> extends ReentrantLock {
    private static final long serialVersionUID = -12432857283423584L;

    private final KeyLocker<K> locker;
    private final K lockId;

    private KeyLock(KeyLocker<K> locker, K lockId) {
      super();
      this.locker = locker;
      this.lockId = lockId;
    }

    @Override
    public void unlock() {
      super.unlock();
      locker.releaseLock(lockId);
    }
  }
}