← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: John Arbash Meinel
  • Date: 2008-08-28 20:47:56 UTC
  • mto: This revision was merged to the branch mainline in revision 3665.
  • Revision ID: john@arbash-meinel.com-20080828204756-iii2npp6ys48xzbo
Change _StatefulDecoder._in_bytes into a _in_bytes_list

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/smart/protocol.py
323
323
 
324
324
    def __init__(self):
325
325
        self.finished_reading = False
326
 
        self._in_buffer = ''
 
326
        # self._in_buffer = None
 
327
        self._in_buffer_list = []
 
328
        self._in_buffer_len = 0
327
329
        self.unused_data = ''
328
330
        self.bytes_left = None
329
331
        self._number_needed_bytes = None
330
332
 
 
333
    def _get_in_buffer(self):
 
334
        if len(self._in_buffer_list) == 1:
 
335
            return self._in_buffer_list[0]
 
336
        in_buffer = ''.join(self._in_buffer_list)
 
337
        assert len(in_buffer) == self._in_buffer_len
 
338
        in_buffer_list = [in_buffer]
 
339
        return in_buffer
 
340
 
 
341
    def _set_in_buffer(self, new_buf):
 
342
        if new_buf is not None:
 
343
            self._in_buffer_list = [new_buf]
 
344
            self._in_buffer_len = len(new_buf)
 
345
        else:
 
346
            self._in_buffer_list = []
 
347
            self._in_buffer_len = 0
 
348
 
331
349
    def accept_bytes(self, bytes):
332
350
        """Decode as much of bytes as possible.
333
351
 
338
356
        data will be appended to self.unused_data.
339
357
        """
340
358
        # accept_bytes is allowed to change the state
341
 
        current_state = self.state_accept
342
359
        self._number_needed_bytes = None
343
 
        self._in_buffer += bytes
 
360
        # lsprof puts a very large amount of time on this specific call for
 
361
        # large readv arrays
 
362
        self._in_buffer_list.append(bytes)
 
363
        self._in_buffer_len += len(bytes)
344
364
        try:
345
365
            # Run the function for the current state.
 
366
            current_state = self.state_accept
346
367
            self.state_accept()
347
368
            while current_state != self.state_accept:
348
369
                # The current state has changed.  Run the function for the new
379
400
            # the rest of this chunk plus an END chunk.
380
401
            return self.bytes_left + 4
381
402
        elif self.state_accept == self._state_accept_expecting_length:
382
 
            if self._in_buffer == '':
 
403
            if self._in_buffer_len == 0:
383
404
                # We're expecting a chunk length.  There's at least two bytes
384
405
                # left: a digit plus '\n'.
385
406
                return 2
390
411
        elif self.state_accept == self._state_accept_reading_unused:
391
412
            return 1
392
413
        elif self.state_accept == self._state_accept_expecting_header:
393
 
            return max(0, len('chunked\n') - len(self._in_buffer))
 
414
            return max(0, len('chunked\n') - self._in_buffer_len)
394
415
        else:
395
416
            raise AssertionError("Impossible state: %r" % (self.state_accept,))
396
417
 
401
422
            return None
402
423
 
403
424
    def _extract_line(self):
404
 
        pos = self._in_buffer.find('\n')
 
425
        in_buf = self._get_in_buffer()
 
426
        pos = in_buf.find('\n')
405
427
        if pos == -1:
406
428
            # We haven't read a complete line yet, so request more bytes before
407
429
            # we continue.
408
430
            raise _NeedMoreBytes(1)
409
 
        line = self._in_buffer[:pos]
 
431
        line = in_buf[:pos]
410
432
        # Trim the prefix (including '\n' delimiter) from the _in_buffer.
411
 
        self._in_buffer = self._in_buffer[pos+1:]
 
433
        self._set_in_buffer(in_buf[pos+1:])
412
434
        return line
413
435
 
414
436
    def _finished(self):
415
 
        self.unused_data = self._in_buffer
416
 
        self._in_buffer = ''
 
437
        self.unused_data = self._get_in_buffer()
 
438
        # self._in_buffer = None
 
439
        self._in_buffer_list = []
 
440
        self._in_buffer_len = 0
417
441
        self.state_accept = self._state_accept_reading_unused
418
442
        if self.error:
419
443
            error_args = tuple(self.error_in_progress)
448
472
            self.state_accept = self._state_accept_reading_chunk
449
473
 
450
474
    def _state_accept_reading_chunk(self):
451
 
        in_buffer_len = len(self._in_buffer)
452
 
        self.chunk_in_progress += self._in_buffer[:self.bytes_left]
453
 
        self._in_buffer = self._in_buffer[self.bytes_left:]
 
475
        in_buf = self._get_in_buffer()
 
476
        in_buffer_len = len(in_buf)
 
477
        self.chunk_in_progress += in_buf[:self.bytes_left]
 
478
        self._set_in_buffer(in_buf[self.bytes_left:])
454
479
        self.bytes_left -= in_buffer_len
455
480
        if self.bytes_left <= 0:
456
481
            # Finished with chunk
463
488
            self.state_accept = self._state_accept_expecting_length
464
489
        
465
490
    def _state_accept_reading_unused(self):
466
 
        self.unused_data += self._in_buffer
467
 
        self._in_buffer = ''
 
491
        self.unused_data += self._get_in_buffer()
 
492
        self._in_buffer_list = []
468
493
 
469
494
 
470
495
class LengthPrefixedBodyDecoder(_StatefulDecoder):
498
523
        return self.state_read()
499
524
 
500
525
    def _state_accept_expecting_length(self):
501
 
        pos = self._in_buffer.find('\n')
 
526
        in_buf = self._get_in_buffer()
 
527
        pos = in_buf.find('\n')
502
528
        if pos == -1:
503
529
            return
504
 
        self.bytes_left = int(self._in_buffer[:pos])
505
 
        self._in_buffer = self._in_buffer[pos+1:]
 
530
        self.bytes_left = int(in_buf[:pos])
 
531
        self._set_in_buffer(in_buf[pos+1:])
506
532
        self.state_accept = self._state_accept_reading_body
507
533
        self.state_read = self._state_read_body_buffer
508
534
 
509
535
    def _state_accept_reading_body(self):
510
 
        self._body += self._in_buffer
511
 
        self.bytes_left -= len(self._in_buffer)
512
 
        self._in_buffer = ''
 
536
        in_buf = self._get_in_buffer()
 
537
        self._body += in_buf
 
538
        self.bytes_left -= len(in_buf)
 
539
        self._set_in_buffer(None)
513
540
        if self.bytes_left <= 0:
514
541
            # Finished with body
515
542
            if self.bytes_left != 0:
519
546
            self.state_accept = self._state_accept_reading_trailer
520
547
        
521
548
    def _state_accept_reading_trailer(self):
522
 
        self._trailer_buffer += self._in_buffer
523
 
        self._in_buffer = ''
 
549
        self._trailer_buffer += self._get_in_buffer()
 
550
        self._set_in_buffer(None)
524
551
        # TODO: what if the trailer does not match "done\n"?  Should this raise
525
552
        # a ProtocolViolation exception?
526
553
        if self._trailer_buffer.startswith('done\n'):
529
556
            self.finished_reading = True
530
557
    
531
558
    def _state_accept_reading_unused(self):
532
 
        self.unused_data += self._in_buffer
533
 
        self._in_buffer = ''
 
559
        self.unused_data += self._get_in_buffer()
 
560
        self._set_in_buffer(None)
534
561
 
535
562
    def _state_read_no_data(self):
536
563
        return ''
865
892
            self.message_handler.protocol_error(exception)
866
893
 
867
894
    def _extract_length_prefixed_bytes(self):
868
 
        if len(self._in_buffer) < 4:
 
895
        if self._in_buffer_len < 4:
869
896
            # A length prefix by itself is 4 bytes, and we don't even have that
870
897
            # many yet.
871
898
            raise _NeedMoreBytes(4)
872
 
        (length,) = struct.unpack('!L', self._in_buffer[:4])
 
899
        in_buf = self._get_in_buffer()
 
900
        (length,) = struct.unpack('!L', in_buf[:4])
873
901
        end_of_bytes = 4 + length
874
 
        if len(self._in_buffer) < end_of_bytes:
 
902
        if self._in_buffer_len < end_of_bytes:
