~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/medium.py

  • Committer: Robert Collins
  • Date: 2007-07-04 08:08:13 UTC
  • mfrom: (2572 +trunk)
  • mto: This revision was merged to the branch mainline in revision 2587.
  • Revision ID: robertc@robertcollins.net-20070704080813-wzebx0r88fvwj5rq
Merge bzr.dev.

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2006 Canonical Ltd
 
2
#
 
3
# This program is free software; you can redistribute it and/or modify
 
4
# it under the terms of the GNU General Public License as published by
 
5
# the Free Software Foundation; either version 2 of the License, or
 
6
# (at your option) any later version.
 
7
#
 
8
# This program is distributed in the hope that it will be useful,
 
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
11
# GNU General Public License for more details.
 
12
#
 
13
# You should have received a copy of the GNU General Public License
 
14
# along with this program; if not, write to the Free Software
 
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
16
 
 
17
"""The 'medium' layer for the smart servers and clients.
 
18
 
 
19
"Medium" here is the noun meaning "a means of transmission", not the adjective
 
20
for "the quality between big and small."
 
21
 
 
22
Media carry the bytes of the requests somehow (e.g. via TCP, wrapped in HTTP, or
 
23
over SSH), and pass them to and from the protocol logic.  See the overview in
 
24
bzrlib/transport/smart/__init__.py.
 
25
"""
 
26
 
 
27
import os
 
28
import socket
 
29
import sys
 
30
 
 
31
from bzrlib import errors
 
32
from bzrlib.smart.protocol import (
 
33
    REQUEST_VERSION_TWO,
 
34
    SmartServerRequestProtocolOne,
 
35
    SmartServerRequestProtocolTwo,
 
36
    )
 
37
 
 
38
try:
 
39
    from bzrlib.transport import ssh
 
40
except errors.ParamikoNotPresent:
 
41
    # no paramiko.  SmartSSHClientMedium will break.
 
42
    pass
 
43
 
 
44
 
 
45
class SmartServerStreamMedium(object):
 
46
    """Handles smart commands coming over a stream.
 
47
 
 
48
    The stream may be a pipe connected to sshd, or a tcp socket, or an
 
49
    in-process fifo for testing.
 
50
 
 
51
    One instance is created for each connected client; it can serve multiple
 
52
    requests in the lifetime of the connection.
 
53
 
 
54
    The server passes requests through to an underlying backing transport, 
 
55
    which will typically be a LocalTransport looking at the server's filesystem.
 
56
    """
 
57
 
 
58
    def __init__(self, backing_transport):
 
59
        """Construct new server.
 
60
 
 
61
        :param backing_transport: Transport for the directory served.
 
62
        """
 
63
        # backing_transport could be passed to serve instead of __init__
 
64
        self.backing_transport = backing_transport
 
65
        self.finished = False
 
66
 
 
67
    def serve(self):
 
68
        """Serve requests until the client disconnects."""
 
69
        # Keep a reference to stderr because the sys module's globals get set to
 
70
        # None during interpreter shutdown.
 
71
        from sys import stderr
 
72
        try:
 
73
            while not self.finished:
 
74
                server_protocol = self._build_protocol()
 
75
                self._serve_one_request(server_protocol)
 
76
        except Exception, e:
 
77
            stderr.write("%s terminating on exception %s\n" % (self, e))
 
78
            raise
 
79
 
 
80
    def _build_protocol(self):
 
81
        """Identifies the version of the incoming request, and returns an
 
82
        a protocol object that can interpret it.
 
83
 
 
84
        If more bytes than the version prefix of the request are read, they will
 
85
        be fed into the protocol before it is returned.
 
86
 
 
87
        :returns: a SmartServerRequestProtocol.
 
88
        """
 
89
        # Identify the protocol version.
 
90
        bytes = self._get_line()
 
91
        if bytes.startswith(REQUEST_VERSION_TWO):
 
92
            protocol_class = SmartServerRequestProtocolTwo
 
93
            bytes = bytes[len(REQUEST_VERSION_TWO):]
 
94
        else:
 
95
            protocol_class = SmartServerRequestProtocolOne
 
