Repository: angr/angr

Issue:
`test_veritesting.TestVeritesting.test_veritesting_a` is flaky
Example failure from CI:
```
Traceback (most recent call last):
  File "/__w/claripy/claripy/build/src/angr/tests/test_veritesting.py", line 56, in test_veritesting_a
    self._run_veritesting_a('x86_64')
  File "/__w/claripy/claripy/build/src/angr/tests/test_veritesting.py", line 31, in _run_veritesting_a
    ex.explore(find=addresses_veritesting_a[arch])
  File "/__w/claripy/claripy/build/src/angr/angr/sim_manager.py", line 268, in explore
    self.run(stash=stash, n=n, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/sim_manager.py", line 298, in run
    self.step(stash=stash, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/misc/hookset.py", line 75, in __call__
    result = current_hook(self.func.__self__, *args, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/exploration_techniques/explorer.py", line 96, in step
    return simgr.step(stash=stash, extra_stop_points=base_extra_stop_points | self._extra_stop_points, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/misc/hookset.py", line 80, in __call__
    return self.func(*args, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/sim_manager.py", line 383, in step
    successors = self.step_state(state, successor_func=successor_func, **run_args)
  File "/__w/claripy/claripy/build/src/angr/angr/misc/hookset.py", line 75, in __call__
    result = current_hook(self.func.__self__, *args, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/exploration_techniques/veritesting.py", line 21, in step_state
    vt = self.project.analyses.Veritesting(state, **self.options)
  File "/__w/claripy/claripy/build/src/angr/angr/analyses/analysis.py", line 138, in __call__
    r = w(*args, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/analyses/analysis.py", line 122, in wrapper
    oself.__init__(*args, **kwargs)
  File "/__w/claripy/claripy/build/src/angr/angr/analyses/veritesting.py", line 227, in __init__
    self.result, self.final_manager = self._veritesting()
  File "/__w/claripy/claripy/build/src/angr/angr/analyses/veritesting.py", line 239, in _veritesting
    new_manager = self._execute_and_merge(s)
  File "/__w/claripy/claripy/build/src/angr/angr/analyses/veritesting.py", line 347, in _execute_and_merge
    manager = self._join_merge_points(manager, merge_points)
  File "/__w/claripy/claripy/build/src/angr/angr/analyses/veritesting.py", line 415, in _join_merge_points
    manager.merge(stash='merge_tmp')
  File "/__w/claripy/claripy/build/src/angr/angr/sim_manager.py", line 690, in merge
    m = self._merge_states(g) if merge_func is None else merge_func(*g)
  File "/__w/claripy/claripy/build/src/angr/angr/sim_manager.py", line 743, in _merge_states
    optimal, common_history, others = self._hierarchy.most_mergeable(states)
  File "/__w/claripy/claripy/build/src/angr/angr/state_hierarchy.py", line 208, in most_mergeable
    for n in networkx.algorithms.dfs_postorder_nodes(self._graph):
  File "/__w/claripy/claripy/build/virtualenv/lib/python3.6/site-packages/networkx/algorithms/traversal/depth_first_search.py", line 290, in <genexpr>
    return (v for u, v, d in edges if d == "reverse")
  File "/__w/claripy/claripy/build/virtualenv/lib/python3.6/site-packages/networkx/algorithms/traversal/depth_first_search.py", line 427, in dfs_labeled_edges
    child = next(children)
RuntimeError: dictionary changed size during iteration
```


Discussion:

Resolving patch:
diff --git a/angr/state_hierarchy.py b/angr/state_hierarchy.py
--- a/angr/state_hierarchy.py
+++ b/angr/state_hierarchy.py
@@ -1,14 +1,21 @@
 import logging
 import weakref
-import networkx
 import itertools
+from contextlib import contextmanager
+import gc
+import networkx
 
 import claripy
 
 l = logging.getLogger(name=__name__)
 
 
-class StateHierarchy(object):
+class StateHierarchy:
+    """
+    The state hierarchy holds weak references to SimStateHistory objects in a directed acyclic graph. It is useful
+    for queries about a state's ancestry, notably "what is the best ancestor state for a merge among these states" and
+    "what is the most recent unsatisfiable state while using LAZY_SOLVES"
+    """
     def __init__(self):
 
         # The New Order
@@ -17,8 +24,11 @@ def __init__(self):
         self._twigs = set() # nodes with one child
         self._weakref_cache = {} # map from object id to weakref
         self._reverse_weakref_cache = {} # map from weakref to object id
