~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/medium.py

  • Committer: Tarmac
  • Author(s): Vincent Ladeuil
  • Date: 2017-01-30 14:42:05 UTC
  • mfrom: (6620.1.1 trunk)
  • Revision ID: tarmac-20170130144205-r8fh2xpmiuxyozpv
Merge  2.7 into trunk including fix for bug #1657238 [r=vila]

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2006 Canonical Ltd
 
1
# Copyright (C) 2006-2011 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
24
24
bzrlib/transport/smart/__init__.py.
25
25
"""
26
26
 
 
27
from __future__ import absolute_import
 
28
 
27
29
import errno
28
30
import os
29
 
import socket
30
31
import sys
31
 
import urllib
 
32
import time
32
33
 
 
34
import bzrlib
33
35
from bzrlib.lazy_import import lazy_import
34
36
lazy_import(globals(), """
35
 
import atexit
 
37
import select
 
38
import socket
 
39
import thread
36
40
import weakref
 
41
 
37
42
from bzrlib import (
38
43
    debug,
39
44
    errors,
40
 
    symbol_versioning,
41
45
    trace,
 
46
    transport,
42
47
    ui,
43
48
    urlutils,
44
49
    )
45
 
from bzrlib.smart import client, protocol, request, vfs
 
50
from bzrlib.i18n import gettext
 
51
from bzrlib.smart import client, protocol, request, signals, vfs
46
52
from bzrlib.transport import ssh
47
53
""")
48
 
#usually already imported, and getting IllegalScoperReplacer on it here.
49
54
from bzrlib import osutils
50
55
 
51
 
# We must not read any more than 64k at a time so we don't risk "no buffer
52
 
# space available" errors on some platforms.  Windows in particular is likely
53
 
# to give error 10053 or 10055 if we read more than 64k from a socket.
54
 
_MAX_READ_SIZE = 64 * 1024
55
 
 
 
56
# Throughout this module buffer size parameters are either limited to be at
 
57
# most _MAX_READ_SIZE, or are ignored and _MAX_READ_SIZE is used instead.
 
58
# For this module's purposes, MAX_SOCKET_CHUNK is a reasonable size for reads
 
59
# from non-sockets as well.
 
60
_MAX_READ_SIZE = osutils.MAX_SOCKET_CHUNK
56
61
 
57
62
def _get_protocol_factory_for_bytes(bytes):
58
63
    """Determine the right protocol factory for 'bytes'.
176
181
        ui.ui_factory.report_transport_activity(self, bytes, direction)
177
182
 
178
183
 
 
184
_bad_file_descriptor = (errno.EBADF,)
 
185
if sys.platform == 'win32':
 
186
    # Given on Windows if you pass a closed socket to select.select. Probably
 
187
    # also given if you pass a file handle to select.
 
188
    WSAENOTSOCK = 10038
 
189
    _bad_file_descriptor += (WSAENOTSOCK,)
 
190
 
 
191
 
179
192
class SmartServerStreamMedium(SmartMedium):
180
193
    """Handles smart commands coming over a stream.
181
194
 
194
207
        the stream.  See also the _push_back method.
195
208
    """
196
209
 
197
 
    def __init__(self, backing_transport, root_client_path='/'):
 
210
    _timer = time.time
 
211
 
 
212
    def __init__(self, backing_transport, root_client_path='/', timeout=None):