875
903
            # We haven't yet read as many bytes as the length-prefix says there
876
904
            # are.
877
905
            raise _NeedMoreBytes(end_of_bytes)
878
906
        # Extract the bytes from the buffer.
879
 
        bytes = self._in_buffer[4:end_of_bytes]
880
 
        self._in_buffer = self._in_buffer[end_of_bytes:]
 
907
        bytes = in_buf[4:end_of_bytes]
 
908
        self._set_in_buffer(in_buf[end_of_bytes:])
881
909
        return bytes
882
910
 
883
911
    def _extract_prefixed_bencoded_data(self):
890
918
        return decoded
891
919
 
892
920
    def _extract_single_byte(self):
893
 
        if self._in_buffer == '':
 
921
        if self._in_buffer_len == 0:
894
922
            # The buffer is empty
895
923
            raise _NeedMoreBytes(1)
896
 
        one_byte = self._in_buffer[0]
897
 
        self._in_buffer = self._in_buffer[1:]
 
924
        in_buf = self._get_in_buffer()
 
925
        one_byte = in_buf[0]
 
926
        self._set_in_buffer(in_buf[1:])
898
927
        return one_byte
899
928
 
900
929
    def _state_accept_expecting_protocol_version(self):
901
 
        needed_bytes = len(MESSAGE_VERSION_THREE) - len(self._in_buffer)
 
930
        needed_bytes = len(MESSAGE_VERSION_THREE) - self._in_buffer_len
 
931
        in_buf = self._get_in_buffer()
902
932
        if needed_bytes > 0:
903
933
            # We don't have enough bytes to check if the protocol version
904
934
            # marker is right.  But we can check if it is already wrong by
908
938
            # len(MESSAGE_VERSION_THREE) bytes.  So if the bytes we have so far
909
939
            # are wrong then we should just raise immediately rather than
910
940
            # stall.]
911
 
            if not MESSAGE_VERSION_THREE.startswith(self._in_buffer):
 
941
            if not MESSAGE_VERSION_THREE.startswith(in_buf):
912
942
                # We have enough bytes to know the protocol version is wrong
913
 
                raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
 
943
                raise errors.UnexpectedProtocolVersionMarker(in_buf)
914
944
            raise _NeedMoreBytes(len(MESSAGE_VERSION_THREE))
915
 
        if not self._in_buffer.startswith(MESSAGE_VERSION_THREE):
916
 
            raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
917
 
        self._in_buffer = self._in_buffer[len(MESSAGE_VERSION_THREE):]
 
945
        if not in_buf.startswith(MESSAGE_VERSION_THREE):
 
946
            raise errors.UnexpectedProtocolVersionMarker(in_buf)
 
947
        self._set_in_buffer(in_buf[len(MESSAGE_VERSION_THREE):])
918
948
        self.state_accept = self._state_accept_expecting_headers
919
949
 
920
950
    def _state_accept_expecting_headers(self):
969
999
            raise errors.SmartMessageHandlerError(sys.exc_info())
970
1000
 
971
1001
    def done(self):
972
 
        self.unused_data = self._in_buffer
973
 
        self._in_buffer = ''
 
1002
        self.unused_data = self._get_in_buffer()
 
1003
        self._set_in_buffer(None)
974
1004
        self.state_accept = self._state_accept_reading_unused
975
1005
        try:
976
1006
            self.message_handler.end_received()
978
1008
            raise errors.SmartMessageHandlerError(sys.exc_info())
979
1009
 
980
1010
    def _state_accept_reading_unused(self):
981
 
        self.unused_data += self._in_buffer
982
 
        self._in_buffer = ''
 
1011
        self.unused_data = self._get_in_buffer()
 
1012
        self._set_in_buffer(None)
983
1013
 
984
1014
    def next_read_size(self):
985
1015
        if self.state_accept == self._state_accept_reading_unused:
992
1022
            return 0
993
1023
        else:
994
1024
            if self._number_needed_bytes is not None:
995
 
                return self._number_needed_bytes - len(self._in_buffer)
 
1025
                return self._number_needed_bytes - self._in_buffer_len
996
1026
            else:
997
1027
                raise AssertionError("don't know how many bytes are expected!")
998
1028