example-tcpstream-deadpool.rs

use std::{io::BufReader, time::Duration};
use deadpool::managed::{Pool, PoolConfig};
use tokio::{
    io::{self, AsyncBufReadExt, AsyncWriteExt},
    net::TcpStream,
    sync::broadcast,
};

#[derive(Debug)]
enum Error {
    Fail,
}

struct Route {
    port: i32,
    sender: broadcast::Sender<Vec<u8>>,
}

impl Route {
    async fn send(&self, msg: String) {
        let _ = self.sender.send(msg.into_bytes());
    }
}

struct RouteManager {
    pub port: i32,
}

impl deadpool::managed::Manager for RouteManager {
    type Type = Route;
    type Error = Error;

    async fn create(&self) -> Result<Route, Error> {
        let (tx, mut rx) = broadcast::channel::<Vec<u8>>(10);
        let port = self.port;

        tokio::spawn(async move {
            match TcpStream::connect(format!("0.0.0.0:{}", port)).await {
                Ok(mut stream) => {
                    while let Ok(msg) = rx.recv().await {
                        let _ = stream.write_all(&msg).await;
                    }
                }
                Err(e) => eprintln!("Failed to connect: {:?}", e),
            }
        });

        Ok(Route { port: self.port, sender: tx })
    }

    async fn recycle(
        &self,
        _obj: &mut Route,
        _metrics: &deadpool::managed::Metrics,
    ) -> deadpool::managed::RecycleResult<Error> {
        Ok(())
    }
}

#[tokio::main]
async fn main() {
    let manager = RouteManager { port: 3000 };
    let pool: Pool<RouteManager> = Pool::builder(manager)
        .config(PoolConfig::new(10)) // Max 10 connections
        .build()
        .unwrap();

    match pool.get().await {
        Ok(route) => {
            let mut reader = BufReader::new(io::stdin()).lines();
            println!("Enter messages (Ctrl+D to exit):");

            while let Ok(Some(line)) = reader.next_line().await {
                route.send(line).await;
            }

            println!("Exiting...");
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
        Err(e) => {
            eprintln!("Failed to get a connection: {:?}", e);
        }
    }
}