1use super::*;
5
6use crate::engine::AsyncEngineContext;
7use crate::metrics::prometheus_names::work_handler;
8use crate::metrics::work_handler_perf::{
9 WORK_HANDLER_NETWORK_TRANSIT_SECONDS, WORK_HANDLER_TIME_TO_FIRST_RESPONSE_SECONDS,
10};
11use crate::pipeline::{ManyIn, RequestStream};
12use futures::StreamExt;
13use prometheus::{Histogram, IntCounter, IntCounterVec, IntGauge};
14use serde::Deserialize;
15use std::sync::Arc;
16use std::time::Instant;
17use tracing::Instrument;
18use tracing::info_span;
19
20#[derive(Clone, Debug)]
22pub struct WorkHandlerMetrics {
23 pub request_counter: IntCounter,
24 pub request_duration: Histogram,
25 pub inflight_requests: IntGauge,
26 pub request_bytes: IntCounter,
27 pub response_bytes: IntCounter,
28 pub error_counter: IntCounterVec,
29 pub cancellation_total: IntCounter,
30}
31
32impl WorkHandlerMetrics {
33 pub fn new(
34 request_counter: IntCounter,
35 request_duration: Histogram,
36 inflight_requests: IntGauge,
37 request_bytes: IntCounter,
38 response_bytes: IntCounter,
39 error_counter: IntCounterVec,
40 cancellation_total: IntCounter,
41 ) -> Self {
42 Self {
43 request_counter,
44 request_duration,
45 inflight_requests,
46 request_bytes,
47 response_bytes,
48 error_counter,
49 cancellation_total,
50 }
51 }
52
53 pub fn from_endpoint(
55 endpoint: &crate::component::Endpoint,
56 metrics_labels: Option<&[(&str, &str)]>,
57 ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
58 let metrics_labels = metrics_labels.unwrap_or(&[]);
59 let metrics = endpoint.metrics();
60 let request_counter = metrics.create_intcounter(
61 work_handler::REQUESTS_TOTAL,
62 "Total number of requests processed by work handler",
63 metrics_labels,
64 )?;
65
66 let request_duration_buckets = vec![
70 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 20.0, 30.0, 60.0, 120.0,
71 300.0, 600.0,
72 ];
73 let request_duration = metrics.create_histogram(
74 work_handler::REQUEST_DURATION_SECONDS,
75 "Time spent processing requests by work handler",
76 metrics_labels,
77 Some(request_duration_buckets),
78 )?;
79
80 let inflight_requests = metrics.create_intgauge(
81 work_handler::INFLIGHT_REQUESTS,
82 "Number of requests currently being processed by work handler",
83 metrics_labels,
84 )?;
85
86 let request_bytes = metrics.create_intcounter(
87 work_handler::REQUEST_BYTES_TOTAL,
88 "Total number of bytes received in requests by work handler",
89 metrics_labels,
90 )?;
91
92 let response_bytes = metrics.create_intcounter(
93 work_handler::RESPONSE_BYTES_TOTAL,
94 "Total number of bytes sent in responses by work handler",
95 metrics_labels,
96 )?;
97
98 let error_counter = metrics.create_intcountervec(
99 work_handler::ERRORS_TOTAL,
100 "Total number of errors in work handler processing",
101 &[work_handler::ERROR_TYPE_LABEL],
102 metrics_labels,
103 )?;
104
105 let cancellation_total = metrics.create_intcounter(
106 work_handler::CANCELLATION_TOTAL,
107 "Total number of requests cancelled by work handler",
108 metrics_labels,
109 )?;
110
111 Ok(Self::new(
112 request_counter,
113 request_duration,
114 inflight_requests,
115 request_bytes,
116 response_bytes,
117 error_counter,
118 cancellation_total,
119 ))
120 }
121}
122
123struct RequestMetricsGuard {
126 inflight_requests: prometheus::IntGauge,
127 request_duration: prometheus::Histogram,
128 start_time: Instant,
129 request_id: Option<String>,
130}
131
132impl Drop for RequestMetricsGuard {
133 fn drop(&mut self) {
134 self.inflight_requests.dec();
135 self.request_duration
136 .observe(self.start_time.elapsed().as_secs_f64());
137 if let Some(request_id) = &self.request_id {
138 tracing::info!(request_id = %request_id, "request completed");
139 }
140 }
141}
142
143impl<Req, Resp, Adapter> Ingress<Req, Resp, Adapter>
144where
145 Req: PipelineIO + Sync,
146 Resp: PipelineIO,
147 Adapter: Send + Sync + 'static,
148{
149 async fn pump_response_stream<U>(
156 &self,
157 mut stream: ManyOut<U>,
158 publisher: &StreamSender,
159 payload_codec: RequestPlanePayloadCodec,
160 ) where
161 U: Data + std::fmt::Debug,
162 Adapter: IngressResponseEncoder<U>,
163 {
164 let context = stream.context();
165
166 let mut send_complete_final = true;
168 let mut saw_error_response = false;
169 while let Some(resp) = stream.next().await {
170 tracing::trace!("Sending response: {:?}", resp);
171 let encoded = match self
172 .payload_adapter
173 .encode_response(payload_codec, Some(resp), false)
174 .await
175 {
176 Ok(encoded) => encoded,
177 Err(err) => {
178 tracing::error!(%err, "failed to encode request-plane response");
179 saw_error_response = true;
180 send_complete_final = false;
181 if let Some(m) = self.metrics() {
182 m.error_counter
183 .with_label_values(&[work_handler::error_types::SERIALIZATION])
184 .inc();
185 }
186 break;
187 }
188 };
189 let is_error = encoded.is_error;
190 saw_error_response |= is_error;
191 let resp_bytes = encoded.bytes;
192 if let Some(m) = self.metrics() {
193 m.response_bytes.inc_by(resp_bytes.len() as u64);
194 }
195 if (publisher.send(resp_bytes).await).is_err() {
196 send_complete_final = false;
197 if context.is_stopped() {
198 tracing::warn!("Failed to publish response for stream {}", context.id());
207 } else {
208 tracing::error!("Failed to publish response for stream {}", context.id());
210 context.stop_generating();
211 }
212 if let Some(m) = self.metrics() {
215 m.error_counter
216 .with_label_values(&[work_handler::error_types::PUBLISH_RESPONSE])
217 .inc();
218 }
219 break;
220 } else if !is_error {
221 if let Some(notifier) = self.endpoint_health_check_notifier.get() {
224 notifier.notify_one();
225 }
226 }
227 if encoded.stop_stream {
228 break;
233 }
234 }
235 if send_complete_final {
236 let encoded = match self
237 .payload_adapter
238 .encode_response(payload_codec, None, true)
239 .await
240 {
241 Ok(encoded) => encoded,
242 Err(err) => {
243 tracing::error!(%err, "failed to encode request-plane final response");
244 if let Some(m) = self.metrics() {
245 m.error_counter
246 .with_label_values(&[work_handler::error_types::PUBLISH_FINAL])
247 .inc();
248 }
249 return;
250 }
251 };
252 let resp_bytes = encoded.bytes;
253 if let Some(m) = self.metrics() {
254 m.response_bytes.inc_by(resp_bytes.len() as u64);
255 }
256 if (publisher.send(resp_bytes).await).is_err() {
257 tracing::error!(
258 "Failed to publish complete final for stream {}",
259 context.id()
260 );
261 if let Some(m) = self.metrics() {
262 m.error_counter
263 .with_label_values(&[work_handler::error_types::PUBLISH_FINAL])
264 .inc();
265 }
266 }
267 if let (false, Some(notifier)) = (
269 saw_error_response,
270 self.endpoint_health_check_notifier.get(),
271 ) {
272 notifier.notify_one();
273 }
274 }
275 }
276
277 fn decode_control_message(
288 &self,
289 payload: Bytes,
290 ) -> Result<(RequestControlMessage, Option<Bytes>), PipelineError> {
291 let msg = TwoPartCodec::default()
292 .decode_message(payload)?
293 .into_message_type();
294
295 let (header, data) = match msg {
296 TwoPartMessageType::HeaderAndData(header, data) => (header, Some(data)),
297 TwoPartMessageType::HeaderOnly(header) => (header, None),
298 _ => {
299 if let Some(m) = self.metrics() {
300 m.error_counter
301 .with_label_values(&[work_handler::error_types::INVALID_MESSAGE])
302 .inc();
303 }
304 return Err(PipelineError::Generic(String::from(
305 "Unexpected message from work queue; expected a header-only or header-and-data TwoPartMessage",
306 )));
307 }
308 };
309
310 let control_msg: RequestControlMessage =
311 serde_json::from_slice(&header).map_err(|err| {
312 if let Some(m) = self.metrics() {
313 m.error_counter
314 .with_label_values(&[work_handler::error_types::DESERIALIZATION])
315 .inc();
316 }
317 let json_str = String::from_utf8_lossy(&header);
318 PipelineError::DeserializationError(format!(
319 "Failed deserializing to RequestControlMessage. err={err}, json_str={json_str}, header_len={}",
320 header.len(),
321 ))
322 })?;
323
324 Ok((control_msg, data))
325 }
326}
327struct ParsedRequest<Req> {
332 request: Req,
333 response_connection_info: ConnectionInfo,
334 frontend_send_ts_ns: Option<u64>,
335 payload_codec: RequestPlanePayloadCodec,
336}
337
338#[async_trait]
345trait IngressDispatch: Send + Sync {
346 type Request: PipelineIO;
347
348 async fn parse_and_build_request(
349 &self,
350 payload: Bytes,
351 ) -> Result<ParsedRequest<Self::Request>, PipelineError>;
352}
353
354#[async_trait]
355impl<T, U, Adapter> IngressDispatch for Ingress<SingleIn<T>, ManyOut<U>, Adapter>
356where
357 T: Data + for<'de> Deserialize<'de> + std::fmt::Debug,
358 U: Data + std::fmt::Debug,
359 Adapter: IngressRequestDecoder<T> + Send + Sync + 'static,
360{
361 type Request = SingleIn<T>;
362
363 async fn parse_and_build_request(
364 &self,
365 payload: Bytes,
366 ) -> Result<ParsedRequest<SingleIn<T>>, PipelineError> {
367 let (control_msg, data) = self.decode_control_message(payload)?;
368
369 let data = data.ok_or_else(|| {
372 if let Some(m) = self.metrics() {
373 m.error_counter
374 .with_label_values(&[work_handler::error_types::INVALID_MESSAGE])
375 .inc();
376 }
377 PipelineError::Generic(String::from(
378 "unary engine received a header-only envelope; expected a request payload",
379 ))
380 })?;
381 let payload_codec = control_msg.payload_codec;
382 let request_t: T = self
383 .payload_adapter
384 .decode_request(payload_codec, data)
385 .await
386 .inspect_err(|_| {
387 if let Some(m) = self.metrics() {
388 m.error_counter
389 .with_label_values(&[work_handler::error_types::DESERIALIZATION])
390 .inc();
391 }
392 })?;
393
394 tracing::trace!(
395 request_id = %control_msg.id,
396 metadata_entries = control_msg.metadata.len(),
397 "received control message"
398 );
399 tracing::trace!("received request: {:?}", request_t);
400
401 let request: context::Context<T> =
402 Context::with_id_and_metadata(request_t, control_msg.id, control_msg.metadata);
403
404 Ok(ParsedRequest {
405 request,
406 response_connection_info: control_msg.connection_info,
407 frontend_send_ts_ns: control_msg.frontend_send_ts_ns,
408 payload_codec,
409 })
410 }
411}
412
413#[async_trait]
414impl<T, U, Adapter> IngressDispatch for Ingress<ManyIn<T>, ManyOut<U>, Adapter>
415where
416 T: Data + for<'de> Deserialize<'de> + std::fmt::Debug,
417 U: Data + std::fmt::Debug,
418 Adapter: IngressRequestDecoder<T> + Send + Sync + 'static,
419{
420 type Request = ManyIn<T>;
421
422 async fn parse_and_build_request(
423 &self,
424 payload: Bytes,
425 ) -> Result<ParsedRequest<ManyIn<T>>, PipelineError> {
426 let (control_msg, data) = self.decode_control_message(payload)?;
427
428 if data.is_some() {
433 if let Some(m) = self.metrics() {
434 m.error_counter
435 .with_label_values(&[work_handler::error_types::INVALID_MESSAGE])
436 .inc();
437 }
438 return Err(PipelineError::Generic(String::from(
439 "bidirectional engine received a non-header-only envelope",
440 )));
441 }
442
443 if !matches!(control_msg.request_type, RequestType::ManyIn) {
444 if let Some(m) = self.metrics() {
445 m.error_counter
446 .with_label_values(&[work_handler::error_types::INVALID_MESSAGE])
447 .inc();
448 }
449 return Err(PipelineError::Generic(String::from(
450 "bidirectional engine received a non-ManyIn request envelope",
451 )));
452 }
453
454 let req_stream_conn_info = control_msg
455 .request_stream_connection_info
456 .clone()
457 .ok_or_else(|| {
458 PipelineError::Generic(String::from(
459 "bidirectional control message missing request_stream_connection_info",
460 ))
461 })?;
462
463 let request_context: context::Context<()> = context::Context::with_id_and_metadata(
464 (),
465 control_msg.id.clone(),
466 control_msg.metadata.clone(),
467 );
468 let payload_codec = control_msg.payload_codec;
469 let context_arc: Arc<dyn AsyncEngineContext> = request_context.context();
470
471 let request_stream_recv = tcp::client::TcpClient::create_request_stream(
477 context_arc.clone(),
478 req_stream_conn_info,
479 None,
480 )
481 .await
482 .map_err(|e| {
483 if let Some(m) = self.metrics() {
484 m.error_counter
485 .with_label_values(&[work_handler::error_types::RESPONSE_STREAM])
486 .inc();
487 }
488 PipelineError::Generic(format!("Failed to create request stream: {e}"))
489 })?;
490
491 let (frame_tx, frame_rx) = tokio::sync::mpsc::channel::<T>(8);
496 let forwarder_ctx = context_arc.clone();
497 let payload_adapter = self.payload_adapter.clone();
498 tokio::spawn(async move {
499 let mut rx = request_stream_recv.rx;
500 while let Some(bytes) = rx.recv().await {
501 if forwarder_ctx.is_killed() || forwarder_ctx.is_stopped() {
506 break;
507 }
508 match payload_adapter.decode_request(payload_codec, bytes).await {
509 Ok(item) => {
510 if frame_tx.send(item).await.is_err() {
511 tracing::debug!(
512 "engine consumer dropped; bidirectional input forwarder exiting"
513 );
514 break;
515 }
516 }
517 Err(e) => {
518 tracing::error!(
519 error = %e,
520 codec = payload_codec.name(),
521 "failed to deserialize bidirectional request frame; killing context"
522 );
523 forwarder_ctx.kill();
524 break;
525 }
526 }
527 }
528 });
529
530 let input_stream: crate::engine::DataStream<T> =
531 Box::pin(tokio_stream::wrappers::ReceiverStream::new(frame_rx));
532 let request: ManyIn<T> = request_context.map(|_| RequestStream::new(input_stream));
533
534 Ok(ParsedRequest {
535 request,
536 response_connection_info: control_msg.connection_info,
537 frontend_send_ts_ns: control_msg.frontend_send_ts_ns,
538 payload_codec,
539 })
540 }
541}
542
543impl<Req, U, Adapter> Ingress<Req, ManyOut<U>, Adapter>
544where
545 Req: PipelineIO + Sync,
546 U: Data + std::fmt::Debug,
547 Adapter: IngressResponseEncoder<U> + Send + Sync + 'static,
548{
549 async fn handle_payload_shared(
557 &self,
558 payload: Bytes,
559 request_id: Option<String>,
560 ) -> Result<(), PipelineError>
561 where
562 Self: IngressDispatch<Request = Req>,
563 {
564 let t2_wallclock_ns = std::time::SystemTime::now()
565 .duration_since(std::time::UNIX_EPOCH)
566 .unwrap_or_default()
567 .as_nanos() as u64;
568 let start_time = std::time::Instant::now();
569
570 let _inflight_guard = self.metrics().map(|m| {
572 m.request_counter.inc();
573 m.inflight_requests.inc();
574 m.request_bytes.inc_by(payload.len() as u64);
575 if let Some(rid) = &request_id {
576 tracing::info!(request_id = %rid, "request received");
577 }
578 RequestMetricsGuard {
579 inflight_requests: m.inflight_requests.clone(),
580 request_duration: m.request_duration.clone(),
581 start_time,
582 request_id: request_id.clone(),
583 }
584 });
585
586 let ParsedRequest {
587 request,
588 response_connection_info,
589 frontend_send_ts_ns,
590 payload_codec,
591 } = self.parse_and_build_request(payload).await?;
592
593 if let Some(t1_ns) = frontend_send_ts_ns {
595 let transit_ns = t2_wallclock_ns.saturating_sub(t1_ns);
596 WORK_HANDLER_NETWORK_TRANSIT_SECONDS.observe(transit_ns as f64 / 1_000_000_000.0);
597 }
598
599 tracing::trace!("creating tcp response stream");
602 let mut publisher = tcp::client::TcpClient::create_response_stream(
603 request.context(),
604 response_connection_info,
605 self.metrics().map(|m| m.cancellation_total.clone()),
606 )
607 .await
608 .map_err(|e| {
609 if let Some(m) = self.metrics() {
610 m.error_counter
611 .with_label_values(&[work_handler::error_types::RESPONSE_STREAM])
612 .inc();
613 }
614 PipelineError::Generic(format!("Failed to create response stream: {e}"))
615 })?;
616
617 tracing::trace!("calling generate");
618 let stream = self
619 .segment
620 .get()
621 .expect("segment not set")
622 .generate(request)
623 .await
624 .map_err(|e| {
625 if let Some(m) = self.metrics() {
626 m.error_counter
627 .with_label_values(&[work_handler::error_types::GENERATE])
628 .inc();
629 }
630 PipelineError::GenerateError(e)
631 });
632
633 let stream = match stream {
636 Ok(stream) => {
637 tracing::trace!("Successfully generated response stream; sending prologue");
638 let _result = publisher.send_prologue(None).await;
639 WORK_HANDLER_TIME_TO_FIRST_RESPONSE_SECONDS
640 .observe(start_time.elapsed().as_secs_f64());
641 stream
642 }
643 Err(e) => {
644 let error_string = e.to_string();
645
646 #[cfg(debug_assertions)]
647 {
648 tracing::debug!(
649 "Failed to generate response stream (with debug backtrace): {:?}",
650 e
651 );
652 }
653 #[cfg(not(debug_assertions))]
654 {
655 tracing::error!("Failed to generate response stream: {error_string}");
656 }
657
658 let _result = publisher.send_prologue(Some(error_string)).await;
659 Err(e)?
660 }
661 };
662
663 self.pump_response_stream(stream, &publisher, payload_codec)
664 .await;
665
666 drop(_inflight_guard);
669
670 Ok(())
671 }
672}
673
674#[async_trait]
675impl<T, U, Adapter> PushWorkHandler for Ingress<SingleIn<T>, ManyOut<U>, Adapter>
676where
677 T: Data + for<'de> Deserialize<'de> + std::fmt::Debug,
678 U: Data + std::fmt::Debug,
679 Adapter: IngressPayloadAdapter<T, U> + Send + Sync + 'static,
680{
681 fn add_metrics(
682 &self,
683 endpoint: &crate::component::Endpoint,
684 metrics_labels: Option<&[(&str, &str)]>,
685 ) -> Result<()> {
686 Ingress::add_metrics(self, endpoint, metrics_labels)
688 }
689
690 fn set_endpoint_health_check_notifier(&self, notifier: Arc<tokio::sync::Notify>) -> Result<()> {
691 self.endpoint_health_check_notifier
692 .set(notifier)
693 .map_err(|_| anyhow::anyhow!("Endpoint health check notifier already set"))?;
694 Ok(())
695 }
696
697 async fn handle_payload(
698 &self,
699 payload: Bytes,
700 request_id: Option<String>,
701 ) -> Result<(), PipelineError> {
702 self.handle_payload_shared(payload, request_id).await
703 }
704}
705
706#[async_trait]
707impl<T, U, Adapter> PushWorkHandler for Ingress<ManyIn<T>, ManyOut<U>, Adapter>
708where
709 T: Data + for<'de> Deserialize<'de> + std::fmt::Debug,
710 U: Data + std::fmt::Debug,
711 Adapter: IngressPayloadAdapter<T, U> + Send + Sync + 'static,
712{
713 fn add_metrics(
714 &self,
715 endpoint: &crate::component::Endpoint,
716 metrics_labels: Option<&[(&str, &str)]>,
717 ) -> Result<()> {
718 Ingress::add_metrics(self, endpoint, metrics_labels)
720 }
721
722 fn set_endpoint_health_check_notifier(&self, notifier: Arc<tokio::sync::Notify>) -> Result<()> {
723 self.endpoint_health_check_notifier
724 .set(notifier)
725 .map_err(|_| anyhow::anyhow!("Endpoint health check notifier already set"))?;
726 Ok(())
727 }
728
729 async fn handle_payload(
730 &self,
731 payload: Bytes,
732 request_id: Option<String>,
733 ) -> Result<(), PipelineError> {
734 self.handle_payload_shared(payload, request_id).await
735 }
736}