[PATCH 1 of 1 hg-bfiles] sshstore: reimplemented adding a new bfserve command

david.douard at logilab.fr david.douard at logilab.fr
Thu Nov 19 14:44:52 CST 2009


# HG changeset patch
# User David Douard <david.douard at logilab.fr>
# Date 1258091865 -3600
# Node ID baea47f3f1139faeda56e61d3f7963cb311da8d3
# Parent  22fb77911970030b6db67d089e02c8b9290b49f6
sshstore: reimplemented adding a new bfserve command

New implementation of the sshstore using a new hg command, bfserve, so we do not need to create a useless hg repo in the remote store anymore, and we do not need to hack hg's sshserver anymore.

diff --git a/bfiles.py b/bfiles.py
--- a/bfiles.py
+++ b/bfiles.py
@@ -1,6 +1,7 @@
 '''track large binary files'''
 
 import os
+import sys
 import errno
 import re
 import hashlib
@@ -13,7 +14,7 @@
 
 from mercurial import \
      util, extensions, dirstate, cmdutil, match, url as url_, node, error
-from mercurial import sshserver
+from mercurial import commands
 from mercurial.i18n import _
 
 # -- Commands ----------------------------------------------------------
@@ -382,6 +383,13 @@
     return store.verify(revs, contents=opts.get('contents'))
 
 
+def bfserve(ui, **opts):
+    """export the bfile store via SSH
+    """
+    s = sshstoreserver(ui)
+    s.serve_forever()
+
+
 # -- Wrappers: modify existing commands --------------------------------
 
 def reposetup(ui, repo):
@@ -421,103 +429,13 @@
                     # someone might delete a big file before committing it
                     pass
                 _move_pending(ui, repo, bfdirstate, ctx, filename, realfile)
-                
+
         bfdirstate.write()
         return node
 
     extensions.wrapfunction(repo, 'status', localrepo_status)
     extensions.wrapfunction(repo, 'commitctx', localrepo_commitctx)
 
-    serversetup()
-
-def serversetup():
-    """
-    Monkeypatch sshserver to add bfiles protocol support.
-    """
-
-    # -- wrap sshserver.do_hello ---------------------------------------
-    # dirty hack to be able to add 'bfstore' capability to sshserve
-    def _sshserver_do_hello_wrapper(origfn, self):
-        respond = self.respond
-        caps = []
-        self.respond = caps.append
-        origfn(self)
-        self.respond = respond
-        caps.append('bfstore')
-        allcaps = ' '.join([x.strip() for x in caps])
-        self.respond("%s\n" % allcaps)
-
-    # much cleaner version, if sshserver have the new getcapabilities method 
-    def _sshserver_getcapabilities_wrapper(origfn, self):
-        caps = origfn(self)
-        caps.append('bfstore')
-        return caps
-
-    if hasattr(sshserver.sshserver, "getcapabilities"):
-        extensions.wrapfunction(sshserver.sshserver, "getcapabilities",
-                                _sshserver_getcapabilities_wrapper)
-    else: # no support for getcapabilities, use good ol' dirty hack
-        extensions.wrapfunction(sshserver.sshserver, "do_hello",
-                                _sshserver_do_hello_wrapper)
-
-
-    # -- add sshserver.do_xxx commands ---------------------------------
-    def _sshserver_do_bfput(self):
-        """Shim this function into the sshserver so that it responds to
-        the bfput command.
-        """
-        key, fname = self.getarg()
-        if os.path.exists(fname):
-            self.respond('file %s already exists' % fname)
-            return
-        destdir = os.path.dirname(fname)
-        util.makedirs(destdir)
-        try:
-            fd = open(fname, "wb")
-        except IOError:
-            self.respond('cannot create file')
-        self.respond("")
-
-        try:
-            count = int(self.fin.readline())
-            while count:
-                fd.write(self.fin.read(count))
-                count = int(self.fin.readline())
-            self.respond('')
-        finally:
-            #
-            fd.close()
-
-    def _sshserver_do_bfget(self):
-        """Shim this function into the sshserver so that it responds to
-        the bfget command.
-        """
-        key, fname = self.getarg()
-        if not os.path.isfile(fname):
-            self.respond('file %s does not exists' % fname)
-            return
-        try:
-            fd = open(fname, "rb")
-        except IOError:
-            self.respond('cannot read file')
-            return
-        size = util.fstat(fd).st_size
-
-        self.respond(str(size))
-
-        try:
-            while 1:
-                d = fd.read(4096)
-                if not d:
-                    break
-                self.respond(d)
-            self.respond('')
-        finally:
-            #
-            fd.close()
-
-    sshserver.sshserver.do_bfput = _sshserver_do_bfput
-    sshserver.sshserver.do_bfget = _sshserver_do_bfget
 
 # -- Private worker functions ------------------------------------------
 
