GH-800: Drain the Pipeline

node/src/proxy_client/stream_establisher.rs

@@ -1,5 +1,6 @@
 // Copyright (c) 2019, MASQ (https://masq.ai) and/or its affiliates. All rights reserved.
 
+use crate::proxy_client::stream_handler_pool::StreamSenders;
 use crate::proxy_client::stream_reader::StreamReader;
 use crate::proxy_client::stream_writer::StreamWriter;
 use crate::sub_lib::channel_wrappers::FuturesChannelFactory;
@@ -14,16 +15,18 @@ use crate::sub_lib::stream_connector::StreamConnectorReal;
 use crate::sub_lib::stream_key::StreamKey;
 use crate::sub_lib::tokio_wrappers::ReadHalfWrapper;
 use actix::Recipient;
-use crossbeam_channel::Sender;
+use crossbeam_channel::{unbounded, Receiver, Sender};
 use masq_lib::logger::Logger;
 use std::io;
 use std::net::IpAddr;
 use std::net::SocketAddr;
+use tokio::sync::mpsc::UnboundedReceiver;
 
 pub struct StreamEstablisher {
     pub cryptde: &'static dyn CryptDE,
-    pub stream_adder_tx: Sender<(StreamKey, Box<dyn SenderWrapper<SequencedPacket>>)>,
+    pub stream_adder_tx: Sender<(StreamKey, StreamSenders)>,
     pub stream_killer_tx: Sender<(StreamKey, u64)>,
+    pub shutdown_signal_rx: Receiver<()>,
     pub stream_connector: Box<dyn StreamConnector>,
     pub proxy_client_sub: Recipient<InboundServerData>,
     pub logger: Logger,
@@ -36,6 +39,7 @@ impl Clone for StreamEstablisher {
             cryptde: self.cryptde,
             stream_adder_tx: self.stream_adder_tx.clone(),
             stream_killer_tx: self.stream_killer_tx.clone(),
+            shutdown_signal_rx: unbounded().1,
             stream_connector: Box::new(StreamConnectorReal {}),
             proxy_client_sub: self.proxy_client_sub.clone(),
             logger: self.logger.clone(),
@@ -57,11 +61,13 @@ impl StreamEstablisher {
             payload.target_port,
             &self.logger,
         )?;
+        let (shutdown_signal_tx, shutdown_signal_rx) = tokio::sync::mpsc::unbounded_channel();
 
         self.spawn_stream_reader(
             &payload.clone(),
             connection_info.reader,
             connection_info.peer_addr,
+            shutdown_signal_rx,
         );
 
         let (tx_to_write, rx_to_write) = self.channel_factory.make(connection_info.peer_addr);
@@ -73,8 +79,13 @@ impl StreamEstablisher {
         );
         tokio::spawn(stream_writer);
 
+        let stream_senders = StreamSenders {
+            writer_data: tx_to_write.clone(),
+            reader_shutdown_tx: shutdown_signal_tx,
+        };
+
         self.stream_adder_tx
-            .send((payload.stream_key, tx_to_write.clone()))
+            .send((payload.stream_key, stream_senders))
             .expect("StreamHandlerPool died");
         Ok(tx_to_write)
     }
@@ -84,12 +95,14 @@ impl StreamEstablisher {
         payload: &ClientRequestPayload_0v1,
         read_stream: Box<dyn ReadHalfWrapper>,
         peer_addr: SocketAddr,
+        shutdown_signal: UnboundedReceiver<()>,
     ) {
         let stream_reader = StreamReader::new(
             payload.stream_key,
             self.proxy_client_sub.clone(),
             read_stream,
             self.stream_killer_tx.clone(),
+            shutdown_signal,
             peer_addr,
         );
         debug!(self.logger, "Spawning StreamReader for {}", peer_addr);
@@ -103,7 +116,7 @@ pub trait StreamEstablisherFactory: Send {
 
 pub struct StreamEstablisherFactoryReal {
     pub cryptde: &'static dyn CryptDE,
-    pub stream_adder_tx: Sender<(StreamKey, Box<dyn SenderWrapper<SequencedPacket>>)>,
+    pub stream_adder_tx: Sender<(StreamKey, StreamSenders)>,
     pub stream_killer_tx: Sender<(StreamKey, u64)>,
     pub proxy_client_subs: ProxyClientSubs,
     pub logger: Logger,
@@ -115,6 +128,7 @@ impl StreamEstablisherFactory for StreamEstablisherFactoryReal {
             cryptde: self.cryptde,
             stream_adder_tx: self.stream_adder_tx.clone(),
             stream_killer_tx: self.stream_killer_tx.clone(),
+            shutdown_signal_rx: unbounded().1,
             stream_connector: Box::new(StreamConnectorReal {}),
             proxy_client_sub: self.proxy_client_subs.inbound_server_data.clone(),
             logger: self.logger.clone(),
@@ -140,6 +154,7 @@ mod tests {
     use std::str::FromStr;
     use std::thread;
     use tokio::prelude::Async;
+    use tokio::sync::mpsc::unbounded_channel;
 
     #[test]
     fn spawn_stream_reader_handles_data() {
@@ -171,6 +186,7 @@ mod tests {
                 cryptde: main_cryptde(),
                 stream_adder_tx,
                 stream_killer_tx,
+                shutdown_signal_rx: unbounded().1,
                 stream_connector: Box::new(StreamConnectorMock::new()), // only used in "establish_stream"
                 proxy_client_sub,
                 logger: Logger::new("ProxyClient"),
@@ -191,6 +207,7 @@ mod tests {
                 },
                 read_stream,
                 SocketAddr::from_str("1.2.3.4:5678").unwrap(),
+                unbounded_channel().1,
             );
 
             proxy_client_awaiter.await_message_count(1);