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 */
018
019package org.apache.oozie.lock;
020
021import java.util.HashMap;
022import java.util.concurrent.TimeUnit;
023import java.util.concurrent.locks.ReentrantReadWriteLock;
024import java.util.concurrent.locks.Lock;
025
026/**
027 * In memory resource locking that provides READ/WRITE lock capabilities.
028 */
029public class MemoryLocks {
030    final private HashMap<String, ReentrantReadWriteLock> locks = new HashMap<String, ReentrantReadWriteLock>();
031
032    private static enum Type {
033        READ, WRITE
034    }
035
036    /**
037     * Implementation of {@link LockToken} for in memory locks.
038     */
039    class MemoryLockToken implements LockToken {
040        private final ReentrantReadWriteLock rwLock;
041        private final java.util.concurrent.locks.Lock lock;
042        private final String resource;
043
044        private MemoryLockToken(ReentrantReadWriteLock rwLock, java.util.concurrent.locks.Lock lock, String resource) {
045            this.rwLock = rwLock;
046            this.lock = lock;
047            this.resource = resource;
048        }
049
050        /**
051         * Release the lock.
052         */
053        @Override
054        public void release() {
055            int val = rwLock.getQueueLength();
056            if (val == 0) {
057                synchronized (locks) {
058                    locks.remove(resource);
059                }
060            }
061            lock.unlock();
062        }
063    }
064
065    /**
066     * Return the number of active locks.
067     *
068     * @return the number of active locks.
069     */
070    public int size() {
071        return locks.size();
072    }
073
074    /**
075     * Obtain a READ lock for a source.
076     *
077     * @param resource resource name.
078     * @param wait time out in milliseconds to wait for the lock, -1 means no timeout and 0 no wait.
079     * @return the lock token for the resource, or <code>null</code> if the lock could not be obtained.
080     * @throws InterruptedException thrown if the thread was interrupted while waiting.
081     */
082    public MemoryLockToken getReadLock(String resource, long wait) throws InterruptedException {
083        return getLock(resource, Type.READ, wait);
084    }
085
086    /**
087     * Obtain a WRITE lock for a source.
088     *
089     * @param resource resource name.
090     * @param wait time out in milliseconds to wait for the lock, -1 means no timeout and 0 no wait.
091     * @return the lock token for the resource, or <code>null</code> if the lock could not be obtained.
092     * @throws InterruptedException thrown if the thread was interrupted while waiting.
093     */
094    public MemoryLockToken getWriteLock(String resource, long wait) throws InterruptedException {
095        return getLock(resource, Type.WRITE, wait);
096    }
097
098    private MemoryLockToken getLock(String resource, Type type, long wait) throws InterruptedException {
099        ReentrantReadWriteLock lockEntry;
100        synchronized (locks) {
101            if (locks.containsKey(resource)) {
102                lockEntry = locks.get(resource);
103            }
104            else {
105                lockEntry = new ReentrantReadWriteLock(true);
106                locks.put(resource, lockEntry);
107            }
108        }
109
110        Lock lock = (type.equals(Type.READ)) ? lockEntry.readLock() : lockEntry.writeLock();
111
112        if (wait == -1) {
113            lock.lock();
114        }
115        else {
116            if (wait > 0) {
117                if (!lock.tryLock(wait, TimeUnit.MILLISECONDS)) {
118                    return null;
119                }
120            }
121            else {
122                if (!lock.tryLock()) {
123                    return null;
124                }
125            }
126        }
127        synchronized (locks) {
128            if (!locks.containsKey(resource)) {
129                locks.put(resource, lockEntry);
130            }
131        }
132        return new MemoryLockToken(lockEntry, lock, resource);
133    }
134}