@@ -607,7 +525,7 @@
 
     (outfd, tmp_name) = tempfile.mkstemp(dir=pending_dir)
     tmpfile = os.fdopen(outfd, 'w')
-   
+
     ui.debug('copying %s to %s\n' % (filename, tmp_name))
     bhash = _copy_and_hash(infile, tmpfile)
     hhash = binascii.hexlify(bhash)
@@ -918,7 +836,7 @@
 
         write(_('searching %d changesets for big files\n') % len(revs))
         verified = set()                # set of (filename, filenode) tuples
-        
+
         for rev in revs:
             cctx = self.repo[rev]
             cset = "%d:%s" % (cctx.rev(), node.short(cctx.node()))
@@ -1000,27 +918,7 @@
         remotecmd = self.ui.config("ui", "remotecmd", "hg")
 
         args = util.sshargs(sshcmd, self.host, self.user, self.port)
-        # -- XXX *snip* --
-
-        try:
-            self.validate_repo(ui, sshcmd, args, remotecmd)
-        except error.RepoError, err:
-            # if validate_repo fails, it means that we must create a
-            # repo in the remote store. This repo is only used to be
-            # able to run a sshserver on the remote store.
-            #
-            # XXX hg's sshserve could be fixed to be able to run
-            # without a remote repo.
-            ui.debug('validate_repo failed: %s\n' % err)
-            self.cleanup()
-            cmd = '%s %s "%s init %s"'
-            cmd = cmd % (sshcmd, args, remotecmd, self.path)
-
-            ui.debug(_('running %s\n') % cmd)
-            res = util.system(cmd)
-            if res != 0:
-                self.abort(error.RepoError(_("could not create remote repo")))
-            self.validate_repo(ui, sshcmd, args, remotecmd)
+        self.validate_connection(ui, sshcmd, args, remotecmd)
 
     def put(self, source, filename, hash):
         destdir = os.path.join(self.path, filename)
@@ -1042,6 +940,7 @@
             try:
                 self.recvfile(dest, storefile)
                 success.append((filename, hash))
+
             except error.RepoError:
                 missing.append((filename, hash))
         return (success, missing)
@@ -1086,45 +985,33 @@
         if wsize != size:
             self.abort(error.RepoError(_('get failed, file size does not match (%s vs %s)') % (size, wsize)))
 
-    # XXX the following methods were copied and slightly modified from
-    # mercurial's sshrepo.sshrepository class!
-
-    # XXX *snip* method copied from sshrepository and modified XXX
-    def validate_repo(self, ui, sshcmd, args, remotecmd):
+    # XXX *snip* method heavily based on sshrepository.validate_repo
+    def validate_connection(self, ui, sshcmd, args, remotecmd):
         # cleanup up previous run
         self.cleanup()
 
-        cmd = '%s %s "%s -R %s serve --stdio"'
-        cmd = cmd % (sshcmd, args, remotecmd, self.path)
+        cmd = '%s %s "%s bfserve"'
+        cmd = cmd % (sshcmd, args, remotecmd)
 
         cmd = util.quotecommand(cmd)
         ui.debug(_('running %s\n') % cmd)
         self.pipeo, self.pipei, self.pipee = util.popen3(cmd)
 
         # skip any noise generated by remote shell
-        self.do_cmd("hello")
-        r = self.do_cmd("between", pairs=("%s-%s" % ("0"*40, "0"*40)))
+        r = self.do_cmd("hello")
         lines = ["", "dummy"]
         max_noise = 500
         while lines[-1] and max_noise:
             l = r.readline()
             self.readerr()
-            if lines[-1] == "1\n" and l == "\n":
+            if lines[-1] == "6\n" and l == "hello\n":
                 break
             if l:
                 ui.debug(_("remote: "), l)
             lines.append(l)
             max_noise -= 1
         else:
-            self.abort(error.RepoError(_("no suitable response from remote hg")))
-
-        self.capabilities = set()
-        for l in reversed(lines):
-            if l.startswith("capabilities:"):
-                self.capabilities.update(l[:-1].split(":")[1].split())
-                break
-        if "bfstore" not in self.capabilities:
-            self.abort(util.Abort(_("remote hg does not support bfiles")))
+            self.abort(util.Abort(_("no suitable response from remote hg; check that remote hg supports bfiles")))
 
     # XXX *snip* method copied from sshrepository
     def readerr(self):
@@ -1137,18 +1024,17 @@
 
     # XXX *snip* method copied from sshrepository
     def abort(self, exception):
-        self.cleanup(reporterr=self.ui.debugflag)
+        self.cleanup()
         raise exception
 
-    # XXX *snip* method copied from sshrepository and modified
-    def cleanup(self, reporterr=True):
+    # XXX *snip* method copied from sshrepository
+    def cleanup(self):
         try:
             self.pipeo.close()
             self.pipei.close()
             # read the error descriptor until EOF
-            if reporterr:
-                for l in self.pipee:
-                    self.ui.status(_("remote: "), l)
+            for l in self.pipee:
+                self.ui.status(_("remote: "), l)
             self.pipee.close()
         except:
             pass
@@ -1196,7 +1082,7 @@
         raise NotImplementedError('sorry, HTTP PUT not implemented yet')
 
     def get(self, files):
-        
+
         success = []
         missing = []
         ui = self.ui
@@ -1250,6 +1136,109 @@
     'https': httpstore,
     }
 
+# -- Private helper store classes --------------------------------------------
+
+# XXX *snip* class heavily based on sshserver.
+
+# We do *not* subclass it here cause it does not have as many 'do_xxx'
+# commands as there are in sshserver. More, several methods here have
+# simpler implementation than sshserver's ones. Last, this
+# sshstoreserver do not require to have a remote hg repository.
+
+class sshstoreserver(object):
+    """A simple ssh server serving files for bfile's sshstore
+    """
+    def __init__(self, ui):
+        self.ui = ui
+        self.fin = sys.stdin
+        self.fout = sys.stdout
+
+        sys.stdout = sys.stderr
+
+        # Prevent insertion/deletion of CRs
+        util.set_binary(self.fin)
+        util.set_binary(self.fout)
+
+    def getarg(self):
+        argline = self.fin.readline()[:-1]
+        arg, l = argline.split()
+        val = self.fin.read(int(l))
+        return arg, val
+
+    def respond(self, v):
+        self.fout.write("%d\n" % len(v))
+        self.fout.write(v)
+        self.fout.flush()
+
+    def serve_forever(self):
+        while self.serve_one(): pass
+        sys.exit(0)
+
+    def serve_one(self):
+        cmd = self.fin.readline()[:-1]
+        if cmd:
+            impl = getattr(self, 'do_' + cmd, None)
+            if impl: impl()
+            else: self.respond("")
+        return cmd != ''
+
+    def do_hello(self):
+        self.respond("hello\n")
+
+    def do_bfput(self):
+        """respond to the bfput command: send a file
+        """
+        key, fname = self.getarg()
+        if os.path.exists(fname):
+            self.respond('file %s already exists' % fname)
+            return
+        destdir = os.path.dirname(fname)
+        util.makedirs(destdir)
+        try:
+            fd = open(fname, "wb")
+        except IOError:
+            self.respond('cannot create file')
+        self.respond("")
+
+        try:
+            count = int(self.fin.readline())
+            while count:
+                fd.write(self.fin.read(count))
+                count = int(self.fin.readline())
+            self.respond('')
+        finally:
+            fd.close()
+
+    def do_bfget(self):
+        """respond to the bfget command: receive a file
+        """
+        key, fname = self.getarg()
+        if not os.path.isfile(fname):
+            self.respond('file %s does not exists' % fname)
+            return
+        try:
+            fd = open(fname, "rb")
+        except IOError:
+            self.respond('cannot read file')
+            return
+        size = util.fstat(fd).st_size
+
+        self.respond(str(size))
+
+        try:
+            while 1:
+                d = fd.read(4096)
+                if not d:
+                    break
+                self.respond(d)
+            self.respond('')
+        finally:
+            fd.close()
+
+
+# -- hg commands declarations ------------------------------------------------
+
+commands.norepo += " bfserve"
 
 cmdtable = {
     'bfadd'    : (bfadd,
@@ -1285,4 +1274,7 @@
                    ('c', 'contents', False,
                     _('check file contents, not just existence'))],
                   _('hg bfverify [--all] [--contents]')),
