amp 0.5.0

data/lib/amp/revlogs/bundle_revlogs.rb

@@ -0,0 +1,246 @@
+module Amp
+  module Bundles
+    ##
+    # This module handles revlogs passed to our client (or server)
+    # through the bundle file format. Thing is, this revision log
+    # spans both a physical filelog, and a bundle (the new revisions),
+    # and we might need to get stuff from both. It's kind of like the
+    # DelayedOpener/FakeAppender for changelogs.
+    module BundleRevlog
+      include RevlogSupport::Node
+      BUNDLED_INDEX_ENTRY_SIZE = 80
+      
+      ##
+      # Initializes a bundle revlog. Takes, in addition to the normal
+      # revlog arguments, a bundle_file. This is any IO we can read
+      # from that will give us additional revisions, aside from the
+      # revisions stored in the real Revlog. It also takes a link_mapper
+      # that will connect things to the changelog revisions (including
+      # changelog revisions in the bundle).
+      #
+      # @param [Opener] opener the opener to use for openinf up the index_file
+      # @param [String] index_file the name of the file containing the revlog's
+      #  index
+      # @param [IO] bundle_file an IO that we can #read from
+      # @param [Proc, #call] link_mapper a function that will give us the link-index
+      #  to connect revisions to changelog revisions based on node_ids
+      def bundle_initialize(opener, index_file, bundle_file, link_mapper = nil)
+        @bundle_file = bundle_file
+        @base_map = {}
+        
+        num_revs = self.index_size
+        previous = nil
+        all_chunk_positions do |chunk, start|
+          chunk_size = chunk.size
+          
+          # each chunk starts with 4 node IDs: the new node's ID, its 2 parent node IDs,
+          # and the node ID of the corresponding revision in the changelog. In that order.
+          # If we have less than 80 bytes (BUNDLED_INDEX_ENTRY_SIZE), then we're fucked.
+          if chunk_size < BUNDLED_INDEX_ENTRY_SIZE
+            raise abort("invalid changegroup")
+          end
+          
+          start      += BUNDLED_INDEX_ENTRY_SIZE
+          chunk_size -= BUNDLED_INDEX_ENTRY_SIZE
+          
+          # Get the aforementioned node IDs
+          node, parent_1, parent_2, changeset = chunk[0..79].unpack("a20a20a20a20")
+          
+          # Do we already have this node? Skip it.
+          if @index.has_node? node
+            previous = node
+            next
+          end
+          
+          # make sure we have the new node's parents, or all of our operations will fail!
+          # at least, the interesting ones.
+          [parent_1, parent_2].each do |parent|
+            unless @index.has_node? parent
+              raise abort("Unknown parent: #{parent}@#{index_file}")
+            end
+          end
+          
+          link_rev = (link_mapper && link_mapper[changeset]) || num_revs
+          previous ||= parent_1
+          
+          @index << [RevlogSupport::Support.offset_version(start, 0), chunk_size, -1, -1, link_rev,
+                     revision_index_for_node(parent_1), revision_index_for_node(parent_2),
+                     node]
+          
+          @index.node_map[node] = num_revs
+          @base_map[num_revs] = previous
+          
+          previous = node
+          num_revs += 1
+        end
+      end
+      alias_method :bundle_revlog_initialize, :initialize
+      ##
+      # Returns whether the revision index is in the bundle part of this revlog,
+      # or if it's in the actual, stored revlog file.
+      #
+      # @param [Fixnum] revision the revision index to lookup
+      # @return [Boolean] is the revision in the bundle section?
+      def bundled_revision?(revision)
+        return false if revision < 0
+        !!@base_map[revision]
+      end
+      
+      ##
+      # Returns the base revision for the revision at the given index, while being
+      # cognizant of the bundle-ness of this revlog.
+      #
+      # @param [Fixnum] revision the revision index to lookup the base-revision of
+      # @return [String] the revision node ID of the base for the requested revision
+      def bundled_base_revision_for_index(revision)
+        @base_map[revision] || base_revision_for_index(revision)
+      end
+      alias_method :bundled_base, :bundled_base_revision_for_index
+      
+      ##
+      # Gets a chunk of data from the datafile (or, if inline, from the index
+      # file). Just give it a revision index and which data file to use. Only difference
+      # is that this will check the bundlefile if necessary.
+      # 
+      # @param [Fixnum] rev the revision index to extract
+      # @param [IO] data_file The IO file descriptor for loading data
+      # @return [String] the raw data from the index (posssibly compressed)
+      def get_chunk(revision, datafile=nil, cache_len = 4096)
+        # Warning: in case of bundle, the diff is against bundlebase, not against
+        # rev - 1
+        # TODO: could use some caching
+        unless bundled_revision?(revision)
+          return super(revision, datafile)
+        end
+        @bundle_file.seek data_start_for_index(revision)
+        @bundle_file.read self[revision].compressed_len
+      end
+      
+      ##
+      # Diffs 2 revisions, based on their indices. They are returned in
+      # BinaryDiff format.
+      # 
+      # @param [Fixnum] rev1 the index of the source revision
+      # @param [Fixnum] rev2 the index of the destination revision
+      # @return [String] The diff of the 2 revisions.
+      def revision_diff(rev1, rev2)
+        both_bundled = bundled_revision?(rev1) && bundled_revision?(rev2)
+        if both_bundled
+          # super-quick path if both are bundled and rev2 == rev1 + 1 diff
+          revision_base = self.revision_index_for_node bundled_base_revision_for_index(rev2)
+          if revision_base == rev1
+            # if rev2 = rev1 + a diff, just get the diff!
+            return get_chunk(revision2)
+          end
+        end
+        # normal style
+        return super(rev1, rev2)
+      end
+      
+      ##
+      # Given a node ID, extracts that revision and decompresses it. What you get
+      # back will the pristine revision data! Checks for bundle-ness when we access
+      # a node.
+      # 
+      # @param [String] node the Node ID of the revision to extract.
+      # @return [String] the pristine revision data.
+      def decompress_revision(node)
+        return "" if node == NULL_ID
+        
+        text = nil
+        chain = []
+        iter_node = node
+        rev = revision_index_for_node(node)
+        # walk backwards down the chain. Every single node is going
+        # to be a diff, because it's from a bundle.
+        while bundled_revision? rev
+          if @index.cache && @index.cache[0] == iter_node
+            text = @index.cache[2]
+            break
+          end
+          chain << rev
+          iter_node = bundled_base_revision_for_index rev
+          rev = revision_index_for_node iter_node
+        end
+        # done walking back, see if we have a stored cache!
+        text = super(iter_node) if text.nil? || text.empty?
+        
+        while chain.any?
+          delta = get_chunk(chain.pop)
+          text = Diffs::MercurialPatch.apply_patches(text, [delta])
+        end
+        p1, p2 = parents_for_node node
+        
+        if node != RevlogSupport::Support.history_hash(text, p1, p2)
+          raise RevlogSupport::RevlogError.new("integrity check failed on %s:%d, data:%s" % 
+                                               [(@index.inline? ? @index_file : @data_file), rev(node), text.inspect])
+        end
+        
+        @index.cache = [node, revision_index_for_node(node), text]
+        text
+      end
+      
+      ##
+      # Give an error to enforce read-only behavior
+      def add_revision(text, transaction, link, p1, p2, d=nil, index_file_handle = nil)
+        raise NotImplementedError.new
+      end
+      
+      ##
+      # Give an error to enforce read-only behavior
+      def add_group(revisions, link_mapper, journal)
+        raise NotImplementedError.new
+      end
+      
+      
+      private
+      
+      def all_chunk_positions
+        results = []
+        RevlogSupport::ChangeGroup.each_chunk(@bundle_file) do |chunk|
+          if block_given?
+            yield(chunk, @bundle_file.tell - chunk.size)
+          else
+            results << [chunk, @bundle_file.tell - chunk.size]
+          end
+        end
+        results unless block_given?
+      end
+    end
+    
+    class BundleChangeLog < ChangeLog
+      include BundleRevlog
+      def initialize(opener, bundle_file)
+        # This is the changelog initializer
+        super(opener)
+        # This is the bundle initializer
+        bundle_initialize(opener, @index_file, bundle_file)
+      end
+    end
+    
+    class BundleManifest < Manifest
+      include BundleRevlog
+      
+      def initialize(opener, bundle_file, link_mapper)
+        # This is the manifest initializer
+        super(opener)
+        # This is the bundle initializer
+        bundle_initialize(opener, @index_file, bundle_file, link_mapper)
+      end
+      
+    end
+    
+    class BundleFileLog < FileLog
+      include BundleRevlog
+      
+      def initialize(opener, path, bundle_file, link_mapper)
+        # This is the manifest initializer
+        super(opener, path)
+        # This is the bundle initializer
+        bundle_initialize(opener, @index_file, bundle_file, link_mapper)
+      end
+      
+    end
+    
+  end
+end

data/lib/amp/revlogs/changegroup.rb

@@ -0,0 +1,217 @@
+require 'tempfile'
+module Amp
+  module RevlogSupport
+    
+    ##
+    # This class handles changegroups - most specifically, packaging up a bunch
+    # of revisions into a single bundled file.
+    #
+    module ChangeGroup
+      BUNDLE_HEADERS = {
+          "" => "",
+          "HG10UN" => "HG10UN",
+          "HG10BZ" => "HG10",
+          "HG10GZ" => "HG10GZ"
+      }
+      FORMAT_PRIORITIES = BUNDLE_TYPES = ["HG10GZ", "HG10BZ", "HG10UN"]
+      class ChangeGroupError < StandardError; end
+      
+      ##
+      # Loads a single chunk of data from a changegroup. Each chunk is stored in the
+      # changegroup as:
+      #   (uint32_t) length    <-- less than 4 if we should terminate
+      #   (length * char) body
+      # Example:
+      #      00 00 00 05 h e l l o
+      #
+      # @param [IO, #read] source the source stream with all the data that we'll be
+      #   working on.
+      # @return [String] the data read in. Will be either nil or the empty string if
+      #   no data was read.
+      def self.get_chunk(source)
+        data = source.read 4
+        return "" if data.nil? || data.empty?
+        l = data.unpack("N")[0]
+        return "" if l <= 4
+        data = source.read(l - 4)
+        if data.size < l - 4
+          raise ChangeGroupError.new("premature EOF when reading changegroup:" +
+                                     "(got #{data.size} bytes, expected #{l-4})")
+        end
+        return data
+      end
+      
+      ##
+      # Loads each chunk, in a row, until we run out of chunks in the changegroup.
+      # Yields the data to the caller. This will run until we hit a terminating chunk.
+      #
+      # @param [IO, #read] source The stream from which we will read the data in sequence
+      # @yield Each chunk of data will be yielded to be processed
+      # @yieldparam [String] chunk the chunk that is yielded
+      def self.each_chunk(source)
+        begin
+          c = self.get_chunk(source) # get a chunk
+          yield c unless c.empty?    # yield if not empty
+        end until c.empty?           # keep going if we have more!
+      end
+      
+      ##
+      # If we have data of size +size+, then return the encoded header for the chunk
+      #
+      # @param [Integer] size the size of the chunk we wish to encode
+      # @return [String] the encoded header for the chunk
+      def self.chunk_header(size)
+        [size + 4].pack("N")
+      end
+      
+      ##
+      # The terminating chunk that indicates the end of a chunk sequence
+      #
+      # @return [String] the encoded terminating chunk
+      def self.closing_chunk
+        "\000\000\000\000" # [0].pack("N")
+      end
+      
+      ##
+      # Returns a compressing stream based on the header for a changegroup bundle.
+      # The bundle header will specify that the contents should be either uncompressed,
+      # BZip compressed, or GZip compressed. This will return a stream that responds
+      # to #<< and #flush, where #flush will return the unread, decompressed data,
+      # and #<< will input uncompressed data.
+      #
+      # @param [String] header the header for the changegroup bundle. Can be either
+      #   HG10UN, HG10GZ, or HG10BZ, or HG10.
+      # @return [IO, #<<, #flush] an IO stream that accepts uncompressed data via #<< or
+      #   #write, and returns compressed data by #flush.
+      def self.compressor_by_type(header)
+        case header
+        when "HG10UN", ""
+          # new StringIO
+          result = StringIO.new "",(ruby_19? ? "w+:ASCII-8BIT" : "w+")
+          # have to fix the fact that StringIO doesn't conform to the other streams,
+          # and give it a #flush method. Kind of hackish.
+          class << result
+            def flush
+              ret = self.string.dup  # get the current read-in string
+              self.string.replace "" # erase our contents
+              self.rewind            # rewind the IO
+              self.tell
+              ret                    # return the string
+            end
+          end
+          #return the altered StringIO
+          result
+        when "HG10GZ"
+          # lazy-load Zlib
+          require 'zlib'
+          # Return a deflating stream (compressor)
+          Zlib::Deflate.new
+        when "HG10BZ", "HG10"
+          # lazy load BZip
+          need { '../../../ext/amp/bz2/bz2' }
+          # Return a compressing BZip stream
+          BZ2::Writer.new
+        end
+      end
+      
+      ##
+      # Returns a stream that will decompress the IO pointed to by file_handle, 
+      # when #read is called upon it. Note: file_handle doesn't have to be a file!
+      #
+      # @param [String] header the header of the stream. Specifies which compression
+      #  to handle
+      # @param [IO, #read] file_handle the input stream that will provide the compressed
+      #  data.
+      # @return [IO, #read] an IO object that we can #read to get uncompressed data
+      def self.unbundle(header, file_handle)
+        # uncompressed? just return the input IO!
+        return file_handle if header == "HG10UN"
+        # if we have no header, we're uncompressed
+        if !header.start_with?("HG")
+          # append the header to it. meh
+          headerio = StringIO.new(header, (ruby_19? ? "w+:ASCII-8BIT" : "w+"))
+          Amp::Support::MultiIO.new(headerio, file_handle)
+          # WOW we have legacy support already
+        elsif header == "HG10GZ"
+          # Get a gzip reader
+          Zlib::GzipReader.new(file_handle)
+        elsif header == "HG10BZ"
+          # get a BZip reader, but it has to decompress "BZ" first. Meh.
+          headerio = StringIO.new("BZ", (ruby_19? ? "w+:ASCII-8BIT" : "w+"))
+          input = Amp::Support::MultiIO.new(headerio, file_handle)
+          BZ2::Reader.new(input)
+        end
+      end
+          
+      ##
+      # Writes a set of changegroups to a bundle. If no IO is specified, a new StringIO
+      # is created, and the bundle is written to that (i.e., memory). the IO used is returned.
+      #
+      # @param [IO, #read, #seek] changegroup A stream that will feed in an uncompressed
+      #   changegroup
+      # @param [String] bundletype A specified compression type - either "HG10UN", "HG10GZ",
+      #   or "HG10BZ". The empty string defaults to "HG10UN".
+      # @param [IO, #write] fh (StringIO.new) An output stream to write to, such as a File
+      #   or a socket. If not specified, a StringIO is created and returned.
+      # @return [IO, #write] the output IO stream is returned, even if a new one is not
+      #   created on the fly.
+      def self.write_bundle(changegroup, bundletype, fh = StringIO.new("", (ruby_19? ? "w+:ASCII-8BIT" : "w+")))
+        # rewind the changegroup to start at the beginning
+        changegroup.rewind
+        # pick out our header
+        header     = BUNDLE_HEADERS[bundletype]
+        # get a compressing stream
+        compressor = compressor_by_type header
+        # output the header (uncompressed)
+        fh.write header
+        
+        # These 2 variables are for checking to see if #changegroup has been fully
+        # read in or not. 
+        empty = false
+        count = 0
+        
+        # Do at least 2 changegroups (changelog + manifest), then go until we're empty
+        while !empty || count <= 2
+          # Set empty to true for this particular file (each iteration of this loop
+          # represents compressing 1 file's changesets into changegroups in the bundle)
+          empty = true
+          # Add 1 to the number of files we've comrpessed
+          count += 1
+          # For each chunk in the changegroup (i.e. each changeset)
+          inner_count = 0
+          self.each_chunk(changegroup) do |chunk|
+            
+            #puts "\t\twrite_bundle inner loop count #{inner_count}"
+            inner_count += 1
+            empty = false
+            # Compress the chunk header
+            compressor << chunk_header(chunk.size)
+            # Write the chunk header
+            fh.write(compressor.flush)
+            
+            # compress the actual chunk 1 megabyte at a time
+            step_amt = 1048576
+            (0..chunk.size).step(step_amt) do |pos|
+              compressor << chunk[pos..(pos+step_amt-1)]
+              fh.write(compressor.flush)
+              fh.flush
+            end
+          end
+          # Compress the terminating chunk - this indicates that there are no more changesets
+          # for the current file
+          compressor << closing_chunk
+          # Write the terminating chunk out!
+          fh.write compressor.flush
+          fh.flush
+        end
+        
+        # Write anything left over in that there compressor
+        fh.write compressor.flush
+        # Kill the compressor
+        compressor.close
+        # Return the IO we wrote to (in case we instantiated it)
+        return fh
+      end
+    end
+  end
+end