Track remote trees via special remote refs file
authorGregory Szorc <gps@mozilla.com>
Sun, 21 Jul 2013 23:02:48 -0700
changeset 10 72299b499208efe068499f72f7d8eacb300b178f
parent 9 916cf879e9a7cb3af1a7f167c4195288fc8bfa2d
child 11 a5568c1aa771bcf04b9f0e0a44e6360f39c908bf
push id2
push usergszorc@mozilla.com
push dateMon, 22 Jul 2013 06:02:53 +0000
Track remote trees via special remote refs file
__init__.py
--- a/__init__.py
+++ b/__init__.py
@@ -32,35 +32,61 @@ This extension provides mechanisms to cr
 repository that contains changesets from multiple "upstream" repositories.
 
 The recommended method to create a unified repository is to run `hg
 cloneunified`.
 
 Once you have a unified repository, you can pull changesets from repositories
 by running `hg pulltree`. e.g. `hg pulltree central fx-team` will pull from
 mozilla-central and fx-team.
+
+Remote References
+=================
+
+When pulling from known Gecko repositories, this extension automatically
+creates references to branches on the remote. These can be referenced via
+the revision <tree>/<name>. e.g. 'central/default'. This makes it possible to
+update to revisions on the remote. e.g. `hg up central/default`.
+
+Remote refs are read-only and are updated automatically during repository pull
+and push operations.
+
+This feature is similar to Git remote refs.
 """
 
+import errno
 import os
 import sys
 
 import mercurial.commands as commands
 
 from mercurial.i18n import _
 from mercurial.commands import (
     bookmark,
     pull,
     push,
 )
+from mercurial.localrepo import (
+    repofilecache,
+)
+from mercurial.node import (
+    bin,
+    hex,
+)
 from mercurial import (
     cmdutil,
+    encoding,
     hg,
     util,
 )
 
+from mozautomation.repository import (
+    resolve_uri_to_tree,
+)
+
 
 commands.norepo += ' cloneunified moztrees treestatus'
 cmdtable = {}
 command = cmdutil.command(cmdtable)
 
 colortable = {
     'buildstatus.success': 'green',
     'buildstatus.failed': 'red',
@@ -143,45 +169,39 @@ def treestatus(ui, *trees, **opts):
         ui.write('%s: %s\n' % (tree.rjust(longest), s.status))
 
 
 @command('pulltree', [], _('hg pulltree [TREE] ...'))
 def pulltree(ui, repo, *trees, **opts):
     """Pull changesets from a Mozilla repository into this repository.
 
     Trees can be specified by their common name or aliases (see |hg moztrees|).
-    When a tree is pulled, a bookmark is created with the common name of the
-    tree. This allows updating to specified trees via e.g. |hg up central|.
+    When a tree is pulled, a reference to the current remote heads is created.
+    This allows updating to revisions of remote trees via e.g.
+    |hg up remote/central|.
 
     If no arguments are specified, the main landing trees (central and inbound)
     will be pulled.
     """
     from mozautomation.repository import resolve_trees_to_uris
 
     if not trees:
         trees = ['central', 'inbound']
 
     uris = resolve_trees_to_uris(trees)
 
-    bms = {}
-
     for tree, uri in uris:
         if uri is None:
             ui.warn('Unknown Mozilla repository: %s\n' % tree)
             continue
 
         if pull(ui, repo, uri):
             ui.warn('Error pulling from %s\n' % uri)
             continue
 
-        peer = hg.peer(repo, {}, uri)
-        default = peer.lookup('default')
-
-        bookmark(ui, repo, rev=default, force=True, mark=tree)
-
 
 @command('pushtree',
     [('r', 'rev', 'tip', _('revision'), _('REV'))],
     _('hg pushtree [-r REV] TREE'))
 def pushtree(ui, repo, tree=None, rev=None, **opts):
     """Push changesets to a Mozilla repository.
 
     If only the tree argument is defined, we will attempt to push the current
@@ -200,8 +220,116 @@ def pushtree(ui, repo, tree=None, rev=No
     from mozautomation.repository import resolve_trees_to_uris
 
     tree, uri = resolve_trees_to_uris([tree], write_access=True)[0]
 
     if not uri:
         raise util.Abort("Don't know about tree %s" % tree)
 
     return push(ui, repo, rev=rev, dest=uri)
+
+
+class remoterefs(dict):
+    """Represents a remote refs file."""
+
+    def __init__(self, repo):
+        dict.__init__(self)
+        self._repo = repo
+
+        try:
+            for line in repo.vfs('remoterefs'):
+                line = line.strip()
+                if not line:
+                    continue
+
+                sha, ref = line.split(None, 1)
+                ref = encoding.tolocal(ref)
+                try:
+                    self[ref] = repo.changelog.lookup(sha)
+                except LookupError:
+                    pass
+
+        except IOError as e:
+            if e.errno != errno.ENOENT:
+                raise
+
+    def write(self):
+        f = self._repo.vfs('remoterefs', 'w', atomictemp=True)
+        for ref in sorted(self):
+            f.write('%s %s\n' % (hex(self[ref]), encoding.fromlocal(ref)))
+        f.close()
+
+
+def reposetup(ui, repo):
+    """Custom repository implementation.
+
+    Our custom repository class tracks remote tree references so users can
+    reference specific revisions on remotes.
+    """
+
+    if not repo.local():
+        return
+
+    orig_findtags = repo._findtags
+    orig_lookup = repo.lookup
+    orig_pull = repo.pull
+    orig_push = repo.push
+
+    class remotestrackingrepo(repo.__class__):
+        @repofilecache('remoterefs')
+        def remoterefs(self):
+            return remoterefs(self)
+
+        # Resolve remote ref symbols. For some reason, we need both lookup
+        # and findtags implemented.
+        def lookup(self, key):
+            try:
+                key = self.remoterefs[key]
+            except KeyError, TypeError:
+                pass
+
+            return orig_lookup(key)
+
+        def _findtags(self):
+            tags, tagtypes = orig_findtags()
+            tags.update(self.remoterefs)
+
+            return tags, tagtypes
+
+        def pull(self, remote, *args, **kwargs):
+            # Pulls from known repositories will automatically update our
+            # remote tracking references.
+            res = orig_pull(remote, *args, **kwargs)
+            lock = self.wlock()
+            try:
+                tree = resolve_uri_to_tree(remote.url())
+
+                if tree:
+                    self._update_remote_refs(remote, tree)
+
+            finally:
+                lock.release()
+
+            return res
+
+        def push(self, remote, *args, **kwargs):
+            res = orig_push(remote, *args, **kwargs)
+            lock = self.wlock()
+            try:
+                tree = resolve_uri_to_tree(remote.url())
+
+                if tree:
+                    self._update_remote_refs(remote, tree)
+
+            finally:
+                lock.release()
+
+            return res
+
+        def _update_remote_refs(self, remote, tree):
+            for branch, nodes in remote.branchmap().items():
+                for node in nodes:
+                    self.remoterefs['%s/%s' % (tree, branch)] = node
+
+            self.remoterefs.write()
+
+    repo.__class__ = remotestrackingrepo
+