198
213
        """Construct new server.
199
214
 
200
215
        :param backing_transport: Transport for the directory served.
203
218
        self.backing_transport = backing_transport
204
219
        self.root_client_path = root_client_path
205
220
        self.finished = False
 
221
        if timeout is None:
 
222
            raise AssertionError('You must supply a timeout.')
 
223
        self._client_timeout = timeout
 
224
        self._client_poll_timeout = min(timeout / 10.0, 1.0)
206
225
        SmartMedium.__init__(self)
207
226
 
208
227
    def serve(self):
214
233
            while not self.finished:
215
234
                server_protocol = self._build_protocol()
216
235
                self._serve_one_request(server_protocol)
 
236
        except errors.ConnectionTimeout, e:
 
237
            trace.note('%s' % (e,))
 
238
            trace.log_exception_quietly()
 
239
            self._disconnect_client()
 
240
            # We reported it, no reason to make a big fuss.
 
241
            return
217
242
        except Exception, e:
218
243
            stderr.write("%s terminating on exception %s\n" % (self, e))
219
244
            raise
 
245
        self._disconnect_client()
 
246
 
 
247
    def _stop_gracefully(self):
 
248
        """When we finish this message, stop looking for more."""
 
249
        trace.mutter('Stopping %s' % (self,))
 
250
        self.finished = True
 
251
 
 
252
    def _disconnect_client(self):
 
253
        """Close the current connection. We stopped due to a timeout/etc."""
 
254
        # The default implementation is a no-op, because that is all we used to
 
255
        # do when disconnecting from a client. I suppose we never had the
 
256
        # *server* initiate a disconnect, before
 
257
 
 
258
    def _wait_for_bytes_with_timeout(self, timeout_seconds):
 
259
        """Wait for more bytes to be read, but timeout if none available.
 
260
 
 
261
        This allows us to detect idle connections, and stop trying to read from
 
262
        them, without setting the socket itself to non-blocking. This also
 
263
        allows us to specify when we watch for idle timeouts.
 
264
 
 
265
        :return: Did we timeout? (True if we timed out, False if there is data
 
266
            to be read)
 
267
        """
 
268
        raise NotImplementedError(self._wait_for_bytes_with_timeout)
220
269
 
221
270
    def _build_protocol(self):
222
271
        """Identifies the version of the incoming request, and returns an
227
276
 
228
277
        :returns: a SmartServerRequestProtocol.
229
278
        """
 
279
        self._wait_for_bytes_with_timeout(self._client_timeout)
 
280
        if self.finished:
 
281
            # We're stopping, so don't try to do any more work
 
282
            return None
230
283
        bytes = self._get_line()
231
284
        protocol_factory, unused_bytes = _get_protocol_factory_for_bytes(bytes)
232
285
        protocol = protocol_factory(
234
287
        protocol.accept_bytes(unused_bytes)
235
288
        return protocol
236
289
 
 
290
    def _wait_on_descriptor(self, fd, timeout_seconds):
 
291
        """select() on a file descriptor, waiting for nonblocking read()
 
292
 
 
293
        This will raise a ConnectionTimeout exception if we do not get a
 
294
        readable handle before timeout_seconds.
 
295
        :return: None
 
296
        """
 
297
        t_end = self._timer() + timeout_seconds
 
298
        poll_timeout = min(timeout_seconds, self._client_poll_timeout)
 
299
        rs = xs = None
 
300
        while not rs and not xs and self._timer() < t_end:
 
301
            if self.finished:
 
302
                return
 
303
            try:
 
304
                rs, _, xs = select.select([fd], [], [fd], poll_timeout)
 
305
            except (select.error, socket.error) as e:
 
306
                err = getattr(e, 'errno', None)
 
307
                if err is None and getattr(e, 'args', None) is not None:
 
308
                    # select.error doesn't have 'errno', it just has args[0]
 
309
                    err = e.args[0]
 
310
                if err in _bad_file_descriptor:
 
311
                    return # Not a socket indicates read() will fail
 
312
                elif err == errno.EINTR:
 
313
                    # Interrupted, keep looping.
 
314
                    continue
 
315
                raise
 
316
        if rs or xs:
 
317
            return
 
318
        raise errors.ConnectionTimeout('disconnecting client after %.1f seconds'
 
319
                                       % (timeout_seconds,))
 
320
 
237
321
    def _serve_one_request(self, protocol):
238
322
        """Read one request from input, process, send back a response.
239
323
 
240
324
        :param protocol: a SmartServerRequestProtocol.
241
325
        """
 
326
        if protocol is None:
 
327
            return
242
328
        try:
243
329
            self._serve_one_request_unguarded(protocol)
244
330
        except KeyboardInterrupt:
260
346
 
261
347
class SmartServerSocketStreamMedium(SmartServerStreamMedium):
262
348
 
263
 
    def __init__(self, sock, backing_transport, root_client_path='/'):
 
349
    def __init__(self, sock, backing_transport, root_client_path='/',
 
350
                 timeout=None):
264
351
        """Constructor.
