~bzr-pqm/bzr/bzr.dev

# Copyright (C) 2008, 2009 Canonical Ltd

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

"""Core compression logic for compressing streams of related files."""

import time

import zlib

try:

    import pylzma

except ImportError:

    pylzma = None

from bzrlib import (

    annotate,

    debug,

    errors,

    graph as _mod_graph,

    knit,

    osutils,

    pack,

    trace,

    )

from bzrlib.graph import Graph

from bzrlib.btree_index import BTreeBuilder

from bzrlib.lru_cache import LRUSizeCache

from bzrlib.tsort import topo_sort

from bzrlib.versionedfile import (

    adapter_registry,

    AbsentContentFactory,

    ChunkedContentFactory,

    FulltextContentFactory,

    VersionedFiles,

    )

_USE_LZMA = False and (pylzma is not None)

# osutils.sha_string('')

_null_sha1 = 'da39a3ee5e6b4b0d3255bfef95601890afd80709'

def sort_gc_optimal(parent_map):

    """Sort and group the keys in parent_map into groupcompress order.

    groupcompress is defined (currently) as reverse-topological order, grouped

    by the key prefix.

    :return: A sorted-list of keys

    """

    # groupcompress ordering is approximately reverse topological,

    # properly grouped by file-id.

    per_prefix_map = {}

    for item in parent_map.iteritems():

        key = item[0]

        if isinstance(key, str) or len(key) == 1:

            prefix = ''

        else:

            prefix = key[0]

        try:

            per_prefix_map[prefix].append(item)

        except KeyError:

            per_prefix_map[prefix] = [item]

    present_keys = []

    for prefix in sorted(per_prefix_map):

        present_keys.extend(reversed(topo_sort(per_prefix_map[prefix])))

    return present_keys

# The max zlib window size is 32kB, so if we set 'max_size' output of the

# decompressor to the requested bytes + 32kB, then we should guarantee

# num_bytes coming out.

_ZLIB_DECOMP_WINDOW = 32*1024