96
        protocol = protocol_class(self.backing_transport, self._write_out)
 
97
        protocol.accept_bytes(bytes)
 
98
        return protocol
 
99
 
 
100
    def _serve_one_request(self, protocol):
 
101
        """Read one request from input, process, send back a response.
 
102
        
 
103
        :param protocol: a SmartServerRequestProtocol.
 
104
        """
 
105
        try:
 
106
            self._serve_one_request_unguarded(protocol)
 
107
        except KeyboardInterrupt:
 
108
            raise
 
109
        except Exception, e:
 
110
            self.terminate_due_to_error()
 
111
 
 
112
    def terminate_due_to_error(self):
 
113
        """Called when an unhandled exception from the protocol occurs."""
 
114
        raise NotImplementedError(self.terminate_due_to_error)
 
115
 
 
116
    def _get_bytes(self, desired_count):
 
117
        """Get some bytes from the medium.
 
118
 
 
119
        :param desired_count: number of bytes we want to read.
 
120
        """
 
121
        raise NotImplementedError(self._get_bytes)
 
122
 
 
123
    def _get_line(self):
 
124
        """Read bytes from this request's response until a newline byte.
 
125
        
 
126
        This isn't particularly efficient, so should only be used when the
 
127
        expected size of the line is quite short.
 
128
 
 
129
        :returns: a string of bytes ending in a newline (byte 0x0A).
 
130
        """
 
131
        # XXX: this duplicates SmartClientRequestProtocolOne._recv_tuple
 
132
        line = ''
 
133
        while not line or line[-1] != '\n':
 
134
            new_char = self._get_bytes(1)
 
135
            line += new_char
 
136
            if new_char == '':
 
137
                # Ran out of bytes before receiving a complete line.
 
138
                break
 
139
        return line
 
140
 
 
141
 
 
142
class SmartServerSocketStreamMedium(SmartServerStreamMedium):
 
143
 
 
144
    def __init__(self, sock, backing_transport):
 
145
        """Constructor.
 
146
 
 
147
        :param sock: the socket the server will read from.  It will be put
 
148
            into blocking mode.
 
149
        """
 
150
        SmartServerStreamMedium.__init__(self, backing_transport)
 
151
        self.push_back = ''
 
152
        sock.setblocking(True)
 
153
        self.socket = sock
 
154
 
 
155
    def _serve_one_request_unguarded(self, protocol):
 
156
        while protocol.next_read_size():
 
157
            if self.push_back:
 
158
                protocol.accept_bytes(self.push_back)
 
159
                self.push_back = ''
 
160
            else:
 
161
                bytes = self._get_bytes(4096)
 
162
                if bytes == '':
 
163
                    self.finished = True
 
164
                    return
 
165
                protocol.accept_bytes(bytes)
 
166
        
 
167
        self.push_back = protocol.excess_buffer
 
168
 
 
169
    def _get_bytes(self, desired_count):
 
170
        # We ignore the desired_count because on sockets it's more efficient to
 
171
        # read 4k at a time.
 
172
        return self.socket.recv(4096)
 
173
    
 
174
    def terminate_due_to_error(self):
 
175
        """Called when an unhandled exception from the protocol occurs."""
 
176
        # TODO: This should log to a server log file, but no such thing
 
177
        # exists yet.  Andrew Bennetts 2006-09-29.
 
178
        self.socket.close()
 
179
        self.finished = True
 
180
 
 
181
    def _write_out(self, bytes):
 
182
        self.socket.sendall(bytes)
 
183
 
 
184
 
 
185
class SmartServerPipeStreamMedium(SmartServerStreamMedium):
 
186
 
 
187
    def __init__(self, in_file, out_file, backing_transport):
 
188
        """Construct new server.
 
189
 
 
190
        :param in_file: Python file from which requests can be read.
 
191
        :param out_file: Python file to write responses.
 
192
        :param backing_transport: Transport for the directory served.
 
193
        """
 
194
        SmartServerStreamMedium.__init__(self, backing_transport)
 
195
        if sys.platform == 'win32':
 
196
            # force binary mode for files
 
