diff --git a/java/src/main/java/org/softwareheritage/graph/Entry.java b/java/src/main/java/org/softwareheritage/graph/Entry.java --- a/java/src/main/java/org/softwareheritage/graph/Entry.java +++ b/java/src/main/java/org/softwareheritage/graph/Entry.java @@ -126,24 +126,24 @@ close(); } - public void visit_nodes(String direction, String edgesFmt, long srcNodeId, long max_edges) { + public void visit_nodes(String direction, String edgesFmt, long srcNodeId, long maxEdges) { open(); - Traversal t = new Traversal(this.graph, direction, edgesFmt); - t.visitNodesVisitor(srcNodeId, this::writeNode, max_edges); + Traversal t = new Traversal(this.graph, direction, edgesFmt, maxEdges); + t.visitNodesVisitor(srcNodeId, this::writeNode); close(); } - public void visit_edges(String direction, String edgesFmt, long srcNodeId, long max_edges) { + public void visit_edges(String direction, String edgesFmt, long srcNodeId, long maxEdges) { open(); - Traversal t = new Traversal(this.graph, direction, edgesFmt); - t.visitNodesVisitor(srcNodeId, null, this::writeEdge, max_edges); + Traversal t = new Traversal(this.graph, direction, edgesFmt, maxEdges); + t.visitNodesVisitor(srcNodeId, null, this::writeEdge); close(); } - public void visit_paths(String direction, String edgesFmt, long srcNodeId, long max_edges) { + public void visit_paths(String direction, String edgesFmt, long srcNodeId, long maxEdges) { open(); - Traversal t = new Traversal(this.graph, direction, edgesFmt); - t.visitPathsVisitor(srcNodeId, this::writePath, max_edges); + Traversal t = new Traversal(this.graph, direction, edgesFmt, maxEdges); + t.visitPathsVisitor(srcNodeId, this::writePath); close(); } diff --git a/java/src/main/java/org/softwareheritage/graph/Traversal.java b/java/src/main/java/org/softwareheritage/graph/Traversal.java --- a/java/src/main/java/org/softwareheritage/graph/Traversal.java +++ b/java/src/main/java/org/softwareheritage/graph/Traversal.java @@ -30,6 +30,9 @@ /** Number of edges accessed during traversal */ long nbEdgesAccessed; + /** The anti Dos limit of edges traversed while a visit */ + long maxEdges; + /** random number generator, for random walks */ Random rng; @@ -43,6 +46,10 @@ * edges */ public Traversal(Graph graph, String direction, String edgesFmt) { + this(graph, direction, edgesFmt, 0); + } + + public Traversal(Graph graph, String direction, String edgesFmt, long maxEdges) { if (!direction.matches("forward|backward")) { throw new IllegalArgumentException("Unknown traversal direction: " + direction); } @@ -57,6 +64,7 @@ this.visited = new HashSet<>(); this.parentNode = new HashMap<>(); this.nbEdgesAccessed = 0; + this.maxEdges = maxEdges; this.rng = new Random(); } @@ -146,7 +154,7 @@ /** * Push version of {@link #visitNodes}: will fire passed callback on each visited node. */ - public void visitNodesVisitor(long srcNodeId, NodeIdConsumer nodeCb, EdgeIdConsumer edgeCb, long max_edges) { + public void visitNodesVisitor(long srcNodeId, NodeIdConsumer nodeCb, EdgeIdConsumer edgeCb) { Stack stack = new Stack<>(); this.nbEdgesAccessed = 0; @@ -159,8 +167,8 @@ nodeCb.accept(currentNodeId); } nbEdgesAccessed += graph.outdegree(currentNodeId); - if (max_edges > 0) { - if (nbEdgesAccessed >= max_edges) { + if (this.maxEdges > 0) { + if (nbEdgesAccessed >= this.maxEdges) { break; } } @@ -177,14 +185,9 @@ } } - /** Two argument version for count_visitor */ - public void visitNodesVisitor(long srcNodeId, NodeIdConsumer cb) { - visitNodesVisitor(srcNodeId, cb, null, 0); - } - /** One-argument version to handle callbacks properly */ - public void visitNodesVisitor(long srcNodeId, NodeIdConsumer cb, long max_edges) { - visitNodesVisitor(srcNodeId, cb, null, max_edges); + public void visitNodesVisitor(long srcNodeId, NodeIdConsumer cb) { + visitNodesVisitor(srcNodeId, cb, null); } /** @@ -195,7 +198,7 @@ */ public ArrayList visitNodes(long srcNodeId) { ArrayList nodeIds = new ArrayList<>(); - visitNodesVisitor(srcNodeId, nodeIds::add, 0); + visitNodesVisitor(srcNodeId, nodeIds::add); return nodeIds; } @@ -203,10 +206,10 @@ * Push version of {@link #visitPaths}: will fire passed callback on each discovered (complete) * path. */ - public void visitPathsVisitor(long srcNodeId, PathConsumer cb, long max_edges) { + public void visitPathsVisitor(long srcNodeId, PathConsumer cb) { Stack currentPath = new Stack<>(); this.nbEdgesAccessed = 0; - visitPathsInternalVisitor(srcNodeId, currentPath, cb, max_edges); + visitPathsInternalVisitor(srcNodeId, currentPath, cb); } /** @@ -217,26 +220,25 @@ */ public ArrayList> visitPaths(long srcNodeId) { ArrayList> paths = new ArrayList<>(); - visitPathsVisitor(srcNodeId, paths::add, 0); + visitPathsVisitor(srcNodeId, paths::add); return paths; } - private void visitPathsInternalVisitor(long currentNodeId, Stack currentPath, PathConsumer cb, - long max_edges) { + private void visitPathsInternalVisitor(long currentNodeId, Stack currentPath, PathConsumer cb) { currentPath.push(currentNodeId); long visitedNeighbors = 0; nbEdgesAccessed += graph.outdegree(currentNodeId); - if (max_edges > 0) { - if (nbEdgesAccessed >= max_edges) { + if (this.maxEdges > 0) { + if (nbEdgesAccessed >= this.maxEdges) { currentPath.pop(); return; } } LazyLongIterator it = graph.successors(currentNodeId, edges); for (long neighborNodeId; (neighborNodeId = it.nextLong()) != -1;) { - visitPathsInternalVisitor(neighborNodeId, currentPath, cb, max_edges); + visitPathsInternalVisitor(neighborNodeId, currentPath, cb); visitedNeighbors++; }