265
352
 
266
353
        :param sock: the socket the server will read from.  It will be put
267
354
            into blocking mode.
268
355
        """
269
356
        SmartServerStreamMedium.__init__(
270
 
            self, backing_transport, root_client_path=root_client_path)
 
357
            self, backing_transport, root_client_path=root_client_path,
 
358
            timeout=timeout)
271
359
        sock.setblocking(True)
272
360
        self.socket = sock
 
361
        # Get the getpeername now, as we might be closed later when we care.
 
362
        try:
 
363
            self._client_info = sock.getpeername()
 
364
        except socket.error:
 
365
            self._client_info = '<unknown>'
 
366
 
 
367
    def __str__(self):
 
368
        return '%s(client=%s)' % (self.__class__.__name__, self._client_info)
 
369
 
 
370
    def __repr__(self):
 
371
        return '%s.%s(client=%s)' % (self.__module__, self.__class__.__name__,
 
372
            self._client_info)
273
373
 
274
374
    def _serve_one_request_unguarded(self, protocol):
275
375
        while protocol.next_read_size():
276
376
            # We can safely try to read large chunks.  If there is less data
277
 
            # than _MAX_READ_SIZE ready, the socket wil just return a short
278
 
            # read immediately rather than block.
279
 
            bytes = self.read_bytes(_MAX_READ_SIZE)
 
377
            # than MAX_SOCKET_CHUNK ready, the socket will just return a
 
378
            # short read immediately rather than block.
 
379
            bytes = self.read_bytes(osutils.MAX_SOCKET_CHUNK)
280
380
            if bytes == '':
281
381
                self.finished = True
282
382
                return
284
384
 
285
385
        self._push_back(protocol.unused_data)
286
386
 
 
387
    def _disconnect_client(self):
 
388
        """Close the current connection. We stopped due to a timeout/etc."""
 
389
        self.socket.close()
 
390
 
 
391
    def _wait_for_bytes_with_timeout(self, timeout_seconds):
 
392
        """Wait for more bytes to be read, but timeout if none available.
 
393
 
 
394
        This allows us to detect idle connections, and stop trying to read from
 
395
        them, without setting the socket itself to non-blocking. This also
 
396
        allows us to specify when we watch for idle timeouts.
 
397
 
 
398
        :return: None, this will raise ConnectionTimeout if we time out before
 
399
            data is available.
 
400
        """
 
401
        return self._wait_on_descriptor(self.socket, timeout_seconds)
 
402
 
287
403
    def _read_bytes(self, desired_count):
288
 
        return _read_bytes_from_socket(
289
 
            self.socket.recv, desired_count, self._report_activity)
 
404
        return osutils.read_bytes_from_socket(
 
405
            self.socket, self._report_activity)
290
406
 
291
407
    def terminate_due_to_error(self):
292
408
        # TODO: This should log to a server log file, but no such thing
295
411
        self.finished = True
296
412
 
297
413
    def _write_out(self, bytes):
 
414
        tstart = osutils.timer_func()
298
415
        osutils.send_all(self.socket, bytes, self._report_activity)
 
416
        if 'hpss' in debug.debug_flags:
 
417
            thread_id = thread.get_ident()
 
418
            trace.mutter('%12s: [%s] %d bytes to the socket in %.3fs'
 
419
                         % ('wrote', thread_id, len(bytes),
 
420
                            osutils.timer_func() - tstart))
299
421
 
300
422
 
301
423
class SmartServerPipeStreamMedium(SmartServerStreamMedium):
302
424
 
303
 
    def __init__(self, in_file, out_file, backing_transport):
 
425
    def __init__(self, in_file, out_file, backing_transport, timeout=None):
304
426
        """Construct new server.
305
427
 
306
428
        :param in_file: Python file from which requests can be read.
307
429
        :param out_file: Python file to write responses.
