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