+    'bfserve' : (bfserve,
+                 [],
+                 ''),
     }
diff --git a/tests/test-sshstore b/tests/test-sshstore
--- a/tests/test-sshstore
+++ b/tests/test-sshstore
@@ -15,9 +15,7 @@
 echo "% setup"
 tar -xf $TESTDIR/test-store.tar
 
-# HACK! the store has to be a repo so "hg serve" can be run from there
 storedir=$PWD/store
-hg init $storedir
 
 hg clone -q -U $TESTDIR/test-repo1.bundle repo1
 cd repo1
@@ -28,19 +26,8 @@
 store = ssh://$user@localhost/$storedir/
 __EOF__
 
-# Try to hit the store repo without configuring bfiles there, and make
-# sure we get a comprehensible error message.
-echo "% unconfigured store"
+echo "% rev 1: get big1, big2"
 hg update 1
-hg bfupdate -v
-
-# Now configure the store repo for actual use.
-cat > $storedir/.hg/hgrc << __EOF__
-[extensions]
-bfiles = $TESTDIR/../bfiles.py
-__EOF__
-
-echo "% rev 1: get big1, big2"
 hg bfstatus
 hg bfupdate -v          # get big1, big2
 hg bfstatus
@@ -60,3 +47,13 @@
 hg bfstatus
 hg bfupdate -v
 hg bfstatus
+
+echo "% new rev: modify, commmit and put big1"
+echo -n x >> big1
+hg bfrefresh -v
+hg commit -v -d "0 0" -m"modify big1"
+
+# Upload (put) all pending committed revisions to the central store.
+hg bfstatus
+hg bfput -v
+hg bfstatus
diff --git a/tests/test-sshstore.out b/tests/test-sshstore.out
--- a/tests/test-sshstore.out
+++ b/tests/test-sshstore.out
@@ -1,12 +1,8 @@
 % setup
-hg init $HGTMP/test-sshstore/store
 hg clone -q -U $TESTDIR/test-repo1.bundle repo1
-% unconfigured store
+% rev 1: get big1, big2
 hg update 1
 4 files updated, 0 files merged, 0 files removed, 0 files unresolved
-hg bfupdate -v
-abort: remote hg does not support bfiles
-% rev 1: get big1, big2
 hg bfstatus
 B-! big1
 B-! big2
@@ -51,3 +47,15 @@
 getting sub/space file
 3 big files updated, 0 removed
 hg bfstatus
+% new rev: modify, commmit and put big1
+hg bfrefresh -v
+big1
+1 big files refreshed
+hg commit -v -d 0 0 -mmodify big1
+.hgbfiles/big1
+committed changeset 5:df9365d7bd33
+hg bfstatus
+BPC big1
+hg bfput -v
+putting big1 (rev 16fded083628dcddefee3e6e966dfac69f395435)
+hg bfstatus


More information about the Mercurial-devel mailing list