handle server not ready when starting OTAP Exporter (open-telemetry#1233)

closes open-telemetry#1096 

Fixes issue where the `start` method of the exporter would exit with
error immediately if the server identified by config.grpc_endpoint is
not ready

Author: albertlockett

rust/otap-dataflow/crates/otap/src/otap_exporter.rs

@@ -114,34 +114,22 @@ impl local::Exporter<OtapPdata> for OTAPExporter {
             ))
             .await;
 
+        let exporter_id = effect_handler.exporter_id();
+        let channel = Channel::from_shared(self.config.grpc_endpoint.clone())
+            .map_err(|e| Error::ExporterError {
+                exporter: exporter_id,
+                error: format!("grpc channel error {e}"),
+            })?
+            .connect_lazy();
+
         let timer_cancel_handle = effect_handler
             .start_periodic_telemetry(Duration::from_secs(1))
             .await?;
 
         // start a grpc client and connect to the server
-        let mut arrow_metrics_client =
-            ArrowMetricsServiceClient::connect(self.config.grpc_endpoint.clone())
-                .await
-                .map_err(|error| Error::ExporterError {
-                    exporter: effect_handler.exporter_id(),
-                    error: error.to_string(),
-                })?;
-
-        let mut arrow_logs_client =
-            ArrowLogsServiceClient::connect(self.config.grpc_endpoint.clone())
-                .await
-                .map_err(|error| Error::ExporterError {
-                    exporter: effect_handler.exporter_id(),
-                    error: error.to_string(),
-                })?;
-
-        let mut arrow_traces_client =
-            ArrowTracesServiceClient::connect(self.config.grpc_endpoint.clone())
-                .await
-                .map_err(|error| Error::ExporterError {
-                    exporter: effect_handler.exporter_id(),
-                    error: error.to_string(),
-                })?;
+        let mut arrow_metrics_client = ArrowMetricsServiceClient::new(channel.clone());
+        let mut arrow_logs_client = ArrowLogsServiceClient::new(channel.clone());
+        let mut arrow_traces_client = ArrowTracesServiceClient::new(channel.clone());
 
         if let Some(ref compression) = self.config.compression_method {
             let encoding = compression.map_to_compression_encoding();
@@ -453,13 +441,24 @@ mod tests {
     use crate::compression::CompressionMethod;
     use otap_df_config::node::NodeUserConfig;
     use otap_df_engine::context::ControllerContext;
+    use otap_df_engine::control::Controllable;
+    use otap_df_engine::control::NodeControlMsg;
+    use otap_df_engine::control::PipelineCtrlMsgSender;
+    use otap_df_engine::control::pipeline_ctrl_msg_channel;
     use otap_df_engine::error::Error;
     use otap_df_engine::exporter::ExporterWrapper;
+    use otap_df_engine::local::message::LocalReceiver;
+    use otap_df_engine::local::message::LocalSender;
+    use otap_df_engine::message::Receiver;
+    use otap_df_engine::message::Sender;
+    use otap_df_engine::node::NodeWithPDataReceiver;
+    use otap_df_engine::testing::create_not_send_channel;
     use otap_df_engine::testing::{
         exporter::{TestContext, TestRuntime},
         test_node,
     };
     use otap_df_telemetry::registry::MetricsRegistryHandle;
+    use otap_df_telemetry::reporter::MetricsReporter;
     use otel_arrow_rust::otap::OtapArrowRecords;
     use otel_arrow_rust::proto::opentelemetry::arrow::v1::{
         ArrowPayloadType, arrow_logs_service_server::ArrowLogsServiceServer,
@@ -736,4 +735,173 @@ mod tests {
             exporter.config.arrow.payload_compression
         );
     }
+
+    #[test]
+    fn test_receiver_not_ready_on_start() {
+        let grpc_addr = "127.0.0.1";
+        let grpc_port = portpicker::pick_unused_port().expect("No free ports");
+        let grpc_endpoint = format!("http://{grpc_addr}:{grpc_port}");
+        let tokio_rt = Runtime::new().unwrap();
+
+        let test_runtime = TestRuntime::<OtapPdata>::new();
+        let node_config = Arc::new(NodeUserConfig::new_exporter_config(OTAP_EXPORTER_URN));
+        let metrics_registry_handle = MetricsRegistryHandle::new();
+        let controller_ctx = ControllerContext::new(metrics_registry_handle);
+        let node_id = test_node(test_runtime.config().name.clone());
+        let pipeline_ctx =
+            controller_ctx.pipeline_context_with("grp".into(), "pipeline".into(), 0, 0);
+
+        let mut exporter = ExporterWrapper::local(
+            OTAPExporter::from_config(
+                pipeline_ctx,
+                &serde_json::json!({
+                    "grpc_endpoint": grpc_endpoint,
+                    "compression_method": "none"
+                }),
+            )
+            .unwrap(),
+            node_id.clone(),
+            node_config,
+            test_runtime.config(),
+        );
+
+        let control_sender = exporter.control_sender();
+        let (pdata_tx, pdata_rx) = create_not_send_channel::<OtapPdata>(1);
+        let pdata_tx = Sender::Local(LocalSender::MpscSender(pdata_tx));
+        let pdata_rx = Receiver::Local(LocalReceiver::MpscReceiver(pdata_rx));
+        let (pipeline_ctrl_msg_tx, _pipeline_ctrl_msg_rx) = pipeline_ctrl_msg_channel(2);
+        exporter
+            .set_pdata_receiver(node_id.clone(), pdata_rx)
+            .expect("Failed to set PData Receiver");
+
+        let (req_sender, req_receiver) = tokio::sync::mpsc::channel(1);
+        let (server_startup_sender, mut server_startup_receiver) = tokio::sync::mpsc::channel(1);
+        let (server_start_ack_sender, server_start_ack_receiver) = tokio::sync::mpsc::channel(1);
+        let (server_shutdown_sender, server_shutdown_signal) = tokio::sync::oneshot::channel();
+
+        async fn start_exporter(
+            exporter: ExporterWrapper<OtapPdata>,
+            pipeline_ctrl_msg_tx: PipelineCtrlMsgSender<OtapPdata>,
+        ) -> Result<(), Error> {
+            _ = exporter.start(pipeline_ctrl_msg_tx).await;
+            Ok(())
+        }
+
+        async fn drive_test(
+            server_startup_sender: tokio::sync::mpsc::Sender<bool>,
+            mut server_startup_ack_receiver: tokio::sync::mpsc::Receiver<bool>,
+            server_shutdown_sender1: tokio::sync::oneshot::Sender<bool>,
+            pdata_tx: Sender<OtapPdata>,
+            control_sender: Sender<NodeControlMsg<OtapPdata>>,
+            mut req_receiver: tokio::sync::mpsc::Receiver<OtapPdata>,
+        ) {
+            // send a request before while the server isn't running and check how we handle it
+            let log_message = create_otap_batch(LOG_BATCH_ID, ArrowPayloadType::Logs);
+            pdata_tx
+                .send(OtapPdata::new_default(log_message.into()))
+                .await
+                .expect("Failed to send log message");
+            // TODO instead of sleeping here, once we handle ACK/NACK we should wait to get a NACK
+            // from the control channel
+            tokio::time::sleep(Duration::from_millis(5)).await;
+
+            // wait a bit before starting the server. This will ensure the exporter no-long exits
+            // when start is called if the endpoint can't be reached
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            server_startup_sender.send(true).await.unwrap();
+            _ = server_startup_ack_receiver.recv().await.unwrap();
+
+            // send another pdata now that the server has started
+            let log_message = create_otap_batch(LOG_BATCH_ID + 1, ArrowPayloadType::Logs);
+            pdata_tx
+                .send(OtapPdata::new_default(log_message.into()))
+                .await
+                .expect("Failed to send log message");
+            _ = req_receiver.recv().await.unwrap(); // ensure we got response
+            // TODO instead of sleeping here, once we handle ACK/NACK we should wait to get a ACK
+            // from the control channel
+            tokio::time::sleep(Duration::from_millis(50)).await;
+
+            // check the metrics:
+            let (metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(32);
+            control_sender
+                .send(NodeControlMsg::CollectTelemetry {
+                    metrics_reporter: metrics_reporter.clone(),
+                })
+                .await
+                .unwrap();
+            let metrics = metrics_rx.recv_async().await.unwrap();
+            let logs_exported_count = metrics.get_metrics()[4]; // logs exported
+            assert_eq!(logs_exported_count, 1);
+            let logs_failed_count = metrics.get_metrics()[5]; // logs failed
+            assert_eq!(logs_failed_count, 1);
+
+            control_sender
+                .send(NodeControlMsg::Shutdown {
+                    deadline: Duration::from_millis(10),
+                    reason: "shutting down".into(),
+                })
+                .await
+                .unwrap();
+
+            server_shutdown_sender1.send(true).unwrap();
+        }
+
+        async fn run_server(
+            listening_addr: String,
+            startup_ack_sender: tokio::sync::mpsc::Sender<bool>,
+            shutdown_signal: tokio::sync::oneshot::Receiver<bool>,
+            req_sender: tokio::sync::mpsc::Sender<OtapPdata>,
+        ) {
+            let listening_addr: SocketAddr = listening_addr.to_string().parse().unwrap();
+            let tcp_listener = TcpListener::bind(listening_addr).await.unwrap();
+            let tcp_stream = TcpListenerStream::new(tcp_listener);
+
+            let logs_service = ArrowLogsServiceServer::new(ArrowLogsServiceMock::new(req_sender));
+
+            Server::builder()
+                .add_service(logs_service)
+                .serve_with_incoming_shutdown(tcp_stream, async {
+                    startup_ack_sender.send(true).await.unwrap();
+                    let _ = shutdown_signal.await;
+                })
+                .await
+                .expect("uh oh server failed");
+        }
+
+        let server_handle = tokio_rt.spawn(async move {
+            let listening_addr = format!("{grpc_addr}:{grpc_port}");
+
+            // wait for signal to start the server
+            _ = server_startup_receiver.recv().await.unwrap();
+            run_server(
+                listening_addr.clone(),
+                server_start_ack_sender.clone(),
+                server_shutdown_signal,
+                req_sender.clone(),
+            )
+            .await;
+        });
+
+        let _ = tokio_rt.block_on(async move {
+            let local_set = tokio::task::LocalSet::new();
+            let _fut = local_set
+                .spawn_local(async move { start_exporter(exporter, pipeline_ctrl_msg_tx).await });
+            tokio::join!(
+                local_set,
+                drive_test(
+                    server_startup_sender,
+                    server_start_ack_receiver,
+                    server_shutdown_sender,
+                    pdata_tx,
+                    control_sender,
+                    req_receiver
+                )
+            )
+        });
+
+        tokio_rt
+            .block_on(server_handle)
+            .expect("server shutdown success");
+    }
 }