+        self._pending_cleanup = set()
+        self._defer_cleanup = False
 
     def __getstate__(self):
+        gc.collect()
         histories = [ h() for h in networkx.algorithms.dfs_postorder_nodes(self._graph) ]
         return (histories,)
 
@@ -35,64 +45,82 @@ def __setstate__(self, s):
 
     def get_ref(self, obj):
         if id(obj) not in self._weakref_cache:
-            ref = weakref.ref(obj, self.clear_ref)
+            ref = weakref.ref(obj, self.dead_ref)
             self._weakref_cache[id(obj)] = ref
             self._reverse_weakref_cache[ref] = id(obj)
             return ref
         else:
             return self._weakref_cache[id(obj)]
 
-    def clear_ref(self, ref):
+    def dead_ref(self, ref):
+        if self._defer_cleanup:
+            self._pending_cleanup.add(ref)
+        else:
+            self._cleanup_ref(ref)
+
+    def _cleanup_ref(self, ref):
         if ref not in self._reverse_weakref_cache:
             l.error("Cleaning mystery weakref %s", ref)
             return
 
         self._remove_history(ref)
 
-        # TODO: this nonsense is very much not thread safe
         del self._weakref_cache[self._reverse_weakref_cache[ref]]
         del self._reverse_weakref_cache[ref]
 
+    @contextmanager
+    def defer_cleanup(self):
+        old_defer, self._defer_cleanup = self._defer_cleanup, True
+        yield
+        self._defer_cleanup = old_defer
+        if not self._defer_cleanup:
+            toclean = list(self._pending_cleanup)
+            self._pending_cleanup.clear()
+            for ref in toclean:
+                self._cleanup_ref(ref)
+
     #
     # Graph management
     #
 
     def _remove_history(self, h):
-        try:
-            predecessors = self._graph.predecessors(h)
-            successors = self._graph.successors(h)
+        with self.defer_cleanup():
+            try:
+                predecessors = self._graph.predecessors(h)
+                successors = self._graph.successors(h)
 
-            for p,s in itertools.product(predecessors, successors):
-                self._graph.add_edge(p, s)
+                for p,s in itertools.product(predecessors, successors):
+                    self._graph.add_edge(p, s)
 
-            self._graph.remove_node(h)
-        except networkx.NetworkXError:
-            pass
+                self._graph.remove_node(h)
+            except networkx.NetworkXError:
+                pass
 
-        self._leaves.discard(h)
-        self._twigs.discard(h)
-        hh = h()
-        if hh is not None:
-            hh.demote()
+            self._leaves.discard(h)
+            self._twigs.discard(h)
+            hh = h()
+            if hh is not None:
+                hh.demote()
 
     def add_state(self, s):
         h = s.history
         self.add_history(h)
 
     def add_history(self, h):
-        cur_node = self.get_ref(h)
-        self._graph.add_node(cur_node)
-        if h.parent is not None:
-            prev_node = self.get_ref(h.parent)
-            self._graph.add_edge(prev_node, cur_node)
-
-            self._leaves.discard(prev_node)
-            if len(list(self._graph.successors(prev_node))) == 1:
-                self._twigs.add(prev_node)
-            else:
-                self._twigs.discard(prev_node)
+        with self.defer_cleanup():
+            cur_node = self.get_ref(h)
+            self._graph.add_node(cur_node)
+            if h.parent is not None:
+                prev_node = self.get_ref(h.parent)
+                self._graph.add_edge(prev_node, cur_node)
 
-        self._leaves.add(cur_node)
+                self._leaves.discard(prev_node)
+                if len(list(self._graph.successors(prev_node))) == 1:
+                    self._twigs.add(prev_node)
+                else:
+                    self._twigs.discard(prev_node)
+
+            self._leaves.add(cur_node)
 
     def simplify(self):
         tw = self._twigs
@@ -101,38 +129,44 @@ def simplify(self):
             self._remove_history(h)
 
     def full_simplify(self):
-        for h in self._graph.nodes():
-            if self._graph.out_degree(h) == 1:
-                self._remove_history(h)
+        with self.defer_cleanup():
+            for h in self._graph.nodes():
+                if self._graph.out_degree(h) == 1:
+                    self._remove_history(h)
 
     def lineage(self, h):
         """
         Returns the lineage of histories leading up to `h`.
         """
 
-        lineage = [ ]
+        with self.defer_cleanup():
+            lineage = [ ]
 
