MCP 서버 성능 최적화: Rust로 구축하는 초고속 처리 파이프라인

[dependencies]
tokio = { version = "1", features = ["full"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"

use tokio::net::TcpListener;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

// MCP 요청 메시지 구조 정의
#[derive(Debug, Deserialize)]
struct McpRequest {
    id: String,
    method: String,
    params: serde_json::Value,
}

#[derive(Debug, Serialize)]
struct McpResponse {
    id: String,
    result: serde_json::Value,
}

// 무거운 작업을 시뮬레이션하는 함수
async fn process_heavy_tool(params: serde_json::Value) -> Result<String, String> {
    // 실제 환경에서는 여기서 DB 조회나 파일 I/O 발생
    // tokio::time::sleep을 통해 비동기 대기 시뮬레이션
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    Ok(format!("Processed: {}", params))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let listener = TcpListener::bind("127.0.0.1:8080").await?;
    println!("MCP Server listening on {}", listener.local_addr()?);

    loop {
        let (mut socket, _) = listener.accept().await?;

        // 연결이 들어올 때마다 새로운 태스크 생성 (병렬 처리)
        tokio::spawn(async move {
            let mut buf = [0; 1024];

            loop {
                let n = match socket.read(&mut buf).await {
                    Ok(n) if n == 0 => return, // 연결 종료
                    Ok(n) => n,
                    Err(e) => {
                        eprintln!("Failed to read from socket; err = {:?}", e);
                        return;
                    }
                };

                let req_str = String::from_utf8_lossy(&buf[0..n]);
                
                // JSON 파싱 및 처리 로직 (에러 처리 생략)
                if let Ok(req) = serde_json::from_str::<McpRequest>(&req_str) {
                    let req_id = req.id.clone();
                    
                    // 핵심: 비동기 함수를 spawn하여 블로킹 없이 처리
                    let handle = tokio::spawn(async move {
                        process_heavy_tool(req.params).await
                    });

                    // 결과 대기 및 응답 (실제 구현에선 채널 사용 권장)
                    if let Ok(Ok(result)) = handle.await {
                        let resp = McpResponse {
                            id: req_id,
                            result: serde_json::json!(result),
                        };
                        
                        if let Ok(serialized) = serde_json::to_string(&resp) {
                            let _ = socket.write_all(serialized.as_bytes()).await;
                        }
                    }
                }
            }
        });
    }
}

use futures::stream::{Stream, StreamExt};
use std::pin::Pin;

// 가상의 로그 데이터 생성 스트림
fn log_stream() -> Pin<Box<dyn Stream<Item = String> + Send>> {
    Box::pin(async_stream::stream! {
        for i in 0..1000 {
            yield format!("Log entry #{}\n", i);
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        }
    })
}

// 스트림 처리 예시
async fn process_logs() {
    let mut stream = log_stream();
    while let Some(log_entry) = stream.next().await {
        // 실시간으로 한 줄씩 처리 (파일 쓰기 or 전송)
        println!("Processing: {}", log_entry);
    }
}