001    /**
002     * Copyright (c) 2010 Yahoo! Inc. All rights reserved.
003     * Licensed under the Apache License, Version 2.0 (the "License");
004     * you may not use this file except in compliance with the License.
005     * You may obtain a copy of the License at
006     *
007     *   http://www.apache.org/licenses/LICENSE-2.0
008     *
009     *  Unless required by applicable law or agreed to in writing, software
010     *  distributed under the License is distributed on an "AS IS" BASIS,
011     *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
012     *  See the License for the specific language governing permissions and
013     *  limitations under the License. See accompanying LICENSE file.
014     */
015    package org.apache.oozie.util;
016    
017    import java.util.HashMap;
018    import java.util.concurrent.TimeUnit;
019    import java.util.concurrent.locks.ReentrantReadWriteLock;
020    import java.util.concurrent.locks.Lock;
021    
022    /**
023     * In memory resource locking that provides READ/WRITE lock capabilities.
024     */
025    public class MemoryLocks {
026        final private HashMap<String, ReentrantReadWriteLock> locks = new HashMap<String, ReentrantReadWriteLock>();
027    
028        private static enum Type {
029            READ, WRITE
030        }
031    
032        /**
033         * Lock token returned when obtaining a lock, the token must be released when the lock is not needed anymore.
034         */
035        public class LockToken {
036            private final ReentrantReadWriteLock rwLock;
037            private final java.util.concurrent.locks.Lock lock;
038            private final String resource;
039    
040            private LockToken(ReentrantReadWriteLock rwLock, java.util.concurrent.locks.Lock lock, String resource) {
041                this.rwLock = rwLock;
042                this.lock = lock;
043                this.resource = resource;
044            }
045    
046            /**
047             * Release the lock.
048             */
049            public void release() {
050                int val = rwLock.getQueueLength();
051                if (val == 0) {
052                    synchronized (locks) {
053                        locks.remove(resource);
054                    }
055                }
056                lock.unlock();
057            }
058        }
059    
060        /**
061         * Return the number of active locks.
062         *
063         * @return the number of active locks.
064         */
065        public int size() {
066            return locks.size();
067        }
068    
069        /**
070         * Obtain a READ lock for a source.
071         *
072         * @param resource resource name.
073         * @param wait time out in milliseconds to wait for the lock, -1 means no timeout and 0 no wait.
074         * @return the lock token for the resource, or <code>null</code> if the lock could not be obtained.
075         * @throws InterruptedException thrown if the thread was interrupted while waiting.
076         */
077        public LockToken getReadLock(String resource, long wait) throws InterruptedException {
078            return getLock(resource, Type.READ, wait);
079        }
080    
081        /**
082         * Obtain a WRITE lock for a source.
083         *
084         * @param resource resource name.
085         * @param wait time out in milliseconds to wait for the lock, -1 means no timeout and 0 no wait.
086         * @return the lock token for the resource, or <code>null</code> if the lock could not be obtained.
087         * @throws InterruptedException thrown if the thread was interrupted while waiting.
088         */
089        public LockToken getWriteLock(String resource, long wait) throws InterruptedException {
090            return getLock(resource, Type.WRITE, wait);
091        }
092    
093        private LockToken getLock(String resource, Type type, long wait) throws InterruptedException {
094            ReentrantReadWriteLock lockEntry;
095            synchronized (locks) {
096                if (locks.containsKey(resource)) {
097                    lockEntry = locks.get(resource);
098                }
099                else {
100                    lockEntry = new ReentrantReadWriteLock(true);
101                    locks.put(resource, lockEntry);
102                }
103            }
104    
105            Lock lock = (type.equals(Type.READ)) ? lockEntry.readLock() : lockEntry.writeLock();
106    
107            if (wait == -1) {
108                lock.lock();
109            }
110            else {
111                if (wait > 0) {
112                    if (!lock.tryLock(wait, TimeUnit.MILLISECONDS)) {
113                        return null;
114                    }
115                }
116                else {
117                    if (!lock.tryLock()) {
118                        return null;
119                    }
120                }
121            }
122            synchronized (locks) {
123                if (!locks.containsKey(resource)) {
124                    locks.put(resource, lockEntry);
125                }
126            }
127            return new LockToken(lockEntry, lock, resource);
128        }
129    
130    }