async-stream-server.rs

// Copyright 2022 Alibaba Cloud. All rights reserved.
//
// SPDX-License-Identifier: Apache-2.0
//

mod protocols;
mod utils;

use std::sync::Arc;

use log::{info, LevelFilter};

use protocols::asynchronous::{empty, streaming, streaming_ttrpc};
use ttrpc::{asynchronous::Server, Error};

use async_trait::async_trait;
use tokio::time::sleep;

struct StreamingService;

#[async_trait]
impl streaming_ttrpc::Streaming for StreamingService {
    async fn echo(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        mut e: streaming::EchoPayload,
    ) -> ::ttrpc::Result<streaming::EchoPayload> {
        e.seq += 1;
        Ok(e)
    }

    async fn echo_stream(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        mut s: ::ttrpc::r#async::ServerStream<streaming::EchoPayload, streaming::EchoPayload>,
    ) -> ::ttrpc::Result<()> {
        while let Some(mut e) = s.recv().await? {
            e.seq += 1;
            s.send(&e).await?;
        }

        Ok(())
    }

    async fn sum_stream(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        mut s: ::ttrpc::r#async::ServerStreamReceiver<streaming::Part>,
    ) -> ::ttrpc::Result<streaming::Sum> {
        let mut sum = streaming::Sum::new();
        while let Some(part) = s.recv().await? {
            sum.sum += part.add;
            sum.num += 1;
        }

        Ok(sum)
    }

    async fn divide_stream(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        sum: streaming::Sum,
        s: ::ttrpc::r#async::ServerStreamSender<streaming::Part>,
    ) -> ::ttrpc::Result<()> {
        let mut parts = vec![streaming::Part::new(); sum.num as usize];

        let mut total = 0i32;
        for i in 1..(sum.num - 2) {
            let add = (rand::random::<u32>() % 1000) as i32 - 500;
            parts[i as usize].add = add;
            total += add;
        }

        parts[sum.num as usize - 2].add = sum.sum - total;

        for part in parts {
            s.send(&part).await.unwrap();
        }

        Ok(())
    }

    async fn echo_null(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        mut s: ::ttrpc::r#async::ServerStreamReceiver<streaming::EchoPayload>,
    ) -> ::ttrpc::Result<empty::Empty> {
        let mut seq = 0;
        while let Some(e) = s.recv().await? {
            assert_eq!(e.seq, seq);
            assert_eq!(e.msg.as_str(), "non-empty empty");
            seq += 1;
        }
        Ok(empty::Empty::new())
    }

    async fn echo_null_stream(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        s: ::ttrpc::r#async::ServerStream<empty::Empty, streaming::EchoPayload>,
    ) -> ::ttrpc::Result<()> {
        let msg = "non-empty empty".to_string();

        let mut tasks = Vec::new();

        let (tx, mut rx) = s.split();
        let mut seq = 0u32;
        while let Some(e) = rx.recv().await? {
            assert_eq!(e.seq, seq);
            assert_eq!(e.msg, msg);
            seq += 1;

            for _i in 0..10 {
                let tx = tx.clone();
                tasks.push(tokio::spawn(
                    async move { tx.send(&empty::Empty::new()).await },
                ));
            }
        }

        for t in tasks {
            t.await.unwrap().map_err(|e| {
                ::ttrpc::Error::RpcStatus(::ttrpc::get_status(
                    ::ttrpc::Code::UNKNOWN,
                    e.to_string(),
                ))
            })?;
        }
        Ok(())
    }

    // It verifies PR #208
    async fn echo_default_value(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        e: streaming::EchoPayload,
        s: ::ttrpc::r#async::ServerStreamSender<streaming::EchoPayload>,
    ) -> ::ttrpc::Result<()> {
        if e.seq != 0 || !e.msg.is_empty() {
            return Err(Error::Others(
                "Expect a request with empty payload to verify #208".to_string(),
            ));
        }

        s.send(&e).await.unwrap();

        Ok(())
    }

    async fn server_send_stream(
        &self,
        _ctx: &::ttrpc::r#async::TtrpcContext,
        _: empty::Empty,
        s: ::ttrpc::r#async::ServerStreamSender<streaming::EchoPayload>,
    ) -> ::ttrpc::Result<()> {
        let mut seq = 0;
        while seq < 10 {
            sleep(std::time::Duration::from_millis(100)).await;
            let mut e = streaming::EchoPayload::new();
            e.seq = seq;
            e.msg = format!("hello");
            s.send(&e).await.unwrap();
            seq += 1;
        }

        Ok(())
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
    simple_logging::log_to_stderr(LevelFilter::Info);
    let service = streaming_ttrpc::create_streaming(Arc::new(StreamingService {}));

    let sock_addr = utils::get_sock_addr();
    utils::remove_if_sock_exist(sock_addr).unwrap();

    let mut server = Server::new()
        .bind(sock_addr)
        .unwrap()
        .register_service(service);

    server.start().await.unwrap();

    tokio::select! {
        _ = utils::hangup() => {
            // test stop_listen -> start
            info!("stop listen");
            server.stop_listen().await;
            info!("start listen");
            server.start().await.unwrap();

            // hold some time for the new test connection.
            sleep(std::time::Duration::from_secs(100)).await;
        }
        _ = utils::interrupt() => {
            // test graceful shutdown
            info!("graceful shutdown");
            server.shutdown().await.unwrap();
        }
    };
}