def __init__(self, server_address, runner_factory, context_lock, handler_class=None, bind_and_activate=True): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param func context_lock: A contextmgr that will be used as a lock during request handling/forking. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket( socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. self._context_lock = context_lock if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise
def __init__(self, server_address, runner_factory, lifecycle_lock, request_complete_callback, handler_class=None, bind_and_activate=True): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param threading.RLock lifecycle_lock: A lock used to guard against abrupt teardown of the servers execution thread during handling. All pailgun request handling will take place under care of this lock, which would be shared with a `PailgunServer`-external lifecycle manager to guard teardown. :param function request_complete_callback: A callback that will be called whenever a pailgun request is completed. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket(socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.lifecycle_lock = lifecycle_lock self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. self.request_complete_callback = request_complete_callback if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise
def setUp(self): self.chunk_size = 512 self.mock_socket = unittest.mock.Mock() self.client_sock, self.server_sock = socket.socketpair() self.buf_sock = RecvBufferedSocket(self.client_sock, chunk_size=self.chunk_size) self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket, chunk_size=self.chunk_size)
class TestRecvBufferedSocket(unittest.TestCase): def setUp(self): self.chunk_size = 512 self.mock_socket = mock.Mock() self.client_sock, self.server_sock = socket.socketpair() self.buf_sock = RecvBufferedSocket(self.client_sock, chunk_size=self.chunk_size) self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket, chunk_size=self.chunk_size) def tearDown(self): self.buf_sock.close() self.server_sock.close() def test_getattr(self): self.assertTrue(inspect.ismethod(self.buf_sock.recv)) self.assertFalse(inspect.isbuiltin(self.buf_sock.recv)) self.assertTrue(inspect.isbuiltin(self.buf_sock.connect)) def test_recv(self): self.server_sock.sendall(b'A' * 300) self.assertEqual(self.buf_sock.recv(1), b'A') self.assertEqual(self.buf_sock.recv(200), b'A' * 200) self.assertEqual(self.buf_sock.recv(99), b'A' * 99) def test_recv_max_larger_than_buf(self): double_chunk = self.chunk_size * 2 self.server_sock.sendall(b'A' * double_chunk) self.assertEqual(self.buf_sock.recv(double_chunk), b'A' * double_chunk) @mock.patch('selectors.DefaultSelector' if PY3 else 'select.select', **PATCH_OPTS) def test_recv_check_calls(self, mock_selector): if PY3: mock_selector = mock_selector.return_value.__enter__.return_value mock_selector.register = mock.Mock() # NB: the return value should actually be List[Tuple[SelectorKey, Events]], but our code only # cares that _some_ event happened so we choose a simpler mock here. See # https://docs.python.org/3/library/selectors.html#selectors.BaseSelector.select. mock_selector.select = mock.Mock(return_value=[(1, "")]) else: mock_selector.return_value = ([1], [], []) self.mock_socket.recv.side_effect = [ b'A' * self.chunk_size, b'B' * self.chunk_size ] self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128) self.mock_socket.recv.assert_called_once_with(self.chunk_size) self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEqual(self.mock_socket.recv.call_count, 1) self.assertEqual(self.mocked_buf_sock.recv(self.chunk_size), b'B' * self.chunk_size) self.assertEqual(self.mock_socket.recv.call_count, 2)
class TestRecvBufferedSocket(unittest.TestCase): def setUp(self): self.chunk_size = 512 self.mock_socket = mock.Mock() self.client_sock, self.server_sock = socket.socketpair() self.buf_sock = RecvBufferedSocket(self.client_sock, chunk_size=self.chunk_size) self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket, chunk_size=self.chunk_size) def tearDown(self): self.buf_sock.close() self.server_sock.close() def test_getattr(self): self.assertTrue(inspect.ismethod(self.buf_sock.recv)) self.assertFalse(inspect.isbuiltin(self.buf_sock.recv)) self.assertTrue(inspect.isbuiltin(self.buf_sock.connect)) def test_recv(self): self.server_sock.sendall(b'A' * 300) self.assertEquals(self.buf_sock.recv(1), b'A') self.assertEquals(self.buf_sock.recv(200), b'A' * 200) self.assertEquals(self.buf_sock.recv(99), b'A' * 99) def test_recv_max_larger_than_buf(self): double_chunk = self.chunk_size * 2 self.server_sock.sendall(b'A' * double_chunk) self.assertEquals(self.buf_sock.recv(double_chunk), b'A' * double_chunk) @mock.patch('pants.util.socket.select.select', **PATCH_OPTS) def test_recv_check_calls(self, mock_select): mock_select.return_value = ([1], [], []) self.mock_socket.recv.side_effect = [ b'A' * self.chunk_size, b'B' * self.chunk_size ] self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.mock_socket.recv.assert_called_once_with(self.chunk_size) self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEquals(self.mock_socket.recv.call_count, 1) self.assertEquals(self.mocked_buf_sock.recv(self.chunk_size), b'B' * self.chunk_size) self.assertEquals(self.mock_socket.recv.call_count, 2)
def __init__(self, server_address, runner_factory, context_lock, handler_class=None, bind_and_activate=True): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param func context_lock: A contextmgr that will be used as a lock during request handling/forking. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket(socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. self._context_lock = context_lock if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise
class PailgunServer(TCPServer): """A (forking) pants nailgun server.""" def __init__(self, server_address, runner_factory, context_lock, handler_class=None, bind_and_activate=True): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param func context_lock: A contextmgr that will be used as a lock during request handling/forking. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket(socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. self._context_lock = context_lock if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise def server_bind(self): """Override of TCPServer.server_bind() that tracks bind-time assigned random ports.""" TCPServer.server_bind(self) _, self.server_port = self.socket.getsockname()[:2] def process_request(self, request, client_address): """Override of TCPServer.process_request() that provides for forking request handlers and delegates error handling to the request handler.""" # Instantiate the request handler. handler = self.RequestHandlerClass(request, client_address, self) try: # Attempt to handle a request with the handler under the context_lock. with self._context_lock(): handler.handle_request() except Exception as e: # If that fails, (synchronously) handle the error with the error handler sans-fork. try: handler.handle_error(e) finally: # Shutdown the socket since we don't expect a fork() in the exception context. self.shutdown_request(request) else: # At this point, we expect a fork() has taken place - the parent side will return, and so we # close the request here from the parent without explicitly shutting down the socket. The # child half of this will perform an os._exit() before it gets to this point and is also # responsible for shutdown and closing of the socket when its execution is complete. self.close_request(request)
class TestRecvBufferedSocket(unittest.TestCase): def setUp(self): self.chunk_size = 512 self.mock_socket = mock.Mock() self.client_sock, self.server_sock = socket.socketpair() self.buf_sock = RecvBufferedSocket(self.client_sock, chunk_size=self.chunk_size) self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket, chunk_size=self.chunk_size) def tearDown(self): self.buf_sock.close() self.server_sock.close() def test_getattr(self): self.assertTrue(inspect.ismethod(self.buf_sock.recv)) self.assertFalse(inspect.isbuiltin(self.buf_sock.recv)) self.assertTrue(inspect.isbuiltin(self.buf_sock.connect)) def test_recv(self): self.server_sock.sendall(b'A' * 300) self.assertEquals(self.buf_sock.recv(1), b'A') self.assertEquals(self.buf_sock.recv(200), b'A' * 200) self.assertEquals(self.buf_sock.recv(99), b'A' * 99) def test_recv_max_larger_than_buf(self): double_chunk = self.chunk_size * 2 self.server_sock.sendall(b'A' * double_chunk) self.assertEquals(self.buf_sock.recv(double_chunk), b'A' * double_chunk) @mock.patch('pants.util.socket.select.select', **PATCH_OPTS) def test_recv_check_calls(self, mock_select): mock_select.return_value = ([1], [], []) self.mock_socket.recv.side_effect = [b'A' * self.chunk_size, b'B' * self.chunk_size] self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.mock_socket.recv.assert_called_once_with(self.chunk_size) self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128) self.assertEquals(self.mock_socket.recv.call_count, 1) self.assertEquals(self.mocked_buf_sock.recv(self.chunk_size), b'B' * self.chunk_size) self.assertEquals(self.mock_socket.recv.call_count, 2)
def try_connect(self): """Creates a socket, connects it to the nailgun and returns the connected socket. :returns: a connected `socket.socket`. :raises: `NailgunClient.NailgunConnectionError` on failure to connect. """ sock = RecvBufferedSocket(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) try: sock.connect(self._address) except (socket.error, socket.gaierror) as e: logger.debug('Encountered socket exception {!r} when attempting connect to nailgun'.format(e)) sock.close() raise self.NailgunConnectionError( address=self._address_string, pid=self._maybe_last_pid(), pgrp=self._maybe_last_pgrp(), wrapped_exc=e, ) else: return sock
def try_connect(self): """Creates a socket, connects it to the nailgun and returns the connected socket. :returns: a connected `socket.socket`. :raises: `NailgunClient.NailgunConnectionError` on failure to connect. """ sock = RecvBufferedSocket(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) try: sock.connect((self._host, self._port)) except (socket.error, socket.gaierror) as e: logger.debug('Encountered socket exception {!r} when attempting connect to nailgun'.format(e)) sock.close() raise self.NailgunConnectionError( 'Problem connecting to nailgun server at {}:{}: {!r}'.format(self._host, self._port, e)) else: return sock
def try_connect(self): """Creates a socket, connects it to the nailgun and returns the connected socket. :returns: a connected `socket.socket`. :raises: `NailgunClient.NailgunConnectionError` on failure to connect. """ sock = RecvBufferedSocket(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) try: sock.connect(self._address) except (socket.error, socket.gaierror) as e: logger.debug('Encountered socket exception {!r} when attempting connect to nailgun'.format(e)) sock.close() raise self.NailgunConnectionError( address=self._address_string, pid=self.pid, wrapped_exc=e, traceback=sys.exc_info()[2] ) else: return sock
def try_connect(self): """Creates a socket, connects it to the nailgun and returns the connected socket. :returns: a connected `socket.socket`. :raises: `NailgunClient.NailgunConnectionError` on failure to connect. """ sock = RecvBufferedSocket( sock=socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)) try: sock.connect(self._address) except (socket.error, socket.gaierror) as e: logger.debug( "Encountered socket exception {!r} when attempting connect to nailgun" .format(e)) sock.close() raise self.NailgunConnectionError( address=self._address_string, pid=self._maybe_last_pid(), pgrp=self._maybe_last_pgrp(), wrapped_exc=e, ) else: return sock
class PailgunServer(TCPServer): """A (forking) pants nailgun server.""" def __init__(self, server_address, runner_factory, context_lock, handler_class=None, bind_and_activate=True): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param func context_lock: A contextmgr that will be used as a lock during request handling/forking. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket( socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. self._context_lock = context_lock if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise def server_bind(self): """Override of TCPServer.server_bind() that tracks bind-time assigned random ports.""" TCPServer.server_bind(self) _, self.server_port = self.socket.getsockname()[:2] def process_request(self, request, client_address): """Override of TCPServer.process_request() that provides for forking request handlers and delegates error handling to the request handler.""" # Instantiate the request handler. handler = self.RequestHandlerClass(request, client_address, self) try: # Attempt to handle a request with the handler under the context_lock. with self._context_lock(): handler.handle_request() except Exception as e: # If that fails, (synchronously) handle the error with the error handler sans-fork. try: handler.handle_error(e) finally: # Shutdown the socket since we don't expect a fork() in the exception context. self.shutdown_request(request) else: # At this point, we expect a fork() has taken place - the parent side will return, and so we # close the request here from the parent without explicitly shutting down the socket. The # child half of this will perform an os._exit() before it gets to this point and is also # responsible for shutdown and closing of the socket when its execution is complete. self.close_request(request)
class PailgunServer(ThreadingMixIn, TCPServer): """A pants nailgun server. This class spawns a thread per request via `ThreadingMixIn`: the thread body runs `process_request_thread`, which we override. """ timeout = 0.05 # Override the ThreadingMixIn default, to minimize the chances of zombie pailgun processes. daemon_threads = True def __init__( self, server_address, runner_factory, lifecycle_lock, request_complete_callback, handler_class=None, bind_and_activate=True, ): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param threading.RLock lifecycle_lock: A lock used to guard against abrupt teardown of the servers execution thread during handling. All pailgun request handling will take place under care of this lock, which would be shared with a `PailgunServer`-external lifecycle manager to guard teardown. :param function request_complete_callback: A callback that will be called whenever a pailgun request is completed. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket( socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.lifecycle_lock = lifecycle_lock self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. self.request_complete_callback = request_complete_callback self.logger = logging.getLogger(__name__) self.free_to_handle_request_lock = PailgunHandleRequestLock() if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise def server_bind(self): """Override of TCPServer.server_bind() that tracks bind-time assigned random ports.""" TCPServer.server_bind(self) _, self.server_port = self.socket.getsockname()[:2] def process_request(self, request, client_address): """Start a new thread to process the request. This is lovingly copied and pasted from ThreadingMixIn, with the addition of setting the name of the thread. It's a shame that ThreadingMixIn doesn't provide a customization hook. """ t = threading.Thread( target=self.process_request_thread, args=(request, client_address), name="PailgunRequestThread", ) t.daemon = self.daemon_threads t.start() def handle_request(self): """Override of TCPServer.handle_request() that provides locking. Calling this method has the effect of "maybe" (if the socket does not time out first) accepting a request and (because we mixin in ThreadingMixIn) spawning it on a thread. It should always return within `min(self.timeout, socket.gettimeout())`. N.B. Most of this is copied verbatim from SocketServer.py in the stdlib. """ timeout = self.socket.gettimeout() if timeout is None: timeout = self.timeout elif self.timeout is not None: timeout = min(timeout, self.timeout) if not is_readable(self, timeout=timeout): self.handle_timeout() return # After select tells us we can safely accept, guard the accept and request # handling with the lifecycle lock to avoid abrupt teardown mid-request. with self.lifecycle_lock(): self._handle_request_noblock() def _should_poll_forever(self, timeout): return timeout < 0 def _should_keep_polling(self, timeout, time_polled): return self._should_poll_forever(timeout) or time_polled < timeout def _send_stderr(self, request, message): NailgunProtocol.send_stderr(request, message) @contextmanager def ensure_request_is_exclusive(self, environment, request): """Ensure that this is the only pants running. We currently don't allow parallel pants runs, so this function blocks a request thread until there are no more requests being handled. """ # TODO add `did_poll` to pantsd metrics timeout = float(environment["PANTSD_REQUEST_TIMEOUT_LIMIT"]) @contextmanager def yield_and_release(time_waited): try: self.logger.debug( f"request lock acquired {('on the first try' if time_waited == 0 else f'in {time_waited} seconds')}." ) yield finally: self.free_to_handle_request_lock.release() self.logger.debug("released request lock.") time_polled = 0.0 user_notification_interval = 5.0 # Stop polling to notify the user every second. self.logger.debug( f"request {request} is trying to acquire the request lock.") # NB: Optimistically try to acquire the lock without blocking, in case we are the only request being handled. # This could be merged into the `while` loop below, but separating this special case for logging helps. if self.free_to_handle_request_lock.acquire(timeout=0): with yield_and_release(time_polled): yield else: self.logger.debug( f"request {request} didn't acquire the lock on the first try, polling..." ) # We have to wait for another request to finish being handled. self._send_stderr( request, "Another pants invocation is running. " "Will wait {} for it to finish before giving up.\n".format( "forever" if self._should_poll_forever( timeout) else "up to {} seconds".format(timeout)), ) self._send_stderr( request, "If you don't want to wait for the first run to finish, please " "press Ctrl-C and run this command with PANTS_CONCURRENT=True " "in the environment.\n", ) while not self.free_to_handle_request_lock.acquire( timeout=user_notification_interval): time_polled += user_notification_interval if self._should_keep_polling(timeout, time_polled): self._send_stderr( request, f"Waiting for invocation to finish (waited for {time_polled}s so far)...\n", ) else: # We have timed out. raise ExclusiveRequestTimeout( "Timed out while waiting for another pants invocation to finish." ) with yield_and_release(time_polled): yield def process_request_thread(self, request, client_address): """Override of ThreadingMixIn.process_request_thread() that delegates to the request handler.""" # Instantiate the request handler. Native().override_thread_logging_destination_to_just_pantsd() handler = self.RequestHandlerClass(request, client_address, self) _, _, _, environment = handler.parsed_request() try: with self.ensure_request_is_exclusive(environment, request): # Attempt to handle a request with the handler. handler.handle_request() self.request_complete_callback() except BrokenPipeError as e: # The client has closed the connection, most likely from a SIGINT self.logger.error( f"Request {request} abruptly closed with {type(e)}, probably because the client crashed or was sent a SIGINT." ) except Exception as e: # If that fails, (synchronously) handle the error with the error handler sans-fork. try: self.logger.error( f"Request {request} errored with {type(e)}({e!r})") handler.handle_error(e) finally: # Shutdown the socket since we don't expect a fork() in the exception context. self.shutdown_request(request) else: # At this point, we expect a fork() has taken place - the parent side will return, and so we # close the request here from the parent without explicitly shutting down the socket. The # child half of this will perform an os._exit() before it gets to this point and is also # responsible for shutdown and closing of the socket when its execution is complete. self.close_request(request)
class PailgunServer(ThreadingMixIn, TCPServer): """A pants nailgun server. This class spawns a thread per request via `ThreadingMixIn`: the thread body runs `process_request_thread`, which we override. """ timeout = 0.05 # Override the ThreadingMixIn default, to minimize the chances of zombie pailgun processes. daemon_threads = True def __init__(self, server_address, runner_factory, lifecycle_lock, handler_class=None, bind_and_activate=True): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param threading.RLock lifecycle_lock: A lock used to guard against abrupt teardown of the servers execution thread during handling. All pailgun request handling will take place under care of this lock, which would be shared with a `PailgunServer`-external lifecycle manager to guard teardown. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket( socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.lifecycle_lock = lifecycle_lock self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise def server_bind(self): """Override of TCPServer.server_bind() that tracks bind-time assigned random ports.""" TCPServer.server_bind(self) _, self.server_port = self.socket.getsockname()[:2] def handle_request(self): """Override of TCPServer.handle_request() that provides locking. Calling this method has the effect of "maybe" (if the socket does not time out first) accepting a request and (because we mixin in ThreadingMixIn) spawning it on a thread. It should always return within `min(self.timeout, socket.gettimeout())`. N.B. Most of this is copied verbatim from SocketServer.py in the stdlib. """ timeout = self.socket.gettimeout() if timeout is None: timeout = self.timeout elif self.timeout is not None: timeout = min(timeout, self.timeout) fd_sets = safe_select([self], [], [], timeout) if not fd_sets[0]: self.handle_timeout() return # After select tells us we can safely accept, guard the accept and request # handling with the lifecycle lock to avoid abrupt teardown mid-request. with self.lifecycle_lock(): self._handle_request_noblock() def process_request_thread(self, request, client_address): """Override of ThreadingMixIn.process_request_thread() that delegates to the request handler.""" # Instantiate the request handler. handler = self.RequestHandlerClass(request, client_address, self) try: # Attempt to handle a request with the handler. handler.handle_request() except Exception as e: # If that fails, (synchronously) handle the error with the error handler sans-fork. try: handler.handle_error(e) finally: # Shutdown the socket since we don't expect a fork() in the exception context. self.shutdown_request(request) else: # At this point, we expect a fork() has taken place - the parent side will return, and so we # close the request here from the parent without explicitly shutting down the socket. The # child half of this will perform an os._exit() before it gets to this point and is also # responsible for shutdown and closing of the socket when its execution is complete. self.close_request(request)
class PailgunServer(ThreadingMixIn, TCPServer): """A pants nailgun server. This class spawns a thread per request via `ThreadingMixIn`: the thread body runs `process_request_thread`, which we override. """ timeout = 0.05 # Override the ThreadingMixIn default, to minimize the chances of zombie pailgun processes. daemon_threads = True def __init__(self, server_address, runner_factory, lifecycle_lock, request_complete_callback, handler_class=None, bind_and_activate=True): """Override of TCPServer.__init__(). N.B. the majority of this function is copied verbatim from TCPServer.__init__(). :param tuple server_address: An address tuple of (hostname, port) for socket.bind(). :param class runner_factory: A factory function for creating a DaemonPantsRunner for each run. :param threading.RLock lifecycle_lock: A lock used to guard against abrupt teardown of the servers execution thread during handling. All pailgun request handling will take place under care of this lock, which would be shared with a `PailgunServer`-external lifecycle manager to guard teardown. :param function request_complete_callback: A callback that will be called whenever a pailgun request is completed. :param class handler_class: The request handler class to use for each request. (Optional) :param bool bind_and_activate: If True, binds and activates networking at __init__ time. (Optional) """ # Old-style class, so we must invoke __init__() this way. BaseServer.__init__(self, server_address, handler_class or PailgunHandler) self.socket = RecvBufferedSocket(socket.socket(self.address_family, self.socket_type)) self.runner_factory = runner_factory self.lifecycle_lock = lifecycle_lock self.allow_reuse_address = True # Allow quick reuse of TCP_WAIT sockets. self.server_port = None # Set during server_bind() once the port is bound. self.request_complete_callback = request_complete_callback if bind_and_activate: try: self.server_bind() self.server_activate() except Exception: self.server_close() raise def server_bind(self): """Override of TCPServer.server_bind() that tracks bind-time assigned random ports.""" TCPServer.server_bind(self) _, self.server_port = self.socket.getsockname()[:2] def handle_request(self): """Override of TCPServer.handle_request() that provides locking. Calling this method has the effect of "maybe" (if the socket does not time out first) accepting a request and (because we mixin in ThreadingMixIn) spawning it on a thread. It should always return within `min(self.timeout, socket.gettimeout())`. N.B. Most of this is copied verbatim from SocketServer.py in the stdlib. """ timeout = self.socket.gettimeout() if timeout is None: timeout = self.timeout elif self.timeout is not None: timeout = min(timeout, self.timeout) fd_sets = safe_select([self], [], [], timeout) if not fd_sets[0]: self.handle_timeout() return # After select tells us we can safely accept, guard the accept and request # handling with the lifecycle lock to avoid abrupt teardown mid-request. with self.lifecycle_lock(): self._handle_request_noblock() def process_request_thread(self, request, client_address): """Override of ThreadingMixIn.process_request_thread() that delegates to the request handler.""" # Instantiate the request handler. handler = self.RequestHandlerClass(request, client_address, self) try: # Attempt to handle a request with the handler. handler.handle_request() self.request_complete_callback() except Exception as e: # If that fails, (synchronously) handle the error with the error handler sans-fork. try: handler.handle_error(e) finally: # Shutdown the socket since we don't expect a fork() in the exception context. self.shutdown_request(request) else: # At this point, we expect a fork() has taken place - the parent side will return, and so we # close the request here from the parent without explicitly shutting down the socket. The # child half of this will perform an os._exit() before it gets to this point and is also # responsible for shutdown and closing of the socket when its execution is complete. self.close_request(request)
def setUp(self): self.chunk_size = 512 self.mock_socket = mock.Mock() self.client_sock, self.server_sock = socket.socketpair() self.buf_sock = RecvBufferedSocket(self.client_sock, chunk_size=self.chunk_size) self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket, chunk_size=self.chunk_size)