308
430
        :param backing_transport: Transport for the directory served.
309
431
        """
310
 
        SmartServerStreamMedium.__init__(self, backing_transport)
 
432
        SmartServerStreamMedium.__init__(self, backing_transport,
 
433
            timeout=timeout)
311
434
        if sys.platform == 'win32':
312
435
            # force binary mode for files
313
436
            import msvcrt
318
441
        self._in = in_file
319
442
        self._out = out_file
320
443
 
 
444
    def serve(self):
 
445
        """See SmartServerStreamMedium.serve"""
 
446
        # This is the regular serve, except it adds signal trapping for soft
 
447
        # shutdown.
 
448
        stop_gracefully = self._stop_gracefully
 
449
        signals.register_on_hangup(id(self), stop_gracefully)
 
450
        try:
 
451
            return super(SmartServerPipeStreamMedium, self).serve()
 
452
        finally:
 
453
            signals.unregister_on_hangup(id(self))
 
454
 
321
455
    def _serve_one_request_unguarded(self, protocol):
322
456
        while True:
323
457
            # We need to be careful not to read past the end of the current
336
470
                return
337
471
            protocol.accept_bytes(bytes)
338
472
 
 
473
    def _disconnect_client(self):
 
474
        self._in.close()
 
475
        self._out.flush()
 
476
        self._out.close()
 
477
 
 
478
    def _wait_for_bytes_with_timeout(self, timeout_seconds):
 
479
        """Wait for more bytes to be read, but timeout if none available.
 
480
 
 
481
        This allows us to detect idle connections, and stop trying to read from
 
482
        them, without setting the socket itself to non-blocking. This also
 
483
        allows us to specify when we watch for idle timeouts.
 
484
 
 
485
        :return: None, this will raise ConnectionTimeout if we time out before
 
486
            data is available.
 
487
        """
 
488
        if (getattr(self._in, 'fileno', None) is None
 
489
            or sys.platform == 'win32'):
 
490
            # You can't select() file descriptors on Windows.
 
491
            return
 
492
        return self._wait_on_descriptor(self._in, timeout_seconds)
 
493
 
339
494
    def _read_bytes(self, desired_count):
340
495
        return self._in.read(desired_count)
341
496
 
485
640
        return self._medium._get_line()
486
641
 
487
642
 
 
643
class _VfsRefuser(object):
 
644
    """An object that refuses all VFS requests.
 
645
 
 
646
    """
 
647
 
 
648
    def __init__(self):
 
649
        client._SmartClient.hooks.install_named_hook(
 
650
            'call', self.check_vfs, 'vfs refuser')
 
651
 
 
652
    def check_vfs(self, params):
 
653
        try:
 
654
            request_method = request.request_handlers.get(params.method)
 
655
        except KeyError:
 
656
            # A method we don't know about doesn't count as a VFS method.
 
657
            return
 
658
        if issubclass(request_method, vfs.VfsRequest):
 
659
            raise errors.HpssVfsRequestNotAllowed(params.method, params.args)
 
660
 
 
661
 
488
662
class _DebugCounter(object):
489
663
    """An object that counts the HPSS calls made to each client medium.
490
664
 
491
 
    When a medium is garbage-collected, or failing that when atexit functions
492
 
    are run, the total number of calls made on that medium are reported via
493
 
    trace.note.
 
665
    When a medium is garbage-collected, or failing that when
 
666
    bzrlib.global_state exits, the total number of calls made on that medium
 
667
    are reported via trace.note.
494
668
    """
495
669
 
496
670
    def __init__(self):
497
671
        self.counts = weakref.WeakKeyDictionary()
498
672
        client._SmartClient.hooks.install_named_hook(
499
673
            'call', self.increment_call_count, 'hpss call counter')
500
 
        atexit.register(self.flush_all)
 
674
        bzrlib.global_state.cleanups.add_cleanup(self.flush_all)
501
675
 
502
676
    def track(self, medium):
503
677
        """Start tracking calls made to a medium.
537
711
        value['count'] = 0
538
712
        value['vfs_count'] = 0
