tests_integration/grpc/
network.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Mock connectors are registered on a shared channel manager.
16#[cfg(test)]
17#[allow(deprecated)]
18mod tests {
19    use std::io;
20    use std::pin::Pin;
21    use std::task::{Context, Poll};
22    use std::time::Duration;
23
24    use client::Client;
25    use common_grpc::channel_manager::ChannelManager;
26    use futures_util::future::BoxFuture;
27    use http::Uri;
28    use hyper_util::rt::TokioIo;
29    use servers::grpc::GrpcServerConfig;
30    use servers::server::Server;
31    use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
32    use tokio::net::TcpStream;
33    use tokio::sync::mpsc;
34    use tower::Service;
35
36    use crate::test_util::{StorageType, setup_grpc_server_with};
37
38    struct NetworkTrafficMonitorableConnector {
39        interested_tx: mpsc::Sender<String>,
40    }
41
42    impl Service<Uri> for NetworkTrafficMonitorableConnector {
43        type Response = TokioIo<CollectGrpcResponseFrameTypeStream>;
44        type Error = String;
45        type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
46
47        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
48            Poll::Ready(Ok(()))
49        }
50
51        fn call(&mut self, uri: Uri) -> Self::Future {
52            let frame_types = self.interested_tx.clone();
53
54            Box::pin(async move {
55                let addr = format!(
56                    "{}:{}",
57                    uri.host().unwrap_or("localhost"),
58                    uri.port_u16().unwrap_or(4001),
59                );
60                let inner = TcpStream::connect(addr).await.map_err(|e| e.to_string())?;
61                Ok(TokioIo::new(CollectGrpcResponseFrameTypeStream {
62                    inner,
63                    frame_types,
64                }))
65            })
66        }
67    }
68
69    struct CollectGrpcResponseFrameTypeStream {
70        inner: TcpStream,
71        frame_types: mpsc::Sender<String>,
72    }
73
74    impl AsyncRead for CollectGrpcResponseFrameTypeStream {
75        fn poll_read(
76            mut self: Pin<&mut Self>,
77            cx: &mut Context<'_>,
78            buf: &mut ReadBuf<'_>,
79        ) -> Poll<io::Result<()>> {
80            let before_len = buf.filled().len();
81
82            let result = Pin::new(&mut self.inner).poll_read(cx, buf);
83            if let Poll::Ready(Ok(())) = &result {
84                let after_len = buf.filled().len();
85
86                let new_data = &buf.filled()[before_len..after_len];
87                if let Some(frame_type) = maybe_decode_frame_type(new_data)
88                    && let Err(_) = self.frame_types.try_send(frame_type.to_string())
89                {
90                    return Poll::Ready(Err(io::Error::other("interested party has gone")));
91                }
92            }
93            result
94        }
95    }
96
97    fn maybe_decode_frame_type(data: &[u8]) -> Option<&str> {
98        (data.len() >= 9).then(|| match data[3] {
99            0x0 => "DATA",
100            0x1 => "HEADERS",
101            0x2 => "PRIORITY",
102            0x3 => "RST_STREAM",
103            0x4 => "SETTINGS",
104            0x5 => "PUSH_PROMISE",
105            0x6 => "PING",
106            0x7 => "GOAWAY",
107            0x8 => "WINDOW_UPDATE",
108            0x9 => "CONTINUATION",
109            _ => "UNKNOWN",
110        })
111    }
112
113    impl AsyncWrite for CollectGrpcResponseFrameTypeStream {
114        fn poll_write(
115            mut self: Pin<&mut Self>,
116            cx: &mut Context<'_>,
117            buf: &[u8],
118        ) -> Poll<Result<usize, io::Error>> {
119            Pin::new(&mut self.inner).poll_write(cx, buf)
120        }
121
122        fn poll_flush(
123            mut self: Pin<&mut Self>,
124            cx: &mut Context<'_>,
125        ) -> Poll<Result<(), io::Error>> {
126            Pin::new(&mut self.inner).poll_flush(cx)
127        }
128
129        fn poll_shutdown(
130            mut self: Pin<&mut Self>,
131            cx: &mut Context<'_>,
132        ) -> Poll<Result<(), io::Error>> {
133            Pin::new(&mut self.inner).poll_shutdown(cx)
134        }
135    }
136
137    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
138    async fn test_grpc_max_connection_age() {
139        let config = GrpcServerConfig {
140            max_connection_age: Some(Duration::from_secs(1)),
141            ..Default::default()
142        };
143        let (_db, server) = setup_grpc_server_with(
144            StorageType::File,
145            "test_grpc_max_connection_age",
146            None,
147            Some(config),
148            None,
149        )
150        .await;
151        let addr = server.bind_addr().unwrap().to_string();
152
153        let channel_manager = ChannelManager::new();
154        let client = Client::with_manager_and_urls(channel_manager.clone(), vec![&addr]);
155
156        let (tx, mut rx) = mpsc::channel(1024);
157        channel_manager
158            .reset_with_connector(
159                &addr,
160                NetworkTrafficMonitorableConnector { interested_tx: tx },
161            )
162            .unwrap();
163
164        let recv = tokio::spawn(async move {
165            let sleep = tokio::time::sleep(Duration::from_secs(3));
166            tokio::pin!(sleep);
167
168            let mut frame_types = vec![];
169            loop {
170                tokio::select! {
171                    x = rx.recv() => {
172                        if let Some(x) = x {
173                            frame_types.push(x);
174                        } else {
175                            break;
176                        }
177                    }
178                    _ = &mut sleep => {
179                        break;
180                    }
181                }
182            }
183            frame_types
184        });
185
186        // Drive the gRPC connection, has no special meaning for this keep-alive test.
187        let _ = client.health_check().await;
188
189        let frame_types = recv.await.unwrap();
190        // If "max_connection_age" has taken effects, server will return a "GOAWAY" message.
191        assert!(
192            frame_types.iter().any(|x| x == "GOAWAY"),
193            "{:?}",
194            frame_types
195        );
196
197        server.shutdown().await.unwrap();
198    }
199}
tests_integration/grpc/network.rs

tests_integration/grpc/
network.rs