class GroupCompressBlock(object):

    """An object which maintains the internal structure of the compressed data.

    This tracks the meta info (start of text, length, type, etc.)

    """

    # Group Compress Block v1 Zlib

    GCB_HEADER = 'gcb1z\n'

    # Group Compress Block v1 Lzma

    GCB_LZ_HEADER = 'gcb1l\n'

    GCB_KNOWN_HEADERS = (GCB_HEADER, GCB_LZ_HEADER)

    def __init__(self):

        # map by key? or just order in file?

        self._compressor_name = None

        self._z_content = None

        self._z_content_decompressor = None

        self._z_content_length = None

        self._content_length = None

        self._content = None

        self._content_chunks = None

    def __len__(self):

        # This is the maximum number of bytes this object will reference if

        # everything is decompressed. However, if we decompress less than

        # everything... (this would cause some problems for LRUSizeCache)

        return self._content_length + self._z_content_length

    def _ensure_content(self, num_bytes=None):

        """Make sure that content has been expanded enough.

        :param num_bytes: Ensure that we have extracted at least num_bytes of

            content. If None, consume everything

        """

        # TODO: If we re-use the same content block at different times during

        #       get_record_stream(), it is possible that the first pass will

        #       get inserted, triggering an extract/_ensure_content() which

        #       will get rid of _z_content. And then the next use of the block

        #       will try to access _z_content (to send it over the wire), and

        #       fail because it is already extracted. Consider never releasing

        #       _z_content because of this.

        if num_bytes is None:

            num_bytes = self._content_length

        elif (self._content_length is not None

              and num_bytes > self._content_length):

            raise AssertionError(

                'requested num_bytes (%d) > content length (%d)'

                % (num_bytes, self._content_length))

        # Expand the content if required

        if self._content is None:

            if self._content_chunks is not None:

                self._content = ''.join(self._content_chunks)

                self._content_chunks = None

        if self._content is None:

            if self._z_content is None:

                raise AssertionError('No content to decompress')

            if self._z_content == '':

                self._content = ''

            elif self._compressor_name == 'lzma':

                # We don't do partial lzma decomp yet

                self._content = pylzma.decompress(self._z_content)

            elif self._compressor_name == 'zlib':

                # Start a zlib decompressor

                if num_bytes is None:

                    self._content = zlib.decompress(self._z_content)

                else:

                    self._z_content_decompressor = zlib.decompressobj()

                    # Seed the decompressor with the uncompressed bytes, so

                    # that the rest of the code is simplified

                    self._content = self._z_content_decompressor.decompress(

                        self._z_content, num_bytes + _ZLIB_DECOMP_WINDOW)

            else:

                raise AssertionError('Unknown compressor: %r'

                                     % self._compressor_name)

        # Any bytes remaining to be decompressed will be in the decompressors

        # 'unconsumed_tail'

        # Do we have enough bytes already?

        if num_bytes is not None and len(self._content) >= num_bytes:

            return

        if num_bytes is None and self._z_content_decompressor is None:

            # We must have already decompressed everything

            return

        # If we got this far, and don't have a decompressor, something is wrong

        if self._z_content_decompressor is None:

            raise AssertionError(

                'No decompressor to decompress %d bytes' % num_bytes)

        remaining_decomp = self._z_content_decompressor.unconsumed_tail

        if num_bytes is None:

            if remaining_decomp:

                # We don't know how much is left, but we'll decompress it all

                self._content += self._z_content_decompressor.decompress(

                    remaining_decomp)

                # Note: There's what I consider a bug in zlib.decompressobj

                #       If you pass back in the entire unconsumed_tail, only

                #       this time you don't pass a max-size, it doesn't

                #       change the unconsumed_tail back to None/''.

                #       However, we know we are done with the whole stream

                self._z_content_decompressor = None

            # XXX: Why is this the only place in this routine we set this?

            self._content_length = len(self._content)

        else:

            if not remaining_decomp:

                raise AssertionError('Nothing left to decompress')

            needed_bytes = num_bytes - len(self._content)

            # We always set max_size to 32kB over the minimum needed, so that

            # zlib will give us as much as we really want.

            # TODO: If this isn't good enough, we could make a loop here,

            #       that keeps expanding the request until we get enough

            self._content += self._z_content_decompressor.decompress(

                remaining_decomp, needed_bytes + _ZLIB_DECOMP_WINDOW)

            if len(self._content) < num_bytes:

                raise AssertionError('%d bytes wanted, only %d available'

                                     % (num_bytes, len(self._content)))

            if not self._z_content_decompressor.unconsumed_tail:

                # The stream is finished

                self._z_content_decompressor = None

    def _parse_bytes(self, bytes, pos):

        """Read the various lengths from the header.

        This also populates the various 'compressed' buffers.

        :return: The position in bytes just after the last newline

        """

        # At present, we have 2 integers for the compressed and uncompressed

        # content. In base10 (ascii) 14 bytes can represent > 1TB, so to avoid

        # checking too far, cap the search to 14 bytes.

        pos2 = bytes.index('\n', pos, pos + 14)

        self._z_content_length = int(bytes[pos:pos2])

        pos = pos2 + 1

        pos2 = bytes.index('\n', pos, pos + 14)

        self._content_length = int(bytes[pos:pos2])

        pos = pos2 + 1

        if len(bytes) != (pos + self._z_content_length):

            # XXX: Define some GCCorrupt error ?

            raise AssertionError('Invalid bytes: (%d) != %d + %d' %

                                 (len(bytes), pos, self._z_content_length))

        self._z_content = bytes[pos:]

    @classmethod

    def from_bytes(cls, bytes):

        out = cls()

        if bytes[:6] not in cls.GCB_KNOWN_HEADERS:

            raise ValueError('bytes did not start with any of %r'

                             % (cls.GCB_KNOWN_HEADERS,))

        # XXX: why not testing the whole header ?

        if bytes[4] == 'z':

            out._compressor_name = 'zlib'

        elif bytes[4] == 'l':

            out._compressor_name = 'lzma'

        else:

            raise ValueError('unknown compressor: %r' % (bytes,))

        out._parse_bytes(bytes, 6)

        return out

    def extract(self, key, start, end, sha1=None):

        """Extract the text for a specific key.

        :param key: The label used for this content

        :param sha1: TODO (should we validate only when sha1 is supplied?)

        :return: The bytes for the content

        """

        if start == end == 0:

            return ''

        self._ensure_content(end)

        # The bytes are 'f' or 'd' for the type, then a variable-length

        # base128 integer for the content size, then the actual content

        # We know that the variable-length integer won't be longer than 5

        # bytes (it takes 5 bytes to encode 2^32)

        c = self._content[start]

        if c == 'f':

            type = 'fulltext'

        else:

            if c != 'd':

                raise ValueError('Unknown content control code: %s'

                                 % (c,))

            type = 'delta'

        content_len, len_len = decode_base128_int(

                            self._content[start + 1:start + 6])

        content_start = start + 1 + len_len

        if end != content_start + content_len:

            raise ValueError('end != len according to field header'

                ' %s != %s' % (end, content_start + content_len))

        if c == 'f':

            bytes = self._content[content_start:end]

        elif c == 'd':

            bytes = apply_delta_to_source(self._content, content_start, end)

        return bytes

    def set_chunked_content(self, content_chunks, length):

        """Set the content of this block to the given chunks."""

        # If we have lots of short lines, it is may be more efficient to join

        # the content ahead of time. If the content is <10MiB, we don't really

        # care about the extra memory consumption, so we can just pack it and

        # be done. However, timing showed 18s => 17.9s for repacking 1k revs of

        # mysql, which is below the noise margin

        self._content_length = length

        self._content_chunks = content_chunks

        self._content = None

        self._z_content = None

    def set_content(self, content):

        """Set the content of this block."""

        self._content_length = len(content)

        self._content = content

        self._z_content = None

    def _create_z_content_using_lzma(self):

        if self._content_chunks is not None:

            self._content = ''.join(self._content_chunks)

            self._content_chunks = None

        if self._content is None:

            raise AssertionError('Nothing to compress')

        self._z_content = pylzma.compress(self._content)

        self._z_content_length = len(self._z_content)

    def _create_z_content_from_chunks(self):

        compressor = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION)

        compressed_chunks = map(compressor.compress, self._content_chunks)

        compressed_chunks.append(compressor.flush())

        self._z_content = ''.join(compressed_chunks)

        self._z_content_length = len(self._z_content)

    def _create_z_content(self):

        if self._z_content is not None:

            return

        if _USE_LZMA:

            self._create_z_content_using_lzma()

            return

        if self._content_chunks is not None:

            self._create_z_content_from_chunks()

            return

        self._z_content = zlib.compress(self._content)

        self._z_content_length = len(self._z_content)

    def to_bytes(self):

        """Encode the information into a byte stream."""

        self._create_z_content()

        if _USE_LZMA:

            header = self.GCB_LZ_HEADER

        else:

            header = self.GCB_HEADER

        chunks = [header,

                  '%d\n%d\n' % (self._z_content_length, self._content_length),

                  self._z_content,

                 ]

        return ''.join(chunks)

    def _dump(self, include_text=False):

        """Take this block, and spit out a human-readable structure.

        :param include_text: Inserts also include text bits, chose whether you

            want this displayed in the dump or not.

        :return: A dump of the given block. The layout is something like:

            [('f', length), ('d', delta_length, text_length, [delta_info])]

            delta_info := [('i', num_bytes, text), ('c', offset, num_bytes),

            ...]

        """

        self._ensure_content()

        result = []

        pos = 0

        while pos < self._content_length:

            kind = self._content[pos]

            pos += 1

            if kind not in ('f', 'd'):

                raise ValueError('invalid kind character: %r' % (kind,))

            content_len, len_len = decode_base128_int(

                                self._content[pos:pos + 5])

            pos += len_len

            if content_len + pos > self._content_length:

                raise ValueError('invalid content_len %d for record @ pos %d'

                                 % (content_len, pos - len_len - 1))

            if kind == 'f': # Fulltext

                if include_text:

                    text = self._content[pos:pos+content_len]

                    result.append(('f', content_len, text))

                else:

                    result.append(('f', content_len))

            elif kind == 'd': # Delta

                delta_content = self._content[pos:pos+content_len]

                delta_info = []

                # The first entry in a delta is the decompressed length

                decomp_len, delta_pos = decode_base128_int(delta_content)

                result.append(('d', content_len, decomp_len, delta_info))

                measured_len = 0

                while delta_pos < content_len:

                    c = ord(delta_content[delta_pos])

                    delta_pos += 1

                    if c & 0x80: # Copy

                        (offset, length,

                         delta_pos) = decode_copy_instruction(delta_content, c,

                                                              delta_pos)

                        if include_text:

                            text = self._content[offset:offset+length]

                            delta_info.append(('c', offset, length, text))

                        else:

                            delta_info.append(('c', offset, length))

                        measured_len += length

                    else: # Insert

                        if include_text:

                            txt = delta_content[delta_pos:delta_pos+c]

                        else:

                            txt = ''

                        delta_info.append(('i', c, txt))

                        measured_len += c

                        delta_pos += c

                if delta_pos != content_len:

                    raise ValueError('Delta consumed a bad number of bytes:'

                                     ' %d != %d' % (delta_pos, content_len))

                if measured_len != decomp_len:

                    raise ValueError('Delta claimed fulltext was %d bytes, but'

                                     ' extraction resulted in %d bytes'

                                     % (decomp_len, measured_len))

            pos += content_len

        return result