539
713
        if count != 0:
540
 
            trace.note('HPSS calls: %d (%d vfs) %s',
541
 
                       count, vfs_count, medium_repr)
 
714
            trace.note(gettext('HPSS calls: {0} ({1} vfs) {2}').format(
 
715
                       count, vfs_count, medium_repr))
542
716
 
543
717
    def flush_all(self):
544
718
        for ref in list(self.counts.keys()):
545
719
            self.done(ref)
546
720
 
547
721
_debug_counter = None
 
722
_vfs_refuser = None
548
723
 
549
724
 
550
725
class SmartClientMedium(SmartMedium):
567
742
            if _debug_counter is None:
568
743
                _debug_counter = _DebugCounter()
569
744
            _debug_counter.track(self)
 
745
        if 'hpss_client_no_vfs' in debug.debug_flags:
 
746
            global _vfs_refuser
 
747
            if _vfs_refuser is None:
 
748
                _vfs_refuser = _VfsRefuser()
570
749
 
571
750
    def _is_remote_before(self, version_tuple):
572
751
        """Is it possible the remote side supports RPCs for a given version?
601
780
            # which is newer than a previously supplied older-than version.
602
781
            # This indicates that some smart verb call is not guarded
603
782
            # appropriately (it should simply not have been tried).
604
 
            raise AssertionError(
 
783
            trace.mutter(
605
784
                "_remember_remote_is_before(%r) called, but "
606
785
                "_remember_remote_is_before(%r) was called previously."
607
 
                % (version_tuple, self._remote_version_is_before))
 
786
                , version_tuple, self._remote_version_is_before)
 
787
            if 'hpss' in debug.debug_flags:
 
788
                ui.ui_factory.show_warning(
 
789
                    "_remember_remote_is_before(%r) called, but "
 
790
                    "_remember_remote_is_before(%r) was called previously."
 
791
                    % (version_tuple, self._remote_version_is_before))
 
792
            return
608
793
        self._remote_version_is_before = version_tuple
609
794
 
610
795
    def protocol_version(self):
657
842
        """
658
843
        medium_base = urlutils.join(self.base, '/')
659
844
        rel_url = urlutils.relative_url(medium_base, transport.base)
660
 
        return urllib.unquote(rel_url)
 
845
        return urlutils.unquote(rel_url)
661
846
 
662
847
 
663
848
class SmartClientStreamMedium(SmartClientMedium):
698
883
        """
699
884
        return SmartClientStreamMediumRequest(self)
700
885
 
 
886
    def reset(self):
 
887
        """We have been disconnected, reset current state.
 
888
 
 
889
        This resets things like _current_request and connected state.
 
890
        """
 
891
        self.disconnect()
 
892
        self._current_request = None
 
893
 
701
894
 
702
895
class SmartSimplePipesClientMedium(SmartClientStreamMedium):
703
896
    """A client medium using simple pipes.
712
905
 
713
906
    def _accept_bytes(self, bytes):
714
907
        """See SmartClientStreamMedium.accept_bytes."""
715
 
        self._writeable_pipe.write(bytes)
 
908
        try:
 
909
            self._writeable_pipe.write(bytes)
 
910
        except IOError, e:
 
911
            if e.errno in (errno.EINVAL, errno.EPIPE):
 
912
                raise errors.ConnectionReset(
 
913
                    "Error trying to write to subprocess", e)
 
914
            raise
716
915
        self._report_activity(len(bytes), 'write')
717
916
 
718
917
    def _flush(self):
719
918
        """See SmartClientStreamMedium._flush()."""
 
919
        # Note: If flush were to fail, we'd like to raise ConnectionReset, etc.
 
920
        #       However, testing shows that even when the child process is
 
921
        #       gone, this doesn't error.
720
922
        self._writeable_pipe.flush()
721
923
 
722
924
    def _read_bytes(self, count):
723
925
        """See SmartClientStreamMedium._read_bytes."""
724
 
        bytes = self._readable_pipe.read(count)
 
926
        bytes_to_read = min(count, _MAX_READ_SIZE)
 
