Skip to content

Instantly share code, notes, and snippets.

@abaybektursun
Created November 23, 2022 22:53
Show Gist options
  • Save abaybektursun/b1ce38c9d6c45f7b09940dc4c8e20f59 to your computer and use it in GitHub Desktop.
Save abaybektursun/b1ce38c9d6c45f7b09940dc4c8e20f59 to your computer and use it in GitHub Desktop.
use pyo3::prelude::*;
use pyo3::buffer::PyBuffer;
use std::thread;
use std::sync::mpsc;
use std::sync::Mutex;
use std::sync::Arc;
use std::net::{TcpListener, TcpStream};
use std::io::prelude::*;
use ctrlc;
use std::ops::Deref;
// Helper function to handle the connections
fn _bridge_client(mut stream: TcpStream, rx: Arc<Mutex<mpsc::Receiver<String>>>) {
let buffer_size = 1024;
let mut buf = vec![0; buffer_size];
loop {
let n = stream.read(&mut buf).unwrap();
//println!("Received {} bytes", n);
let signal = rx.lock().unwrap().try_recv();
let mut signal = match signal {
Ok(signal) => signal,
Err(_) => String::from(""),
};
if n == 0 || signal == "stop" {
return;
}
}
}
#[pyfunction]
fn start(server_address: &str) -> PyResult<()> {
// One way communication from the main thread to listener thread to stop the children threads
let (tx_signal, rx_signal) = mpsc::channel();
// Handle Ctrl+C. Shut down the threads
//ctrlc::set_handler(move || tx_signal.send("stop").expect("Could not send signal on channel."))
// .expect("Error setting Ctrl-C handler");
tx_signal.send("test").expect("Could not send signal on channel.");
// Mutex to pass IP address to the listener thread
let mut_server_address = Mutex::new(server_address.to_string());
// -- Listener thread ------------------------------------------------------
let handle = thread::spawn(move || {
// Is this the right way to get the value out of the mutex?
let server_address = mut_server_address.lock().unwrap().clone();
// Create a TCP listener
let listener = TcpListener::bind(server_address).unwrap();
// Accept connections and process them in new threads
let mut handle = None;
let (tx_client, rx_client) = mpsc::channel();
let rx_client = Arc::new(Mutex::new(rx_client));
for stream in listener.incoming() {
match stream {
Ok(stream) => {
let receiver = Arc::clone(&rx_client);
println!("Connection established!");
handle = Some(thread::spawn(move || {
_bridge_client(stream, receiver);
}));
}
Err(e) => { /* connection failed */ }
}
let signal = rx_signal.try_recv();
let signal = match signal {
Ok(signal) => signal,
Err(_) => "",
};
// TODO: currently listener.incoming() is blocking. Find a way to iterate over the incoming connections with timeout
if signal == "stop" {
println!("Stopping listener thread");
break;
}
}
});
// -------------------------------------------------------------------------
Ok(())
}
/// A Python module implemented in Rust.
#[pymodule]
fn tcp_server_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(start, m)?)?;
Ok(())
}
@tuncatunc
Copy link

Keep it simple and create another udp port between python process and rust process.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment