comparison mercurial/wireprotoframing.py @ 40138:b5bf3dd6ec5b

wireprotov2: send content encoded frames from server

Now that we have support for negotiating encodings and configuring
an encoder, we can start sending content encoded frames from the
server.

This commit teaches the wireprotov2 server code to send content
encoded frames.

On the mozilla-unified repository with zstd enabled peers, this change
reduces the total amount of data transferred from server to client
drastically:

befor: 7,190,995,812 bytes
after: 1,605,508,691 bytes

Differential Revision: https://phab.mercurial-scm.org/D4927

40137:ed4ebbb98ca0

                break

def createcommandresponseokframe(stream, requestid):
    overall = b''.join(cborutil.streamencode({b'status': b'ok'}))

    return stream.makeframe(requestid=requestid,
                            typeid=FRAME_TYPE_COMMAND_RESPONSE,
                            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
                            payload=overall)

def createcommandresponseeosframe(stream, requestid):
    """Create an empty payload frame representing command end-of-stream."""
    return stream.makeframe(requestid=requestid,
                            typeid=FRAME_TYPE_COMMAND_RESPONSE,
                            flags=FLAG_COMMAND_RESPONSE_EOS,
                            payload=b'')

def createalternatelocationresponseframe(stream, requestid, location):
    data = {
        b'status': b'redirect',
        b'location': {

              r'servercadercerts'):
        value = getattr(location, a)
        if value is not None:
            data[b'location'][pycompat.bytestr(a)] = value

    return stream.makeframe(requestid=requestid,
                            typeid=FRAME_TYPE_COMMAND_RESPONSE,
                            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
                            payload=b''.join(cborutil.streamencode(data)))

def createcommanderrorresponse(stream, requestid, message, args=None):
    # TODO should this be using a list of {'msg': ..., 'args': {}} so atom
    # formatting works consistently?
    m = {

        if data is None:
            for frame in self._flush():
                yield frame
            return

        # There is a ton of potential to do more complicated things here.
        # Our immediate goal is to coalesce small chunks into big frames,
        # not achieve the fewest number of frames possible. So we go with
        # a simple implementation:

                yield self._stream.makeframe(
                    self._requestid,
                    typeid=FRAME_TYPE_COMMAND_RESPONSE,
                    flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
                    payload=chunk)

                if offset == len(data):
                    return

        # If we don't have enough to constitute a full frame, buffer and

        yield self._stream.makeframe(
            self._requestid,
            typeid=FRAME_TYPE_COMMAND_RESPONSE,
            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
            payload=payload)

# TODO consider defining encoders/decoders using the util.compressionengine
# mechanism.

class identityencoder(object):

class outputstream(stream):
    """Represents a stream used for sending data."""

    def __init__(self, streamid, active=False):
        super(outputstream, self).__init__(streamid, active=active)
        self._encoder = None

    def setencoder(self, ui, name):
        """Set the encoder for this stream.

        Receives the stream profile name.
        """
        if name not in STREAM_ENCODERS:
            raise error.Abort(_('unknown stream encoder: %s') % name)

        self._encoder = STREAM_ENCODERS[name][0](ui)

    def encode(self, data):
        if not self._encoder:
            return data

    def finish(self):
        if not self._encoder:
            return b''

        self._encoder.finish()

def ensureserverstream(stream):
    if stream.streamid % 2:
        raise error.ProgrammingError('server should only write to even '
                                     'numbered streams; %d is not even' %

                except StopIteration:
                    for frame in emitter.send(None):
                        yield frame

                    if emitted:
                        yield createcommandresponseeosframe(stream, requestid)
                    break

                except error.WireprotoCommandError as e:
                    for frame in createcommanderrorresponse(
                        stream, requestid, e.message, e.messageargs):

                        if emitted:
                            raise error.ProgrammingError(
                                'alternatelocationresponse seen after initial '
                                'output object')

                        yield createalternatelocationresponseframe(
                            stream, requestid, o)

                        alternatelocationsent = True
                        emitted = True

                    if alternatelocationsent:
                        raise error.ProgrammingError(
                            'object follows alternatelocationresponse')

                    if not emitted:
                        yield createcommandresponseokframe(stream, requestid)
                        emitted = True

                    # Objects emitted by command functions can be serializable
                    # data structures or special types.
                    # TODO consider extracting the content normalization to a

            self._activecommands.remove(requestid)

        return self._handlesendframes(sendframes())

    def makeoutputstream(self):
        """Create a stream to be used for sending data to the client."""
        streamid = self._nextoutgoingstreamid
        self._nextoutgoingstreamid += 2

        s = outputstream(streamid)
        self._outgoingstreams[streamid] = s

        return s

    def _makeerrorresult(self, msg):
        return 'error', {

40138:b5bf3dd6ec5b

                break

def createcommandresponseokframe(stream, requestid):
    overall = b''.join(cborutil.streamencode({b'status': b'ok'}))

    if stream.streamsettingssent:
        overall = stream.encode(overall)
        encoded = True

        if not overall:
            return None
    else:
        encoded = False

    return stream.makeframe(requestid=requestid,
                            typeid=FRAME_TYPE_COMMAND_RESPONSE,
                            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
                            payload=overall,
                            encoded=encoded)

def createcommandresponseeosframes(stream, requestid,
                                   maxframesize=DEFAULT_MAX_FRAME_SIZE):
    """Create an empty payload frame representing command end-of-stream."""
    payload = stream.flush()

    offset = 0
    while True:
        chunk = payload[offset:offset + maxframesize]
        offset += len(chunk)

        done = offset == len(payload)

        if done:
            flags = FLAG_COMMAND_RESPONSE_EOS
        else:
            flags = FLAG_COMMAND_RESPONSE_CONTINUATION

        yield stream.makeframe(requestid=requestid,
                               typeid=FRAME_TYPE_COMMAND_RESPONSE,
                               flags=flags,
                               payload=chunk,
                               encoded=payload != b'')

        if done:
            break

def createalternatelocationresponseframe(stream, requestid, location):
    data = {
        b'status': b'redirect',
        b'location': {

              r'servercadercerts'):
        value = getattr(location, a)
        if value is not None:
            data[b'location'][pycompat.bytestr(a)] = value

    payload = b''.join(cborutil.streamencode(data))

    if stream.streamsettingssent:
        payload = stream.encode(payload)
        encoded = True
    else:
        encoded = False

    return stream.makeframe(requestid=requestid,
                            typeid=FRAME_TYPE_COMMAND_RESPONSE,
                            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
                            payload=payload,
                            encoded=encoded)

def createcommanderrorresponse(stream, requestid, message, args=None):
    # TODO should this be using a list of {'msg': ..., 'args': {}} so atom
    # formatting works consistently?
    m = {

        if data is None:
            for frame in self._flush():
                yield frame
            return

        data = self._stream.encode(data)

        # There is a ton of potential to do more complicated things here.
        # Our immediate goal is to coalesce small chunks into big frames,
        # not achieve the fewest number of frames possible. So we go with
        # a simple implementation:

                yield self._stream.makeframe(
                    self._requestid,
                    typeid=FRAME_TYPE_COMMAND_RESPONSE,
                    flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
                    payload=chunk,
                    encoded=True)

                if offset == len(data):
                    return

        # If we don't have enough to constitute a full frame, buffer and

        yield self._stream.makeframe(
            self._requestid,
            typeid=FRAME_TYPE_COMMAND_RESPONSE,
            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,
            payload=payload,
            encoded=True)

# TODO consider defining encoders/decoders using the util.compressionengine
# mechanism.

class identityencoder(object):

class outputstream(stream):
    """Represents a stream used for sending data."""

    def __init__(self, streamid, active=False):
        super(outputstream, self).__init__(streamid, active=active)
        self.streamsettingssent = False
        self._encoder = None
        self._encodername = None

    def setencoder(self, ui, name):
        """Set the encoder for this stream.

        Receives the stream profile name.
        """
        if name not in STREAM_ENCODERS:
            raise error.Abort(_('unknown stream encoder: %s') % name)

        self._encoder = STREAM_ENCODERS[name][0](ui)
        self._encodername = name

    def encode(self, data):
        if not self._encoder:
            return data

    def finish(self):
        if not self._encoder:
            return b''

        self._encoder.finish()

    def makeframe(self, requestid, typeid, flags, payload,
                  encoded=False):
        """Create a frame to be sent out over this stream.

        Only returns the frame instance. Does not actually send it.
        """
        streamflags = 0
        if not self._active:
            streamflags |= STREAM_FLAG_BEGIN_STREAM
            self._active = True

        if encoded:
            if not self.streamsettingssent:
                raise error.ProgrammingError(
                    b'attempting to send encoded frame without sending stream '
                    b'settings')

            streamflags |= STREAM_FLAG_ENCODING_APPLIED

        if (typeid == FRAME_TYPE_STREAM_SETTINGS
            and flags & FLAG_STREAM_ENCODING_SETTINGS_EOS):
            self.streamsettingssent = True

        return makeframe(requestid, self.streamid, streamflags, typeid, flags,
                         payload)

    def makestreamsettingsframe(self, requestid):
        """Create a stream settings frame for this stream.

        Returns frame data or None if no stream settings frame is needed or has
        already been sent.
        """
        if not self._encoder or self.streamsettingssent:
            return None

        payload = b''.join(cborutil.streamencode(self._encodername))
        return self.makeframe(requestid, FRAME_TYPE_STREAM_SETTINGS,
                              FLAG_STREAM_ENCODING_SETTINGS_EOS, payload)

def ensureserverstream(stream):
    if stream.streamid % 2:
        raise error.ProgrammingError('server should only write to even '
                                     'numbered streams; %d is not even' %

                except StopIteration:
                    for frame in emitter.send(None):
                        yield frame

                    if emitted:
                        for frame in createcommandresponseeosframes(
                            stream, requestid):
                            yield frame
                    break

                except error.WireprotoCommandError as e:
                    for frame in createcommanderrorresponse(
                        stream, requestid, e.message, e.messageargs):

                        if emitted:
                            raise error.ProgrammingError(
                                'alternatelocationresponse seen after initial '
                                'output object')

                        frame = stream.makestreamsettingsframe(requestid)
                        if frame:
                            yield frame

                        yield createalternatelocationresponseframe(
                            stream, requestid, o)

                        alternatelocationsent = True
                        emitted = True

                    if alternatelocationsent:
                        raise error.ProgrammingError(
                            'object follows alternatelocationresponse')

                    if not emitted:
                        # Frame is optional.
                        frame = stream.makestreamsettingsframe(requestid)
                        if frame:
                            yield frame

                        # May be None if empty frame (due to encoding).
                        frame = createcommandresponseokframe(stream, requestid)
                        if frame:
                            yield frame

                        emitted = True

                    # Objects emitted by command functions can be serializable
                    # data structures or special types.
                    # TODO consider extracting the content normalization to a

            self._activecommands.remove(requestid)

        return self._handlesendframes(sendframes())

    def makeoutputstream(self):
        """Create a stream to be used for sending data to the client.

        If this is called before protocol settings frames are received, we
        don't know what stream encodings are supported by the client and
        we will default to identity.
        """
        streamid = self._nextoutgoingstreamid
        self._nextoutgoingstreamid += 2

        s = outputstream(streamid)
        self._outgoingstreams[streamid] = s

        # Always use the *server's* preferred encoder over the client's,
        # as servers have more to lose from sub-optimal encoders being used.
        for name in STREAM_ENCODERS_ORDER:
            if name in self._sendersettings['contentencodings']:
                s.setencoder(self._ui, name)
                break

        return s

    def _makeerrorresult(self, msg):
        return 'error', {