197
            import msvcrt
 
198
            for f in (in_file, out_file):
 
199
                fileno = getattr(f, 'fileno', None)
 
200
                if fileno:
 
201
                    msvcrt.setmode(fileno(), os.O_BINARY)
 
202
        self._in = in_file
 
203
        self._out = out_file
 
204
 
 
205
    def _serve_one_request_unguarded(self, protocol):
 
206
        while True:
 
207
            bytes_to_read = protocol.next_read_size()
 
208
            if bytes_to_read == 0:
 
209
                # Finished serving this request.
 
210
                self._out.flush()
 
211
                return
 
212
            bytes = self._get_bytes(bytes_to_read)
 
213
            if bytes == '':
 
214
                # Connection has been closed.
 
215
                self.finished = True
 
216
                self._out.flush()
 
217
                return
 
218
            protocol.accept_bytes(bytes)
 
219
 
 
220
    def _get_bytes(self, desired_count):
 
221
        return self._in.read(desired_count)
 
222
 
 
223
    def terminate_due_to_error(self):
 
224
        # TODO: This should log to a server log file, but no such thing
 
225
        # exists yet.  Andrew Bennetts 2006-09-29.
 
226
        self._out.close()
 
227
        self.finished = True
 
228
 
 
229
    def _write_out(self, bytes):
 
230
        self._out.write(bytes)
 
231
 
 
232
 
 
233
class SmartClientMediumRequest(object):
 
234
    """A request on a SmartClientMedium.
 
235
 
 
236
    Each request allows bytes to be provided to it via accept_bytes, and then
 
237
    the response bytes to be read via read_bytes.
 
238
 
 
239
    For instance:
 
240
    request.accept_bytes('123')
 
241
    request.finished_writing()
 
242
    result = request.read_bytes(3)
 
243
    request.finished_reading()
 
244
 
 
245
    It is up to the individual SmartClientMedium whether multiple concurrent
 
246
    requests can exist. See SmartClientMedium.get_request to obtain instances 
 
247
    of SmartClientMediumRequest, and the concrete Medium you are using for 
 
248
    details on concurrency and pipelining.
 
249
    """
 
250
 
 
251
    def __init__(self, medium):
 
252
        """Construct a SmartClientMediumRequest for the medium medium."""
 
253
        self._medium = medium
 
254
        # we track state by constants - we may want to use the same
 
255
        # pattern as BodyReader if it gets more complex.
 
256
        # valid states are: "writing", "reading", "done"
 
257
        self._state = "writing"
 
258
 
 
259
    def accept_bytes(self, bytes):
 
260
        """Accept bytes for inclusion in this request.
 
261
 
 
262
        This method may not be be called after finished_writing() has been
 
263
        called.  It depends upon the Medium whether or not the bytes will be
 
264
        immediately transmitted. Message based Mediums will tend to buffer the
 
265
        bytes until finished_writing() is called.
 
266
 
 
267
        :param bytes: A bytestring.
 
268
        """
 
269
        if self._state != "writing":
 
270
            raise errors.WritingCompleted(self)
 
271
        self._accept_bytes(bytes)
 
272
 
 
273
    def _accept_bytes(self, bytes):
 
274
        """Helper for accept_bytes.
 
275
 
 
276
        Accept_bytes checks the state of the request to determing if bytes
 
277
        should be accepted. After that it hands off to _accept_bytes to do the
 
278
        actual acceptance.
 
279
        """
 
280
        raise NotImplementedError(self._accept_bytes)
 
281
 
 
282
    def finished_reading(self):
 
283
        """Inform the request that all desired data has been read.
 
284
 
 
285
        This will remove the request from the pipeline for its medium (if the
 
286
        medium supports pipelining) and any further calls to methods on the
 
287
        request will raise ReadingCompleted.
 
288
        """
 
289
        if self._state == "writing":
 
290
            raise errors.WritingNotComplete(self)
 
291
        if self._state != "reading":
 
292
            raise errors.ReadingCompleted(self)
 
293
        self._state = "done"
 
294
        self._finished_reading()
 
