001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase.util;
019
020import java.util.Arrays;
021import java.util.LinkedHashMap;
022import java.util.Map;
023import java.util.Set;
024import java.util.concurrent.locks.Lock;
025import java.util.concurrent.locks.ReentrantLock;
026import org.apache.yetus.audience.InterfaceAudience;
027
028/**
029 * A utility class to manage a set of locks. Each lock is identified by a String which serves as a
030 * key. Typical usage is:
031 *
032 * <pre>
033 * class Example {
034 *   private final static KeyLocker&lt;String&gt; locker = new Locker&lt;String&gt;();
035 *
036 *   public void foo(String s) {
037 *     Lock lock = locker.acquireLock(s);
038 *     try {
039 *       // whatever
040 *     } finally {
041 *       lock.unlock();
042 *     }
043 *   }
044 * }
045 * </pre>
046 */
047@InterfaceAudience.Private
048public class KeyLocker<K> {
049  // The number of lock we want to easily support. It's not a maximum.
050  private static final int NB_CONCURRENT_LOCKS = 1000;
051
052  private final WeakObjectPool<K, ReentrantLock> lockPool =
053    new WeakObjectPool<>(new ObjectPool.ObjectFactory<K, ReentrantLock>() {
054      @Override
055      public ReentrantLock createObject(K key) {
056        return new ReentrantLock();
057      }
058    }, NB_CONCURRENT_LOCKS);
059
060  /**
061   * Return a lock for the given key. The lock is already locked. n
062   */
063  public ReentrantLock acquireLock(K key) {
064    if (key == null) throw new IllegalArgumentException("key must not be null");
065
066    lockPool.purge();
067    ReentrantLock lock = lockPool.get(key);
068
069    lock.lock();
070    return lock;
071  }
072
073  /**
074   * Acquire locks for a set of keys. The keys will be sorted internally to avoid possible deadlock.
075   * @throws ClassCastException if the given {@code keys} contains elements that are not mutually
076   *                            comparable
077   */
078  public Map<K, Lock> acquireLocks(Set<? extends K> keys) {
079    Object[] keyArray = keys.toArray();
080    Arrays.sort(keyArray);
081
082    lockPool.purge();
083    Map<K, Lock> locks = new LinkedHashMap<>(keyArray.length);
084    for (Object o : keyArray) {
085      @SuppressWarnings("unchecked")
086      K key = (K) o;
087      ReentrantLock lock = lockPool.get(key);
088      locks.put(key, lock);
089    }
090
091    for (Lock lock : locks.values()) {
092      lock.lock();
093    }
094    return locks;
095  }
096}