buffered_pipe.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. # Copyright (C) 2006-2007 Robey Pointer <robeypointer@gmail.com>
  2. #
  3. # This file is part of paramiko.
  4. #
  5. # Paramiko is free software; you can redistribute it and/or modify it under the
  6. # terms of the GNU Lesser General Public License as published by the Free
  7. # Software Foundation; either version 2.1 of the License, or (at your option)
  8. # any later version.
  9. #
  10. # Paramiko is distributed in the hope that it will be useful, but WITHOUT ANY
  11. # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  12. # A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
  13. # details.
  14. #
  15. # You should have received a copy of the GNU Lesser General Public License
  16. # along with Paramiko; if not, write to the Free Software Foundation, Inc.,
  17. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  18. """
  19. Attempt to generalize the "feeder" part of a `.Channel`: an object which can be
  20. read from and closed, but is reading from a buffer fed by another thread. The
  21. read operations are blocking and can have a timeout set.
  22. """
  23. import array
  24. import threading
  25. import time
  26. from paramiko.py3compat import PY2, b
  27. class PipeTimeout(IOError):
  28. """
  29. Indicates that a timeout was reached on a read from a `.BufferedPipe`.
  30. """
  31. pass
  32. class BufferedPipe(object):
  33. """
  34. A buffer that obeys normal read (with timeout) & close semantics for a
  35. file or socket, but is fed data from another thread. This is used by
  36. `.Channel`.
  37. """
  38. def __init__(self):
  39. self._lock = threading.Lock()
  40. self._cv = threading.Condition(self._lock)
  41. self._event = None
  42. self._buffer = array.array("B")
  43. self._closed = False
  44. if PY2:
  45. def _buffer_frombytes(self, data):
  46. self._buffer.fromstring(data)
  47. def _buffer_tobytes(self, limit=None):
  48. return self._buffer[:limit].tostring()
  49. else:
  50. def _buffer_frombytes(self, data):
  51. self._buffer.frombytes(data)
  52. def _buffer_tobytes(self, limit=None):
  53. return self._buffer[:limit].tobytes()
  54. def set_event(self, event):
  55. """
  56. Set an event on this buffer. When data is ready to be read (or the
  57. buffer has been closed), the event will be set. When no data is
  58. ready, the event will be cleared.
  59. :param threading.Event event: the event to set/clear
  60. """
  61. self._lock.acquire()
  62. try:
  63. self._event = event
  64. # Make sure the event starts in `set` state if we appear to already
  65. # be closed; otherwise, if we start in `clear` state & are closed,
  66. # nothing will ever call `.feed` and the event (& OS pipe, if we're
  67. # wrapping one - see `Channel.fileno`) will permanently stay in
  68. # `clear`, causing deadlock if e.g. `select`ed upon.
  69. if self._closed or len(self._buffer) > 0:
  70. event.set()
  71. else:
  72. event.clear()
  73. finally:
  74. self._lock.release()
  75. def feed(self, data):
  76. """
  77. Feed new data into this pipe. This method is assumed to be called
  78. from a separate thread, so synchronization is done.
  79. :param data: the data to add, as a ``str`` or ``bytes``
  80. """
  81. self._lock.acquire()
  82. try:
  83. if self._event is not None:
  84. self._event.set()
  85. self._buffer_frombytes(b(data))
  86. self._cv.notify_all()
  87. finally:
  88. self._lock.release()
  89. def read_ready(self):
  90. """
  91. Returns true if data is buffered and ready to be read from this
  92. feeder. A ``False`` result does not mean that the feeder has closed;
  93. it means you may need to wait before more data arrives.
  94. :return:
  95. ``True`` if a `read` call would immediately return at least one
  96. byte; ``False`` otherwise.
  97. """
  98. self._lock.acquire()
  99. try:
  100. if len(self._buffer) == 0:
  101. return False
  102. return True
  103. finally:
  104. self._lock.release()
  105. def read(self, nbytes, timeout=None):
  106. """
  107. Read data from the pipe. The return value is a string representing
  108. the data received. The maximum amount of data to be received at once
  109. is specified by ``nbytes``. If a string of length zero is returned,
  110. the pipe has been closed.
  111. The optional ``timeout`` argument can be a nonnegative float expressing
  112. seconds, or ``None`` for no timeout. If a float is given, a
  113. `.PipeTimeout` will be raised if the timeout period value has elapsed
  114. before any data arrives.
  115. :param int nbytes: maximum number of bytes to read
  116. :param float timeout:
  117. maximum seconds to wait (or ``None``, the default, to wait forever)
  118. :return: the read data, as a ``str`` or ``bytes``
  119. :raises:
  120. `.PipeTimeout` -- if a timeout was specified and no data was ready
  121. before that timeout
  122. """
  123. out = bytes()
  124. self._lock.acquire()
  125. try:
  126. if len(self._buffer) == 0:
  127. if self._closed:
  128. return out
  129. # should we block?
  130. if timeout == 0.0:
  131. raise PipeTimeout()
  132. # loop here in case we get woken up but a different thread has
  133. # grabbed everything in the buffer.
  134. while (len(self._buffer) == 0) and not self._closed:
  135. then = time.time()
  136. self._cv.wait(timeout)
  137. if timeout is not None:
  138. timeout -= time.time() - then
  139. if timeout <= 0.0:
  140. raise PipeTimeout()
  141. # something's in the buffer and we have the lock!
  142. if len(self._buffer) <= nbytes:
  143. out = self._buffer_tobytes()
  144. del self._buffer[:]
  145. if (self._event is not None) and not self._closed:
  146. self._event.clear()
  147. else:
  148. out = self._buffer_tobytes(nbytes)
  149. del self._buffer[:nbytes]
  150. finally:
  151. self._lock.release()
  152. return out
  153. def empty(self):
  154. """
  155. Clear out the buffer and return all data that was in it.
  156. :return:
  157. any data that was in the buffer prior to clearing it out, as a
  158. `str`
  159. """
  160. self._lock.acquire()
  161. try:
  162. out = self._buffer_tobytes()
  163. del self._buffer[:]
  164. if (self._event is not None) and not self._closed:
  165. self._event.clear()
  166. return out
  167. finally:
  168. self._lock.release()
  169. def close(self):
  170. """
  171. Close this pipe object. Future calls to `read` after the buffer
  172. has been emptied will return immediately with an empty string.
  173. """
  174. self._lock.acquire()
  175. try:
  176. self._closed = True
  177. self._cv.notify_all()
  178. if self._event is not None:
  179. self._event.set()
  180. finally:
  181. self._lock.release()
  182. def __len__(self):
  183. """
  184. Return the number of bytes buffered.
  185. :return: number (`int`) of bytes buffered
  186. """
  187. self._lock.acquire()
  188. try:
  189. return len(self._buffer)
  190. finally:
  191. self._lock.release()