-        predecessors = list(self._graph.predecessors(h))
-        while len(predecessors):
-            lineage.append(predecessors[0])
-            predecessors = list(self._graph.predecessors(predecessors[0]))
+            predecessors = list(self._graph.predecessors(h))
+            while len(predecessors):
+                lineage.append(predecessors[0])
+                predecessors = list(self._graph.predecessors(predecessors[0]))
 
-        lineage.reverse()
-        return lineage
+            lineage.reverse()
+            return lineage
 
     def all_successors(self, h):
-        nodes = list(networkx.algorithms.dfs_postorder_nodes(self._graph, h))[:-1]
-        nodes.reverse()
-        return nodes
+        with self.defer_cleanup():
+            nodes = list(networkx.algorithms.dfs_postorder_nodes(self._graph, h))[:-1]
+            nodes.reverse()
+            return nodes
 
     def history_successors(self, h):
-        return [ ref() for ref in self._graph.successors(self.get_ref(h)) ]
+        with self.defer_cleanup():
+            return [ ref() for ref in self._graph.successors(self.get_ref(h)) ]
 
     def history_predecessors(self, h):
-        return [ ref() for ref in self._graph.predecessors(self.get_ref(h)) ]
+        with self.defer_cleanup():
+            return [ ref() for ref in self._graph.predecessors(self.get_ref(h)) ]
 
     def history_contains(self, h):
-        return self.get_ref(h) in self._graph
+        with self.defer_cleanup():
+            return self.get_ref(h) in self._graph
 
     #
     # LAZY_SOLVES support
@@ -150,7 +184,7 @@ def _find_root_unreachable(self, h):
             l.debug("... looking between %d and %d in %d states", good, bad, len(lineage))
             cur = (bad+good)//2
 
-            if cur == good or cur == bad:
+            if cur in (good, bad):
                 if lineage[bad]().reachable():
                     bad += 1
 
@@ -165,16 +199,17 @@ def _find_root_unreachable(self, h):
                 bad = cur
 
     def _prune_subtree(self, h):
-        ph = list(self._graph.predecessors(h))
-        if len(ph) == 1 and len(list(self._graph.successors(ph[0]))) == 1:
-            self._twigs.add(ph[0])
-
-        all_children = list(networkx.algorithms.dfs_postorder_nodes(self._graph, h))
-        for n in all_children:
-            n()._satisfiable = False
-            if n().state is not None:
-                n().state.add_constraints(claripy.false)
-        self._graph.remove_nodes_from(all_children)
+        with self.defer_cleanup():
+            ph = list(self._graph.predecessors(h))
+            if len(ph) == 1 and len(list(self._graph.successors(ph[0]))) == 1:
+                self._twigs.add(ph[0])
+
+            all_children = list(networkx.algorithms.dfs_postorder_nodes(self._graph, h))
+            for n in all_children:
+                n()._satisfiable = False
+                if n().state is not None:
+                    n().state.add_constraints(claripy.false)
+            self._graph.remove_nodes_from(all_children)
 
     def unreachable_state(self, state):
         self.unreachable_history(state.history)
@@ -200,19 +235,20 @@ def most_mergeable(self, states):
         Find the "most mergeable" set of states from those provided.
 
         :param states: a list of states
-        :returns: a tuple of: (a list of states to merge, those states' common history, a list of states to not merge yet)
+        :returns: a tuple of: (list of states to merge, those states' common history, list of states to not merge yet)
         """
 
-        histories = set(self.get_ref(s.history) for s in states)
+        with self.defer_cleanup():
+            histories = set(self.get_ref(s.history) for s in states)
 
-        for n in networkx.algorithms.dfs_postorder_nodes(self._graph):
-            intersection = histories.intersection(self.all_successors(n))
-            if len(intersection) > 1:
-                return (
-                    [ s for s in states if self.get_ref(s.history) in intersection ],
-                    n(),
-                    [ s for s in states if self.get_ref(s.history) not in intersection ]
-                )
+            for n in networkx.algorithms.dfs_postorder_nodes(self._graph):
+                intersection = histories.intersection(self.all_successors(n))
+                if len(intersection) > 1:
+                    return (
+                        [ s for s in states if self.get_ref(s.history) in intersection ],
+                        n(),
+                        [ s for s in states if self.get_ref(s.history) not in intersection ]
+                    )
 
-        # didn't find any?
-        return set(), None, states
+            # didn't find any?
+            return set(), None, states


Test patch: