treemanifest: server - write trees at the same time as flat manifests

Summary:
When a server is a treemanifest enabled server, we want to ensure that every
commit that comes in will be added to the treemanifest as well as the flat
manifest. This patch accomplishes that by hooking into the addmanifest logic and
applying the new manifest to the tree as well.

Test Plan: Added a test

Reviewers: #mercurial, simonfar

Reviewed By: simonfar

Subscribers: mjpieters

Differential Revision: https://phabricator.intern.facebook.com/D4093738

Signature: t1:4093738:1478260442:58f77bd88ab09ccda4ac238a065388be8f4a271d

tests/test-treemanifest-server.t

@@ -0,0 +1,43 @@
+  $ . "$TESTDIR/library.sh"
+
+  $ PYTHONPATH=$TESTDIR/..:$PYTHONPATH
+  $ export PYTHONPATH
+
+  $ hginit master
+  $ cd master
+  $ cat >> .hg/hgrc <<EOF
+  > [extensions]
+  > bundle2hooks=
+  > pushrebase=
+  > treemanifest=
+  > [treemanifest]
+  > server=True
+  > [remotefilelog]
+  > server=True
+  > EOF
+  $ mkdir subdir
+  $ echo x > subdir/x
+  $ hg commit -qAm 'add subdir/x'
+  $ hg book mybook
+  $ hg debugdata .hg/store/00manifesttree.i 0
+  subdir\x00bc0c2c938b929f98b1c31a8c5994396ebb096bf0t (esc)
+  $ cd ..
+
+  $ hgcloneshallow ssh://user@dummy/master client -q
+  1 files fetched over 1 fetches - (1 misses, 0.00% hit ratio) over * (glob)
+  $ cd client
+  $ cat >> .hg/hgrc <<EOF
+  > [extensions]
+  > pushrebase=
+  > EOF
+  $ mkdir subdir2
+  $ echo z >> subdir2/z
+  $ hg commit -qAm "add subdir2/z"
+  $ hg push --to mybook
+  pushing to ssh://user@dummy/master
+  searching for changes
+  remote: pushing 1 changset:
+  remote:     15486e46ccf6  add subdir2/z
+  $ ls ../master/.hg/store/meta
+  subdir
+  subdir2

treemanifest/__init__.py

@@ -68,6 +68,8 @@ def extsetup(ui):
 
     wrappropertycache(localrepo.localrepository, 'manifestlog', getmanifestlog)
 
+    extensions.wrapfunction(manifest.memmanifestctx, 'write', _writemanifest)
+
 def reposetup(ui, repo):
     wraprepo(repo)
 
@@ -75,7 +77,8 @@ def wraprepo(repo):
     if not isinstance(repo, localrepo.localrepository):
         return
 
-    if not repo.ui.configbool('treemanifest', 'server'):
+    repo.svfs.treemanifestserver = repo.ui.configbool('treemanifest', 'server')
+    if not repo.svfs.treemanifestserver:
         clientreposetup(repo)
 
 def clientreposetup(repo):
@@ -155,6 +158,52 @@ def getmanifestlog(orig, self):
 
     return mfl
 
+def _writemanifest(orig, self, transaction, link, p1, p2, added, removed):
+    n = orig(self, transaction, link, p1, p2, added, removed)
+
+    if not self._manifestlog._revlog.opener.treemanifestserver:
+        return n
+
+    # Since we're adding the root flat manifest, let's add the corresponding
+    # root tree manifest.
+    mfl = self._manifestlog
+    treemfl = mfl.treemanifestlog
+
+    m = self._manifestdict
+
+    parentflat = mfl[p1].read()
+    diff = parentflat.diff(m)
+
+    newtree = treemfl[p1].read().copy()
+    added = []
+    removed = []
+    for filename, (old, new) in diff.iteritems():
+        if new is not None and new[0] is not None:
+            added.append(filename)
+            newtree[filename] = new[0]
+            newtree.setflag(filename, new[1])
+        else:
+            removed.append(filename)
+            del newtree[filename]
+
+    try:
+        treemfrevlog = treemfl._revlog
+        oldaddrevision = treemfrevlog.addrevision
+        def addusingnode(*args, **kwargs):
+            newkwargs = kwargs.copy()
+            newkwargs['node'] = n
+            return oldaddrevision(*args, **newkwargs)
+        treemfrevlog.addrevision = addusingnode
+
+        def readtree(dir, node):
+            return treemfl.get(dir, node).read()
+        treemfrevlog.add(newtree, transaction, link, p1, p2, added, removed,
+                         readtree=readtree)
+    finally:
+        del treemfrevlog.__dict__['addrevision']
+
+    return n
+
 @command('backfilltree', [
     ('l', 'limit', '10000000', _(''))
     ], _('hg backfilltree [OPTIONS]'))