comparison mercurial/bundle2.py @ 21019:3dc09f831a2e

bundle2: lazy unbundle of part payload

The `unbundle` part gains a `read` method to retrieve payload content.
This method behaves as a python file-like read method.

The bundle-processing code is updated to make sure a part is fully consumed before
another one is extracted.

Test output changes because the debug output is even more interleaved now.

21018:c848bfd02366

    # todo:
    # - replace this is a init function soon.
    # - exception catching
    unbundler.params
    iterparts = iter(unbundler)
    try:
        for part in iterparts:
            parttype = part.type
            # part key are matched lower case
            key = parttype.lower()

                if key != parttype: # mandatory parts
                    # todo:
                    # - use a more precise exception
                    raise
                op.ui.debug('ignoring unknown advisory part %r\n' % key)
                # todo:
                # - consume the part once we use streaming
                continue

            # handler is called outside the above try block so that we don't
            # risk catching KeyErrors from anything other than the
            # parthandlermapping lookup (any KeyError raised by handler()
            # itself represents a defect of a different variety).
            handler(op, part)
    except Exception:
        for part in iterparts:
            pass # consume the bundle content
        raise
    return op

class bundle20(object):
    """represent an outgoing bundle2 container

        super(unbundlepart, self).__init__(fp)
        self.ui = ui
        # unbundle state attr
        self._headerdata = header
        self._headeroffset = 0
        # part data
        self.id = None
        self.type = None
        self.mandatoryparams = None
        self.advisoryparams = None
        self.data = None
        self._readdata()

    def _fromheader(self, size):
        """return the next <size> byte from the header"""
        offset = self._headeroffset
        data = self._headerdata[offset:(offset + size)]
        self._headeroffset += size
        return data

    def _unpackheader(self, format):
        """read given format from header

        This automatically compute the size of the format to read."""
        data = self._fromheader(struct.calcsize(format))
        return _unpack(format, data)

    def _readdata(self):
        """read the header and setup the object"""
        # some utility to help reading from the header block

        typesize = self._unpackheader(_fparttypesize)[0]
        self.type = self._fromheader(typesize)
        self.ui.debug('part type: "%s"\n' % self.type)
        self.id = self._unpackheader(_fpartid)[0]
        self.ui.debug('part id: "%s"\n' % self.id)

        for key, value in advsizes:
            advparams.append((self._fromheader(key), self._fromheader(value)))
        self.mandatoryparams = manparams
        self.advisoryparams  = advparams
        ## part payload
        payload = []
        payloadsize = self._unpack(_fpayloadsize)[0]
        self.ui.debug('payload chunk size: %i\n' % payloadsize)
        while payloadsize:
            payload.append(self._readexact(payloadsize))
            payloadsize = self._unpack(_fpayloadsize)[0]
            self.ui.debug('payload chunk size: %i\n' % payloadsize)
        self.data = ''.join(payload)

@parthandler('changegroup')
def handlechangegroup(op, inpart):
    """apply a changegroup part on the repo

    #
    # The addchangegroup function will get a transaction object by itself, but
    # we need to make sure we trigger the creation of a transaction object used
    # for the whole processing scope.
    op.gettransaction()
    data = StringIO.StringIO(inpart.data)
    data.seek(0)
    cg = changegroup.readbundle(data, 'bundle2part')
    ret = changegroup.addchangegroup(op.repo, cg, 'bundle2', 'bundle2')
    op.records.add('changegroup', {'return': ret})
    if op.reply is not None:

        # return. But one need to start somewhere.
        part = bundlepart('reply:changegroup', (),
                           [('in-reply-to', str(inpart.id)),
                            ('return', '%i' % ret)])
        op.reply.addpart(part)

@parthandler('reply:changegroup')
def handlechangegroup(op, inpart):
    p = dict(inpart.advisoryparams)
    ret = int(p['return'])

