aws_smithy_http_server_python/
socket.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Socket implementation that can be shared between multiple Python processes.
7
8use pyo3::prelude::*;
9
10use socket2::{Domain, Protocol, Socket, Type};
11use std::net::SocketAddr;
12
13/// Socket implementation that can be shared between multiple Python processes.
14///
15/// Python cannot handle true multi-threaded applications due to the [GIL],
16/// often resulting in reduced performance and only one core used by the application.
17/// To work around this, Python web applications usually create a socket with
18/// SO_REUSEADDR and SO_REUSEPORT enabled that can be shared between multiple
19/// Python processes, allowing you to maximize performance and use all available
20/// computing capacity of the host.
21///
22/// [GIL]: https://wiki.python.org/moin/GlobalInterpreterLock
23///
24/// :param address str:
25/// :param port int:
26/// :param backlog typing.Optional\[int\]:
27/// :rtype None:
28#[pyclass]
29#[derive(Debug)]
30pub struct PySocket {
31    pub(crate) inner: Socket,
32}
33
34#[pymethods]
35impl PySocket {
36    /// Create a new UNIX `SharedSocket` from an address, port and backlog.
37    /// If not specified, the backlog defaults to 1024 connections.
38    #[pyo3(text_signature = "($self, address, port, backlog=None)")]
39    #[new]
40    pub fn new(address: String, port: i32, backlog: Option<i32>) -> PyResult<Self> {
41        let address: SocketAddr = format!("{}:{}", address, port).parse()?;
42        let (domain, ip_version) = PySocket::socket_domain(address);
43        tracing::trace!(address = %address, ip_version, "shared socket listening");
44        let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;
45        // Set value for the `SO_REUSEPORT` and `SO_REUSEADDR` options on this socket.
46        // This indicates that further calls to `bind` may allow reuse of local
47        // addresses. For IPv4 sockets this means that a socket may bind even when
48        // there's a socket already listening on this port.
49        socket.set_reuse_port(true)?;
50        socket.set_reuse_address(true)?;
51        socket.bind(&address.into())?;
52        socket.listen(backlog.unwrap_or(1024))?;
53        Ok(PySocket { inner: socket })
54    }
55
56    /// Clone the inner socket allowing it to be shared between multiple
57    /// Python processes.
58    ///
59    /// :rtype PySocket:
60    pub fn try_clone(&self) -> PyResult<PySocket> {
61        let copied = self.inner.try_clone()?;
62        Ok(PySocket { inner: copied })
63    }
64}
65
66impl PySocket {
67    /// Get a cloned inner socket.
68    pub fn get_socket(&self) -> Result<Socket, std::io::Error> {
69        self.inner.try_clone()
70    }
71
72    /// Find the socket domain
73    fn socket_domain(address: SocketAddr) -> (Domain, &'static str) {
74        if address.is_ipv6() {
75            (Domain::IPV6, "6")
76        } else {
77            (Domain::IPV4, "4")
78        }
79    }
80}
81
82#[cfg(test)]
83// `is_listener` on `Socket` is only available on certain platforms.
84// In particular, this fails to compile on MacOS.
85#[cfg(any(
86    target_os = "android",
87    target_os = "freebsd",
88    target_os = "fuchsia",
89    target_os = "linux",
90))]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn socket_can_bind_on_random_port() {
96        let socket = PySocket::new("127.0.0.1".to_owned(), 0, None).unwrap();
97        assert!(socket.inner.is_listener().is_ok());
98    }
99
100    #[test]
101    fn socket_can_be_cloned() {
102        let socket = PySocket::new("127.0.0.1".to_owned(), 0, None).unwrap();
103        let cloned_socket = socket.try_clone().unwrap();
104        assert!(cloned_socket.inner.is_listener().is_ok());
105    }
106}