← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/repository.py

  • Committer: Martin Pool
  • Date: 2010-08-13 08:09:53 UTC
  • mto: (5050.17.6 2.2)
  • mto: This revision was merged to the branch mainline in revision 5379.
  • Revision ID: mbp@sourcefrog.net-20100813080953-c00cm9l3qgu2flj9
Remove spuriously-resurrected test

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/smart/repository.py
1
 
# Copyright (C) 2006, 2007 Canonical Ltd
 
1
# Copyright (C) 2006-2010 Canonical Ltd
2
2
#
3
3
# This program is free software; you can redistribute it and/or modify
4
4
# it under the terms of the GNU General Public License as published by
20
20
import os
21
21
import Queue
22
22
import sys
23
 
import tarfile
24
23
import tempfile
25
24
import threading
26
25
 
30
29
    graph,
31
30
    osutils,
32
31
    pack,
 
32
    ui,
33
33
    versionedfile,
34
34
    )
35
35
from bzrlib.bzrdir import BzrDir
39
39
    SuccessfulSmartServerResponse,
40
40
    )
41
41
from bzrlib.repository import _strip_NULL_ghosts, network_format_registry
 
42
from bzrlib.recordcounter import RecordCounter
42
43
from bzrlib import revision as _mod_revision
43
44
from bzrlib.versionedfile import (
44
45
    NetworkRecordStream,
392
393
        if token == '':
393
394
            token = None
394
395
        try:
395
 
            token = repository.lock_write(token=token)
 
396
            token = repository.lock_write(token=token).repository_token
396
397
        except errors.LockContention, e:
397
398
            return FailedSmartServerResponse(('LockContention',))
398
399
        except errors.UnlockableTransport:
519
520
    yield pack_writer.end()
520
521
 
521
522
 
522
 
def _byte_stream_to_stream(byte_stream):
 
523
class _ByteStreamDecoder(object):
 
524
    """Helper for _byte_stream_to_stream.
 
525
 
 
526
    The expected usage of this class is via the function _byte_stream_to_stream
 
527
    which creates a _ByteStreamDecoder, pops off the stream format and then
 
528
    yields the output of record_stream(), the main entry point to
 
529
    _ByteStreamDecoder.
 
530
 
 
531
    Broadly this class has to unwrap two layers of iterators:
 
532
    (type, substream)
 
533
    (substream details)
 
534
 
 
535
    This is complicated by wishing to return type, iterator_for_type, but
 
536
    getting the data for iterator_for_type when we find out type: we can't
 
537
    simply pass a generator down to the NetworkRecordStream parser, instead
 
538
    we have a little local state to seed each NetworkRecordStream instance,
 
539
    and gather the type that we'll be yielding.
 
540
 
 
541
    :ivar byte_stream: The byte stream being decoded.
 
542
    :ivar stream_decoder: A pack parser used to decode the bytestream
 
543
    :ivar current_type: The current type, used to join adjacent records of the
 
544
        same type into a single stream.
 
545
    :ivar first_bytes: The first bytes to give the next NetworkRecordStream.
 
546
    """
 
547
 
 
548
    def __init__(self, byte_stream, record_counter):
 
549
        """Create a _ByteStreamDecoder."""
 
550
        self.stream_decoder = pack.ContainerPushParser()
 
551
        self.current_type = None
 
552
        self.first_bytes = None
 
553
        self.byte_stream = byte_stream
 
554
        self._record_counter = record_counter
 
555
        self.key_count = 0
 
556
 
 
557
    def iter_stream_decoder(self):
 
558
        """Iterate the contents of the pack from stream_decoder."""
 
559
        # dequeue pending items
 
560
        for record in self.stream_decoder.read_pending_records():
 
561
            yield record
 
562
        # Pull bytes of the wire, decode them to records, yield those records.
 
563
        for bytes in self.byte_stream:
 
564
            self.stream_decoder.accept_bytes(bytes)
 
565
            for record in self.stream_decoder.read_pending_records():
 
566
                yield record
 
567
 
 
568
    def iter_substream_bytes(self):
 
569
        if self.first_bytes is not None:
 
570
            yield self.first_bytes
 
571
            # If we run out of pack records, single the outer layer to stop.
 
572
            self.first_bytes = None
 
573
        for record in self.iter_pack_records:
 
574
            record_names, record_bytes = record
 
575
            record_name, = record_names
 
576
            substream_type = record_name[0]
 
577
            if substream_type != self.current_type:
 
578
                # end of a substream, seed the next substream.
 
579
                self.current_type = substream_type
 
580
                self.first_bytes = record_bytes
 
581
                return
 
582
            yield record_bytes
 
583
 
 
584
    def record_stream(self):
 
585
        """Yield substream_type, substream from the byte stream."""
 
586
        def wrap_and_count(pb, rc, substream):
 
587
            """Yield records from stream while showing progress."""
 
588
            counter = 0
 
589
            if rc:
 
590
                if self.current_type != 'revisions' and self.key_count != 0:
 
591
                    # As we know the number of revisions now (in self.key_count)
 
592
                    # we can setup and use record_counter (rc).
 
593
                    if not rc.is_initialized():
 
594
                        rc.setup(self.key_count, self.key_count)
 
595
            for record in substream.read():
 
596
                if rc:
 
597
                    if rc.is_initialized() and counter == rc.STEP:
 
598
                        rc.increment(counter)
 
599
                        pb.update('Estimate', rc.current, rc.max)
 
600
                        counter = 0
 
601
                    if self.current_type == 'revisions':
 
602
                        # Total records is proportional to number of revs
 
603
                        # to fetch. With remote, we used self.key_count to
 
604
                        # track the number of revs. Once we have the revs
 
605
                        # counts in self.key_count, the progress bar changes
 
606
                        # from 'Estimating..' to 'Estimate' above.
 
607
                        self.key_count += 1
 
608
                        if counter == rc.STEP:
 
609
                            pb.update('Estimating..', self.key_count)
 
610
                            counter = 0
 
611
                counter += 1
 
612
                yield record
 
613
 
 
614
        self.seed_state()
 
615
        pb = ui.ui_factory.nested_progress_bar()
 
616
        rc = self._record_counter
 
617
        # Make and consume sub generators, one per substream type:
 
618
        while self.first_bytes is not None:
 
619
            substream = NetworkRecordStream(self.iter_substream_bytes())
 
620
            # after substream is fully consumed, self.current_type is set to
 
621
            # the next type, and self.first_bytes is set to the matching bytes.
 
622
            yield self.current_type, wrap_and_count(pb, rc, substream)
 
623
        if rc:
 
624
            pb.update('Done', rc.max, rc.max)
 
625
        pb.finished()
 
626
 
 
627
    def seed_state(self):
 
628
        """Prepare the _ByteStreamDecoder to decode from the pack stream."""
 
629
        # Set a single generator we can use to get data from the pack stream.
 
630
        self.iter_pack_records = self.iter_stream_decoder()
 
631
        # Seed the very first subiterator with content; after this each one
 
632
        # seeds the next.
 
633
        list(self.iter_substream_bytes())
 
634
 
 
635
 
 
636
def _byte_stream_to_stream(byte_stream, record_counter=None):
523
637
    """Convert a byte stream into a format and a stream.
524
638
 
525
639
    :param byte_stream: A bytes iterator, as output by _stream_to_byte_stream.
526
640
    :return: (RepositoryFormat, stream_generator)
527
641
    """
528
 
    stream_decoder = pack.ContainerPushParser()
529
 
    def record_stream():
530
 
        """Closure to return the substreams."""
531
 
        # May have fully parsed records already.
532
 
        for record in stream_decoder.read_pending_records():
533
 
            record_names, record_bytes = record
534
 
            record_name, = record_names
535
 
            substream_type = record_name[0]
536
 
            substream = NetworkRecordStream([record_bytes])
537
 
            yield substream_type, substream.read()
538
 
        for bytes in byte_stream:
539
 
            stream_decoder.accept_bytes(bytes)
540
 
            for record in stream_decoder.read_pending_records():
541
 
                record_names, record_bytes = record
542
 
                record_name, = record_names
543
 
                substream_type = record_name[0]
544
 
                substream = NetworkRecordStream([record_bytes])
545
 
                yield substream_type, substream.read()
 
642
    decoder = _ByteStreamDecoder(byte_stream, record_counter)
546
643
    for bytes in byte_stream:
547
 
        stream_decoder.accept_bytes(bytes)
548
 
        for record in stream_decoder.read_pending_records(max=1):
 
644
        decoder.stream_decoder.accept_bytes(bytes)
 
645
        for record in decoder.stream_decoder.read_pending_records(max=1):
549
646
            record_names, src_format_name = record
550
647
            src_format = network_format_registry.get(src_format_name)
551
 
            return src_format, record_stream()
 
648
            return src_format, decoder.record_stream()
552
649
 
553
650
 
554
651
class SmartServerRepositoryUnlock(SmartServerRepositoryRequest):
614
711
            temp.close()
615
712
 
616
713
    def _tarball_of_dir(self, dirname, compression, ofile):
 
714
        import tarfile
617
715
        filename = os.path.basename(ofile.name)
618
716
        tarball = tarfile.open(fileobj=ofile, name=filename,
619
717
            mode='w|' + compression)