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 }