sapling/treemanifest/__init__.py

# __init__.py
#
# Copyright 2016 Facebook, Inc.
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.

"""allows using and migrating to tree manifests

When autocreatetrees is enabled, you can limit which bookmarks are initially
converted to trees during pull by specifying `treemanifest.allowedtreeroots`.

    [treemanifest]
    allowedtreeroots = master,stable

"""

from mercurial import (
    changegroup,
    cmdutil,
    error,
    extensions,
    localrepo,
    scmutil,
    util,
)
from mercurial.i18n import _
from mercurial.node import bin, nullid

from remotefilelog.contentstore import unioncontentstore
from remotefilelog.datapack import datapackstore, mutabledatapack
from remotefilelog import shallowutil
import ctreemanifest

import struct

cmdtable = {}
command = cmdutil.command(cmdtable)

PACK_CATEGORY='manifests'

def extsetup(ui):
    extensions.wrapfunction(changegroup.cg1unpacker, '_unpackmanifests',
                            _unpackmanifests)

def reposetup(ui, repo):
    wraprepo(repo)

def wraprepo(repo):
    if not isinstance(repo, localrepo.localrepository):
        return

    repo.name = repo.ui.config('remotefilelog', 'reponame')
    if not repo.name:
        raise error.Abort(_("remotefilelog.reponame must be configured"))

    usecdatapack = repo.ui.configbool('remotefilelog', 'fastdatapack')

    packpath = shallowutil.getcachepackpath(repo, PACK_CATEGORY)
    datastore = datapackstore(packpath, usecdatapack=usecdatapack)

    localpackpath = shallowutil.getlocalpackpath(repo.svfs.vfs.base,
                                                 PACK_CATEGORY)
    localdatastore = datapackstore(localpackpath, usecdatapack=usecdatapack)

    repo.svfs.sharedmanifestdatastores = [datastore]
    repo.svfs.localmanifestdatastores = [localdatastore]
    repo.svfs.manifestdatastore = unioncontentstore(localdatastore, datastore,
        writestore=localdatastore)

def _unpackmanifests(orig, self, repo, *args, **kwargs):
    mfrevlog = repo.manifestlog._revlog
    oldtip = len(mfrevlog)

    orig(self, repo, *args, **kwargs)

    if (util.safehasattr(repo.svfs, "manifestdatastore") and
        repo.ui.configbool('treemanifest', 'autocreatetrees')):

        # TODO: only put in cache if pulling from main server
        packpath = shallowutil.getcachepackpath(repo, PACK_CATEGORY)
        opener = scmutil.vfs(packpath)
        with mutabledatapack(repo.ui, opener) as dpack:
            recordmanifest(dpack, repo, oldtip, len(mfrevlog))
            dpack.close()

        # Alert the store that there may be new packs
        repo.svfs.manifestdatastore.markforrefresh()

class InterceptedMutablePack(object):
    """This classes intercepts pack writes and replaces the node for the root
    with the provided node. This is useful for forcing a tree manifest to be
    referencable via its flat hash.
    """
    def __init__(self, pack, node, p1node):
        self._pack = pack
        self._node = node
        self._p1node = p1node

    def add(self, name, node, deltabasenode, delta):
        # For the root node, provide the flat manifest as the key
        if name == "":
            node = self._node
            if deltabasenode != nullid:
                deltabasenode = self._p1node
        return self._pack.add(name, node, deltabasenode, delta)

def recordmanifest(pack, repo, oldtip, newtip):
    mfl = repo.manifestlog
    mfrevlog = mfl._revlog
    total = newtip - oldtip
    ui = repo.ui
    builttrees = {}
    message = _('priming tree cache')
    ui.progress(message, 0, total=total)

    refcount = {}
    for rev in xrange(oldtip, newtip):
        p1 = mfrevlog.parentrevs(rev)[0]
        p1node = mfrevlog.node(p1)
        refcount[p1node] = refcount.get(p1node, 0) + 1

    allowedtreeroots = set()
    for name in repo.ui.configlist('treemanifest', 'allowedtreeroots'):
        if name in repo:
            allowedtreeroots.add(repo[name].manifestnode())

    for rev in xrange(oldtip, newtip):
        ui.progress(message, rev - oldtip, total=total)
        p1 = mfrevlog.parentrevs(rev)[0]
        p1node = mfrevlog.node(p1)

        if p1node == nullid:
            origtree = ctreemanifest.treemanifest(repo.svfs.manifestdatastore)
        elif p1node in builttrees:
            origtree = builttrees[p1node]
        else:
            origtree = mfl[p1node].read()._treemanifest()

        if origtree is None:
            if allowedtreeroots and p1node not in allowedtreeroots:
                continue

            p1mf = mfl[p1node].read()
            origtree = ctreemanifest.treemanifest(repo.svfs.manifestdatastore)
            for filename, node, flag in p1mf.iterentries():
                origtree.set(filename, node, flag)
            origtree.write(InterceptedMutablePack(pack, p1node, nullid))
            builttrees[p1node] = origtree

        # Remove the tree from the cache once we've processed its final use.
        # Otherwise memory explodes
        p1refcount = refcount[p1node] - 1
        if p1refcount == 0:
            builttrees.pop(p1node, None)
        refcount[p1node] = p1refcount

        # This will generally be very quick, since p1 == deltabase
        delta = mfrevlog.revdiff(p1, rev)

        deletes = []
        adds = []

        # Inspect the delta and read the added files from it
        current = 0
        end = len(delta)
        while current < end:
            try:
                block = ''
                # Deltas are of the form:
                #   <start><end><datalen><data>
                # Where start and end say what bytes to delete, and data says
                # what bytes to insert in their place. So we can just read
                # <data> to figure out all the added files.
                byte1, byte2, blocklen = struct.unpack(">lll",
                        delta[current:current + 12])
                current += 12
                if blocklen:
                    block = delta[current:current + blocklen]
                    current += blocklen
            except struct.error:
                raise RuntimeError("patch cannot be decoded")

            # An individual delta block may contain multiple newline delimited
            # entries.
            for line in block.split('\n'):
                if not line:
                    continue
                fname, rest = line.split('\0')
                fnode = rest[:40]
                fflag = rest[40:]
                adds.append((fname, bin(fnode), fflag))

        allfiles = set(repo.changelog.readfiles(mfrevlog.linkrev(rev)))
        deletes = allfiles.difference(fname for fname, fnode, fflag in adds)

        # Apply the changes on top of the parent tree
        newtree = origtree.copy()
        for fname in deletes:
            newtree.set(fname, None, None)

        for fname, fnode, fflags in adds:
            newtree.set(fname, fnode, fflags)

        newtree.write(InterceptedMutablePack(pack, mfrevlog.node(rev), p1node),
                      origtree if p1node != nullid else None)

        if ui.configbool('treemanifest', 'verifyautocreate', True):
            diff = newtree.diff(origtree)
            if len(diff) != len(adds) + len(deletes):
                import pdb
                pdb.set_trace()

            for fname in deletes:
                l, r = diff[fname]
                if l != (None, ''):
                    import pdb
                    pdb.set_trace()
                    pass

            for fname, fnode, fflags in adds:
                l, r = diff[fname]
                if l != (fnode, fflags):
                    import pdb
                    pdb.set_trace()
                    pass
        builttrees[mfrevlog.node(rev)] = newtree

        mfnode = mfrevlog.node(rev)
        if refcount.get(mfnode) > 0:
            builttrees[mfnode] = newtree

    ui.progress(message, None)