203 

   204 

   205 auto-generation of clone bundles

   206 --------------------------------

   207 

   208 It is possible to set Mercurial to automatically re-generate clone bundles when

   209 new content is available.

   210 

   211 Mercurial will take care of the process asynchronously. The defined list of

   212 bundle type will be generated, uploaded, and advertised.

   213 

   214 Bundles Generation:

   215 ...................

   216 

   217 The extension can generate multiple variants of the clone bundle. Each

   218 different variant will be defined by the "bundle-spec" they use::

   219 

   220     [clone-bundles]

   221     auto-generate.formats= zstd-v2, gzip-v2

   222 

   223 See `hg help bundlespec` for details about available options.

   224 

   225 Bundles Upload and Serving:

   226 ...........................

   227 

   228 The generated bundles need to be made available to users through a "public" URL.

   229 This should be donne through `clone-bundles.upload-command` configuration. The

   230 value of this command should be a shell command. It will have access to the

   231 bundle file path through the `$HGCB_BUNDLE_PATH` variable. And the expected

   232 basename in the "public" URL is accessible at::

   233 

   234   [clone-bundles]

   235   upload-command=sftp put $HGCB_BUNDLE_PATH \

   236       sftp://bundles.host/clone-bundles/$HGCB_BUNDLE_BASENAME

   237 

   238 After upload, the file should be available at an url defined by

   239 `clone-bundles.url-template`.

   240 

   241   [clone-bundles]

   242   url-template=https://bundles.host/cache/clone-bundles/{basename}

   245 

   246 import os

   247 import weakref

   248 

   249 from mercurial.i18n import _

   253     commands,

   254     error,

   256     localrepo,

   257     lock,

   258     node,

   259     registrar,

   260     util,

   262 )

   263 

   264 

   265 from mercurial.utils import (

   266     procutil,

   286 

   287 

   288 # logic for bundle auto-generation

   289 

   290 

   291 configtable = {}

   292 configitem = registrar.configitem(configtable)

   293 

   294 cmdtable = {}

   295 command = registrar.command(cmdtable)

   296 

   297 configitem(b'clone-bundles', b'auto-generate.formats', default=list)

   298 

   299 

   300 configitem(b'clone-bundles', b'upload-command', default=None)

   301 

   302 configitem(b'clone-bundles', b'url-template', default=None)

   303 

   304 configitem(b'devel', b'debug.clonebundles', default=False)

   305 

   306 

   307 # category for the post-close transaction hooks

   308 CAT_POSTCLOSE = b"clonebundles-autobundles"

   309 

   310 # template for bundle file names

   311 BUNDLE_MASK = (

   312     b"full-%(bundle_type)s-%(revs)d_revs-%(tip_short)s_tip-%(op_id)s.hg"

   313 )

   314 

   315 

   316 # file in .hg/ use to track clonebundles being auto-generated

   317 AUTO_GEN_FILE = b'clonebundles.auto-gen'

   318 

   319 

   320 class BundleBase(object):

   321     """represents the core of properties that matters for us in a bundle

   322 

   323     :bundle_type: the bundlespec (see hg help bundlespec)

   324     :revs:        the number of revisions in the repo at bundle creation time

   325     :tip_rev:     the rev-num of the tip revision

   326     :tip_node:    the node id of the tip-most revision in the bundle

   327 

   328     :ready:       True if the bundle is ready to be served

   329     """

   330 

   331     ready = False

   332 

   333     def __init__(self, bundle_type, revs, tip_rev, tip_node):

   334         self.bundle_type = bundle_type

   335         self.revs = revs

   336         self.tip_rev = tip_rev

   337         self.tip_node = tip_node

   338 

   339     def valid_for(self, repo):

   340         """is this bundle applicable to the current repository

   341 

   342         This is useful for detecting bundles made irrelevant by stripping.

   343         """

   344         tip_node = node.bin(self.tip_node)

   345         return repo.changelog.index.get_rev(tip_node) == self.tip_rev

   346 

   347     def __eq__(self, other):

   348         left = (self.ready, self.bundle_type, self.tip_rev, self.tip_node)

   349         right = (other.ready, other.bundle_type, other.tip_rev, other.tip_node)

   350         return left == right

   351 

   352     def __neq__(self, other):

   353         return not self == other

   354 

   355     def __cmp__(self, other):

   356         if self == other:

   357             return 0

   358         return -1

   359 

   360 

   361 class RequestedBundle(BundleBase):

   362     """A bundle that should be generated.

   363 

   364     Additional attributes compared to BundleBase

   365     :heads:       list of head revisions (as rev-num)

   366     :op_id:       a "unique" identifier for the operation triggering the change

   367     """

   368 

   369     def __init__(self, bundle_type, revs, tip_rev, tip_node, head_revs, op_id):

   370         self.head_revs = head_revs

   371         self.op_id = op_id

   372         super(RequestedBundle, self).__init__(

   373             bundle_type,

   374             revs,

   375             tip_rev,

   376             tip_node,

   377         )

   378 

   379     @property

   380     def suggested_filename(self):

   381         """A filename that can be used for the generated bundle"""

   382         data = {

   383             b'bundle_type': self.bundle_type,

   384             b'revs': self.revs,

   385             b'heads': self.head_revs,

   386             b'tip_rev': self.tip_rev,

   387             b'tip_node': self.tip_node,

   388             b'tip_short': self.tip_node[:12],

   389             b'op_id': self.op_id,

   390         }

   391         return BUNDLE_MASK % data

   392 

   393     def generate_bundle(self, repo, file_path):

   394         """generate the bundle at `filepath`"""

   395         commands.bundle(

   396             repo.ui,

   397             repo,

   398             file_path,

   399             base=[b"null"],

   400             rev=self.head_revs,

   401             type=self.bundle_type,

   402             quiet=True,

   403         )

   404 

   405     def generating(self, file_path, hostname=None, pid=None):

   406         """return a GeneratingBundle object from this object"""

   407         if pid is None:

   408             pid = os.getpid()

   409         if hostname is None:

   410             hostname = lock._getlockprefix()

   411         return GeneratingBundle(

   412             self.bundle_type,

   413             self.revs,

   414             self.tip_rev,

   415             self.tip_node,

   416             hostname,

   417             pid,

   418             file_path,

   419         )

   420 

   421 

   422 class GeneratingBundle(BundleBase):

   423     """A bundle being generated

   424 

   425     extra attributes compared to BundleBase:

   426 

   427     :hostname: the hostname of the machine generating the bundle

   428     :pid:      the pid of the process generating the bundle

   429     :filepath: the target filename of the bundle

   430 

   431     These attributes exist to help detect stalled generation processes.

   432     """

   433 

   434     ready = False

   435 

   436     def __init__(

   437         self, bundle_type, revs, tip_rev, tip_node, hostname, pid, filepath

   438     ):

   439         self.hostname = hostname

   440         self.pid = pid

   441         self.filepath = filepath

   442         super(GeneratingBundle, self).__init__(

   443             bundle_type, revs, tip_rev, tip_node

   444         )

   445 

   446     @classmethod

   447     def from_line(cls, line):

   448         """create an object by deserializing a line from AUTO_GEN_FILE"""

   449         assert line.startswith(b'PENDING-v1 ')

   450         (

   451             __,

   452             bundle_type,

   453             revs,

   454             tip_rev,

   455             tip_node,

   456             hostname,

   457             pid,

   458             filepath,

   459         ) = line.split()

   460         hostname = util.urlreq.unquote(hostname)

   461         filepath = util.urlreq.unquote(filepath)

   462         revs = int(revs)

   463         tip_rev = int(tip_rev)

   464         pid = int(pid)

   465         return cls(

   466             bundle_type, revs, tip_rev, tip_node, hostname, pid, filepath

   467         )

   468 

   469     def to_line(self):

   470         """serialize the object to include as a line in AUTO_GEN_FILE"""

   471         templ = b"PENDING-v1 %s %d %d %s %s %d %s"

   472         data = (

   473             self.bundle_type,

   474             self.revs,

   475             self.tip_rev,

   476             self.tip_node,

   477             util.urlreq.quote(self.hostname),

   478             self.pid,

   479             util.urlreq.quote(self.filepath),

   480         )

   481         return templ % data

   482 

   483     def __eq__(self, other):

   484         if not super(GeneratingBundle, self).__eq__(other):

   485             return False

   486         left = (self.hostname, self.pid, self.filepath)

   487         right = (other.hostname, other.pid, other.filepath)

   488         return left == right

   489 

   490     def uploaded(self, url, basename):

   491         """return a GeneratedBundle from this object"""

   492         return GeneratedBundle(

   493             self.bundle_type,

   494             self.revs,

   495             self.tip_rev,

   496             self.tip_node,

   497             url,

   498             basename,

   499         )

   500 

   501 

   502 class GeneratedBundle(BundleBase):

   503     """A bundle that is done being generated and can be served

   504 

   505     extra attributes compared to BundleBase:

   506 

   507     :file_url: the url where the bundle is available.

   508     :basename: the "basename" used to upload (useful for deletion)

   509 

   510     These attributes exist to generate a bundle manifest

   511     (.hg/pullbundles.manifest)

   512     """

   513 

   514     ready = True

   515 

   516     def __init__(

   517         self, bundle_type, revs, tip_rev, tip_node, file_url, basename

   518     ):

   519         self.file_url = file_url

   520         self.basename = basename

   521         super(GeneratedBundle, self).__init__(

   522             bundle_type, revs, tip_rev, tip_node

   523         )

   524 

   525     @classmethod

   526     def from_line(cls, line):

   527         """create an object by deserializing a line from AUTO_GEN_FILE"""

   528         assert line.startswith(b'DONE-v1 ')

   529         (

   530             __,

   531             bundle_type,

   532             revs,

   533             tip_rev,

   534             tip_node,

   535             file_url,

   536             basename,

   537         ) = line.split()

   538         revs = int(revs)

   539         tip_rev = int(tip_rev)

   540         file_url = util.urlreq.unquote(file_url)

   541         return cls(bundle_type, revs, tip_rev, tip_node, file_url, basename)

   542 

   543     def to_line(self):

   544         """serialize the object to include as a line in AUTO_GEN_FILE"""

   545         templ = b"DONE-v1 %s %d %d %s %s %s"

   546         data = (

   547             self.bundle_type,

   548             self.revs,

   549             self.tip_rev,

   550             self.tip_node,

   551             util.urlreq.quote(self.file_url),

   552             self.basename,

   553         )

   554         return templ % data

   555 

   556     def manifest_line(self):

   557         """serialize the object to include as a line in pullbundles.manifest"""

   558         templ = b"%s BUNDLESPEC=%s REQUIRESNI=true"

   559         return templ % (self.file_url, self.bundle_type)

   560 

   561     def __eq__(self, other):

   562         if not super(GeneratedBundle, self).__eq__(other):

   563             return False

   564         return self.file_url == other.file_url

   565 

   566 

   567 def parse_auto_gen(content):

   568     """parse the AUTO_GEN_FILE to return a list of Bundle object"""

   569     bundles = []

   570     for line in content.splitlines():

   571         if line.startswith(b'PENDING-v1 '):

   572             bundles.append(GeneratingBundle.from_line(line))

   573         elif line.startswith(b'DONE-v1 '):

   574             bundles.append(GeneratedBundle.from_line(line))

   575     return bundles

   576 

   577 

   578 def dumps_auto_gen(bundles):

   579     """serialize a list of Bundle as a AUTO_GEN_FILE content"""

   580     lines = []

   581     for b in bundles:

   582         lines.append(b"%s\n" % b.to_line())

   583     lines.sort()

   584     return b"".join(lines)

   585 

   586 

   587 def read_auto_gen(repo):

   588     """read the AUTO_GEN_FILE for the <repo> a list of Bundle object"""

   589     data = repo.vfs.tryread(AUTO_GEN_FILE)

   590     if not data:

   591         return []

   592     return parse_auto_gen(data)

   593 

   594 

   595 def write_auto_gen(repo, bundles):

   596     """write a list of Bundle objects into the repo's AUTO_GEN_FILE"""

   597     assert repo._cb_lock_ref is not None

   598     data = dumps_auto_gen(bundles)

   599     with repo.vfs(AUTO_GEN_FILE, mode=b'wb', atomictemp=True) as f:

   600         f.write(data)

   601 

   602 

   603 def generate_manifest(bundles):

   604     """write a list of Bundle objects into the repo's AUTO_GEN_FILE"""

   605     bundles = list(bundles)

   606     bundles.sort(key=lambda b: b.bundle_type)

   607     lines = []

   608     for b in bundles:

   609         lines.append(b"%s\n" % b.manifest_line())

   610     return b"".join(lines)

   611 

   612 

   613 def update_ondisk_manifest(repo):

   614     """update the clonebundle manifest with latest url"""

   615     with repo.clonebundles_lock():

   616         bundles = read_auto_gen(repo)

   617 

   618         per_types = {}

   619         for b in bundles:

   620             if not (b.ready and b.valid_for(repo)):

   621                 continue

   622             current = per_types.get(b.bundle_type)

   623             if current is not None and current.revs >= b.revs:

   624                 continue

   625             per_types[b.bundle_type] = b

   626         manifest = generate_manifest(per_types.values())

   627         with repo.vfs(

   628             bundlecaches.CB_MANIFEST_FILE, mode=b"wb", atomictemp=True

   629         ) as f:

   630             f.write(manifest)

   631 

   632 

   633 def update_bundle_list(repo, new_bundles=(), del_bundles=()):

   634     """modify the repo's AUTO_GEN_FILE

   635 

   636     This method also regenerates the clone bundle manifest when needed"""

   637     with repo.clonebundles_lock():

   638         bundles = read_auto_gen(repo)

   639         if del_bundles:

   640             bundles = [b for b in bundles if b not in del_bundles]

   641         new_bundles = [b for b in new_bundles if b not in bundles]

   642         bundles.extend(new_bundles)

   643         write_auto_gen(repo, bundles)

   644         all_changed = []

   645         all_changed.extend(new_bundles)

   646         all_changed.extend(del_bundles)

   647         if any(b.ready for b in all_changed):

   648             update_ondisk_manifest(repo)

   649 

   650 

   651 def cleanup_tmp_bundle(repo, target):

   652     """remove a GeneratingBundle file and entry"""

   653     assert not target.ready

   654     with repo.clonebundles_lock():

   655         repo.vfs.tryunlink(target.filepath)

   656         update_bundle_list(repo, del_bundles=[target])

   657 

   658 

   659 def finalize_one_bundle(repo, target):

   660     """upload a generated bundle and advertise it in the clonebundles.manifest"""

   661     with repo.clonebundles_lock():

   662         bundles = read_auto_gen(repo)

   663         if target in bundles and target.valid_for(repo):

   664             result = upload_bundle(repo, target)

   665             update_bundle_list(repo, new_bundles=[result])

   666     cleanup_tmp_bundle(repo, target)

   667 

   668 

   669 def upload_bundle(repo, bundle):

   670     """upload the result of a GeneratingBundle and return a GeneratedBundle

   671 

   672     The upload is done using the `clone-bundles.upload-command`

   673     """

   674     cmd = repo.ui.config(b'clone-bundles', b'upload-command')

   675     url = repo.ui.config(b'clone-bundles', b'url-template')

   676     basename = repo.vfs.basename(bundle.filepath)

   677     filepath = procutil.shellquote(bundle.filepath)

   678     variables = {

   679         b'HGCB_BUNDLE_PATH': filepath,

   680         b'HGCB_BUNDLE_BASENAME': basename,

   681     }

   682     env = procutil.shellenviron(environ=variables)

   683     ret = repo.ui.system(cmd, environ=env)

   684     if ret:

   685         raise error.Abort(b"command returned status %d: %s" % (ret, cmd))

   686     url = (

   687         url.decode('utf8')

   688         .format(basename=basename.decode('utf8'))

   689         .encode('utf8')

   690     )

   691     return bundle.uploaded(url, basename)

   692 

   693 

   694 def auto_bundle_needed_actions(repo, bundles, op_id):

   695     """find the list of bundles that need action

   696 

   697     returns a list of RequestedBundle objects that need to be generated and

   698     uploaded."""

   699     create_bundles = []

   700     repo = repo.filtered(b"immutable")

   701     targets = repo.ui.configlist(b'clone-bundles', b'auto-generate.formats')

   702     revs = len(repo.changelog)

   703     generic_data = {

   704         'revs': revs,

   705         'head_revs': repo.changelog.headrevs(),

   706         'tip_rev': repo.changelog.tiprev(),

   707         'tip_node': node.hex(repo.changelog.tip()),

   708         'op_id': op_id,

   709     }

   710     for t in targets:

   711         data = generic_data.copy()

   712         data['bundle_type'] = t

   713         b = RequestedBundle(**data)

   714         create_bundles.append(b)

   715     return create_bundles

   716 

   717 

   718 def start_one_bundle(repo, bundle):

   719     """start the generation of a single bundle file

   720 

   721     the `bundle` argument should be a RequestedBundle object.

   722 

   723     This data is passed to the `debugmakeclonebundles` "as is".

   724     """

   725     data = util.pickle.dumps(bundle)

   726     cmd = [procutil.hgexecutable(), b'--cwd', repo.path, INTERNAL_CMD]

   727     env = procutil.shellenviron()

   728     msg = b'clone-bundles: starting bundle generation: %s\n'

   729     stdout = None

   730     stderr = None

   731     waits = []

   732     record_wait = None

   733     if repo.ui.configbool(b'devel', b'debug.clonebundles'):

   734         stdout = procutil.stdout

   735         stderr = procutil.stderr

   736         repo.ui.write(msg % bundle.bundle_type)

   737         record_wait = waits.append

   738     else:

   739         repo.ui.debug(msg % bundle.bundle_type)

   740     bg = procutil.runbgcommand

   741     bg(

   742         cmd,

   743         env,

   744         stdin_bytes=data,

   745         stdout=stdout,

   746         stderr=stderr,

   747         record_wait=record_wait,

   748     )

   749     for f in waits:

   750         f()

   751 

   752 

   753 INTERNAL_CMD = b'debug::internal-make-clone-bundles'

   754 

   755 

   756 @command(INTERNAL_CMD, [], b'')

   757 def debugmakeclonebundles(ui, repo):

   758     """Internal command to auto-generate debug bundles"""

   759     requested_bundle = util.pickle.load(procutil.stdin)

   760     procutil.stdin.close()

   761 

   762     fname = requested_bundle.suggested_filename

   763     fpath = repo.vfs.makedirs(b'tmp-bundles')

   764     fpath = repo.vfs.join(b'tmp-bundles', fname)

   765     bundle = requested_bundle.generating(fpath)

   766     update_bundle_list(repo, new_bundles=[bundle])

   767 

   768     requested_bundle.generate_bundle(repo, fpath)

   769 

   770     repo.invalidate()

   771     finalize_one_bundle(repo, bundle)

   772 

   773 

   774 def make_auto_bundler(source_repo):

   775     reporef = weakref.ref(source_repo)

   776 

   777     def autobundle(tr):

   778         repo = reporef()

   779         assert repo is not None

   780         bundles = read_auto_gen(repo)

   781         new = auto_bundle_needed_actions(repo, bundles, b"%d_txn" % id(tr))

   782         for data in new:

   783             start_one_bundle(repo, data)

   784         return None

   785 

   786     return autobundle

   787 

   788 

   789 def reposetup(ui, repo):

   790     """install the two pieces needed for automatic clonebundle generation

   791 

   792     - add a "post-close" hook that fires bundling when needed

   793     - introduce a clone-bundle lock to let multiple processes meddle with the

   794       state files.

   795     """

   796     if not repo.local():

   797         return

   798 

   799     class autobundlesrepo(repo.__class__):

   800         def transaction(self, *args, **kwargs):

   801             tr = super(autobundlesrepo, self).transaction(*args, **kwargs)

   802             targets = repo.ui.configlist(

   803                 b'clone-bundles', b'auto-generate.formats'

   804             )

   805             if targets:

   806                 tr.addpostclose(CAT_POSTCLOSE, make_auto_bundler(self))

   807             return tr

   808 

   809         @localrepo.unfilteredmethod

   810         def clonebundles_lock(self, wait=True):

   811             '''Lock the repository file related to clone bundles'''

   812             if not util.safehasattr(self, '_cb_lock_ref'):

   813                 self._cb_lock_ref = None

   814             l = self._currentlock(self._cb_lock_ref)

   815             if l is not None:

   816                 l.lock()

   817                 return l

   818 

   819             l = self._lock(

   820                 vfs=self.vfs,

   821                 lockname=b"clonebundleslock",

   822                 wait=wait,

   823                 releasefn=None,

   824                 acquirefn=None,

   825                 desc=_(b'repository %s') % self.origroot,

   826             )

   827             self._cb_lock_ref = weakref.ref(l)

   828             return l

   829 

   830     repo._wlockfreeprefix.add(AUTO_GEN_FILE)

   831     repo._wlockfreeprefix.add(bundlecaches.CB_MANIFEST_FILE)

   832     repo.__class__ = autobundlesrepo