class _LazyGroupCompressFactory(object):

    """Yield content from a GroupCompressBlock on demand."""

    def __init__(self, key, parents, manager, start, end, first):

        """Create a _LazyGroupCompressFactory

        :param key: The key of just this record

        :param parents: The parents of this key (possibly None)

        :param gc_block: A GroupCompressBlock object

        :param start: Offset of the first byte for this record in the

            uncompressd content

        :param end: Offset of the byte just after the end of this record

            (ie, bytes = content[start:end])

        :param first: Is this the first Factory for the given block?

        """

        self.key = key

        self.parents = parents

        self.sha1 = None

        # Note: This attribute coupled with Manager._factories creates a

        #       reference cycle. Perhaps we would rather use a weakref(), or

        #       find an appropriate time to release the ref. After the first

        #       get_bytes_as call? After Manager.get_record_stream() returns

        #       the object?

        self._manager = manager

        self._bytes = None

        self.storage_kind = 'groupcompress-block'

        if not first:

            self.storage_kind = 'groupcompress-block-ref'

        self._first = first

        self._start = start

        self._end = end

    def __repr__(self):

        return '%s(%s, first=%s)' % (self.__class__.__name__,

            self.key, self._first)

    def get_bytes_as(self, storage_kind):

        if storage_kind == self.storage_kind:

            if self._first:

                # wire bytes, something...

                return self._manager._wire_bytes()

            else:

                return ''

        if storage_kind in ('fulltext', 'chunked'):

            if self._bytes is None:

                # Grab and cache the raw bytes for this entry

                # and break the ref-cycle with _manager since we don't need it

                # anymore

                self._manager._prepare_for_extract()

                block = self._manager._block

                self._bytes = block.extract(self.key, self._start, self._end)

                # There are code paths that first extract as fulltext, and then

                # extract as storage_kind (smart fetch). So we don't break the

                # refcycle here, but instead in manager.get_record_stream()

                # self._manager = None

            if storage_kind == 'fulltext':

                return self._bytes

            else:

                return [self._bytes]

        raise errors.UnavailableRepresentation(self.key, storage_kind,

                                               self.storage_kind)

