remotefilelog: move stream_out wireprotocol to core

Reviewed By: DurhamG

Differential Revision: D21011398

fbshipit-source-id: c47307d06517e145837854db8e6bbec97b17e6bd

eden/scm/edenscm/hgext/hggit/gitrepo.py

@@ -97,7 +97,7 @@ class gitrepo(peerrepository):
         def peer(self):
             return self
 
-        def stream_out(self):
+        def stream_out(self, shallow=False):
             raise NotImplementedError
 
         def unbundle(self):

eden/scm/edenscm/hgext/remotefilelog/__init__.py

@@ -368,31 +368,8 @@ def cloneshallow(orig, ui, repo, *args, **opts):
 
         wrapfunction(exchange, "pull", pull_shallow)
 
-        # Wrap the stream logic to add requirements and to pass include/exclude
-        # patterns around.
-        def setup_streamout(repo, remote):
-            # Replace remote.stream_out with a version that sends file
-            # patterns.
-            def stream_out_shallow(orig):
-                caps = shallowutil.peercapabilities(remote)
-                if shallowrepo.requirement in caps:
-                    opts = {}
-                    if repo.ui.configbool("treemanifest", "treeonly"):
-                        opts["noflatmanifest"] = "True"
-                    return remote._callstream("stream_out_shallow", **opts)
-                else:
-                    return orig()
-
-            wrapfunction(remote, "stream_out", stream_out_shallow)
-
         if hasstreamclone:
 
-            def stream_wrap(orig, op):
-                setup_streamout(op.repo, op.remote)
-                return orig(op)
-
-            wrapfunction(streamclone, "maybeperformlegacystreamclone", stream_wrap)
-
             def canperformstreamclone(orig, *args, **kwargs):
                 supported, requirements = orig(*args, **kwargs)
                 if requirements is not None:
@@ -403,7 +380,6 @@ def cloneshallow(orig, ui, repo, *args, **opts):
         else:
 
             def stream_in_shallow(orig, repo, remote, requirements):
-                setup_streamout(repo, remote)
                 requirements.add(shallowrepo.requirement)
                 return orig(repo, remote, requirements)
 

eden/scm/edenscm/hgext/remotefilelog/remotefilelogserver.py

@@ -64,37 +64,7 @@ def onetimesetup(ui):
     wireproto.commands["getpackv1"] = (getpackv1, "*")
     wireproto.commands["getpackv2"] = (getpackv2, "*")
 
-    class streamstate(object):
-        shallowremote = False
-        noflatmf = False
-
-    state = streamstate()
-
-    def stream_out_shallow(repo, proto, other):
-        oldshallow = state.shallowremote
-        oldnoflatmf = state.noflatmf
-        try:
-            state.shallowremote = True
-            state.noflatmf = other.get("noflatmanifest") == "True"
-            streamres = wireproto.stream(repo, proto)
-
-            # Force the first value to execute, so the file list is computed
-            # within the try/finally scope
-            first = next(streamres.gen)
-            second = next(streamres.gen)
-
-            def gen():
-                yield first
-                yield second
-                for value in streamres.gen:
-                    yield value
-
-            return wireproto.streamres(gen())
-        finally:
-            state.shallowremote = oldshallow
-            state.noflatmf = oldnoflatmf
-
-    wireproto.commands["stream_out_shallow"] = (stream_out_shallow, "*")
+    state = wireproto._streamstate
 
     # don't clone filelogs to shallow clients
     def _walkstreamfiles(orig, repo):

eden/scm/edenscm/mercurial/localrepo.py

@@ -264,7 +264,7 @@ class localpeer(repository.peer):
     def pushkey(self, namespace, key, old, new):
         return self._repo.pushkey(namespace, key, old, new)
 
-    def stream_out(self):
+    def stream_out(self, shallow=False):
         raise errormod.Abort(_("cannot perform stream clone against local " "peer"))
 
     def unbundle(self, cg, heads, url):

eden/scm/edenscm/mercurial/repository.py

@@ -172,7 +172,7 @@ class _basewirecommands(pycompat.ABC):
         """
 
     @abc.abstractmethod
-    def stream_out(self):
+    def stream_out(self, shallow=False):
         """Obtain streaming clone data.
 
         Successful result should be a generator of data chunks.

eden/scm/edenscm/mercurial/streamclone.py

@@ -134,7 +134,9 @@ def maybeperformlegacystreamclone(pullop):
 
     repo.ui.status(_("streaming all changes\n"))
 
-    fp = remote.stream_out()
+    shallow = pullop.extras.get("shallow", False)
+
+    fp = remote.stream_out(shallow=shallow)
     l = fp.readline()
     try:
         resp = int(l)

eden/scm/edenscm/mercurial/wireproto.py

@@ -364,12 +364,22 @@ class wirepeer(repository.legacypeer):
             self.ui.status(_("remote: "), l)
         yield d
 
-    def stream_out(self):
+    def stream_out(self, shallow=False):
+        fullclone = self.ui.configbool("clone", "requestfullclone")
+        if shallow and fullclone:
+            raise error.Abort(
+                _("--shallow is incompatible with clone.requestfullclone")
+            )
+        if shallow and self.capable("remotefilelog"):
+            opts = {}
+            if self.ui.configbool("treemanifest", "treeonly"):
+                opts["noflatmanifest"] = "True"
+            return self._callstream("stream_out_shallow", **opts)
         if self.capable("stream_option"):
-            fullclone = self.ui.configbool("clone", "requestfullclone")
             args = {"fullclone": str(fullclone)}
             return self._callstream("stream_out_option", **args)
-        return self._callstream("stream_out")
+        else:
+            return self._callstream("stream_out")
 
     def getbundle(self, source, **kwargs):
         kwargs = kwargs
@@ -1181,6 +1191,41 @@ def stream(repo, proto):
         return b"2\n"
 
 
+class streamstate(object):
+    shallowremote = False
+    noflatmf = False
+
+
+# XXX: Consider move this to another state.
+_streamstate = streamstate()
+
+
+@wireprotocommand("stream_out_shallow", "*")
+def stream_out_shallow(repo, proto, other):
+    oldshallow = _streamstate.shallowremote
+    oldnoflatmf = _streamstate.noflatmf
+    try:
+        _streamstate.shallowremote = True
+        _streamstate.noflatmf = other.get("noflatmanifest") == "True"
+        s = stream(repo, proto)
+
+        # Force the first value to execute, so the file list is computed
+        # within the try/finally scope
+        first = next(s.gen)
+        second = next(s.gen)
+
+        def gen():
+            yield first
+            yield second
+            for value in s.gen:
+                yield value
+
+        return streamres(gen())
+    finally:
+        _streamstate.shallowremote = oldshallow
+        _streamstate.noflatmf = oldnoflatmf
+
+
 @wireprotocommand("stream_out_option", "*")
 def streamoption(repo, proto, options):
     if (