wasm_sandbox/communication/
streaming.rs

1//! Streaming API implementation for large data handling
2//! 
3//! This module provides streaming capabilities for processing large data
4//! that doesn't fit entirely in memory.
5
6use std::io;
7use std::pin::Pin;
8use std::sync::Arc;
9use std::path::Path;
10
11use async_trait::async_trait;
12use futures::{Stream, StreamExt};
13use tokio::io::{AsyncReadExt, AsyncWriteExt};
14
15use crate::error::{Error, Result, SandboxError, ResourceKind};
16
17/// Stream direction
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum StreamDirection {
20    /// Host to guest (input to WebAssembly)
21    HostToGuest,
22    
23    /// Guest to host (output from WebAssembly)
24    GuestToHost,
25    
26    /// Bidirectional
27    Bidirectional,
28}
29
30/// Stream type
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum StreamType {
33    /// Raw bytes
34    Bytes,
35    
36    /// UTF-8 text
37    Text,
38    
39    /// JSON data
40    Json,
41    
42    /// MessagePack data
43    MessagePack,
44}
45
46/// Streaming channel configuration
47#[derive(Debug, Clone)]
48pub struct StreamingChannelConfig {
49    /// Stream direction
50    pub direction: StreamDirection,
51    
52    /// Stream type
53    pub stream_type: StreamType,
54    
55    /// Buffer size for the channel
56    pub buffer_size: usize,
57    
58    /// Maximum chunk size
59    pub max_chunk_size: usize,
60    
61    /// Whether to validate UTF-8 for text streams
62    pub validate_utf8: bool,
63}
64
65impl Default for StreamingChannelConfig {
66    fn default() -> Self {
67        Self {
68            direction: StreamDirection::Bidirectional,
69            stream_type: StreamType::Bytes,
70            buffer_size: 64 * 1024, // 64KB
71            max_chunk_size: 16 * 1024, // 16KB
72            validate_utf8: true,
73        }
74    }
75}
76
77/// Streaming statistics
78#[derive(Debug, Clone)]
79pub struct StreamingStats {
80    /// Number of bytes sent
81    pub bytes_sent: u64,
82    
83    /// Number of bytes received
84    pub bytes_received: u64,
85    
86    /// Number of chunks sent
87    pub chunks_sent: u64,
88    
89    /// Number of chunks received
90    pub chunks_received: u64,
91    
92    /// Average chunk size
93    pub average_chunk_size: f64,
94    
95    /// Maximum chunk size seen
96    pub max_chunk_size_seen: usize,
97    
98    /// Number of errors
99    pub error_count: u64,
100}
101
102/// Streaming channel trait
103pub trait StreamingChannel: Send + Sync {
104    /// Get the channel ID
105    fn id(&self) -> &str;
106    
107    /// Get the channel configuration
108    fn config(&self) -> &StreamingChannelConfig;
109    
110    /// Get the streaming statistics
111    fn stats(&self) -> StreamingStats;
112    
113    /// Check if the channel is open
114    fn is_open(&self) -> bool;
115    
116    /// Close the channel
117    fn close(&self) -> Result<()>;
118}
119
120/// Byte stream chunk with metadata
121#[derive(Debug, Clone)]
122pub struct StreamChunk {
123    /// Chunk data
124    pub data: Vec<u8>,
125    
126    /// Whether this is the final chunk
127    pub is_final: bool,
128    
129    /// Sequence number
130    pub sequence: u64,
131    
132    /// Chunk metadata
133    pub metadata: Option<serde_json::Value>,
134}
135
136/// Streaming input channel
137#[async_trait]
138pub trait StreamingInput: StreamingChannel {
139    /// Send a chunk to the guest
140    async fn send_chunk(&self, chunk: StreamChunk) -> Result<()>;
141    
142    /// Send raw bytes to the guest
143    async fn send_bytes(&self, data: &[u8], is_final: bool) -> Result<()>;
144    
145    /// Send a stream of chunks to the guest
146    async fn send_stream<S>(&self, stream: &mut S) -> Result<u64>
147    where
148        S: Stream<Item = Result<StreamChunk>> + Unpin + Send;
149}
150
151/// Streaming output channel
152#[async_trait]
153pub trait StreamingOutput: StreamingChannel {
154    /// Receive a chunk from the guest
155    async fn receive_chunk(&self) -> Result<StreamChunk>;
156    
157    /// Receive raw bytes from the guest
158    async fn receive_bytes(&self, max_size: Option<usize>) -> Result<Vec<u8>>;
159    
160    /// Receive a stream of chunks from the guest
161    fn receive_stream(&self) -> Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send>>;
162}
163
164/// Bidirectional streaming channel
165pub trait StreamingChannel2Way: StreamingInput + StreamingOutput {}
166
167/// Concrete implementation for dynamic dispatch
168#[derive(Clone)]
169pub enum StreamingChannel2WayImpl {
170    Memory(MemoryStreamingChannel),
171    File(FileStreamingChannel),
172}
173
174impl StreamingChannel for StreamingChannel2WayImpl {
175    fn id(&self) -> &str {
176        match self {
177            StreamingChannel2WayImpl::Memory(ch) => ch.id(),
178            StreamingChannel2WayImpl::File(ch) => ch.id(),
179        }
180    }
181    
182    fn config(&self) -> &StreamingChannelConfig {
183        match self {
184            StreamingChannel2WayImpl::Memory(ch) => ch.config(),
185            StreamingChannel2WayImpl::File(ch) => ch.config(),
186        }
187    }
188    
189    fn stats(&self) -> StreamingStats {
190        match self {
191            StreamingChannel2WayImpl::Memory(ch) => ch.stats(),
192            StreamingChannel2WayImpl::File(ch) => ch.stats(),
193        }
194    }
195    
196    fn is_open(&self) -> bool {
197        match self {
198            StreamingChannel2WayImpl::Memory(ch) => ch.is_open(),
199            StreamingChannel2WayImpl::File(ch) => ch.is_open(),
200        }
201    }
202    
203    fn close(&self) -> Result<()> {
204        match self {
205            StreamingChannel2WayImpl::Memory(ch) => ch.close(),
206            StreamingChannel2WayImpl::File(ch) => ch.close(),
207        }
208    }
209}
210
211#[async_trait]
212impl StreamingInput for StreamingChannel2WayImpl {
213    async fn send_chunk(&self, chunk: StreamChunk) -> Result<()> {
214        match self {
215            StreamingChannel2WayImpl::Memory(ch) => ch.send_chunk(chunk).await,
216            StreamingChannel2WayImpl::File(ch) => ch.send_chunk(chunk).await,
217        }
218    }
219    
220    async fn send_bytes(&self, data: &[u8], is_final: bool) -> Result<()> {
221        match self {
222            StreamingChannel2WayImpl::Memory(ch) => ch.send_bytes(data, is_final).await,
223            StreamingChannel2WayImpl::File(ch) => ch.send_bytes(data, is_final).await,
224        }
225    }
226    
227    async fn send_stream<S>(&self, stream: &mut S) -> Result<u64>
228    where
229        S: Stream<Item = Result<StreamChunk>> + Unpin + Send
230    {
231        match self {
232            StreamingChannel2WayImpl::Memory(ch) => ch.send_stream(stream).await,
233            StreamingChannel2WayImpl::File(ch) => ch.send_stream(stream).await,
234        }
235    }
236}
237
238#[async_trait]
239impl StreamingOutput for StreamingChannel2WayImpl {
240    async fn receive_chunk(&self) -> Result<StreamChunk> {
241        match self {
242            StreamingChannel2WayImpl::Memory(ch) => ch.receive_chunk().await,
243            StreamingChannel2WayImpl::File(ch) => ch.receive_chunk().await,
244        }
245    }
246    
247    async fn receive_bytes(&self, max_size: Option<usize>) -> Result<Vec<u8>> {
248        match self {
249            StreamingChannel2WayImpl::Memory(ch) => ch.receive_bytes(max_size).await,
250            StreamingChannel2WayImpl::File(ch) => ch.receive_bytes(max_size).await,
251        }
252    }
253    
254    fn receive_stream(&self) -> Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send>> {
255        match self {
256            StreamingChannel2WayImpl::Memory(ch) => ch.receive_stream(),
257            StreamingChannel2WayImpl::File(ch) => ch.receive_stream(),
258        }
259    }
260}
261
262impl StreamingChannel2Way for StreamingChannel2WayImpl {}
263
264/// Concrete streaming channel types
265#[derive(Clone)]
266pub enum StreamingChannelType {
267    Memory(Arc<MemoryStreamingChannel>),
268    File(Arc<FileStreamingChannel>),
269}
270
271impl StreamingChannelType {
272    pub fn id(&self) -> &str {
273        match self {
274            StreamingChannelType::Memory(channel) => &channel.id,
275            StreamingChannelType::File(channel) => &channel.id,
276        }
277    }
278    
279    pub async fn send_chunk(&self, chunk: StreamChunk) -> Result<()> {
280        match self {
281            StreamingChannelType::Memory(channel) => channel.send_chunk(chunk).await,
282            StreamingChannelType::File(channel) => channel.send_chunk(chunk).await,
283        }
284    }
285    
286    pub async fn receive_chunk(&self) -> Result<StreamChunk> {
287        match self {
288            StreamingChannelType::Memory(channel) => channel.receive_chunk().await,
289            StreamingChannelType::File(channel) => channel.receive_chunk().await,
290        }
291    }
292    
293    pub async fn close(&self) -> Result<()> {
294        match self {
295            StreamingChannelType::Memory(channel) => channel.close(),
296            StreamingChannelType::File(channel) => channel.close(),
297        }
298    }
299}
300
301/// Memory-based streaming channel implementation
302#[derive(Clone)]
303pub struct MemoryStreamingChannel {
304    /// Channel ID
305    id: String,
306    
307    /// Channel configuration
308    config: StreamingChannelConfig,
309    
310    /// Internal state
311    state: Arc<tokio::sync::RwLock<MemoryStreamingState>>,
312}
313
314/// Internal state for memory streaming channel
315struct MemoryStreamingState {
316    /// Whether the channel is open
317    is_open: bool,
318    
319    /// Host to guest queue
320    h2g_queue: Vec<StreamChunk>,
321    
322    /// Guest to host queue
323    g2h_queue: Vec<StreamChunk>,
324    
325    /// Stream statistics
326    stats: StreamingStats,
327}
328
329impl MemoryStreamingChannel {
330    /// Create a new memory-based streaming channel
331    pub fn new(id: impl Into<String>, config: StreamingChannelConfig) -> Self {
332        Self {
333            id: id.into(),
334            config,
335            state: Arc::new(tokio::sync::RwLock::new(MemoryStreamingState {
336                is_open: true,
337                h2g_queue: Vec::new(),
338                g2h_queue: Vec::new(),
339                stats: StreamingStats {
340                    bytes_sent: 0,
341                    bytes_received: 0,
342                    chunks_sent: 0,
343                    chunks_received: 0,
344                    average_chunk_size: 0.0,
345                    max_chunk_size_seen: 0,
346                    error_count: 0,
347                },
348            })),
349        }
350    }
351}
352
353impl StreamingChannel for MemoryStreamingChannel {
354    fn id(&self) -> &str {
355        &self.id
356    }
357    
358    fn config(&self) -> &StreamingChannelConfig {
359        &self.config
360    }
361    
362    fn stats(&self) -> StreamingStats {
363        tokio::task::block_in_place(|| {
364            tokio::runtime::Handle::current().block_on(async {
365                self.state.read().await.stats.clone()
366            })
367        })
368    }
369    
370    fn is_open(&self) -> bool {
371        tokio::task::block_in_place(|| {
372            tokio::runtime::Handle::current().block_on(async {
373                self.state.read().await.is_open
374            })
375        })
376    }
377    
378    fn close(&self) -> Result<()> {
379        tokio::task::block_in_place(|| {
380            tokio::runtime::Handle::current().block_on(async {
381                let mut state = self.state.write().await;
382                state.is_open = false;
383                Ok(())
384            })
385        })
386    }
387}
388
389#[async_trait]
390impl StreamingInput for MemoryStreamingChannel {
391    async fn send_chunk(&self, chunk: StreamChunk) -> Result<()> {
392        let mut state = self.state.write().await;
393        
394        if !state.is_open {
395            return Err(Error::Communication {
396                channel: "memory_streaming".to_string(),
397                reason: "Streaming channel is closed".to_string(),
398                instance_id: None,
399            });
400        }
401        
402        if chunk.data.len() > self.config.max_chunk_size {
403            state.stats.error_count += 1;
404            return Err(Error::ResourceExhausted {
405                kind: ResourceKind::Memory,
406                limit: self.config.max_chunk_size as u64,
407                used: chunk.data.len() as u64,
408                instance_id: None,
409                suggestion: Some(format!("Consider reducing chunk size to {} bytes or less", self.config.max_chunk_size)),
410            });
411        }
412        
413        // Update statistics
414        state.stats.bytes_sent += chunk.data.len() as u64;
415        state.stats.chunks_sent += 1;
416        
417        // Update average chunk size
418        let new_avg = ((state.stats.average_chunk_size * (state.stats.chunks_sent - 1) as f64) 
419            + chunk.data.len() as f64) / state.stats.chunks_sent as f64;
420        state.stats.average_chunk_size = new_avg;
421        
422        // Update max chunk size
423        state.stats.max_chunk_size_seen = state.stats.max_chunk_size_seen.max(chunk.data.len());
424        
425        // Add to queue
426        state.h2g_queue.push(chunk);
427        
428        Ok(())
429    }
430    
431    async fn send_bytes(&self, data: &[u8], is_final: bool) -> Result<()> {
432        // Create a chunk from raw bytes
433        let chunk = StreamChunk {
434            data: data.to_vec(),
435            is_final,
436            sequence: 0, // Will be assigned in send_chunk
437            metadata: None,
438        };
439        
440        self.send_chunk(chunk).await
441    }
442    
443    async fn send_stream<S>(&self, stream: &mut S) -> Result<u64>
444    where
445        S: Stream<Item = Result<StreamChunk>> + Unpin + Send,
446    {
447        let mut bytes_sent = 0;
448        
449        while let Some(chunk_result) = stream.next().await {
450            let chunk = chunk_result?;
451            bytes_sent += chunk.data.len() as u64;
452            self.send_chunk(chunk).await?;
453        }
454        
455        Ok(bytes_sent)
456    }
457}
458
459#[async_trait]
460impl StreamingOutput for MemoryStreamingChannel {
461    async fn receive_chunk(&self) -> Result<StreamChunk> {
462        let mut state = self.state.write().await;
463        
464        if !state.is_open {
465            return Err(Error::Communication {
466                channel: "memory_streaming".to_string(),
467                reason: "Streaming channel is closed".to_string(),
468                instance_id: None,
469            });
470        }
471        
472        if state.g2h_queue.is_empty() {
473            return Err(Error::NotFound {
474                resource_type: "stream_data".to_string(),
475                identifier: "No data available in streaming channel".to_string(),
476            });
477        }
478        
479        // Get the first chunk
480        let chunk = state.g2h_queue.remove(0);
481        
482        // Update statistics
483        state.stats.bytes_received += chunk.data.len() as u64;
484        state.stats.chunks_received += 1;
485        
486        Ok(chunk)
487    }
488    
489    async fn receive_bytes(&self, max_size: Option<usize>) -> Result<Vec<u8>> {
490        let chunk = self.receive_chunk().await?;
491        
492        // Check size limit if specified
493        if let Some(max) = max_size {
494            if chunk.data.len() > max {
495                return Err(Error::ResourceExhausted {
496                    kind: ResourceKind::Memory,
497                    limit: max as u64,
498                    used: chunk.data.len() as u64,
499                    instance_id: None,
500                    suggestion: Some(format!("Consider increasing max_size to {} bytes or more", chunk.data.len())),
501                });
502            }
503        }
504        
505        Ok(chunk.data)
506    }
507    
508    fn receive_stream(&self) -> Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send>> {
509        // Create a clone of the state for the stream
510        let state = self.state.clone();
511        
512        // Create a stream that yields chunks from the g2h_queue
513        let stream = async_stream::stream! {
514            let mut last_seen_index = 0;
515            
516            loop {
517                // Get the current state
518                let state_guard = state.read().await;
519                
520                // Check if closed
521                if !state_guard.is_open {
522                    break;
523                }
524                
525                // Get all new chunks
526                let new_chunks: Vec<StreamChunk> = state_guard.g2h_queue
527                    .iter()
528                    .skip(last_seen_index)
529                    .cloned()
530                    .collect();
531                
532                // Update last seen index
533                last_seen_index = state_guard.g2h_queue.len();
534                
535                // Drop the guard
536                drop(state_guard);
537                
538                // Yield all new chunks
539                for chunk in new_chunks {
540                    let is_final = chunk.is_final;
541                    yield Ok(chunk);
542                    
543                    // If this is the final chunk, break
544                    if is_final {
545                        return;
546                    }
547                }
548                
549                // Wait a bit before polling again
550                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
551            }
552            
553            // If we get here, the channel was closed
554            yield Err(Error::Communication {
555                channel: "memory_streaming".to_string(),
556                reason: "Streaming channel was closed".to_string(),
557                instance_id: None,
558            });
559        };
560        
561        Box::pin(stream)
562    }
563}
564
565impl StreamingChannel2Way for MemoryStreamingChannel {}
566
567/// File-based streaming channel
568#[derive(Clone)]
569pub struct FileStreamingChannel {
570    /// Channel ID
571    id: String,
572    
573    /// Channel configuration
574    config: StreamingChannelConfig,
575    
576    /// Input file path
577    input_path: Option<std::path::PathBuf>,
578    
579    /// Output file path
580    output_path: Option<std::path::PathBuf>,
581    
582    /// Internal state
583    state: Arc<tokio::sync::RwLock<FileStreamingState>>,
584}
585
586/// Internal state for file streaming channel
587struct FileStreamingState {
588    /// Whether the channel is open
589    is_open: bool,
590    
591    /// Stream statistics
592    stats: StreamingStats,
593}
594
595impl FileStreamingChannel {
596    /// Create a new file-based streaming channel
597    pub fn new(
598        id: impl Into<String>,
599        config: StreamingChannelConfig,
600        input_path: Option<impl AsRef<Path>>,
601        output_path: Option<impl AsRef<Path>>,
602    ) -> Self {
603        Self {
604            id: id.into(),
605            config,
606            input_path: input_path.map(|p| p.as_ref().to_path_buf()),
607            output_path: output_path.map(|p| p.as_ref().to_path_buf()),
608            state: Arc::new(tokio::sync::RwLock::new(FileStreamingState {
609                is_open: true,
610                stats: StreamingStats {
611                    bytes_sent: 0,
612                    bytes_received: 0,
613                    chunks_sent: 0,
614                    chunks_received: 0,
615                    average_chunk_size: 0.0,
616                    max_chunk_size_seen: 0,
617                    error_count: 0,
618                },
619            })),
620        }
621    }
622}
623
624impl StreamingChannel for FileStreamingChannel {
625    fn id(&self) -> &str {
626        &self.id
627    }
628    
629    fn config(&self) -> &StreamingChannelConfig {
630        &self.config
631    }
632    
633    fn stats(&self) -> StreamingStats {
634        tokio::task::block_in_place(|| {
635            tokio::runtime::Handle::current().block_on(async {
636                self.state.read().await.stats.clone()
637            })
638        })
639    }
640    
641    fn is_open(&self) -> bool {
642        tokio::task::block_in_place(|| {
643            tokio::runtime::Handle::current().block_on(async {
644                self.state.read().await.is_open
645            })
646        })
647    }
648    
649    fn close(&self) -> Result<()> {
650        tokio::task::block_in_place(|| {
651            tokio::runtime::Handle::current().block_on(async {
652                let mut state = self.state.write().await;
653                state.is_open = false;
654                Ok(())
655            })
656        })
657    }
658}
659
660#[async_trait]
661impl StreamingInput for FileStreamingChannel {
662    async fn send_chunk(&self, chunk: StreamChunk) -> Result<()> {
663        let mut state = self.state.write().await;
664        
665        if !state.is_open {
666            return Err(Error::Communication {
667                channel: "file_streaming".to_string(),
668                reason: "Streaming channel is closed".to_string(),
669                instance_id: None,
670            });
671        }
672        
673        let output_path = match &self.output_path {
674            Some(path) => path,
675            None => return Err(Error::Configuration {
676                message: "No output file configured".to_string(),
677                suggestion: Some("Configure an output file path using with_output_file()".to_string()),
678                field: Some("output_path".to_string()),
679            }),
680        };
681        
682        // Open the file for appending
683        let mut file = tokio::fs::OpenOptions::new()
684            .write(true)
685            .create(true)
686            .append(true)
687            .open(output_path)
688            .await
689            .map_err(|e| Error::Filesystem {
690                operation: "open".to_string(),
691                path: output_path.clone(),
692                reason: format!("Failed to open output file: {}", e),
693            })?;
694        
695        // Write the chunk size and data
696        let size = chunk.data.len() as u32;
697        let size_bytes = size.to_le_bytes();
698        
699        file.write_all(size_bytes.as_slice()).await
700            .map_err(|e| Error::Filesystem {
701                operation: "write".to_string(),
702                path: output_path.clone(),
703                reason: format!("Failed to write chunk size: {}", e),
704            })?;
705            
706        file.write_all(&chunk.data).await
707            .map_err(|e| Error::Filesystem {
708                operation: "write".to_string(),
709                path: output_path.clone(),
710                reason: format!("Failed to write chunk data: {}", e),
711            })?;
712        
713        // Write is_final flag
714        let final_flag = if chunk.is_final { 1u8 } else { 0u8 };
715        file.write_all([final_flag].as_slice()).await
716            .map_err(|e| Error::Filesystem {
717                operation: "write".to_string(),
718                path: output_path.clone(),
719                reason: format!("Failed to write final flag: {}", e),
720            })?;
721        
722        // Update statistics
723        state.stats.bytes_sent += chunk.data.len() as u64;
724        state.stats.chunks_sent += 1;
725        
726        // Update average chunk size
727        let new_avg = ((state.stats.average_chunk_size * (state.stats.chunks_sent - 1) as f64) 
728            + chunk.data.len() as f64) / state.stats.chunks_sent as f64;
729        state.stats.average_chunk_size = new_avg;
730        
731        // Update max chunk size
732        state.stats.max_chunk_size_seen = state.stats.max_chunk_size_seen.max(chunk.data.len());
733        
734        Ok(())
735    }
736    
737    async fn send_bytes(&self, data: &[u8], is_final: bool) -> Result<()> {
738        // Create a chunk from raw bytes
739        let chunk = StreamChunk {
740            data: data.to_vec(),
741            is_final,
742            sequence: 0,
743            metadata: None,
744        };
745        
746        self.send_chunk(chunk).await
747    }
748    
749    async fn send_stream<S>(&self, stream: &mut S) -> Result<u64>
750    where
751        S: Stream<Item = Result<StreamChunk>> + Unpin + Send,
752    {
753        let mut bytes_sent = 0;
754        
755        while let Some(chunk_result) = stream.next().await {
756            let chunk = chunk_result?;
757            bytes_sent += chunk.data.len() as u64;
758            self.send_chunk(chunk).await?;
759        }
760        
761        Ok(bytes_sent)
762    }
763}
764
765#[async_trait]
766impl StreamingOutput for FileStreamingChannel {
767    async fn receive_chunk(&self) -> Result<StreamChunk> {
768        let mut state = self.state.write().await;
769        
770        if !state.is_open {
771            return Err(Error::Communication {
772                channel: "file_streaming".to_string(),
773                reason: "Streaming channel is closed".to_string(),
774                instance_id: None,
775            });
776        }
777        
778        let input_path = match &self.input_path {
779            Some(path) => path,
780            None => return Err(Error::Configuration {
781                message: "No input file configured".to_string(),
782                suggestion: Some("Configure an input file path using with_input_file()".to_string()),
783                field: Some("input_path".to_string()),
784            }),
785        };
786        
787        // Check if file exists
788        if !tokio::fs::try_exists(input_path).await.unwrap_or(false) {
789            return Err(Error::NotFound {
790                resource_type: "input_file".to_string(),
791                identifier: "Input file does not exist".to_string(),
792            });
793        }
794        
795        // Open the file for reading
796        let mut file = tokio::fs::File::open(input_path).await
797            .map_err(|e| Error::Filesystem {
798                operation: "open".to_string(),
799                path: input_path.clone(),
800                reason: format!("Failed to open input file: {}", e),
801            })?;
802        
803        // Get file metadata
804        let metadata = file.metadata().await
805            .map_err(|e| Error::Filesystem {
806                operation: "metadata".to_string(),
807                path: input_path.clone(),
808                reason: format!("Failed to get file metadata: {}", e),
809            })?;
810            
811        // Check if file is empty
812        if metadata.len() == 0 {
813            return Err(Error::NotFound {
814                resource_type: "stream_data".to_string(),
815                identifier: "Input file is empty".to_string(),
816            });
817        }
818        
819        // Read the chunk size
820        let mut size_bytes = [0u8; 4];
821        file.read_exact(&mut size_bytes).await
822            .map_err(|e| Error::Filesystem {
823                operation: "read".to_string(),
824                path: input_path.clone(),
825                reason: format!("Failed to read chunk size: {}", e),
826            })?;
827            
828        let size = u32::from_le_bytes(size_bytes) as usize;
829        
830        // Read the chunk data
831        let mut data = vec![0u8; size];
832        file.read_exact(&mut data).await
833            .map_err(|e| Error::Filesystem {
834                operation: "read".to_string(),
835                path: input_path.clone(),
836                reason: format!("Failed to read chunk data: {}", e),
837            })?;
838        
839        // Read is_final flag
840        let mut final_flag = [0u8; 1];
841        file.read_exact(&mut final_flag).await
842            .map_err(|e| Error::Filesystem {
843                operation: "read".to_string(),
844                path: input_path.clone(),
845                reason: format!("Failed to read final flag: {}", e),
846            })?;
847            
848        let is_final = final_flag[0] != 0;
849        
850        // Create the chunk
851        let chunk = StreamChunk {
852            data,
853            is_final,
854            sequence: state.stats.chunks_received,
855            metadata: None,
856        };
857        
858        // Update statistics
859        state.stats.bytes_received += chunk.data.len() as u64;
860        state.stats.chunks_received += 1;
861        
862        Ok(chunk)
863    }
864    
865    async fn receive_bytes(&self, max_size: Option<usize>) -> Result<Vec<u8>> {
866        let chunk = self.receive_chunk().await?;
867        
868        // Check size limit if specified
869        if let Some(max) = max_size {
870            if chunk.data.len() > max {
871                return Err(Error::ResourceExhausted {
872                    kind: ResourceKind::Memory,
873                    limit: max as u64,
874                    used: chunk.data.len() as u64,
875                    instance_id: None,
876                    suggestion: Some(format!("Consider increasing max_size to {} bytes or more", chunk.data.len())),
877                });
878            }
879        }
880        
881        Ok(chunk.data)
882    }
883    
884    fn receive_stream(&self) -> Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send>> {
885        // Create clones for the stream
886        let state = self.state.clone();
887        let input_path = self.input_path.clone();
888        
889        // Create a stream that reads chunks from the file
890        let stream = async_stream::stream! {
891            // Check if input path is configured
892            let input_path = match &input_path {
893                Some(path) => path.clone(),
894                None => {
895                    yield Err(Error::Configuration {
896                        message: "No input file configured".to_string(),
897                        suggestion: Some("Configure an input file path using with_input_file()".to_string()),
898                        field: Some("input_path".to_string()),
899                    });
900                    return;
901                }
902            };
903            
904            // Try to open the file
905            let file = match tokio::fs::File::open(&input_path).await {
906                Ok(file) => file,
907                Err(e) => {
908                    yield Err(Error::Filesystem {
909                        operation: "open".to_string(),
910                        path: input_path.clone(),
911                        reason: format!("Failed to open input file: {}", e),
912                    });
913                    return;
914                }
915            };
916            
917            // Create a buffered reader
918            let mut reader = tokio::io::BufReader::new(file);
919            let mut sequence = 0;
920            
921            loop {
922                // Check if channel is still open
923                {
924                    let state_guard = state.read().await;
925                    if !state_guard.is_open {
926                        yield Err(Error::Communication {
927                            channel: "file_streaming".to_string(),
928                            reason: "Streaming channel was closed".to_string(),
929                            instance_id: None,
930                        });
931                        return;
932                    }
933                }
934                
935                // Read chunk size
936                let mut size_bytes = [0u8; 4];
937                match reader.read_exact(&mut size_bytes).await {
938                    Ok(_) => {}
939                    Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
940                        // End of file, wait for more data
941                        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
942                        continue;
943                    }
944                    Err(e) => {
945                        yield Err(Error::Filesystem {
946                            operation: "read".to_string(),
947                            path: input_path.clone(),
948                            reason: format!("Failed to read chunk size: {}", e),
949                        });
950                        return;
951                    }
952                }
953                
954                let size = u32::from_le_bytes(size_bytes) as usize;
955                
956                // Read chunk data
957                let mut data = vec![0u8; size];
958                if let Err(e) = reader.read_exact(&mut data).await {
959                    yield Err(Error::Filesystem {
960                        operation: "read".to_string(),
961                        path: input_path.clone(),
962                        reason: format!("Failed to read chunk data: {}", e),
963                    });
964                    return;
965                }
966                
967                // Read is_final flag
968                let mut final_flag = [0u8; 1];
969                if let Err(e) = reader.read_exact(&mut final_flag).await {
970                    yield Err(Error::Filesystem {
971                        operation: "read".to_string(),
972                        path: input_path.clone(),
973                        reason: format!("Failed to read final flag: {}", e),
974                    });
975                    return;
976                }
977                
978                let is_final = final_flag[0] != 0;
979                
980                // Create the chunk
981                let chunk = StreamChunk {
982                    data,
983                    is_final,
984                    sequence,
985                    metadata: None,
986                };
987                
988                // Update statistics
989                {
990                    let mut state_guard = state.write().await;
991                    state_guard.stats.bytes_received += chunk.data.len() as u64;
992                    state_guard.stats.chunks_received += 1;
993                }
994                
995                // Yield the chunk
996                yield Ok(chunk);
997                sequence += 1;
998                
999                // If this is the final chunk, break
1000                if is_final {
1001                    break;
1002                }
1003            }
1004        };
1005        
1006        Box::pin(stream)
1007    }
1008}
1009
1010impl StreamingChannel2Way for FileStreamingChannel {}
1011
1012/// Streaming factory
1013pub struct StreamingFactory;
1014
1015impl StreamingFactory {
1016    /// Create a new memory-based streaming channel
1017    pub fn create_memory_channel(
1018        id: impl Into<String>,
1019        config: Option<StreamingChannelConfig>,
1020    ) -> StreamingChannelType {
1021        let config = config.unwrap_or_default();
1022        StreamingChannelType::Memory(Arc::new(MemoryStreamingChannel::new(id, config)))
1023    }
1024    
1025    /// Create a new file-based streaming channel
1026    pub fn create_file_channel(
1027        id: impl Into<String>,
1028        config: Option<StreamingChannelConfig>,
1029        input_path: Option<impl AsRef<Path>>,
1030        output_path: Option<impl AsRef<Path>>,
1031    ) -> StreamingChannelType {
1032        let config = config.unwrap_or_default();
1033        StreamingChannelType::File(Arc::new(FileStreamingChannel::new(id, config, input_path, output_path)))
1034    }
1035}
1036
1037/// Host function wrapper for handling streams
1038pub struct StreamingFunction<T, R> {
1039    /// Function to process the stream
1040    processor: Box<dyn Fn(T, Arc<StreamingChannel2WayImpl>) -> Result<R> + Send + Sync>,
1041    
1042    /// Input stream
1043    input_stream: Arc<StreamingChannel2WayImpl>,
1044}
1045
1046impl<T, R> StreamingFunction<T, R>
1047where
1048    T: serde::de::DeserializeOwned + Send + 'static,
1049    R: serde::Serialize + Send + 'static,
1050{
1051    /// Create a new streaming function
1052    pub fn new<F>(processor: F, input_stream: Arc<StreamingChannel2WayImpl>) -> Self
1053    where
1054        F: Fn(T, Arc<StreamingChannel2WayImpl>) -> Result<R> + Send + Sync + 'static,
1055    {
1056        Self {
1057            processor: Box::new(processor),
1058            input_stream,
1059        }
1060    }
1061    
1062    /// Process a request
1063    pub fn process(&self, request: T) -> Result<R> {
1064        (self.processor)(request, self.input_stream.clone())
1065    }
1066}
1067
1068/// Extension methods for streaming with WebAssembly instances
1069#[async_trait]
1070pub trait StreamingInstanceExt {
1071    /// Call a function with streaming data
1072    async fn call_streaming_function<Params, Return>(
1073        &self,
1074        function_name: &str,
1075        params: &Params,
1076        stream: Arc<StreamingChannel2WayImpl>,
1077    ) -> Result<Return>
1078    where
1079        Params: serde::Serialize + ?Sized + Send + Sync,
1080        Return: serde::de::DeserializeOwned + 'static;
1081}
1082
1083/// Streaming manager for coordinating streams between host and guest
1084pub struct StreamingManager {
1085    /// Active streaming channels
1086    channels: tokio::sync::RwLock<std::collections::HashMap<String, StreamingChannelType>>,
1087}
1088
1089impl StreamingManager {
1090    /// Create a new streaming manager
1091    pub fn new() -> Self {
1092        Self {
1093            channels: tokio::sync::RwLock::new(std::collections::HashMap::new()),
1094        }
1095    }
1096    
1097    /// Create a new streaming channel
1098    pub async fn create_channel(
1099        &self,
1100        id: Option<String>,
1101        config: Option<StreamingChannelConfig>,
1102    ) -> Result<StreamingChannelType> {
1103        let id = id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
1104        let channel = StreamingFactory::create_memory_channel(id.clone(), config);
1105        
1106        // Register the channel
1107        let mut channels = self.channels.write().await;
1108        channels.insert(id, channel.clone());
1109        
1110        Ok(channel)
1111    }
1112    
1113    /// Get a streaming channel by ID
1114    pub async fn get_channel(&self, id: &str) -> Option<StreamingChannelType> {
1115        let channels = self.channels.read().await;
1116        channels.get(id).cloned()
1117    }
1118    
1119    /// Close a streaming channel
1120    pub async fn close_channel(&self, id: &str) -> Result<()> {
1121        let mut channels = self.channels.write().await;
1122        
1123        if let Some(channel) = channels.remove(id) {
1124            channel.close().await?;
1125            Ok(())
1126        } else {
1127            Err(SandboxError::NotFound { 
1128                resource_type: "streaming_channel".to_string(), 
1129                identifier: id.to_string() 
1130            })
1131        }
1132    }
1133    
1134    /// List all active streaming channels
1135    pub async fn list_channels(&self) -> Vec<String> {
1136        let channels = self.channels.read().await;
1137        channels.keys().cloned().collect()
1138    }
1139}
1140
1141impl Default for StreamingManager {
1142    fn default() -> Self {
1143        Self::new()
1144    }
1145}
1146
1147/// Helper for transforming data in streams
1148pub struct StreamTransformer;
1149
1150impl StreamTransformer {
1151    /// Transform a byte stream to a text stream
1152    pub fn bytes_to_text_stream<S>(
1153        stream: S,
1154    ) -> impl Stream<Item = Result<String>>
1155    where
1156        S: Stream<Item = Result<StreamChunk>> + Send,
1157    {
1158        stream.map(|chunk_result| {
1159            let chunk = chunk_result?;
1160            let text = String::from_utf8(chunk.data)
1161                .map_err(|e| SandboxError::Serialization {
1162                    format: "utf-8".to_string(),
1163                    operation: "decode".to_string(),
1164                    reason: e.to_string(),
1165                })?;
1166            Ok(text)
1167        })
1168    }
1169    
1170    /// Transform a byte stream to a JSON stream
1171    pub fn bytes_to_json_stream<S, T>(
1172        stream: S,
1173    ) -> impl Stream<Item = Result<T>>
1174    where
1175        S: Stream<Item = Result<StreamChunk>> + Send,
1176        T: serde::de::DeserializeOwned,
1177    {
1178        stream.map(|chunk_result| {
1179            let chunk = chunk_result?;
1180            let value = serde_json::from_slice(&chunk.data)
1181                .map_err(|e| SandboxError::Serialization {
1182                    format: "json".to_string(),
1183                    operation: "decode".to_string(),
1184                    reason: e.to_string(),
1185                })?;
1186            Ok(value)
1187        })
1188    }
1189    
1190    /// Transform objects to a JSON byte stream
1191    pub fn objects_to_json_byte_stream<S, T>(
1192        stream: S,
1193    ) -> impl Stream<Item = Result<StreamChunk>>
1194    where
1195        S: Stream<Item = Result<T>> + Send,
1196        T: serde::Serialize,
1197    {
1198        stream.map(|item_result| {
1199            let item = item_result?;
1200            let json = serde_json::to_vec(&item)
1201                .map_err(|e| SandboxError::Serialization {
1202                    format: "json".to_string(),
1203                    operation: "encode".to_string(),
1204                    reason: e.to_string(),
1205                })?;
1206                
1207            Ok(StreamChunk {
1208                data: json,
1209                is_final: false, // Set to true for last chunk
1210                sequence: 0,     // Will be assigned by the channel
1211                metadata: None,
1212            })
1213        })
1214    }
1215}
wasm_sandbox/communication/streaming.rs

wasm_sandbox/communication/
streaming.rs