295
 
 
296
    def _finished_reading(self):
 
297
        """Helper for finished_reading.
 
298
 
 
299
        finished_reading checks the state of the request to determine if 
 
300
        finished_reading is allowed, and if it is hands off to _finished_reading
 
301
        to perform the action.
 
302
        """
 
303
        raise NotImplementedError(self._finished_reading)
 
304
 
 
305
    def finished_writing(self):
 
306
        """Finish the writing phase of this request.
 
307
 
 
308
        This will flush all pending data for this request along the medium.
 
309
        After calling finished_writing, you may not call accept_bytes anymore.
 
310
        """
 
311
        if self._state != "writing":
 
312
            raise errors.WritingCompleted(self)
 
313
        self._state = "reading"
 
314
        self._finished_writing()
 
315
 
 
316
    def _finished_writing(self):
 
317
        """Helper for finished_writing.
 
318
 
 
319
        finished_writing checks the state of the request to determine if 
 
320
        finished_writing is allowed, and if it is hands off to _finished_writing
 
321
        to perform the action.
 
322
        """
 
323
        raise NotImplementedError(self._finished_writing)
 
324
 
 
325
    def read_bytes(self, count):
 
326
        """Read bytes from this requests response.
 
327
 
 
328
        This method will block and wait for count bytes to be read. It may not
 
329
        be invoked until finished_writing() has been called - this is to ensure
 
330
        a message-based approach to requests, for compatibility with message
 
331
        based mediums like HTTP.
 
332
        """
 
333
        if self._state == "writing":
 
334
            raise errors.WritingNotComplete(self)
 
335
        if self._state != "reading":
 
336
            raise errors.ReadingCompleted(self)
 
337
        return self._read_bytes(count)
 
338
 
 
339
    def _read_bytes(self, count):
 
340
        """Helper for read_bytes.
 
341
 
 
342
        read_bytes checks the state of the request to determing if bytes
 
343
        should be read. After that it hands off to _read_bytes to do the
 
344
        actual read.
 
345
        """
 
346
        raise NotImplementedError(self._read_bytes)
 
347
 
 
348
    def read_line(self):
 
349
        """Read bytes from this request's response until a newline byte.
 
350
        
 
351
        This isn't particularly efficient, so should only be used when the
 
352
        expected size of the line is quite short.
 
353
 
 
354
        :returns: a string of bytes ending in a newline (byte 0x0A).
 
355
        """
 
356
        # XXX: this duplicates SmartClientRequestProtocolOne._recv_tuple
 
357
        line = ''
 
358
        while not line or line[-1] != '\n':
 
359
            new_char = self.read_bytes(1)
 
360
            line += new_char
 
361
            if new_char == '':
 
362
                raise errors.SmartProtocolError(
 
363
                    'unexpected end of file reading from server')
 
364
        return line
 
365
 
 
366
 
 
367
class SmartClientMedium(object):
 
368
    """Smart client is a medium for sending smart protocol requests over."""
 
369
 
 
370
    def disconnect(self):
 
371
        """If this medium maintains a persistent connection, close it.
 
372
        
 
373
        The default implementation does nothing.
 
374
        """
 
375
        
 
376
 
 
377
class SmartClientStreamMedium(SmartClientMedium):
 
378
    """Stream based medium common class.
 
379
 
 
380
    SmartClientStreamMediums operate on a stream. All subclasses use a common
 
381
    SmartClientStreamMediumRequest for their requests, and should implement
 
382
    _accept_bytes and _read_bytes to allow the request objects to send and
 
383
    receive bytes.
 
384
    """
 
385
 
 
386
    def __init__(self):
 
387
        self._current_request = None
 
388
 
 
389
    def accept_bytes(self, bytes):
 
390
        self._accept_bytes(bytes)
 
391
 
 
392
    def __del__(self):
 
393
        """The SmartClientStreamMedium knows how to close the stream when it is
 
394
        finished with it.
 
395
        """
 
396
        self.disconnect()
 
397
 
 
398
    def _flush(self):
 
