From a684e53a311abc1b2a2bd19768bf35fc4d34aaf9 Mon Sep 17 00:00:00 2001
From: Jun Wu <quark@fb.com>
Date: Fri, 2 Oct 2020 16:49:04 -0700
Subject: [PATCH] phabstatus: wrap all logrevs with PeekaheadRevsetIter

Summary:
Previously, only the first `__iter__` gets wrapped. With D23095468 (https://github.com/facebookexperimental/eden/commit/34df76813652c73c9ad683700d15d1bce9f4a3e0), the first
`__iter__` is used by the "simplify graph" feature, not the main iteration
loop rendering the graph log output, causing the prefetch feature to fail.

  File "edenscm/mercurial/commands/__init__.py", line 4196, in log
    return cmdutil.graphlog(ui, repo, pats, opts)
    # pats = ()
  File "edenscm/mercurial/cmdutil.py", line 3250, in graphlog
    ui, repo, revdag, displayer, graphmod.asciiedges, getrenamed, filematcher
  File "edenscm/mercurial/cmdutil.py", line 3106, in displaygraph
    rustdisplaygraph(ui, repo, dag, displayer, getrenamed, filematcher, props)
  File "edenscm/mercurial/cmdutil.py", line 3208, in rustdisplaygraph
    for (rev, _type, ctx, parents) in dag:
  File "edenscm/mercurial/graphmod.py", line 63, in dagwalker
    rootnodes = cl.tonodes(revs)
    # revs = <baseset- [7408158, 72057594037927936, ...]
  File "edenscm/mercurial/changelog2.py", line 196, in tonodes
    return self.inner.tonodes(revs)
    # revs = <baseset- [7408158, 72057594037927936, ...]
  File "edenscm/hgext/phabstatus.py", line 281, in next
    return self.__next__()

Lift the "first time" limitation for wrapping `__iter__` to solve the problem.

Reviewed By: simpkins

Differential Revision: D24066918

fbshipit-source-id: 6bbd244e729724e5143147bde60bcb4c8ee4bc80
---
 eden/scm/edenscm/hgext/phabstatus.py | 23 ++++++++---------------
 1 file changed, 8 insertions(+), 15 deletions(-)

diff --git a/eden/scm/edenscm/hgext/phabstatus.py b/eden/scm/edenscm/hgext/phabstatus.py
index e9a239fad2..7ebe426378 100644
--- a/eden/scm/edenscm/hgext/phabstatus.py
+++ b/eden/scm/edenscm/hgext/phabstatus.py
@@ -263,8 +263,8 @@ class PeekaheadRevsetIter(Iterator):
     processing them one at a time as the logging code requests them.
     """
 
-    def __init__(self, revs, chunksize=30):
-        self.mainiter = iter(revs)
+    def __init__(self, revsiter, chunksize=30):
+        self.mainiter = revsiter
         # done is set to true once mainiter has thrown StopIteration
         self.done = False
 
@@ -313,25 +313,18 @@ def _getlogrevs(orig, repo, pats, opts):
     # Call the original function
     revs, expr, filematcher = orig(repo, pats, opts)
 
-    # Wrap the revs result so that iter(revs) returns a PeekaheadRevsetIter()
-    # the first time it is invoked, and sets repo._phabstatusrevs so that the
-    # phabstatus code will be able to peek ahead at the revs to be logged.
-    orig_type = revs.__class__
-
+    # Wrap the revs result so that iter(revs) returns a PeekaheadRevsetIter().
     class wrapped_class(type(revs)):
         def __iter__(self):
-            # The first time __iter__() is called, return a
-            # PeekaheadRevsetIter(), and assign it to repo._phabstatusrevs
-            revs.__class__ = orig_type
             # By default, peek ahead 30 revisions at a time
             peekahead = repo.ui.configint("phabstatus", "logpeekahead", 30)
-            repo._phabstatusrevs = PeekaheadRevsetIter(revs, peekahead)
+            revsiter = super(wrapped_class, self).__iter__()
+            # In case __iter__ is called multiple times, keep the last
+            # wrapped iter accessible for populateresponseforphab (via repo).
+            repo._phabstatusrevs = PeekaheadRevsetIter(revsiter, peekahead)
             return repo._phabstatusrevs
 
-        _is_phabstatus_wrapped = True
-
-    if not hgutil.safehasattr(revs, "_is_phabstatus_wrapped"):
-        revs.__class__ = wrapped_class
+    revs.__class__ = wrapped_class
 
     return revs, expr, filematcher