21019:3dc09f831a2e

    # todo:
    # - replace this is a init function soon.
    # - exception catching
    unbundler.params
    iterparts = iter(unbundler)
    part = None
    try:
        for part in iterparts:
            parttype = part.type
            # part key are matched lower case
            key = parttype.lower()

                if key != parttype: # mandatory parts
                    # todo:
                    # - use a more precise exception
                    raise
                op.ui.debug('ignoring unknown advisory part %r\n' % key)
                # consuming the part
                part.read()
                continue

            # handler is called outside the above try block so that we don't
            # risk catching KeyErrors from anything other than the
            # parthandlermapping lookup (any KeyError raised by handler()
            # itself represents a defect of a different variety).
            handler(op, part)
            part.read()
    except Exception:
        if part is not None:
            # consume the bundle content
            part.read()
        for part in iterparts:
            # consume the bundle content
            part.read()
        raise
    return op

class bundle20(object):
    """represent an outgoing bundle2 container

        super(unbundlepart, self).__init__(fp)
        self.ui = ui
        # unbundle state attr
        self._headerdata = header
        self._headeroffset = 0
        self._initialized = False
        self.consumed = False
        # part data
        self.id = None
        self.type = None
        self.mandatoryparams = None
        self.advisoryparams = None
        self._payloadstream = None
        self._readheader()

    def _fromheader(self, size):
        """return the next <size> byte from the header"""
        offset = self._headeroffset
        data = self._headerdata[offset:(offset + size)]
        self._headeroffset = offset + size
        return data

    def _unpackheader(self, format):
        """read given format from header

        This automatically compute the size of the format to read."""
        data = self._fromheader(struct.calcsize(format))
        return _unpack(format, data)

    def _readheader(self):
        """read the header and setup the object"""
        typesize = self._unpackheader(_fparttypesize)[0]
        self.type = self._fromheader(typesize)
        self.ui.debug('part type: "%s"\n' % self.type)
        self.id = self._unpackheader(_fpartid)[0]
        self.ui.debug('part id: "%s"\n' % self.id)

        for key, value in advsizes:
            advparams.append((self._fromheader(key), self._fromheader(value)))
        self.mandatoryparams = manparams
        self.advisoryparams  = advparams
        ## part payload
        def payloadchunks():
            payloadsize = self._unpack(_fpayloadsize)[0]
            self.ui.debug('payload chunk size: %i\n' % payloadsize)
            while payloadsize:
                yield self._readexact(payloadsize)
                payloadsize = self._unpack(_fpayloadsize)[0]
                self.ui.debug('payload chunk size: %i\n' % payloadsize)
        self._payloadstream = util.chunkbuffer(payloadchunks())
        # we read the data, tell it
        self._initialized = True

    def read(self, size=None):
        """read payload data"""
        if not self._initialized:
            self._readheader()
        if size is None:
            data = self._payloadstream.read()
        else:
            data = self._payloadstream.read(size)
        if size is None or len(data) < size:
            self.consumed = True
        return data


@parthandler('changegroup')
def handlechangegroup(op, inpart):
    """apply a changegroup part on the repo

    #
    # The addchangegroup function will get a transaction object by itself, but
    # we need to make sure we trigger the creation of a transaction object used
    # for the whole processing scope.
    op.gettransaction()
    data = StringIO.StringIO(inpart.read())
    data.seek(0)
    cg = changegroup.readbundle(data, 'bundle2part')
    ret = changegroup.addchangegroup(op.repo, cg, 'bundle2', 'bundle2')
    op.records.add('changegroup', {'return': ret})
    if op.reply is not None:

        # return. But one need to start somewhere.
        part = bundlepart('reply:changegroup', (),
                           [('in-reply-to', str(inpart.id)),
                            ('return', '%i' % ret)])
        op.reply.addpart(part)
    assert not inpart.read()

@parthandler('reply:changegroup')
def handlechangegroup(op, inpart):
    p = dict(inpart.advisoryparams)
    ret = int(p['return'])