399
        """Flush the output stream.
 
400
        
 
401
        This method is used by the SmartClientStreamMediumRequest to ensure that
 
402
        all data for a request is sent, to avoid long timeouts or deadlocks.
 
403
        """
 
404
        raise NotImplementedError(self._flush)
 
405
 
 
406
    def get_request(self):
 
407
        """See SmartClientMedium.get_request().
 
408
 
 
409
        SmartClientStreamMedium always returns a SmartClientStreamMediumRequest
 
410
        for get_request.
 
411
        """
 
412
        return SmartClientStreamMediumRequest(self)
 
413
 
 
414
    def read_bytes(self, count):
 
415
        return self._read_bytes(count)
 
416
 
 
417
 
 
418
class SmartSimplePipesClientMedium(SmartClientStreamMedium):
 
419
    """A client medium using simple pipes.
 
420
    
 
421
    This client does not manage the pipes: it assumes they will always be open.
 
422
    """
 
423
 
 
424
    def __init__(self, readable_pipe, writeable_pipe):
 
425
        SmartClientStreamMedium.__init__(self)
 
426
        self._readable_pipe = readable_pipe
 
427
        self._writeable_pipe = writeable_pipe
 
428
 
 
429
    def _accept_bytes(self, bytes):
 
430
        """See SmartClientStreamMedium.accept_bytes."""
 
431
        self._writeable_pipe.write(bytes)
 
432
 
 
433
    def _flush(self):
 
434
        """See SmartClientStreamMedium._flush()."""
 
435
        self._writeable_pipe.flush()
 
436
 
 
437
    def _read_bytes(self, count):
 
438
        """See SmartClientStreamMedium._read_bytes."""
 
439
        return self._readable_pipe.read(count)
 
440
 
 
441
 
 
442
class SmartSSHClientMedium(SmartClientStreamMedium):
 
443
    """A client medium using SSH."""
 
444
    
 
445
    def __init__(self, host, port=None, username=None, password=None,
 
446
            vendor=None):
 
447
        """Creates a client that will connect on the first use.
 
448
        
 
449
        :param vendor: An optional override for the ssh vendor to use. See
 
450
            bzrlib.transport.ssh for details on ssh vendors.
 
451
        """
 
452
        SmartClientStreamMedium.__init__(self)
 
453
        self._connected = False
 
454
        self._host = host
 
455
        self._password = password
 
456
        self._port = port
 
457
        self._username = username
 
458
        self._read_from = None
 
459
        self._ssh_connection = None
 
460
        self._vendor = vendor
 
461
        self._write_to = None
 
462
 
 
463
    def _accept_bytes(self, bytes):
 
464
        """See SmartClientStreamMedium.accept_bytes."""
 
465
        self._ensure_connection()
 
466
        self._write_to.write(bytes)
 
467
 
 
468
    def disconnect(self):
 
469
        """See SmartClientMedium.disconnect()."""
 
470
        if not self._connected:
 
471
            return
 
472
        self._read_from.close()
 
473
        self._write_to.close()
 
474
        self._ssh_connection.close()
 
475
        self._connected = False
 
476
 
 
477
    def _ensure_connection(self):
 
478
        """Connect this medium if not already connected."""
 
479
        if self._connected:
 
480
            return
 
481
        executable = os.environ.get('BZR_REMOTE_PATH', 'bzr')
 
482
        if self._vendor is None:
 
483
            vendor = ssh._get_ssh_vendor()
 
484
        else:
 
485
            vendor = self._vendor
 
486
        self._ssh_connection = vendor.connect_ssh(self._username,
 
487
                self._password, self._host, self._port,
 
488
                command=[executable, 'serve', '--inet', '--directory=/',
 
489
                         '--allow-writes'])
 
490
        self._read_from, self._write_to = \
 
491
            self._ssh_connection.get_filelike_channels()
 
492
        self._connected = True
 
493
 
 
494
    def _flush(self):
 
495
        """See SmartClientStreamMedium._flush()."""
 
496
        self._write_to.flush()
 
497
 
 
498
    def _read_bytes(self, count):
 
499
        """See SmartClientStreamMedium.read_bytes."""
 
500
        if not self._connected:
 
