1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19 package org.apache.hadoop.hbase.util;
20
21
22 import java.util.HashMap;
23 import java.util.Map;
24 import java.util.Set;
25 import java.util.SortedSet;
26 import java.util.TreeSet;
27 import java.util.concurrent.atomic.AtomicInteger;
28 import java.util.concurrent.locks.Lock;
29 import java.util.concurrent.locks.ReentrantLock;
30
31 import org.apache.commons.logging.Log;
32 import org.apache.commons.logging.LogFactory;
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51 public class KeyLocker<K extends Comparable<? super K>> {
52 private static final Log LOG = LogFactory.getLog(KeyLocker.class);
53
54
55 private static final int NB_CONCURRENT_LOCKS = 1000;
56
57
58
59 private final Map<K, Pair<KeyLock<K>, AtomicInteger>> locks =
60 new HashMap<K, Pair<KeyLock<K>, AtomicInteger>>(NB_CONCURRENT_LOCKS);
61
62
63
64
65
66
67 public ReentrantLock acquireLock(K key) {
68 if (key == null) throw new IllegalArgumentException("key must not be null");
69
70 Pair<KeyLock<K>, AtomicInteger> lock;
71 synchronized (this) {
72 lock = locks.get(key);
73 if (lock == null) {
74 lock = new Pair<KeyLock<K>, AtomicInteger>(
75 new KeyLock<K>(this, key), new AtomicInteger(1));
76 locks.put(key, lock);
77 } else {
78 lock.getSecond().incrementAndGet();
79 }
80 }
81 lock.getFirst().lock();
82 return lock.getFirst();
83 }
84
85
86
87
88
89 public Map<K, Lock> acquireLocks(final Set<K> keys) {
90 Map<K, Lock> locks = new HashMap<K, Lock>(keys.size());
91 SortedSet<K> sortedKeys = new TreeSet<K>(keys);
92 for (K key : sortedKeys) {
93 locks.put(key, acquireLock(key));
94 }
95 return locks;
96 }
97
98
99
100
101 private synchronized void releaseLock(K key) {
102 Pair<KeyLock<K>, AtomicInteger> lock = locks.get(key);
103 if (lock != null) {
104 if (lock.getSecond().decrementAndGet() == 0) {
105 locks.remove(key);
106 }
107 } else {
108 String message = "Can't release the lock for " + key+", this key is not in the key list." +
109 " known keys are: "+ locks.keySet();
110 LOG.error(message);
111 throw new RuntimeException(message);
112 }
113 }
114
115 static class KeyLock<K extends Comparable<? super K>> extends ReentrantLock {
116 private static final long serialVersionUID = -12432857283423584L;
117
118 private final KeyLocker<K> locker;
119 private final K lockId;
120
121 private KeyLock(KeyLocker<K> locker, K lockId) {
122 super();
123 this.locker = locker;
124 this.lockId = lockId;
125 }
126
127 @Override
128 public void unlock() {
129 super.unlock();
130 locker.releaseLock(lockId);
131 }
132 }
133 }