diff --git a/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/GraphCache.java b/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/GraphCache.java index 07c7c332d9..8a576e1b6b 100644 --- a/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/GraphCache.java +++ b/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/GraphCache.java @@ -22,6 +22,8 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.ReentrantReadWriteLock; +import com.google.common.collect.Range; + import org.apache.hugegraph.pd.grpc.Metapb.Graph; import org.apache.hugegraph.pd.grpc.Metapb.Partition; @@ -39,7 +41,7 @@ public class GraphCache { private ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); private Map state = new ConcurrentHashMap<>(); private Map partitions = new ConcurrentHashMap<>(); - private RangeMap range = TreeRangeMap.create(); + private RangeMap range = new SynchronizedRangeMap().rangeMap; public GraphCache(Graph graph) { this.graph = graph; @@ -59,4 +61,56 @@ public Partition addPartition(Integer id, Partition p) { public Partition removePartition(Integer id) { return partitions.remove(id); } + + public class SynchronizedRangeMap, V> { + + private final RangeMap rangeMap = TreeRangeMap.create(); + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + + public void put(Range range, V value) { + lock.writeLock().lock(); + try { + rangeMap.put(range, value); + } finally { + lock.writeLock().unlock(); + } + } + + public V get(K key) { + lock.readLock().lock(); + try { + return rangeMap.get(key); + } finally { + lock.readLock().unlock(); + } + } + + public void remove(Range range) { + lock.writeLock().lock(); + try { + rangeMap.remove(range); + } finally { + lock.writeLock().unlock(); + } + } + + public Map.Entry, V> getEntry(K key) { + lock.readLock().lock(); + try { + return rangeMap.getEntry(key); + } finally { + lock.readLock().unlock(); + } + } + + public void clear() { + lock.writeLock().lock(); + try { + rangeMap.clear(); + } finally { + lock.writeLock().unlock(); + } + } + } + } diff --git a/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/PartitionCache.java b/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/PartitionCache.java index 9b98cbbb03..31cc29deed 100644 --- a/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/PartitionCache.java +++ b/hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/PartitionCache.java @@ -85,7 +85,7 @@ public void waitGraphLock(String graphName) { public void lockGraph(String graphName) { var lock = getOrCreateGraphLock(graphName); - while (lock.compareAndSet(false, true)) { + while (!lock.compareAndSet(false, true)) { Thread.onSpinWait(); } }