tcp_server_pyo3.rs

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(())
}

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