927
        bytes = self._readable_pipe.read(bytes_to_read)
725
928
        self._report_activity(len(bytes), 'read')
726
929
        return bytes
727
930
 
728
931
 
 
932
class SSHParams(object):
 
933
    """A set of parameters for starting a remote bzr via SSH."""
 
934
 
 
935
    def __init__(self, host, port=None, username=None, password=None,
 
936
            bzr_remote_path='bzr'):
 
937
        self.host = host
 
938
        self.port = port
 
939
        self.username = username
 
940
        self.password = password
 
941
        self.bzr_remote_path = bzr_remote_path
 
942
 
 
943
 
729
944
class SmartSSHClientMedium(SmartClientStreamMedium):
730
 
    """A client medium using SSH."""
731
 
 
732
 
    def __init__(self, host, port=None, username=None, password=None,
733
 
            base=None, vendor=None, bzr_remote_path=None):
 
945
    """A client medium using SSH.
 
946
 
 
947
    It delegates IO to a SmartSimplePipesClientMedium or
 
948
    SmartClientAlreadyConnectedSocketMedium (depending on platform).
 
949
    """
 
950
 
 
951
    def __init__(self, base, ssh_params, vendor=None):
734
952
        """Creates a client that will connect on the first use.
735
953
 
 
954
        :param ssh_params: A SSHParams instance.
736
955
        :param vendor: An optional override for the ssh vendor to use. See
737
956
            bzrlib.transport.ssh for details on ssh vendors.
738
957
        """
739
 
        self._connected = False
740
 
        self._host = host
741
 
        self._password = password
742
 
        self._port = port
743
 
        self._username = username
 
958
        self._real_medium = None
 
959
        self._ssh_params = ssh_params
 
960
        # for the benefit of progress making a short description of this
 
961
        # transport
 
962
        self._scheme = 'bzr+ssh'
744
963
        # SmartClientStreamMedium stores the repr of this object in its
745
964
        # _DebugCounter so we have to store all the values used in our repr
746
965
        # method before calling the super init.
747
966
        SmartClientStreamMedium.__init__(self, base)
748
 
        self._read_from = None
 
967
        self._vendor = vendor
749
968
        self._ssh_connection = None
750
 
        self._vendor = vendor
751
 
        self._write_to = None
752
 
        self._bzr_remote_path = bzr_remote_path
753
 
        # for the benefit of progress making a short description of this
754
 
        # transport
755
 
        self._scheme = 'bzr+ssh'
756
969
 
757
970
    def __repr__(self):
758
 
        return "%s(connected=%r, username=%r, host=%r, port=%r)" % (
 
971
        if self._ssh_params.port is None:
 
972
            maybe_port = ''
 
973
        else:
 
974
            maybe_port = ':%s' % self._ssh_params.port
 
975
        if self._ssh_params.username is None:
 
976
            maybe_user = ''
 
977
        else:
 
978
            maybe_user = '%s@' % self._ssh_params.username
 
979
        return "%s(%s://%s%s%s/)" % (
759
980
            self.__class__.__name__,
760
 
            self._connected,
761
 
            self._username,
762
 
            self._host,
763
 
            self._port)
 
981
            self._scheme,
 
982
            maybe_user,
 
983
            self._ssh_params.host,
 
984
            maybe_port)
764
985
 
765
986
    def _accept_bytes(self, bytes):
766
987
        """See SmartClientStreamMedium.accept_bytes."""
767
988
        self._ensure_connection()
768
 
        self._write_to.write(bytes)
769
 
        self._report_activity(len(bytes), 'write')
 
989
        self._real_medium.accept_bytes(bytes)
770
990
 
771
991
    def disconnect(self):
772
992
        """See SmartClientMedium.disconnect()."""
773
 
        if not self._connected:
774
 
            return
775
 
        self._read_from.close()
776
 
        self._write_to.close()
777
 
        self._ssh_connection.close()
778
 
        self._connected = False
 
993
        if self._real_medium is not None:
 
994
            self._real_medium.disconnect()
 
995
            self._real_medium = None
 
996
        if self._ssh_connection is not None:
 
