PageRenderTime 41ms CodeModel.GetById 18ms RepoModel.GetById 0ms app.codeStats 0ms

/lib/concurrence/io/socket.py

https://bitbucket.org/geiregaj/pylabs-core-5.1
Python | 354 lines | 339 code | 11 blank | 4 comment | 4 complexity | 85b906bb3bdbc40734ca77ae3b794c37 MD5 | raw file
Possible License(s): MIT, MPL-2.0-no-copyleft-exception
  1. # Copyright (C) 2009, Hyves (Startphone Ltd.)
  2. #
  3. # This module is part of the Concurrence Framework and is released under
  4. # the New BSD License: http://www.opensource.org/licenses/bsd-license.php
  5. import logging
  6. import _socket
  7. import types
  8. import os
  9. from errno import EALREADY, EINPROGRESS, EWOULDBLOCK, ECONNRESET, ENOTCONN, ESHUTDOWN, EINTR, EISCONN, ENOENT, EAGAIN
  10. import _io
  11. from concurrence import Tasklet, FileDescriptorEvent, TIMEOUT_CURRENT
  12. from concurrence.io import IOStream
  13. DEFAULT_BACKLOG = 512
  14. XMOD = 8
  15. _interceptor = None
  16. class Socket(IOStream):
  17. log = logging.getLogger('Socket')
  18. __slots__ = ['socket', 'fd', '_readable', '_writable', 'state']
  19. STATE_INIT = 0
  20. STATE_LISTENING = 1
  21. STATE_CONNECTING = 2
  22. STATE_CONNECTED = 3
  23. STATE_CLOSING = 4
  24. STATE_CLOSED = 5
  25. _x = 0
  26. def __init__(self, socket, state = STATE_INIT):
  27. """don't call directly pls use one of the provided classmethod to create a socket"""
  28. self.socket = socket
  29. if _socket.AF_INET == socket.family:
  30. #always set the nodelay option on tcp sockets. This turns off the Nagle algorithm
  31. #we don't need this because in concurrence we are always buffering ourselves
  32. #before sending out data, so no need to let the tcp stack do it again and possibly delay
  33. #sending
  34. try:
  35. self.socket.setsockopt(_socket.IPPROTO_TCP, _socket.TCP_NODELAY, 1)
  36. except:
  37. self.log.warn("could not set TCP_NODELAY")
  38. #concurrence sockets are always non-blocking, this is the whole idea :-) :
  39. self.socket.setblocking(0)
  40. self.fd = self.socket.fileno()
  41. self._readable = None #will be created lazily
  42. self._writable = None #will be created lazily
  43. self.state = state
  44. @classmethod
  45. def set_interceptor(cls, interceptor):
  46. global _interceptor
  47. _interceptor = interceptor
  48. @classmethod
  49. def from_address(cls, addr):
  50. """Creates a new socket from the given address. If the addr is a tuple (host, port)
  51. a normal tcp socket is assumed. if addr is a string, a UNIX Domain socket is assumed"""
  52. if _interceptor is not None:
  53. return _interceptor(addr)
  54. elif type(addr) == types.StringType:
  55. return cls(_socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM))
  56. else:
  57. return cls(_socket.socket(_socket.AF_INET, _socket.SOCK_STREAM))
  58. @classmethod
  59. def new(cls):
  60. return cls(_socket.socket(_socket.AF_INET, _socket.SOCK_STREAM))
  61. @classmethod
  62. def server(cls, addr, backlog = DEFAULT_BACKLOG, reuse_address = True):
  63. s = cls.from_address(addr)
  64. s.set_reuse_address(reuse_address)
  65. s.bind(addr)
  66. s.listen(backlog)
  67. return s
  68. @classmethod
  69. def connect(cls, addr, timeout = TIMEOUT_CURRENT):
  70. """creates a new socket and connects it to the given address.
  71. returns the connected socket"""
  72. socket = cls.from_address(addr)
  73. socket._connect(addr, timeout)
  74. return socket
  75. @classmethod
  76. def from_file_descriptor(cls, fd, socket_family = _socket.AF_UNIX, socket_type = _socket.SOCK_STREAM, socket_state = STATE_INIT):
  77. return cls(_socket.fromfd(fd, socket_family, socket_type), socket_state)
  78. def _get_readable(self):
  79. if self._readable is None:
  80. self._readable = FileDescriptorEvent(self.fd, 'r')
  81. return self._readable
  82. def _set_readable(self, readable):
  83. self._readable = readable
  84. readable = property(_get_readable, _set_readable)
  85. def _get_writable(self):
  86. if self._writable is None:
  87. self._writable = FileDescriptorEvent(self.fd, 'w')
  88. return self._writable
  89. def _set_writable(self, writable):
  90. self._writable = writable
  91. writable = property(_get_writable, _set_writable)
  92. def fileno(self):
  93. return self.fd
  94. def set_reuse_address(self, reuse_address):
  95. self.socket.setsockopt(_socket.SOL_SOCKET, _socket.SO_REUSEADDR, int(reuse_address))
  96. def set_send_buffer_size(self, n):
  97. self.socket.setsockopt(_socket.SOL_SOCKET, _socket.SO_SNDBUF, n)
  98. def set_recv_buffer_size(self, n):
  99. self.socket.setsockopt(_socket.SOL_SOCKET, _socket.SO_RCVBUF, n)
  100. def bind(self, addr):
  101. self.socket.bind(addr)
  102. def listen(self, backlog = DEFAULT_BACKLOG):
  103. self.socket.listen(backlog)
  104. self.state = self.STATE_LISTENING
  105. def accept(self):
  106. """waits on a listening socket, returns a new socket_class instance
  107. for the incoming connection"""
  108. assert self.state == self.STATE_LISTENING, "make sure socket is listening before calling accept"
  109. while True:
  110. #we need a loop because sometimes we become readable and still not a valid
  111. #connection was accepted, in which case we return here and wait some more.
  112. self.readable.wait()
  113. try:
  114. s, _ = self.socket.accept()
  115. except _socket.error, (errno, _):
  116. if errno in [EAGAIN, EWOULDBLOCK]:
  117. #this can happen when more than one process received readability on the same socket (forked/cloned/dupped)
  118. #in that case 1 process will do the accept, the others receive this error, and should continue waiting for
  119. #readability
  120. continue
  121. else:
  122. raise
  123. return self.__class__(s, self.STATE_CONNECTED)
  124. def accept_iter(self):
  125. while True:
  126. try:
  127. yield self.accept()
  128. except Exception:
  129. self.log.exception("in accept_iter")
  130. Tasklet.sleep(1.0) #prevent hogging
  131. def _connect(self, addr, timeout = TIMEOUT_CURRENT):
  132. assert self.state == self.STATE_INIT, "make sure socket is not already connected or closed"
  133. try:
  134. err = self.socket.connect_ex(addr)
  135. serr = self.socket.getsockopt(_socket.SOL_SOCKET, _socket.SO_ERROR)
  136. except:
  137. self.log.exception("unexpected exception thrown by connect_ex")
  138. raise
  139. if err == 0 and serr == 0:
  140. self.state = self.STATE_CONNECTED
  141. elif err == EINPROGRESS and serr != 0:
  142. raise IOError(serr, os.strerror(serr))
  143. elif err == EINPROGRESS and serr == 0:
  144. self.state = self.STATE_CONNECTING
  145. try:
  146. self.writable.wait(timeout = timeout)
  147. self.state = self.STATE_CONNECTED
  148. except:
  149. self.state = self.STATE_INIT
  150. raise
  151. else:
  152. #some other error,
  153. #unix domain socket that does not exist, Cannot assign requested address etc etc
  154. raise _io.error_from_errno(IOError)
  155. def write(self, buffer, timeout = TIMEOUT_CURRENT, assume_writable = True):
  156. """Writes as many bytes as possible from the given buffer to this socket.
  157. The buffer position is updated according to the number of bytes succesfully written to the socket.
  158. This method returns the total number of bytes written. This method could possible write 0 bytes"""
  159. assert self.state == self.STATE_CONNECTED, "socket must be connected in order to write to it"
  160. Socket._x += 1
  161. if Socket._x % XMOD == 0:
  162. assume_writable = False
  163. #by default assume that we can write to the socket without blocking
  164. if assume_writable:
  165. bytes_written, _ = buffer.send(self.fd) #write to fd from buffer
  166. if bytes_written < 0 and _io.get_errno() == EAGAIN:
  167. #nope, need to wait before sending our data
  168. assume_writable = False
  169. #else if error != EAGAIN, assume_writable will stay True, and we fall trough and raise error below
  170. #if we cannot assume write-ability we will wait until data can be written again
  171. if not assume_writable:
  172. self.writable.wait(timeout = timeout)
  173. bytes_written, _ = buffer.send(self.fd) #write to fd from buffer
  174. #print 'bw', bytes_written, buffer.capacity
  175. #
  176. if bytes_written < 0:
  177. raise _io.error_from_errno(IOError)
  178. else:
  179. return bytes_written
  180. def read(self, buffer, timeout = TIMEOUT_CURRENT, assume_readable = True):
  181. """Reads as many bytes as possible the socket into the given buffer.
  182. The buffer position is updated according to the number of bytes read from the socket.
  183. This method could possible read 0 bytes. The method returns the total number of bytes read"""
  184. assert self.state == self.STATE_CONNECTED, "socket must be connected in order to read from it"
  185. Socket._x += 1
  186. if Socket._x % XMOD == 0:
  187. assume_readable = False
  188. #by default assume that we can read from the socket without blocking
  189. if assume_readable:
  190. bytes_read, _ = buffer.recv(self.fd) #read from fd to
  191. if bytes_read < 0 and _io.get_errno() == EAGAIN:
  192. #nope, need to wait before reading our data
  193. assume_readable = False
  194. #else if error != EAGAIN, assume_readable will stay True, and we fall trough and raise error below
  195. #if we cannot assume readability we will wait until data can be read again
  196. if not assume_readable:
  197. self.readable.wait(timeout = timeout)
  198. bytes_read, _ = buffer.recv(self.fd) #read from fd to
  199. #print 'br', bytes_read, buffer.capacity
  200. #
  201. if bytes_read < 0:
  202. raise _io.error_from_errno(IOError)
  203. else:
  204. return bytes_read
  205. def write_socket(self, socket, timeout = TIMEOUT_CURRENT):
  206. """writes a socket trough this socket"""
  207. self.writable.wait(timeout = timeout)
  208. _io.msgsendfd(self.fd, socket.fd)
  209. def read_socket(self, socket_class = None, socket_family = _socket.AF_INET, socket_type = _socket.SOCK_STREAM, socket_state = STATE_INIT, timeout = TIMEOUT_CURRENT):
  210. """reads a socket from this socket"""
  211. self.readable.wait(timeout = timeout)
  212. fd = _io.msgrecvfd(self.fd)
  213. return (socket_class or self.__class__).from_file_descriptor(fd, socket_family, socket_type, socket_state)
  214. def is_closed(self):
  215. return self.state == self.STATE_CLOSED
  216. def close(self):
  217. assert self.state in [self.STATE_CONNECTED, self.STATE_LISTENING]
  218. self.state = self.STATE_CLOSING
  219. if self._readable is not None:
  220. self._readable.close()
  221. if self._writable is not None:
  222. self._writable.close()
  223. self.socket.close()
  224. del self.socket
  225. del self._readable
  226. del self._writable
  227. self.state = self.STATE_CLOSED
  228. class SocketServer(object):
  229. log = logging.getLogger('SocketServer')
  230. def __init__(self, endpoint, handler = None):
  231. self._addr = None
  232. self._socket = None
  233. if isinstance(endpoint, Socket):
  234. self._socket = endpoint
  235. else:
  236. self._addr = endpoint
  237. self._handler = handler
  238. self._reuseaddress = True
  239. self._handler_task_name = 'socket_handler'
  240. self._accept_task = None
  241. self._accept_task_name = 'socket_acceptor'
  242. @property
  243. def socket(self):
  244. return self._socket
  245. def _handle_accept(self, accepted_socket):
  246. result = None
  247. try:
  248. result = self._handler(accepted_socket)
  249. except TaskletExit:
  250. raise
  251. except:
  252. self.log.exception("unhandled exception in socket handler")
  253. finally:
  254. if result is None and not accepted_socket.is_closed():
  255. try:
  256. accepted_socket.close()
  257. except TaskletExit:
  258. raise
  259. except:
  260. self.log.exception("unhandled exception while forcefully closing client")
  261. def _create_socket(self):
  262. if self._socket is None:
  263. if self._addr is None:
  264. assert False, "address must be set or accepting socket must be explicitly set"
  265. self._socket = Socket.from_address(self._addr)
  266. self._socket.set_reuse_address(self._reuseaddress)
  267. return self._socket
  268. def _accept_task_loop(self):
  269. accepted_socket = self._socket.accept()
  270. Tasklet.new(self._handle_accept, self._handler_task_name)(accepted_socket)
  271. def bind(self):
  272. """creates socket if needed, and binds it"""
  273. socket = self._create_socket()
  274. socket.bind(self._addr)
  275. def listen(self, backlog = DEFAULT_BACKLOG):
  276. """creates socket if needed, and listens it"""
  277. socket = self._create_socket()
  278. socket.listen(backlog)
  279. def serve(self):
  280. """listens and starts a new tasks accepting incoming connections on the configured address"""
  281. if self._socket is None:
  282. self.bind()
  283. self.listen()
  284. if not callable(self._handler):
  285. assert False, "handler not set or not callable"
  286. self._accept_task = Tasklet.loop(self._accept_task_loop, name = self._accept_task_name, daemon = True)()
  287. def close(self):
  288. self._accept_task.kill()
  289. self._socket.close()