From e1325615dc5f4849c4cee45de89af3544e3d9883 Mon Sep 17 00:00:00 2001
From: wu-sheng <wu.sheng@foxmail.com>
Date: Tue, 7 Nov 2017 12:12:46 +0800
Subject: [PATCH] Support to find node by id and output type. Node<?, Integer>
 foundNode = GraphManager.INSTANCE.findGraph(6).toBuilder().findNode(2,
 Integer.class);

---
 .../apm/collector/core/graph/Graph.java       |  9 +++-
 .../collector/core/graph/GraphBuilder.java    | 43 +++++++++++++++++++
 .../collector/core/graph/GraphManager.java    |  4 ++
 .../core/graph/GraphManagerTest.java          | 14 ++++++
 4 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphBuilder.java

diff --git a/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/Graph.java b/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/Graph.java
index 13d4ea729..d241b41ec 100644
--- a/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/Graph.java
+++ b/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/Graph.java
@@ -48,7 +48,7 @@ public final class Graph<INPUT> {
         if (node == null) {
             throw new NodeNotFoundException("Can't find node with handlerId="
                 + handlerId
-                + " in graph[" + id + "】");
+                + " in graph[" + id + "]");
         }
         return node.getNext();
     }
@@ -63,4 +63,11 @@ public final class Graph<INPUT> {
         nodeIndex.put(nodeId, node);
     }
 
+    public GraphBuilder toBuilder(){
+        return new GraphBuilder(this);
+    }
+
+    ConcurrentHashMap<Integer, Node> getNodeIndex() {
+        return nodeIndex;
+    }
 }
diff --git a/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphBuilder.java b/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphBuilder.java
new file mode 100644
index 000000000..eb671e0e8
--- /dev/null
+++ b/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphBuilder.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2017, OpenSkywalking Organization All rights reserved.
+ *
+ * Licensed 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.
+ *
+ * Project repository: https://github.com/OpenSkywalking/skywalking
+ */
+
+package org.skywalking.apm.collector.core.graph;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * @author wu-sheng
+ */
+public class GraphBuilder {
+    private Graph graph;
+
+    GraphBuilder(Graph graph) {
+        this.graph = graph;
+    }
+
+    public <NODE_OUTPUT> Node<?, NODE_OUTPUT> findNode(int handlerId, Class<NODE_OUTPUT> outputClass) {
+        ConcurrentHashMap<Integer, Node> graphNodeIndex = graph.getNodeIndex();
+        Node node = graphNodeIndex.get(handlerId);
+        if (node == null) {
+            throw new NodeNotFoundException("Can't find node with handlerId="
+                + handlerId
+                + " in graph[" + handlerId + "]");
+        }
+        return node;
+    }
+}
diff --git a/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphManager.java b/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphManager.java
index a77d9de9b..b0803a951 100644
--- a/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphManager.java
+++ b/apm-collector/apm-collector-core/src/main/java/org/skywalking/apm/collector/core/graph/GraphManager.java
@@ -52,4 +52,8 @@ public enum GraphManager {
         }
         return graph;
     }
+
+    public void reset() {
+        allGraphs.clear();
+    }
 }
diff --git a/apm-collector/apm-collector-core/src/test/java/org/skywalking/apm/collector/core/graph/GraphManagerTest.java b/apm-collector/apm-collector-core/src/test/java/org/skywalking/apm/collector/core/graph/GraphManagerTest.java
index c78624333..1d6ed6d58 100644
--- a/apm-collector/apm-collector-core/src/test/java/org/skywalking/apm/collector/core/graph/GraphManagerTest.java
+++ b/apm-collector/apm-collector-core/src/test/java/org/skywalking/apm/collector/core/graph/GraphManagerTest.java
@@ -104,4 +104,18 @@ public class GraphManagerTest {
 
         Next next = GraphManager.INSTANCE.findGraph(5).findNext(3);
     }
+
+    @Test
+    public void testFindNode() {
+        Graph<String> graph = GraphManager.INSTANCE.createIfAbsent(6, String.class);
+        graph.addNode(new Node1Processor()).addNext(new Node2Processor());
+
+        Node<?, Integer> foundNode = GraphManager.INSTANCE.findGraph(6).toBuilder().findNode(2, Integer.class);
+        foundNode.addNext(new Node4Processor());
+    }
+
+    @After
+    public void tearDown() {
+        GraphManager.INSTANCE.reset();
+    }
 }