997
            self._ssh_connection.close()
 
998
            self._ssh_connection = None
779
999
 
780
1000
    def _ensure_connection(self):
781
1001
        """Connect this medium if not already connected."""
782
 
        if self._connected:
 
1002
        if self._real_medium is not None:
783
1003
            return
784
1004
        if self._vendor is None:
785
1005
            vendor = ssh._get_ssh_vendor()
786
1006
        else:
787
1007
            vendor = self._vendor
788
 
        self._ssh_connection = vendor.connect_ssh(self._username,
789
 
                self._password, self._host, self._port,
790
 
                command=[self._bzr_remote_path, 'serve', '--inet',
 
1008
        self._ssh_connection = vendor.connect_ssh(self._ssh_params.username,
 
1009
                self._ssh_params.password, self._ssh_params.host,
 
1010
                self._ssh_params.port,
 
1011
                command=[self._ssh_params.bzr_remote_path, 'serve', '--inet',
791
1012
                         '--directory=/', '--allow-writes'])
792
 
        self._read_from, self._write_to = \
793
 
            self._ssh_connection.get_filelike_channels()
794
 
        self._connected = True
 
1013
        io_kind, io_object = self._ssh_connection.get_sock_or_pipes()
 
1014
        if io_kind == 'socket':
 
1015
            self._real_medium = SmartClientAlreadyConnectedSocketMedium(
 
1016
                self.base, io_object)
 
1017
        elif io_kind == 'pipes':
 
1018
            read_from, write_to = io_object
 
1019
            self._real_medium = SmartSimplePipesClientMedium(
 
1020
                read_from, write_to, self.base)
 
1021
        else:
 
1022
            raise AssertionError(
 
1023
                "Unexpected io_kind %r from %r"
 
1024
                % (io_kind, self._ssh_connection))
 
1025
        for hook in transport.Transport.hooks["post_connect"]:
 
1026
            hook(self)
795
1027
 
796
1028
    def _flush(self):
797
1029
        """See SmartClientStreamMedium._flush()."""
798
 
        self._write_to.flush()
 
1030
        self._real_medium._flush()
799
1031
 
800
1032
    def _read_bytes(self, count):
801
1033
        """See SmartClientStreamMedium.read_bytes."""
802
 
        if not self._connected:
 
1034
        if self._real_medium is None:
803
1035
            raise errors.MediumNotConnected(self)
804
 
        bytes_to_read = min(count, _MAX_READ_SIZE)
805
 
        bytes = self._read_from.read(bytes_to_read)
806
 
        self._report_activity(len(bytes), 'read')
807
 
        return bytes
 
1036
        return self._real_medium.read_bytes(count)
808
1037
 
809
1038
 
810
1039
# Port 4155 is the default port for bzr://, registered with IANA.
812
1041
BZR_DEFAULT_PORT = 4155
813
1042
 
814
1043
 
815
 
class SmartTCPClientMedium(SmartClientStreamMedium):
816
 
    """A client medium using TCP."""
817
 
 
818
 
    def __init__(self, host, port, base):
819
 
        """Creates a client that will connect on the first use."""
 
1044
class SmartClientSocketMedium(SmartClientStreamMedium):
 
1045
    """A client medium using a socket.
 
1046
 
 
1047
    This class isn't usable directly.  Use one of its subclasses instead.
 
1048
    """
 
1049
 
 
1050
    def __init__(self, base):
820
1051
        SmartClientStreamMedium.__init__(self, base)
 
1052
        self._socket = None
821
1053
        self._connected = False
822
 
        self._host = host
823
 
        self._port = port
824
 
        self._socket = None
825
1054
 
826
1055
    def _accept_bytes(self, bytes):
827
1056
        """See SmartClientMedium.accept_bytes."""
828
1057
        self._ensure_connection()
829
1058
        osutils.send_all(self._socket, bytes, self._report_activity)
830
1059
 
 
1060
    def _ensure_connection(self):
 
1061
        """Connect this medium if not already connected."""
 
1062
        raise NotImplementedError(self._ensure_connection)
 
1063
 
 
1064
    def _flush(self):
 
1065
        """See SmartClientStreamMedium._flush().
 
1066
 
 
1067
        For sockets we do no flushing. For TCP sockets we may want to turn off
 
1068
        TCP_NODELAY and add a means to do a flush, but that can be done in the
 
1069
        future.
 
1070
        """
 
1071
 
 
1072
    def _read_bytes(self, count):
 
1073
        """See SmartClientMedium.read_bytes."""
 
1074
        if not self._connected:
 
1075
            raise errors.MediumNotConnected(self)
 
1076
        return osutils.read_bytes_from_socket(
 
1077
            self._socket, self._report_activity)
 
1078
 
831
1079
    def disconnect(self):
832
1080
        """See SmartClientMedium.disconnect()."""
833
1081
        if not self._connected:
836
1084
        self._socket = None
837
1085
        self._connected = False
838
1086
 
 
1087
 
 
1088
class SmartTCPClientMedium(SmartClientSocketMedium):
 
1089
    """A client medium that creates a TCP connection."""
 
1090
 
 
1091
    def __init__(self, host, port, base):
 
1092
        """Creates a client that will connect on the first use."""
 
1093
        SmartClientSocketMedium.__init__(self, base)
 
1094
        self._host = host
 
1095
        self._port = port
 
1096
 
839
1097
    def _ensure_connection(self):
840
1098
        """Connect this medium if not already connected."""
841
1099
        if self._connected:
874
1132
            raise errors.ConnectionError("failed to connect to %s:%d: %s" %
875
1133
                    (self._host, port, err_msg))
876
1134
        self._connected = True
877
 
 
878
 
    def _flush(self):
879
 
        """See SmartClientStreamMedium._flush().
880
 
 
881
 
        For TCP we do no flushing. We may want to turn off TCP_NODELAY and
882
 
        add a means to do a flush, but that can be done in the future.
883
 
        """
884
 
 
885
 
    def _read_bytes(self, count):
886
 
        """See SmartClientMedium.read_bytes."""
887
 
        if not self._connected:
888
 
            raise errors.MediumNotConnected(self)
889
 
        return _read_bytes_from_socket(
890
 
            self._socket.recv, count, self._report_activity)
 
1135
        for hook in transport.Transport.hooks["post_connect"]:
 
1136
            hook(self)
 
1137
 
 
1138
 
 
1139
class SmartClientAlreadyConnectedSocketMedium(SmartClientSocketMedium):
 
1140
    """A client medium for an already connected socket.
 
1141
    
 
1142
    Note that this class will assume it "owns" the socket, so it will close it
 
1143
    when its disconnect method is called.
 
1144
    """
 
1145
 
 
1146
    def __init__(self, base, sock):
 
1147
        SmartClientSocketMedium.__init__(self, base)
 
1148
        self._socket = sock
 
1149
        self._connected = True
 
1150
 
 
1151
    def _ensure_connection(self):
 
1152
        # Already connected, by definition!  So nothing to do.
 
1153
        pass
891
1154
 
892
1155
 
893
1156
class SmartClientStreamMediumRequest(SmartClientMediumRequest):
928
1191
        This invokes self._medium._flush to ensure all bytes are transmitted.
929
1192
        """
930
1193
        self._medium._flush()
931
 
 
932
 
 
933
 
def _read_bytes_from_socket(sock, desired_count, report_activity):
934
 
    # We ignore the desired_count because on sockets it's more efficient to
935
 
    # read large chunks (of _MAX_READ_SIZE bytes) at a time.
936
 
    try:
937
 
        bytes = osutils.until_no_eintr(sock, _MAX_READ_SIZE)
938
 
    except socket.error, e:
939
 
        if len(e.args) and e.args[0] in (errno.ECONNRESET, 10054):
940
 
            # The connection was closed by the other side.  Callers expect an
941
 
            # empty string to signal end-of-stream.
942
 
            bytes = ''
943
 
        else:
944
 
            raise
945
 
    else:
946
 
        report_activity(len(bytes), 'read')
947
 
    return bytes
948