class _LazyGroupContentManager(object):

    """This manages a group of _LazyGroupCompressFactory objects."""

    def __init__(self, block):

        self._block = block

        # We need to preserve the ordering

        self._factories = []

        self._last_byte = 0

    def add_factory(self, key, parents, start, end):

        if not self._factories:

            first = True

        else:

            first = False

        # Note that this creates a reference cycle....

        factory = _LazyGroupCompressFactory(key, parents, self,

            start, end, first=first)

        # max() works here, but as a function call, doing a compare seems to be

        # significantly faster, timeit says 250ms for max() and 100ms for the

        # comparison

        if end > self._last_byte:

            self._last_byte = end

        self._factories.append(factory)

    def get_record_stream(self):

        """Get a record for all keys added so far."""

        for factory in self._factories:

            yield factory

            # Break the ref-cycle

            factory._bytes = None

            factory._manager = None

        # TODO: Consider setting self._factories = None after the above loop,

        #       as it will break the reference cycle

    def _trim_block(self, last_byte):

        """Create a new GroupCompressBlock, with just some of the content."""

        # None of the factories need to be adjusted, because the content is

        # located in an identical place. Just that some of the unreferenced

        # trailing bytes are stripped

        trace.mutter('stripping trailing bytes from groupcompress block'

                     ' %d => %d', self._block._content_length, last_byte)

        new_block = GroupCompressBlock()

        self._block._ensure_content(last_byte)

        new_block.set_content(self._block._content[:last_byte])

        self._block = new_block

    def _rebuild_block(self):

        """Create a new GroupCompressBlock with only the referenced texts."""

        compressor = GroupCompressor()

        tstart = time.time()

        old_length = self._block._content_length

        end_point = 0

        for factory in self._factories:

            bytes = factory.get_bytes_as('fulltext')

            (found_sha1, start_point, end_point,

             type) = compressor.compress(factory.key, bytes, factory.sha1)

            # Now update this factory with the new offsets, etc

            factory.sha1 = found_sha1

            factory._start = start_point

            factory._end = end_point

        self._last_byte = end_point

        new_block = compressor.flush()

        # TODO: Should we check that new_block really *is* smaller than the old

        #       block? It seems hard to come up with a method that it would

        #       expand, since we do full compression again. Perhaps based on a

        #       request that ends up poorly ordered?

        delta = time.time() - tstart

        self._block = new_block

        trace.mutter('creating new compressed block on-the-fly in %.3fs'

                     ' %d bytes => %d bytes', delta, old_length,

                     self._block._content_length)

    def _prepare_for_extract(self):

        """A _LazyGroupCompressFactory is about to extract to fulltext."""

        # We expect that if one child is going to fulltext, all will be. This

        # helps prevent all of them from extracting a small amount at a time.

        # Which in itself isn't terribly expensive, but resizing 2MB 32kB at a

        # time (self._block._content) is a little expensive.

        self._block._ensure_content(self._last_byte)

    def _check_rebuild_block(self):

        """Check to see if our block should be repacked."""

        total_bytes_used = 0

        last_byte_used = 0

        for factory in self._factories:

            total_bytes_used += factory._end - factory._start

            last_byte_used = max(last_byte_used, factory._end)

        # If we are using most of the bytes from the block, we have nothing

        # else to check (currently more that 1/2)

        if total_bytes_used * 2 >= self._block._content_length:

            return

        # Can we just strip off the trailing bytes? If we are going to be

        # transmitting more than 50% of the front of the content, go ahead

        if total_bytes_used * 2 > last_byte_used:

            self._trim_block(last_byte_used)

            return

        # We are using a small amount of the data, and it isn't just packed

        # nicely at the front, so rebuild the content.

        # Note: This would be *nicer* as a strip-data-from-group, rather than

        #       building it up again from scratch

        #       It might be reasonable to consider the fulltext sizes for

        #       different bits when deciding this, too. As you may have a small

        #       fulltext, and a trivial delta, and you are just trading around

        #       for another fulltext. If we do a simple 'prune' you may end up

        #       expanding many deltas into fulltexts, as well.

        #       If we build a cheap enough 'strip', then we could try a strip,

        #       if that expands the content, we then rebuild.

        self._rebuild_block()

    def _wire_bytes(self):

        """Return a byte stream suitable for transmitting over the wire."""

        self._check_rebuild_block()

        # The outer block starts with:

        #   'groupcompress-block\n'

        #   <length of compressed key info>\n

        #   <length of uncompressed info>\n

        #   <length of gc block>\n

        #   <header bytes>

        #   <gc-block>

        lines = ['groupcompress-block\n']

        # The minimal info we need is the key, the start offset, and the

        # parents. The length and type are encoded in the record itself.

        # However, passing in the other bits makes it easier.  The list of

        # keys, and the start offset, the length

        # 1 line key

        # 1 line with parents, '' for ()

        # 1 line for start offset

        # 1 line for end byte

        header_lines = []

        for factory in self._factories:

            key_bytes = '\x00'.join(factory.key)

            parents = factory.parents

            if parents is None:

                parent_bytes = 'None:'

            else:

                parent_bytes = '\t'.join('\x00'.join(key) for key in parents)

            record_header = '%s\n%s\n%d\n%d\n' % (

                key_bytes, parent_bytes, factory._start, factory._end)

            header_lines.append(record_header)

            # TODO: Can we break the refcycle at this point and set

            #       factory._manager = None?

        header_bytes = ''.join(header_lines)

        del header_lines

        header_bytes_len = len(header_bytes)

        z_header_bytes = zlib.compress(header_bytes)

        del header_bytes

        z_header_bytes_len = len(z_header_bytes)

        block_bytes = self._block.to_bytes()

        lines.append('%d\n%d\n%d\n' % (z_header_bytes_len, header_bytes_len,

                                       len(block_bytes)))

        lines.append(z_header_bytes)

        lines.append(block_bytes)

        del z_header_bytes, block_bytes

        return ''.join(lines)

    @classmethod

    def from_bytes(cls, bytes):

        # TODO: This does extra string copying, probably better to do it a

        #       different way

        (storage_kind, z_header_len, header_len,

         block_len, rest) = bytes.split('\n', 4)

        del bytes

        if storage_kind != 'groupcompress-block':

            raise ValueError('Unknown storage kind: %s' % (storage_kind,))

        z_header_len = int(z_header_len)

        if len(rest) < z_header_len:

            raise ValueError('Compressed header len shorter than all bytes')

        z_header = rest[:z_header_len]

        header_len = int(header_len)

        header = zlib.decompress(z_header)

        if len(header) != header_len:

            raise ValueError('invalid length for decompressed bytes')

        del z_header

        block_len = int(block_len)

        if len(rest) != z_header_len + block_len:

            raise ValueError('Invalid length for block')

        block_bytes = rest[z_header_len:]

        del rest

        # So now we have a valid GCB, we just need to parse the factories that

        # were sent to us

        header_lines = header.split('\n')

        del header

        last = header_lines.pop()

        if last != '':

            raise ValueError('header lines did not end with a trailing'

                             ' newline')

        if len(header_lines) % 4 != 0:

            raise ValueError('The header was not an even multiple of 4 lines')

        block = GroupCompressBlock.from_bytes(block_bytes)

        del block_bytes

        result = cls(block)

        for start in xrange(0, len(header_lines), 4):

            # intern()?

            key = tuple(header_lines[start].split('\x00'))

            parents_line = header_lines[start+1]

            if parents_line == 'None:':

                parents = None

            else:

                parents = tuple([tuple(segment.split('\x00'))

                                 for segment in parents_line.split('\t')