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(text_signature = "($self, address, port, backlog=None)")]
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 #[new]
39 pub fn new(address: String, port: i32, backlog: Option<i32>) -> PyResult<Self> {
40 let address: SocketAddr = format!("{}:{}", address, port).parse()?;
41 let (domain, ip_version) = PySocket::socket_domain(address);
42 tracing::trace!(address = %address, ip_version, "shared socket listening");
43 let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;
44 // Set value for the `SO_REUSEPORT` and `SO_REUSEADDR` options on this socket.
45 // This indicates that further calls to `bind` may allow reuse of local
46 // addresses. For IPv4 sockets this means that a socket may bind even when
47 // there's a socket already listening on this port.
48 socket.set_reuse_port(true)?;
49 socket.set_reuse_address(true)?;
50 socket.bind(&address.into())?;
51 socket.listen(backlog.unwrap_or(1024))?;
52 Ok(PySocket { inner: socket })
53 }
54
55 /// Clone the inner socket allowing it to be shared between multiple
56 /// Python processes.
57 ///
58 /// :rtype PySocket:
59 pub fn try_clone(&self) -> PyResult<PySocket> {
60 let copied = self.inner.try_clone()?;
61 Ok(PySocket { inner: copied })
62 }
63}
64
65impl PySocket {
66 /// Get a cloned inner socket.
67 pub fn get_socket(&self) -> Result<Socket, std::io::Error> {
68 self.inner.try_clone()
69 }
70
71 /// Find the socket domain
72 fn socket_domain(address: SocketAddr) -> (Domain, &'static str) {
73 if address.is_ipv6() {
74 (Domain::IPV6, "6")
75 } else {
76 (Domain::IPV4, "4")
77 }
78 }
79}
80
81#[cfg(test)]
82// `is_listener` on `Socket` is only available on certain platforms.
83// In particular, this fails to compile on MacOS.
84#[cfg(any(
85 target_os = "android",
86 target_os = "freebsd",
87 target_os = "fuchsia",
88 target_os = "linux",
89))]
90mod tests {
91 use super::*;
92
93 #[test]
94 fn socket_can_bind_on_random_port() {
95 let socket = PySocket::new("127.0.0.1".to_owned(), 0, None).unwrap();
96 assert!(socket.inner.is_listener().is_ok());
97 }
98
99 #[test]
100 fn socket_can_be_cloned() {
101 let socket = PySocket::new("127.0.0.1".to_owned(), 0, None).unwrap();
102 let cloned_socket = socket.try_clone().unwrap();
103 assert!(cloned_socket.inner.is_listener().is_ok());
104 }
105}