501
            raise errors.MediumNotConnected(self)
 
502
        return self._read_from.read(count)
 
503
 
 
504
 
 
505
class SmartTCPClientMedium(SmartClientStreamMedium):
 
506
    """A client medium using TCP."""
 
507
    
 
508
    def __init__(self, host, port):
 
509
        """Creates a client that will connect on the first use."""
 
510
        SmartClientStreamMedium.__init__(self)
 
511
        self._connected = False
 
512
        self._host = host
 
513
        self._port = port
 
514
        self._socket = None
 
515
 
 
516
    def _accept_bytes(self, bytes):
 
517
        """See SmartClientMedium.accept_bytes."""
 
518
        self._ensure_connection()
 
519
        self._socket.sendall(bytes)
 
520
 
 
521
    def disconnect(self):
 
522
        """See SmartClientMedium.disconnect()."""
 
523
        if not self._connected:
 
524
            return
 
525
        self._socket.close()
 
526
        self._socket = None
 
527
        self._connected = False
 
528
 
 
529
    def _ensure_connection(self):
 
530
        """Connect this medium if not already connected."""
 
531
        if self._connected:
 
532
            return
 
533
        self._socket = socket.socket()
 
534
        self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
535
        result = self._socket.connect_ex((self._host, int(self._port)))
 
536
        if result:
 
537
            raise errors.ConnectionError("failed to connect to %s:%d: %s" %
 
538
                    (self._host, self._port, os.strerror(result)))
 
539
        self._connected = True
 
540
 
 
541
    def _flush(self):
 
542
        """See SmartClientStreamMedium._flush().
 
543
        
 
544
        For TCP we do no flushing. We may want to turn off TCP_NODELAY and 
 
545
        add a means to do a flush, but that can be done in the future.
 
546
        """
 
547
 
 
548
    def _read_bytes(self, count):
 
549
        """See SmartClientMedium.read_bytes."""
 
550
        if not self._connected:
 
551
            raise errors.MediumNotConnected(self)
 
552
        return self._socket.recv(count)
 
553
 
 
554
 
 
555
class SmartClientStreamMediumRequest(SmartClientMediumRequest):
 
556
    """A SmartClientMediumRequest that works with an SmartClientStreamMedium."""
 
557
 
 
558
    def __init__(self, medium):
 
559
        SmartClientMediumRequest.__init__(self, medium)
 
560
        # check that we are safe concurrency wise. If some streams start
 
561
        # allowing concurrent requests - i.e. via multiplexing - then this
 
562
        # assert should be moved to SmartClientStreamMedium.get_request,
 
563
        # and the setting/unsetting of _current_request likewise moved into
 
564
        # that class : but its unneeded overhead for now. RBC 20060922
 
565
        if self._medium._current_request is not None:
 
566
            raise errors.TooManyConcurrentRequests(self._medium)
 
567
        self._medium._current_request = self
 
568
 
 
569
    def _accept_bytes(self, bytes):
 
570
        """See SmartClientMediumRequest._accept_bytes.
 
571
        
 
572
        This forwards to self._medium._accept_bytes because we are operating
 
573
        on the mediums stream.
 
574
        """
 
575
        self._medium._accept_bytes(bytes)
 
576
 
 
577
    def _finished_reading(self):
 
578
        """See SmartClientMediumRequest._finished_reading.
 
579
 
 
580
        This clears the _current_request on self._medium to allow a new 
 
581
        request to be created.
 
582
        """
 
583
        assert self._medium._current_request is self
 
584
        self._medium._current_request = None
 
585
        
 
586
    def _finished_writing(self):
 
587
        """See SmartClientMediumRequest._finished_writing.
 
588
 
 
589
        This invokes self._medium._flush to ensure all bytes are transmitted.
 
590
        """
 
591
        self._medium._flush()
 
592
 
 
593
    def _read_bytes(self, count):
 
594
        """See SmartClientMediumRequest._read_bytes.
 
595
        
 
596
        This forwards to self._medium._read_bytes because we are operating
 
597
        on the mediums stream.
 
598
        """
 
599
        return self._medium._read_bytes(count)
 
600