com.alibaba.wasp.util.KeyLocker.java Source code

Introduction

Here is the source code for com.alibaba.wasp.util.KeyLocker.java

Source

/**
 * 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 com.alibaba.wasp.util;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.hbase.util.Pair;

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;

/**
 * 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);
        }
    }

}