refactor(pd): refactor common module

Author: koi2000

hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/ClientCache.java

@@ -61,7 +61,9 @@ private GraphCache getGraphCache(String graphName) {
         if ((graph = caches.get(graphName)) == null) {
             synchronized (caches) {
                 if ((graph = caches.get(graphName)) == null) {
-                    graph = new GraphCache();
+                    Metapb.Graph.Builder builder = Metapb.Graph.newBuilder().setGraphName(graphName);
+                    Metapb.Graph g = builder.build();
+                    graph = new GraphCache(g);
                     caches.put(graphName, graph);
                 }
             }

hugegraph-pd/hg-pd-common/pom.xml

@@ -44,5 +44,10 @@
             <artifactId>commons-collections4</artifactId>
             <version>4.4</version>
         </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-slf4j-impl</artifactId>
+            <version>${log4j2.version}</version>
+        </dependency>
     </dependencies>
 </project>

hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/Cache.java

@@ -0,0 +1,109 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.pd.common;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+
+public class Cache<T> implements Closeable {
+
+    ScheduledExecutorService ex =
+            Executors.newSingleThreadScheduledExecutor(new DefaultThreadFactory("hg-cache"));
+    private ConcurrentMap<String, CacheValue> map = new ConcurrentHashMap();
+    private ScheduledFuture<?> future;
+    private Runnable checker = () -> {
+        for (Map.Entry<String, CacheValue> e : map.entrySet()) {
+            if (e.getValue().getValue() == null) {
+                map.remove(e.getKey());
+            }
+        }
+    };
+
+    public Cache() {
+        future = ex.scheduleWithFixedDelay(checker, 1, 1, TimeUnit.SECONDS);
+    }
+
+    public CacheValue put(String key, T value, long ttl) {
+        return map.put(key, new CacheValue(value, ttl));
+    }
+
+    public T get(String key) {
+        CacheValue value = map.get(key);
+        if (value == null) {
+            return null;
+        }
+        T t = value.getValue();
+        if (t == null) {
+            map.remove(key);
+        }
+        return t;
+    }
+
+    public boolean keepAlive(String key, long ttl) {
+        CacheValue value = map.get(key);
+        if (value == null) {
+            return false;
+        }
+        value.keepAlive(ttl);
+        return true;
+    }
+
+    @Override
+    public void close() throws IOException {
+        try {
+            future.cancel(true);
+            ex.shutdownNow();
+        } catch (Exception e) {
+            try {
+                ex.shutdownNow();
+            } catch (Exception ex) {
+
+            }
+        }
+    }
+
+    private class CacheValue {
+
+        private final T value;
+        long outTime;
+
+        protected CacheValue(T value, long ttl) {
+            this.value = value;
+            this.outTime = System.currentTimeMillis() + ttl;
+        }
+
+        protected T getValue() {
+            if (System.currentTimeMillis() >= outTime) {
+                return null;
+            }
+            return value;
+        }
+
+        protected void keepAlive(long ttl) {
+            this.outTime = System.currentTimeMillis() + ttl;
+        }
+
+    }
+}

hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/Consts.java

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.pd.common;
+
+import io.grpc.Metadata;
+
+public class Consts {
+
+    public static final Metadata.Key<String> CREDENTIAL_KEY = Metadata.Key.of("credential",
+                                                                              Metadata.ASCII_STRING_MARSHALLER);
+    public static final Metadata.Key<String> TOKEN_KEY = Metadata.Key.of("Pd-Token",
+                                                                         Metadata.ASCII_STRING_MARSHALLER);
+}

hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/DefaultThreadFactory.java

@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.pd.common;
+
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class DefaultThreadFactory implements ThreadFactory {
+
+    private final AtomicInteger number = new AtomicInteger(1);
+    private final String prefix;
+
+    public DefaultThreadFactory(String prefix) {
+        this.prefix = prefix + "-";
+    }
+
+    @Override
+    public Thread newThread(Runnable r) {
+        Thread t = new Thread(r, prefix + number.getAndIncrement());
+        t.setDaemon(true);
+        t.setPriority(Thread.NORM_PRIORITY);
+        return t;
+    }
+}

hugegraph-pd/hg-pd-common/src/main/java/org/apache/hugegraph/pd/common/GraphCache.java

@@ -17,22 +17,27 @@
 
 package org.apache.hugegraph.pd.common;
 
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock.WriteLock;
 
-import com.google.common.collect.Range;
-
+import org.apache.commons.collections4.CollectionUtils;
 import org.apache.hugegraph.pd.grpc.Metapb.Graph;
 import org.apache.hugegraph.pd.grpc.Metapb.Partition;
 
+import com.google.common.collect.Range;
 import com.google.common.collect.RangeMap;
 import com.google.common.collect.TreeRangeMap;
 
 import lombok.Data;
+import lombok.extern.slf4j.Slf4j;
 
 @Data
+@Slf4j
 public class GraphCache {
 
     private Graph graph;
@@ -41,13 +46,30 @@ public class GraphCache {
     private ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
     private Map<Integer, AtomicBoolean> state = new ConcurrentHashMap<>();
     private Map<Integer, Partition> partitions = new ConcurrentHashMap<>();
-    private RangeMap<Long, Integer> range = new SynchronizedRangeMap<Long, Integer>().rangeMap;
+    private volatile RangeMap<Long, Integer> range = TreeRangeMap.create();
 
     public GraphCache(Graph graph) {
         this.graph = graph;
     }
 
-    public GraphCache() {
+    public void init(List<Partition> ps) {
+        Map<Integer, Partition> gps = new ConcurrentHashMap<>(ps.size(), 1);
+        if (!CollectionUtils.isEmpty(ps)) {
+            WriteLock lock = getLock().writeLock();
+            try {
+                lock.lock();
+                for (Partition p : ps) {
+                    gps.put(p.getId(), p);
+                    range.put(Range.closedOpen(p.getStartKey(), p.getEndKey()), p.getId());
+                }
+            } catch (Exception e) {
+                log.warn("init graph with error:", e);
+            } finally {
+                lock.unlock();
+            }
+        }
+        setPartitions(gps);
+
     }
 
     public Partition getPartition(Integer id) {
@@ -59,58 +81,87 @@ public Partition addPartition(Integer id, Partition p) {
     }
 
     public Partition removePartition(Integer id) {
+        Partition p = partitions.get(id);
+        if (p != null) {
+            RangeMap<Long, Integer> range = getRange();
+            if (Objects.equals(p.getId(), range.get(p.getStartKey())) &&
+                Objects.equals(p.getId(), range.get(p.getEndKey() - 1))) {
+                WriteLock lock = getLock().writeLock();
+                lock.lock();
+                try {
+                    range.remove(range.getEntry(p.getStartKey()).getKey());
+                } catch (Exception e) {
+                    log.warn("remove partition with error:", e);
+                } finally {
+                    lock.unlock();
+                }
+            }
+        }
         return partitions.remove(id);
     }
 
-    public class SynchronizedRangeMap<K extends Comparable<K>, V> {
-
-        private final RangeMap<K, V> rangeMap = TreeRangeMap.create();
-        private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
-
-        public void put(Range<K> range, V value) {
-            lock.writeLock().lock();
-            try {
-                rangeMap.put(range, value);
-            } finally {
-                lock.writeLock().unlock();
+    public void removePartitions() {
+        getState().clear();
+        RangeMap<Long, Integer> range = getRange();
+        WriteLock lock = getLock().writeLock();
+        try {
+            lock.lock();
+            if (range != null) {
+                range.clear();
             }
+        } catch (Exception e) {
+            log.warn("remove partition with error:", e);
+        } finally {
+            lock.unlock();
         }
+        getPartitions().clear();
+        getInitialized().set(false);
+    }
 
-        public V get(K key) {
-            lock.readLock().lock();
-            try {
-                return rangeMap.get(key);
-            } finally {
-                lock.readLock().unlock();
-            }
-        }
+    /*
+     * 需要外部加写锁
+     * */
+    public void reset() {
+        partitions.clear();
+        try {
+            range.clear();
+        } catch (Exception e) {
 
-        public void remove(Range<K> range) {
-            lock.writeLock().lock();
-            try {
-                rangeMap.remove(range);
-            } finally {
-                lock.writeLock().unlock();
-            }
         }
+    }
 
-        public Map.Entry<Range<K>, V> getEntry(K key) {
-            lock.readLock().lock();
-            try {
-                return rangeMap.getEntry(key);
-            } finally {
-                lock.readLock().unlock();
-            }
+    public boolean updatePartition(Partition partition) {
+        int partId = partition.getId();
+        Partition p = getPartition(partId);
+        if (p != null && p.equals(partition)) {
+            return false;
         }
-
-        public void clear() {
-            lock.writeLock().lock();
+        WriteLock lock = getLock().writeLock();
+        try {
+            lock.lock();
+            RangeMap<Long, Integer> range = getRange();
+            addPartition(partId, partition);
             try {
-                rangeMap.clear();
-            } finally {
-                lock.writeLock().unlock();
+                if (p != null) {
+                    // The old [1-3) is overwritten by [2-3). When [1-3) becomes [1-2), the
+                    // original [1-3) should not be deleted.
+                    // Only when it is confirmed that the old start and end are both your own can
+                    // the old be deleted (i.e., before it is overwritten).
+                    if (Objects.equals(partId, range.get(partition.getStartKey())) &&
+                        Objects.equals(partId, range.get(partition.getEndKey() - 1))) {
+                        range.remove(range.getEntry(partition.getStartKey()).getKey());
+                    }
+                }
+                range.put(Range.closedOpen(partition.getStartKey(), partition.getEndKey()), partId);
+            } catch (Exception e) {
+                log.warn("update partition with error:", e);
             }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        } finally {
+            lock.unlock();
         }
+        return true;
     }
 
 }