Line data Source code
1 : use tokio::time::timeout;
2 : use std::sync::atomic::AtomicUsize;
3 : use std::time::Instant;
4 : use std::collections::HashMap;
5 : use std::sync::Arc;
6 : use std::sync::RwLock;
7 : use std::time::Duration;
8 : use tokio::sync::Mutex;
9 : use tonic::metadata::MetadataMap;
10 :
11 : use crate::status::GrpcStreamElement;
12 : use crate::status::Status;
13 :
14 : #[derive(Debug)]
15 : pub enum GrpcServerStreamEvent {
16 : Message(Vec<u8>),
17 : End(MetadataMap),
18 : ReadFailed(tonic::Status, GrpcStreamElement),
19 : }
20 :
21 : #[derive(Debug, PartialEq)]
22 : pub enum GrpcServerStreamBatchEvent {
23 : StreamFailed,
24 : StreamFinished,
25 : FlushAfterClose,
26 : FlushAfterTimeout,
27 : FlushAfterBytes,
28 : }
29 :
30 : impl GrpcServerStreamBatchEvent {
31 4 : pub fn to_str(&self) -> &'static str {
32 4 : match self {
33 0 : GrpcServerStreamBatchEvent::StreamFailed => "StreamFailed",
34 2 : GrpcServerStreamBatchEvent::StreamFinished => "StreamFinished",
35 0 : GrpcServerStreamBatchEvent::FlushAfterClose => "FlushAfterClose",
36 1 : GrpcServerStreamBatchEvent::FlushAfterTimeout => "FlushAfterTimeout",
37 1 : GrpcServerStreamBatchEvent::FlushAfterBytes => "FlushAfterBytes",
38 : }
39 4 : }
40 : }
41 :
42 : /// A batch of messages read from a gRPC server stream in one read call
43 : #[derive(Debug)]
44 : pub struct GrpcServerStreamBatch {
45 : /// The event that emitted the batch
46 : pub event: GrpcServerStreamBatchEvent,
47 : /// The messages in the batch
48 : pub messages: Vec<Vec<u8>>,
49 : /// The total message bytes in the batch
50 : pub total_message_bytes: usize,
51 : /// The trailing headers (if any)
52 : pub trailers: Option<MetadataMap>,
53 : }
54 :
55 : impl Default for GrpcServerStreamBatch {
56 7 : fn default() -> Self {
57 7 : Self {
58 7 : event: GrpcServerStreamBatchEvent::StreamFailed,
59 7 : messages: Vec::default(),
60 7 : total_message_bytes: 0,
61 7 : trailers: None,
62 7 : }
63 7 : }
64 : }
65 :
66 : /// A gRPC server stream whose messages are buffered by a background reader task
67 : pub struct GRPCServerStream {
68 : /// The channel id
69 : #[allow(dead_code)]
70 : channel_id: usize,
71 : /// The response queue reading. Dropping it stops the reader task, which cancels the call.
72 : response_receiver: Mutex<tokio::sync::mpsc::Receiver<GrpcServerStreamEvent>>,
73 : }
74 :
75 : /// The registry of open gRPC server streams, keyed by stream id
76 : #[derive(Default)]
77 : pub struct GrpcStreamManager {
78 : /// The next message id
79 : pub next_stream_id: AtomicUsize,
80 : /// The server streams
81 : pub server_streams: RwLock<HashMap<usize, Arc<GRPCServerStream>>>,
82 : }
83 :
84 : impl GrpcStreamManager {
85 : /// Start a server stream
86 : #[allow(dead_code)]
87 6 : pub fn start_server_stream<T: Into<Vec<u8>> + Send + 'static>(
88 6 : self: &Arc<Self>,
89 6 : channel_id: usize,
90 6 : mut streaming: tonic::Streaming<T>,
91 6 : ) -> Result<usize, Status> {
92 : // Allocate an identifier for the stream
93 6 : let reg = self.clone();
94 6 : let stream_id = reg
95 6 : .next_stream_id
96 6 : .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
97 :
98 : // Register the receiver
99 : const STREAM_CAPACITY: usize = 10;
100 6 : let (sender, receiver) = tokio::sync::mpsc::channel(STREAM_CAPACITY);
101 6 : let stream_entry = Arc::new(GRPCServerStream {
102 6 : channel_id,
103 6 : response_receiver: Mutex::new(receiver),
104 6 : });
105 6 : if let Ok(mut streams) = self.server_streams.write() {
106 6 : streams.insert(stream_id, stream_entry);
107 6 : }
108 :
109 : // Spawn the async reader
110 6 : tokio::spawn(async move {
111 6 : let reader = async {
112 6 : let mut stream_alive = true;
113 28 : while stream_alive {
114 : // Read the next message from the gRPC output stream
115 23 : let response = match streaming.message().await {
116 : // Received bytes, forward the message as-is to receivers
117 18 : Ok(Some(m)) => {
118 18 : let buffer: Vec<u8> = m.into();
119 18 : log::debug!(
120 : "Receiving message from server stream, bytes={}",
121 0 : buffer.len()
122 : );
123 18 : GrpcServerStreamEvent::Message(buffer)
124 : }
125 : // Stream was closed, delete stream and return trailers
126 : Ok(None) => {
127 3 : stream_alive = false;
128 3 : log::debug!("Server stream ended");
129 3 : match streaming.trailers().await {
130 3 : Ok(Some(trailer)) => GrpcServerStreamEvent::End(trailer),
131 0 : Ok(None) => GrpcServerStreamEvent::End(MetadataMap::new()),
132 0 : Err(e) => {
133 0 : GrpcServerStreamEvent::ReadFailed(e, GrpcStreamElement::Trailers)
134 : }
135 : }
136 : }
137 : // Stream failed, send error to receivers and close the stream
138 1 : Err(e) => {
139 1 : stream_alive = false;
140 1 : log::warn!("Failed reading from server stream: {}", e);
141 1 : GrpcServerStreamEvent::ReadFailed(e, GrpcStreamElement::Message)
142 : }
143 : };
144 :
145 : // Push the response into the queue
146 22 : if let Err(e) = sender
147 22 : .send(response)
148 22 : .await
149 : {
150 1 : stream_alive = false;
151 1 : log::warn!("Failed writing response to stream: {}", e);
152 21 : }
153 : }
154 5 : };
155 : // A deleted stream drops the receiver, dropping the call then tells the server to stop
156 6 : tokio::select! {
157 6 : _ = sender.closed() => log::debug!("Server stream {} was deleted", stream_id),
158 6 : _ = reader => {}
159 : }
160 6 : });
161 :
162 : // Return initial metadata
163 6 : Ok(stream_id)
164 6 : }
165 :
166 : /// Read from a server stream
167 : #[allow(dead_code)]
168 10 : pub async fn read_server_stream(
169 10 : self: &Arc<Self>,
170 10 : channel_id: usize,
171 10 : stream_id: usize,
172 10 : timeout_read_after: Duration,
173 10 : flush_batch_after: Duration,
174 10 : flush_batch_bytes: usize,
175 10 : ) -> Result<GrpcServerStreamBatch, Status> {
176 10 : let started_at = Instant::now();
177 :
178 : // Try to find the stream.
179 : // If there is none, the request might have finished already.
180 10 : let reg = self.clone();
181 10 : let stream = if let Some(streams) = reg.server_streams.read().unwrap().get(&stream_id) {
182 7 : streams.clone()
183 : } else {
184 3 : return Err(Status::GrpcStreamIsUnknown { channel_id, stream_id });
185 : };
186 :
187 : // Acquire the receiver lock.
188 : // There can only be one active reader at a time.
189 7 : let mut receiver = stream.response_receiver.lock().await;
190 :
191 : // Receive messages and return them in batches
192 7 : let mut batch = GrpcServerStreamBatch::default();
193 : loop {
194 : // Receive the next message from the reader
195 13 : let elapsed_since_start = started_at.elapsed();
196 13 : let response = if batch.messages.len() == 0 {
197 : // Initially, we check the `timeout_read_after` timeout.
198 7 : let receive_timeout = timeout_read_after.checked_sub(elapsed_since_start).unwrap_or_default();
199 7 : match timeout(receive_timeout, receiver.recv()).await {
200 6 : Ok(Some(response)) => response,
201 : Ok(None) => {
202 0 : return Err(Status::GrpcStreamClosed { channel_id, stream_id });
203 : },
204 : Err(_) => {
205 1 : return Err(Status::GrpcStreamReadTimedOut { channel_id, stream_id });
206 : }
207 : }
208 : } else {
209 : // After receiving one batch, we check the `flush_batch_after` timeout.
210 6 : let receive_timeout = match flush_batch_after.checked_sub(elapsed_since_start) {
211 6 : Some(timeout) => timeout,
212 : None => {
213 0 : batch.event = GrpcServerStreamBatchEvent::FlushAfterTimeout;
214 0 : return Ok(batch);
215 : }
216 : };
217 6 : match timeout(receive_timeout, receiver.recv()).await {
218 5 : Ok(Some(response)) => response,
219 : Ok(None) => {
220 0 : batch.event = GrpcServerStreamBatchEvent::FlushAfterClose;
221 0 : return Ok(batch);
222 : },
223 : Err(_) => {
224 1 : batch.event = GrpcServerStreamBatchEvent::FlushAfterTimeout;
225 1 : return Ok(batch);
226 : }
227 : }
228 : };
229 :
230 11 : match response {
231 : // We reached the end of the stream.
232 : // Flush the current batch with trailers and an EOS marker.
233 3 : GrpcServerStreamEvent::End(trailers) => {
234 3 : if let Ok(mut streams) = self.server_streams.write() {
235 3 : streams.remove(&stream_id);
236 3 : }
237 3 : batch.trailers = Some(trailers);
238 3 : batch.event = GrpcServerStreamBatchEvent::StreamFinished;
239 3 : return Ok(batch);
240 : }
241 : // An error occurred,.
242 : // Throw away any intermediate messages that we held back.
243 1 : GrpcServerStreamEvent::ReadFailed(status, element) => {
244 1 : if let Ok(mut streams) = self.server_streams.write() {
245 1 : streams.remove(&stream_id);
246 1 : }
247 1 : return Err(Status::GrpcStreamReadFailed{
248 1 : channel_id,
249 1 : stream_id,
250 1 : element,
251 1 : status,
252 1 : });
253 : }
254 : // Return the message
255 7 : GrpcServerStreamEvent::Message(m) => {
256 : // Add the message to the current batch and check if we should flush
257 7 : batch.total_message_bytes += m.len();
258 7 : batch.messages.push(m);
259 :
260 : // Flush if we hit the message size
261 7 : if batch.total_message_bytes > flush_batch_bytes {
262 1 : batch.event = GrpcServerStreamBatchEvent::FlushAfterBytes;
263 1 : return Ok(batch);
264 6 : }
265 : },
266 : }
267 : }
268 10 : }
269 :
270 : /// Cancellation is done by just dropping everything, see `start_server_stream`
271 : #[allow(dead_code)]
272 2 : pub async fn destroy_server_stream(self: &Arc<Self>, stream_id: usize) {
273 2 : if let Ok(mut streams) = self.server_streams.write() {
274 2 : streams.remove(&stream_id);
275 2 : }
276 2 : }
277 : }
278 :
279 : #[cfg(test)]
280 : mod test {
281 : use super::*;
282 : use std::collections::HashMap;
283 : use tokio_stream::StreamExt;
284 : use prost::Message;
285 :
286 : use crate::status::Status;
287 : use crate::test::hyper_service_mock::{spawn_hyper_service_mock, HyperServiceMock};
288 : use crate::proto::salesforce_hyperdb_grpc_v1::hyper_service_client::HyperServiceClient;
289 : use crate::proto::salesforce_hyperdb_grpc_v1::query_param;
290 : use crate::proto::salesforce_hyperdb_grpc_v1::query_result;
291 : use crate::proto::salesforce_hyperdb_grpc_v1::query_result_header::Header;
292 : use crate::proto::salesforce_hyperdb_grpc_v1::QueryBinaryResultChunk;
293 : use crate::proto::salesforce_hyperdb_grpc_v1::QueryParam;
294 : use crate::proto::salesforce_hyperdb_grpc_v1::QueryResultSchema;
295 : use crate::proto::salesforce_hyperdb_grpc_v1::SqlType;
296 : use crate::proto::salesforce_hyperdb_grpc_v1::{sql_type, QueryResultHeader};
297 : use crate::proto::salesforce_hyperdb_grpc_v1::{ColumnDescription, QueryResult};
298 :
299 : impl Into<Vec<u8>> for QueryResult {
300 2 : fn into(self) -> Vec<u8> {
301 2 : self.encode_to_vec()
302 2 : }
303 : }
304 :
305 : #[tokio::test]
306 1 : async fn test_execute_query_mock() -> Result<(), Status> {
307 1 : let (mock, mut setup_execute_query) = HyperServiceMock::new();
308 1 : let (addr, shutdown) = spawn_hyper_service_mock(mock).await;
309 1 : let mut client = HyperServiceClient::connect(format!("http://{}", addr))
310 1 : .await
311 1 : .unwrap();
312 :
313 1 : let param = QueryParam {
314 1 : query: "select 1".to_string(),
315 1 : output_format: query_param::OutputFormat::ArrowStream.into(),
316 1 : database: Vec::new(),
317 1 : params: HashMap::new(),
318 1 : };
319 1 : let result = client.execute_query(param).await.unwrap();
320 1 : let mut stream = result.into_inner();
321 :
322 : // Send the results
323 1 : let (_params, result_sender) = setup_execute_query.recv().await.unwrap();
324 1 : result_sender
325 1 : .send(Ok(QueryResult {
326 1 : result: Some(query_result::Result::Header(QueryResultHeader {
327 1 : header: Some(Header::Schema(QueryResultSchema {
328 1 : column: vec![ColumnDescription {
329 1 : name: "foo".to_string(),
330 1 : r#type: Some(SqlType {
331 1 : tag: sql_type::TypeTag::HyperBool.into(),
332 1 : modifier: None,
333 1 : }),
334 1 : }],
335 1 : })),
336 1 : })),
337 1 : }))
338 1 : .await.unwrap();
339 1 : result_sender
340 1 : .send(Ok(QueryResult {
341 1 : result: Some(query_result::Result::ArrowChunk(QueryBinaryResultChunk {
342 1 : data: vec![0x1, 0x2, 0x3, 0x4],
343 1 : })),
344 1 : }))
345 1 : .await.unwrap();
346 1 : drop(result_sender);
347 :
348 1 : let mut results = Vec::new();
349 3 : while let Some(item) = stream.next().await {
350 2 : assert!(item.is_ok());
351 2 : results.push(item.unwrap());
352 : }
353 1 : assert_eq!(results.len(), 2);
354 1 : assert!(results[0].result.is_some());
355 1 : assert_eq!(
356 1 : &results[0],
357 1 : &QueryResult {
358 1 : result: Some(query_result::Result::Header(QueryResultHeader {
359 1 : header: Some(Header::Schema(QueryResultSchema {
360 1 : column: vec![ColumnDescription {
361 1 : name: "foo".to_string(),
362 1 : r#type: Some(SqlType {
363 1 : tag: sql_type::TypeTag::HyperBool.into(),
364 1 : modifier: None,
365 1 : }),
366 1 : }],
367 1 : })),
368 1 : })),
369 1 : }
370 : );
371 1 : assert!(results[1].result.is_some());
372 1 : assert_eq!(
373 1 : &results[1],
374 1 : &QueryResult {
375 1 : result: Some(query_result::Result::ArrowChunk(QueryBinaryResultChunk {
376 1 : data: vec![0x1, 0x2, 0x3, 0x4],
377 1 : })),
378 1 : }
379 : );
380 1 : shutdown.send(()).unwrap();
381 :
382 2 : Ok(())
383 1 : }
384 :
385 : #[tokio::test]
386 1 : async fn test_stream_reader() -> Result<(), Status> {
387 : // Spawn a hyper service mock
388 1 : let (mock, mut setup_execute_query) = HyperServiceMock::new();
389 1 : let (addr, shutdown) = spawn_hyper_service_mock(mock).await;
390 1 : let mut client = HyperServiceClient::connect(format!("http://{}", addr))
391 1 : .await
392 1 : .unwrap();
393 :
394 : // Exeute the query
395 1 : let param = QueryParam {
396 1 : query: "select 1".to_string(),
397 1 : output_format: query_param::OutputFormat::ArrowStream.into(),
398 1 : database: Vec::new(),
399 1 : params: HashMap::new(),
400 1 : };
401 1 : let res = client.execute_query(param).await.unwrap();
402 1 : let (_params, result_sender) = setup_execute_query.recv().await.unwrap();
403 :
404 : // Send the results
405 1 : let messages = vec![
406 1 : Ok(QueryResult {
407 1 : result: Some(query_result::Result::ArrowChunk(QueryBinaryResultChunk {
408 1 : data: vec![0x1, 0x2, 0x3, 0x4],
409 1 : })),
410 1 : }),
411 1 : Ok(QueryResult {
412 1 : result: Some(query_result::Result::ArrowChunk(QueryBinaryResultChunk {
413 1 : data: vec![0x5, 0x6, 0x7, 0x8],
414 1 : })),
415 1 : }),
416 : ];
417 2 : for msg in messages.iter() {
418 2 : result_sender.send(msg.clone()).await.unwrap();
419 : }
420 1 : drop(result_sender);
421 :
422 : // Setup the stream manager
423 1 : let reg = Arc::new(GrpcStreamManager::default());
424 1 : reg.next_stream_id
425 1 : .fetch_add(42, std::sync::atomic::Ordering::SeqCst);
426 1 : let stream_id = reg.start_server_stream(0, res.into_inner()).unwrap();
427 1 : assert_eq!(stream_id, 42);
428 :
429 : // Read batched messages
430 1 : let batch = reg.read_server_stream(0, stream_id, Duration::from_secs(10), Duration::from_secs(10), 1000000).await?;
431 1 : assert_eq!(batch.messages.len(), 2);
432 1 : assert_eq!(batch.messages[0], messages[0].as_ref().unwrap().encode_to_vec());
433 1 : assert_eq!(batch.messages[1], messages[1].as_ref().unwrap().encode_to_vec());
434 1 : assert_eq!(batch.event, GrpcServerStreamBatchEvent::StreamFinished);
435 1 : assert!(!reg.server_streams.read().unwrap().contains_key(&stream_id));
436 :
437 1 : shutdown.send(()).unwrap();
438 2 : Ok(())
439 1 : }
440 : }
|
