src/diagnostics/archivist/src/accessor.rs - fuchsia - Git at Google

 // Copyright 2019 The Fuchsia Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 use {
     crate::{
         constants::{self, FORMATTED_CONTENT_CHUNK_SIZE_TARGET},
         diagnostics::{AccessorStats, BatchIteratorConnectionStats},
         error::AccessorError,
         formatter::{new_batcher, FormattedStream, JsonPacketSerializer, JsonString},
         inspect,
         lifecycle::LifecycleServer,
         moniker_rewriter::MonikerRewriter,
         pipeline::Pipeline,
         ImmutableString,
     },
     diagnostics_data::{Data, DiagnosticsData},
     fidl::endpoints::RequestStream,
     fidl_fuchsia_diagnostics::{
         self, ArchiveAccessorRequest, ArchiveAccessorRequestStream, BatchIteratorRequest,
         BatchIteratorRequestStream, ClientSelectorConfiguration, DataType, Format,
         PerformanceConfiguration, Selector, SelectorArgument, StreamMode, StreamParameters,
         StringSelector, TreeSelector, TreeSelectorUnknown,
     },
     fuchsia_async::{self as fasync, Task},
     fuchsia_inspect::NumericProperty,
     fuchsia_zircon as zx,
     futures::{
         channel::mpsc::UnboundedSender,
         future::{select, Either},
         prelude::*,
     },
     parking_lot::RwLock,
     selectors::{self, FastError},
     serde::Serialize,
     std::collections::HashMap,
     std::convert::TryFrom,
     std::sync::{Arc, Mutex},
     tracing::warn,
 };

 /// ArchiveAccessor represents an incoming connection from a client to an Archivist
 /// instance, through which the client may make Reader requests to the various data
 /// sources the Archivist offers.
 pub struct ArchiveAccessor {
     // The inspect repository containing read-only inspect data shared across
     // all inspect reader instances.
     pipeline: Arc<RwLock<Pipeline>>,
     archive_accessor_stats: Arc<AccessorStats>,
     moniker_rewriter: Option<Arc<MonikerRewriter>>,
 }

 fn validate_and_parse_selectors(
     selector_args: Vec<SelectorArgument>,
 ) -> Result<Vec<Selector>, AccessorError> {
     let mut selectors = vec![];
     let mut errors = vec![];

     if selector_args.is_empty() {
         return Err(AccessorError::EmptySelectors);
     }

     for selector_arg in selector_args {
         match selectors::take_from_argument::<FastError>(selector_arg) {
             Ok(s) => selectors.push(s),
             Err(e) => errors.push(e),
         }
     }

     if !errors.is_empty() {
         warn!(?errors, "Found errors in selector arguments");
     }

     Ok(selectors)
 }

 fn validate_and_parse_log_selectors(
     selector_args: Vec<SelectorArgument>,
 ) -> Result<Vec<Selector>, AccessorError> {
     // Only accept selectors of the type: `component:root` for logs for now.
     let selectors = validate_and_parse_selectors(selector_args)?;
     for selector in &selectors {
         // Unwrap safe: Previous validation discards any selector without a node.
         let tree_selector = selector.tree_selector.as_ref().unwrap();
         match tree_selector {
             TreeSelector::PropertySelector(_) => {
                 return Err(AccessorError::InvalidLogSelector);
             }
             TreeSelector::SubtreeSelector(subtree_selector) => {
                 if subtree_selector.node_path.len() != 1 {
                     return Err(AccessorError::InvalidLogSelector);
                 }
                 match &subtree_selector.node_path[0] {
                     StringSelector::ExactMatch(val) if val == "root" => {}
                     StringSelector::StringPattern(val) if val == "root" => {}
                     _ => {
                         return Err(AccessorError::InvalidLogSelector);
                     }
                 }
             }
             TreeSelectorUnknown!() => {}
         }
     }
     Ok(selectors)
 }

 impl ArchiveAccessor {
     /// Create a new accessor for interacting with the archivist's data. The pipeline
     /// parameter determines which static configurations scope/restrict the visibility of
     /// data accessed by readers spawned by this accessor.
     pub fn new(
         pipeline: Arc<RwLock<Pipeline>>,
         archive_accessor_stats: Arc<AccessorStats>,
     ) -> Self {
         ArchiveAccessor { pipeline, archive_accessor_stats, moniker_rewriter: None }
     }

     pub fn add_moniker_rewriter(&mut self, rewriter: Arc<MonikerRewriter>) -> &mut Self {
         self.moniker_rewriter = Some(rewriter);
         self
     }

     async fn run_server(
         pipeline: Arc<RwLock<Pipeline>>,
         requests: BatchIteratorRequestStream,
         params: StreamParameters,
         rewriter: Option<Arc<MonikerRewriter>>,
         accessor_stats: Arc<AccessorStats>,
     ) -> Result<(), AccessorError> {
         let format = params.format.ok_or(AccessorError::MissingFormat)?;
         if !matches!(format, Format::Json) {
             return Err(AccessorError::UnsupportedFormat);
         }
         let mode = params.stream_mode.ok_or(AccessorError::MissingMode)?;

         let performance_config: PerformanceConfig = PerformanceConfig::try_from(&params)?;

         match params.data_type.ok_or(AccessorError::MissingDataType)? {
             DataType::Inspect => {
                 if !matches!(mode, StreamMode::Snapshot) {
                     return Err(AccessorError::UnsupportedMode);
                 }
                 let stats = Arc::new(accessor_stats.new_inspect_batch_iterator());

                 let selectors =
                     params.client_selector_configuration.ok_or(AccessorError::MissingSelectors)?;

                 let selectors = match selectors {
                     ClientSelectorConfiguration::Selectors(selectors) => {
                         Some(validate_and_parse_selectors(selectors)?)
                     }
                     ClientSelectorConfiguration::SelectAll(_) => None,
                     _ => return Err(AccessorError::InvalidSelectors("unrecognized selectors")),
                 };

                 let (selectors, output_rewriter) = match (selectors, rewriter) {
                     (Some(selectors), Some(rewriter)) => rewriter.rewrite_selectors(selectors),
                     // behaves correctly whether selectors is Some(_) or None
                     (selectors, _) => (selectors, None),
                 };

                 let unpopulated_container_vec = pipeline.read().fetch_inspect_data(&selectors);

                 let per_component_budget_opt = if unpopulated_container_vec.is_empty() {
                     None
                 } else {
                     performance_config
                         .aggregated_content_limit_bytes
                         .map(|limit| (limit as usize) / unpopulated_container_vec.len())
                 };

                 if let Some(max_snapshot_size) = performance_config.aggregated_content_limit_bytes {
                     stats.global_stats().record_max_snapshot_size_config(max_snapshot_size);
                 }

                 BatchIterator::new(
                     inspect::ReaderServer::stream(
                         unpopulated_container_vec,
                         performance_config,
                         selectors,
                         output_rewriter,
                         stats.clone(),
                     ),
                     requests,
                     mode,
                     stats,
                     per_component_budget_opt,
                 )?
                 .run()
                 .await
             }
             DataType::Lifecycle => {
                 // TODO(fxbug.dev/61350) support other modes
                 if !matches!(mode, StreamMode::Snapshot) {
                     return Err(AccessorError::UnsupportedMode);
                 }
                 let stats = Arc::new(accessor_stats.new_lifecycle_batch_iterator());

                 let selectors =
                     params.client_selector_configuration.ok_or(AccessorError::MissingSelectors)?;
                 if !matches!(selectors, ClientSelectorConfiguration::SelectAll(_)) {
                     return Err(AccessorError::InvalidSelectors(
                         "lifecycle only supports SelectAll at the moment",
                     ));
                 }

                 let events = LifecycleServer::new(pipeline);

                 BatchIterator::new(events, requests, mode, stats, None)?.run().await
             }
             DataType::Logs => {
                 let stats = Arc::new(accessor_stats.new_logs_batch_iterator());
                 let selectors = match params.client_selector_configuration {
                     Some(ClientSelectorConfiguration::Selectors(selectors)) => {
                         Some(validate_and_parse_log_selectors(selectors)?)
                     }
                     Some(ClientSelectorConfiguration::SelectAll(_)) => None,
                     _ => return Err(AccessorError::InvalidSelectors("unrecognized selectors")),
                 };
                 let logs =
                     pipeline.read().logs(mode, selectors).map(move |inner: _| (*inner).clone());
                 BatchIterator::new_serving_arrays(logs, requests, mode, stats)?.run().await
             }
         }
     }

     /// Spawn an instance `fidl_fuchsia_diagnostics/Archive` that allows clients to open
     /// reader session to diagnostics data.
     pub fn spawn_server(
         self,
         mut stream: ArchiveAccessorRequestStream,
         task_sender: UnboundedSender<Task<()>>,
     ) {
         // Self isn't guaranteed to live into the exception handling of the async block. We need to clone self
         // to have a version that can be referenced in the exception handling.
         let batch_iterator_task_sender = task_sender.clone();
         task_sender
             .unbounded_send(fasync::Task::spawn(async move {
                 self.archive_accessor_stats.global_stats.connections_opened.add(1);
                 while let Ok(Some(ArchiveAccessorRequest::StreamDiagnostics {
                     result_stream,
                     stream_parameters,
                     control_handle: _,
                 })) = stream.try_next().await
                 {
                     let (requests, control) = match result_stream.into_stream_and_control_handle() {
                         Ok(r) => r,
                         Err(e) => {
                             warn!(?e, "Couldn't bind results channel to executor.");
                             continue;
                         }
                     };

                     self.archive_accessor_stats.global_stats.stream_diagnostics_requests.add(1);
                     let pipeline = self.pipeline.clone();
                     let accessor_stats = self.archive_accessor_stats.clone();
                     let moniker_rewriter = self.moniker_rewriter.clone();
                     // Store the batch iterator task so that we can ensure that the client finishes
                     // draining items through it when a Controller#Stop call happens. For example,
                     // this allows tests to fetch all isolated logs before finishing.
                     batch_iterator_task_sender
                         .unbounded_send(Task::spawn(async move {
                             if let Err(e) = Self::run_server(
                                 pipeline,
                                 requests,
                                 stream_parameters,
                                 moniker_rewriter,
                                 accessor_stats,
                             )
                             .await
                             {
                                 e.close(control);
                             }
                         }))
                         .ok();
                 }
                 self.archive_accessor_stats.global_stats.connections_closed.add(1);
             }))
             .ok();
     }
 }

 struct SchemaTruncationCounter {
     truncated_schemas: u64,
     total_schemas: u64,
 }

 impl SchemaTruncationCounter {
     pub fn new() -> Arc<Mutex<Self>> {
         Arc::new(Mutex::new(Self { truncated_schemas: 0, total_schemas: 0 }))
     }
 }

 pub struct BatchIterator {
     /// requests is always populated on construction and is removed in run().
     /// This is an option as run() needs to consume it, but the Drop impl prevents us
     /// from unpacking BatchIterator.
     requests: Option<BatchIteratorRequestStream>,
     stats: Arc<BatchIteratorConnectionStats>,
     data: FormattedStream,
     truncation_counter: Option<Arc<Mutex<SchemaTruncationCounter>>>,
 }

 // Checks if a given schema is within a components budget, and if it is, updates the budget,
 // then returns true. Otherwise, if the schema is not within budget, returns false.
 fn maybe_update_budget(
     budget_map: &mut HashMap<ImmutableString, usize>,
     moniker: &str,
     bytes: usize,
     byte_limit: usize,
 ) -> bool {
     if let Some(remaining_budget) = budget_map.get_mut(moniker) {
         if *remaining_budget + bytes > byte_limit {
             false
         } else {
             *remaining_budget += bytes;
             true
         }
     } else if bytes > byte_limit {
         budget_map.insert(moniker.to_string().into_boxed_str(), 0);
         false
     } else {
         budget_map.insert(moniker.to_string().into_boxed_str(), bytes);
         true
     }
 }

 impl BatchIterator {
     pub fn new<Items, D>(
         data: Items,
         requests: BatchIteratorRequestStream,
         mode: StreamMode,
         stats: Arc<BatchIteratorConnectionStats>,
         per_component_byte_limit_opt: Option<usize>,
     ) -> Result<Self, AccessorError>
     where
         Items: Stream<Item = Data<D>> + Send + 'static,
         D: DiagnosticsData,
     {
         let result_stats = stats.clone();

         let mut budget_tracker: HashMap<ImmutableString, usize> = HashMap::new();

         let truncation_counter = SchemaTruncationCounter::new();
         let stream_owned_counter = truncation_counter.clone();

         let data = data.map(move |d| {
             let mut unlocked_counter = stream_owned_counter.lock().unwrap();
             unlocked_counter.total_schemas += 1;
             if D::has_errors(&d.metadata) {
                 result_stats.add_result_error();
             }

             match JsonString::serialize(&d, D::DATA_TYPE) {
                 Err(e) => {
                     result_stats.add_result_error();
                     Err(e)
                 }
                 Ok(contents) => {
                     result_stats.add_result();
                     match per_component_byte_limit_opt {
                         Some(x) => {
                             if maybe_update_budget(
                                 &mut budget_tracker,
                                 &d.moniker,
                                 contents.size as usize,
                                 x,
                             ) {
                                 Ok(contents)
                             } else {
                                 result_stats.add_schema_truncated();
                                 unlocked_counter.truncated_schemas += 1;

                                 let new_data = d.dropped_payload_schema(
                                     "Schema failed to fit component budget.".to_string(),
                                 );

                                 // TODO(66085): If a payload is truncated, cache the
                                 // new schema so that we can reuse if other schemas from
                                 // the same component get dropped.
                                 JsonString::serialize(&new_data, D::DATA_TYPE)
                             }
                         }
                         None => Ok(contents),
                     }
                 }
             }
         });

         Self::new_inner(
             new_batcher(data, stats.clone(), mode),
             requests,
             stats,
             Some(truncation_counter),
         )
     }

     pub fn new_serving_arrays<D, S>(
         data: S,
         requests: BatchIteratorRequestStream,
         mode: StreamMode,
         stats: Arc<BatchIteratorConnectionStats>,
     ) -> Result<Self, AccessorError>
     where
         D: Serialize + Send + 'static,
         S: Stream<Item = D> + Send + Unpin + 'static,
     {
         let data =
             JsonPacketSerializer::new(stats.clone(), FORMATTED_CONTENT_CHUNK_SIZE_TARGET, data);
         Self::new_inner(new_batcher(data, stats.clone(), mode), requests, stats, None)
     }

     fn new_inner(
         data: FormattedStream,
         requests: BatchIteratorRequestStream,
         stats: Arc<BatchIteratorConnectionStats>,
         truncation_counter: Option<Arc<Mutex<SchemaTruncationCounter>>>,
     ) -> Result<Self, AccessorError> {
         stats.open_connection();
         Ok(Self { data, requests: Some(requests), stats, truncation_counter })
     }

     pub async fn run(mut self) -> Result<(), AccessorError> {
         let (serve_inner, terminated) =
             self.requests.take().expect("request stream should be present").into_inner();
         let serve_inner_clone = serve_inner.clone();
         let channel_closed_fut =
             fasync::OnSignals::new(serve_inner_clone.channel(), zx::Signals::CHANNEL_PEER_CLOSED)
                 .shared();
         let mut requests = BatchIteratorRequestStream::from_inner(serve_inner, terminated);

         while let Some(res) = requests.next().await {
             let BatchIteratorRequest::GetNext { responder } = res?;
             self.stats.add_request();
             let start_time = zx::Time::get_monotonic();
             let batch = match select(self.data.next(), channel_closed_fut.clone()).await {
                 // if we get None back, treat that as a terminal batch with an empty vec
                 Either::Left((batch_option, _)) => batch_option.unwrap_or_default(),
                 // if the client closes the channel, stop waiting and terminate.
                 Either::Right(_) => break,
             };

             // turn errors into epitaphs -- we drop intermediate items if there was an error midway
             let batch = batch.into_iter().collect::<Result<Vec<_>, _>>()?;

             // increment counters
             self.stats.add_response();
             if batch.is_empty() {
                 if let Some(truncation_count) = &self.truncation_counter {
                     let unlocked_count = truncation_count.lock().unwrap();
                     if unlocked_count.total_schemas > 0 {
                         self.stats.global_stats().record_percent_truncated_schemas(
                             ((unlocked_count.truncated_schemas as f32
                                 / unlocked_count.total_schemas as f32)
                                 * 100.0)
                                 .round() as u64,
                         );
                     }
                 }
                 self.stats.add_terminal();
             }
             self.stats.global_stats().record_batch_duration(zx::Time::get_monotonic() - start_time);

             let mut response = Ok(batch);
             responder.send(&mut response)?;
         }
         Ok(())
     }
 }

 impl Drop for BatchIterator {
     fn drop(&mut self) {
         self.stats.close_connection();
     }
 }

 pub struct PerformanceConfig {
     pub batch_timeout_sec: i64,
     pub aggregated_content_limit_bytes: Option<u64>,
 }

 impl TryFrom<&StreamParameters> for PerformanceConfig {
     type Error = AccessorError;
     fn try_from(params: &StreamParameters) -> Result<PerformanceConfig, Self::Error> {
         let batch_timeout_sec_opt = match params {
             // If only nested batch retrieval timeout is definitely not set,
             // use the optional outer field.
             StreamParameters {
                 batch_retrieval_timeout_seconds,
                 performance_configuration: None,
                 ..
             }
             | StreamParameters {
                 batch_retrieval_timeout_seconds,
                 performance_configuration:
                     Some(PerformanceConfiguration { batch_retrieval_timeout_seconds: None, .. }),
                 ..
             } => batch_retrieval_timeout_seconds,
             // If the outer field is definitely not set, and the inner field might be,
             // use the inner field.
             StreamParameters {
                 batch_retrieval_timeout_seconds: None,
                 performance_configuration:
                     Some(PerformanceConfiguration { batch_retrieval_timeout_seconds, .. }),
                 ..
             } => batch_retrieval_timeout_seconds,
             // Both the inner and outer fields are set, which is an error.
             _ => return Err(AccessorError::DuplicateBatchTimeout),
         };

         let aggregated_content_limit_bytes = match params {
             StreamParameters {
                 performance_configuration:
                     Some(PerformanceConfiguration { max_aggregate_content_size_bytes, .. }),
                 ..
             } => *max_aggregate_content_size_bytes,
             _ => None,
         };

         Ok(PerformanceConfig {
             batch_timeout_sec: batch_timeout_sec_opt
                 .unwrap_or(constants::PER_COMPONENT_ASYNC_TIMEOUT_SECONDS),
             aggregated_content_limit_bytes,
         })
     }
 }

 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::{pipeline::Pipeline, repository::DataRepo};
     use assert_matches::assert_matches;
     use fidl_fuchsia_diagnostics::{ArchiveAccessorMarker, BatchIteratorMarker};
     use fuchsia_inspect::Node;
     use fuchsia_zircon_status as zx_status;
     use futures::channel::mpsc;
     use parking_lot::RwLock;

     #[fuchsia::test]
     async fn logs_only_accept_basic_component_selectors() {
         let (accessor, stream) =
             fidl::endpoints::create_proxy_and_stream::<ArchiveAccessorMarker>().unwrap();
         let (snd, _rcv) = mpsc::unbounded();
         fasync::Task::spawn(async move {
             let pipeline = Arc::new(RwLock::new(Pipeline::for_test(None, DataRepo::default())));
             let accessor =
                 ArchiveAccessor::new(pipeline, Arc::new(AccessorStats::new(Node::default())));
             accessor.spawn_server(stream, snd);
         })
         .detach();

         // A selector of the form `component:node/path:property` is rejected.
         let (batch_iterator, server_end) =
             fidl::endpoints::create_proxy::<BatchIteratorMarker>().unwrap();
         assert!(accessor
             .r#stream_diagnostics(
                 StreamParameters {
                     data_type: Some(DataType::Logs),
                     stream_mode: Some(StreamMode::SnapshotThenSubscribe),
                     format: Some(Format::Json),
                     client_selector_configuration: Some(ClientSelectorConfiguration::Selectors(
                         vec![SelectorArgument::RawSelector("foo:root/bar:baz".to_string()),]
                     )),
                     ..StreamParameters::EMPTY
                 },
                 server_end
             )
             .is_ok());
         assert_matches!(
             batch_iterator.get_next().await,
             Err(fidl::Error::ClientChannelClosed { status: zx_status::Status::INVALID_ARGS, .. })
         );

         // A selector of the form `component:root` is accepted.
         let (batch_iterator, server_end) =
             fidl::endpoints::create_proxy::<BatchIteratorMarker>().unwrap();
         assert!(accessor
             .r#stream_diagnostics(
                 StreamParameters {
                     data_type: Some(DataType::Logs),
                     stream_mode: Some(StreamMode::Snapshot),
                     format: Some(Format::Json),
                     client_selector_configuration: Some(ClientSelectorConfiguration::Selectors(
                         vec![SelectorArgument::RawSelector("foo:root".to_string()),]
                     )),
                     ..StreamParameters::EMPTY
                 },
                 server_end
             )
             .is_ok());

         assert!(batch_iterator.get_next().await.is_ok());
     }

     #[fuchsia::test]
     async fn accessor_skips_invalid_selectors() {
         let (accessor, stream) =
             fidl::endpoints::create_proxy_and_stream::<ArchiveAccessorMarker>().unwrap();
         let (snd, _rcv) = mpsc::unbounded();
         fasync::Task::spawn(async move {
             let pipeline = Arc::new(RwLock::new(Pipeline::for_test(None, DataRepo::default())));
             let accessor =
                 ArchiveAccessor::new(pipeline, Arc::new(AccessorStats::new(Node::default())));
             accessor.spawn_server(stream, snd);
         })
         .detach();

         // A selector of the form `component:node/path:property` is rejected.
         let (batch_iterator, server_end) =
             fidl::endpoints::create_proxy::<BatchIteratorMarker>().unwrap();

         assert!(accessor
             .r#stream_diagnostics(
                 StreamParameters {
                     data_type: Some(DataType::Inspect),
                     stream_mode: Some(StreamMode::Snapshot),
                     format: Some(Format::Json),
                     client_selector_configuration: Some(ClientSelectorConfiguration::Selectors(
                         vec![
                             SelectorArgument::RawSelector("invalid".to_string()),
                             SelectorArgument::RawSelector("valid:root".to_string()),
                         ]
                     )),
                     ..StreamParameters::EMPTY
                 },
                 server_end
             )
             .is_ok());

         // The batch iterator proxy should remain valid and providing responses regardless of the
         // invalid selectors that were given.
         assert!(batch_iterator.get_next().await.is_ok());
     }

     #[fuchsia::test]
     fn batch_iterator_terminates_on_client_disconnect() {
         let mut executor = fasync::TestExecutor::new().expect("create executor");
         let (batch_iterator_proxy, stream) =
             fidl::endpoints::create_proxy_and_stream::<BatchIteratorMarker>().unwrap();
         // Create a batch iterator that uses a hung stream to serve logs.
         let batch_iterator = BatchIterator::new(
             futures::stream::pending::<diagnostics_data::Data<diagnostics_data::Logs>>(),
             stream,
             StreamMode::Subscribe,
             Arc::new(AccessorStats::new(Node::default()).new_inspect_batch_iterator()),
             None,
         )
         .expect("create batch iterator");

         let mut batch_iterator_fut = batch_iterator.run().boxed();
         assert!(executor.run_until_stalled(&mut batch_iterator_fut).is_pending());

         // After sending a request, the request should be unfulfilled.
         let mut iterator_request_fut = batch_iterator_proxy.get_next();
         assert!(executor.run_until_stalled(&mut iterator_request_fut).is_pending());
         assert!(executor.run_until_stalled(&mut batch_iterator_fut).is_pending());
         assert!(executor.run_until_stalled(&mut iterator_request_fut).is_pending());

         // After closing the client end of the channel, the server should terminate and release
         // resources.
         drop(iterator_request_fut);
         drop(batch_iterator_proxy);
         assert_matches!(
             executor.run_until_stalled(&mut batch_iterator_fut),
             core::task::Poll::Ready(Ok(()))
         );
     }
 }
	// Copyright 2019 The Fuchsia Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	use {
	crate::{
	constants::{self, FORMATTED_CONTENT_CHUNK_SIZE_TARGET},
	diagnostics::{AccessorStats, BatchIteratorConnectionStats},
	error::AccessorError,
	formatter::{new_batcher, FormattedStream, JsonPacketSerializer, JsonString},
	inspect,
	lifecycle::LifecycleServer,
	moniker_rewriter::MonikerRewriter,
	pipeline::Pipeline,
	ImmutableString,
	},
	diagnostics_data::{Data, DiagnosticsData},
	fidl::endpoints::RequestStream,
	fidl_fuchsia_diagnostics::{
	self, ArchiveAccessorRequest, ArchiveAccessorRequestStream, BatchIteratorRequest,
	BatchIteratorRequestStream, ClientSelectorConfiguration, DataType, Format,
	PerformanceConfiguration, Selector, SelectorArgument, StreamMode, StreamParameters,
	StringSelector, TreeSelector, TreeSelectorUnknown,
	},
	fuchsia_async::{self as fasync, Task},
	fuchsia_inspect::NumericProperty,
	fuchsia_zircon as zx,
	futures::{
	channel::mpsc::UnboundedSender,
	future::{select, Either},
	prelude::*,
	},
	parking_lot::RwLock,
	selectors::{self, FastError},
	serde::Serialize,
	std::collections::HashMap,
	std::convert::TryFrom,
	std::sync::{Arc, Mutex},
	tracing::warn,
	};

	/// ArchiveAccessor represents an incoming connection from a client to an Archivist
	/// instance, through which the client may make Reader requests to the various data
	/// sources the Archivist offers.
	pub struct ArchiveAccessor {
	// The inspect repository containing read-only inspect data shared across
	// all inspect reader instances.
	pipeline: Arc<RwLock<Pipeline>>,
	archive_accessor_stats: Arc<AccessorStats>,
	moniker_rewriter: Option<Arc<MonikerRewriter>>,
	}

	fn validate_and_parse_selectors(
	selector_args: Vec<SelectorArgument>,
	) -> Result<Vec<Selector>, AccessorError> {
	let mut selectors = vec![];
	let mut errors = vec![];

	if selector_args.is_empty() {
	return Err(AccessorError::EmptySelectors);
	}

	for selector_arg in selector_args {
	match selectors::take_from_argument::<FastError>(selector_arg) {
	Ok(s) => selectors.push(s),
	Err(e) => errors.push(e),
	}
	}

	if !errors.is_empty() {
	warn!(?errors, "Found errors in selector arguments");
	}

	Ok(selectors)
	}

	fn validate_and_parse_log_selectors(
	selector_args: Vec<SelectorArgument>,
	) -> Result<Vec<Selector>, AccessorError> {
	// Only accept selectors of the type: `component:root` for logs for now.
	let selectors = validate_and_parse_selectors(selector_args)?;
	for selector in &selectors {
	// Unwrap safe: Previous validation discards any selector without a node.
	let tree_selector = selector.tree_selector.as_ref().unwrap();
	match tree_selector {
	TreeSelector::PropertySelector(_) => {
	return Err(AccessorError::InvalidLogSelector);
	}
	TreeSelector::SubtreeSelector(subtree_selector) => {
	if subtree_selector.node_path.len() != 1 {
	return Err(AccessorError::InvalidLogSelector);
	}
	match &subtree_selector.node_path[0] {
	StringSelector::ExactMatch(val) if val == "root" => {}
	StringSelector::StringPattern(val) if val == "root" => {}
	_ => {
	return Err(AccessorError::InvalidLogSelector);
	}
	}
	}
	TreeSelectorUnknown!() => {}
	}
	}
	Ok(selectors)
	}

	impl ArchiveAccessor {
	/// Create a new accessor for interacting with the archivist's data. The pipeline
	/// parameter determines which static configurations scope/restrict the visibility of
	/// data accessed by readers spawned by this accessor.
	pub fn new(
	pipeline: Arc<RwLock<Pipeline>>,
	archive_accessor_stats: Arc<AccessorStats>,
	) -> Self {
	ArchiveAccessor { pipeline, archive_accessor_stats, moniker_rewriter: None }
	}

	pub fn add_moniker_rewriter(&mut self, rewriter: Arc<MonikerRewriter>) -> &mut Self {
	self.moniker_rewriter = Some(rewriter);
	self
	}

	async fn run_server(
	pipeline: Arc<RwLock<Pipeline>>,
	requests: BatchIteratorRequestStream,
	params: StreamParameters,
	rewriter: Option<Arc<MonikerRewriter>>,
	accessor_stats: Arc<AccessorStats>,
	) -> Result<(), AccessorError> {
	let format = params.format.ok_or(AccessorError::MissingFormat)?;
	if !matches!(format, Format::Json) {
	return Err(AccessorError::UnsupportedFormat);
	}
	let mode = params.stream_mode.ok_or(AccessorError::MissingMode)?;

	let performance_config: PerformanceConfig = PerformanceConfig::try_from(&params)?;

	match params.data_type.ok_or(AccessorError::MissingDataType)? {
	DataType::Inspect => {
	if !matches!(mode, StreamMode::Snapshot) {
	return Err(AccessorError::UnsupportedMode);
	}
	let stats = Arc::new(accessor_stats.new_inspect_batch_iterator());

	let selectors =
	params.client_selector_configuration.ok_or(AccessorError::MissingSelectors)?;

	let selectors = match selectors {
	ClientSelectorConfiguration::Selectors(selectors) => {
	Some(validate_and_parse_selectors(selectors)?)
	}
	ClientSelectorConfiguration::SelectAll(_) => None,
	_ => return Err(AccessorError::InvalidSelectors("unrecognized selectors")),
	};

	let (selectors, output_rewriter) = match (selectors, rewriter) {
	(Some(selectors), Some(rewriter)) => rewriter.rewrite_selectors(selectors),
	// behaves correctly whether selectors is Some(_) or None
	(selectors, _) => (selectors, None),
	};

	let unpopulated_container_vec = pipeline.read().fetch_inspect_data(&selectors);

	let per_component_budget_opt = if unpopulated_container_vec.is_empty() {
	None
	} else {
	performance_config
	.aggregated_content_limit_bytes
	.map(\|limit\| (limit as usize) / unpopulated_container_vec.len())
	};

	if let Some(max_snapshot_size) = performance_config.aggregated_content_limit_bytes {
	stats.global_stats().record_max_snapshot_size_config(max_snapshot_size);
	}

	BatchIterator::new(
	inspect::ReaderServer::stream(
	unpopulated_container_vec,
	performance_config,
	selectors,
	output_rewriter,
	stats.clone(),
	),
	requests,
	mode,
	stats,
	per_component_budget_opt,
	)?
	.run()
	.await
	}
	DataType::Lifecycle => {
	// TODO(fxbug.dev/61350) support other modes
	if !matches!(mode, StreamMode::Snapshot) {
	return Err(AccessorError::UnsupportedMode);
	}
	let stats = Arc::new(accessor_stats.new_lifecycle_batch_iterator());

	let selectors =
	params.client_selector_configuration.ok_or(AccessorError::MissingSelectors)?;
	if !matches!(selectors, ClientSelectorConfiguration::SelectAll(_)) {
	return Err(AccessorError::InvalidSelectors(
	"lifecycle only supports SelectAll at the moment",
	));
	}

	let events = LifecycleServer::new(pipeline);

	BatchIterator::new(events, requests, mode, stats, None)?.run().await
	}
	DataType::Logs => {
	let stats = Arc::new(accessor_stats.new_logs_batch_iterator());
	let selectors = match params.client_selector_configuration {
	Some(ClientSelectorConfiguration::Selectors(selectors)) => {
	Some(validate_and_parse_log_selectors(selectors)?)
	}
	Some(ClientSelectorConfiguration::SelectAll(_)) => None,
	_ => return Err(AccessorError::InvalidSelectors("unrecognized selectors")),
	};
	let logs =
	pipeline.read().logs(mode, selectors).map(move \|inner: _\| (*inner).clone());
	BatchIterator::new_serving_arrays(logs, requests, mode, stats)?.run().await
	}
	}
	}

	/// Spawn an instance `fidl_fuchsia_diagnostics/Archive` that allows clients to open
	/// reader session to diagnostics data.
	pub fn spawn_server(
	self,
	mut stream: ArchiveAccessorRequestStream,
	task_sender: UnboundedSender<Task<()>>,
	) {
	// Self isn't guaranteed to live into the exception handling of the async block. We need to clone self
	// to have a version that can be referenced in the exception handling.
	let batch_iterator_task_sender = task_sender.clone();
	task_sender
	.unbounded_send(fasync::Task::spawn(async move {
	self.archive_accessor_stats.global_stats.connections_opened.add(1);
	while let Ok(Some(ArchiveAccessorRequest::StreamDiagnostics {
	result_stream,
	stream_parameters,
	control_handle: _,
	})) = stream.try_next().await
	{
	let (requests, control) = match result_stream.into_stream_and_control_handle() {
	Ok(r) => r,
	Err(e) => {
	warn!(?e, "Couldn't bind results channel to executor.");
	continue;
	}
	};

	self.archive_accessor_stats.global_stats.stream_diagnostics_requests.add(1);
	let pipeline = self.pipeline.clone();
	let accessor_stats = self.archive_accessor_stats.clone();
	let moniker_rewriter = self.moniker_rewriter.clone();
	// Store the batch iterator task so that we can ensure that the client finishes
	// draining items through it when a Controller#Stop call happens. For example,
	// this allows tests to fetch all isolated logs before finishing.
	batch_iterator_task_sender
	.unbounded_send(Task::spawn(async move {
	if let Err(e) = Self::run_server(
	pipeline,
	requests,
	stream_parameters,
	moniker_rewriter,
	accessor_stats,
	)
	.await
	{
	e.close(control);
	}
	}))
	.ok();
	}
	self.archive_accessor_stats.global_stats.connections_closed.add(1);
	}))
	.ok();
	}
	}

	struct SchemaTruncationCounter {
	truncated_schemas: u64,
	total_schemas: u64,
	}

	impl SchemaTruncationCounter {
	pub fn new() -> Arc<Mutex<Self>> {
	Arc::new(Mutex::new(Self { truncated_schemas: 0, total_schemas: 0 }))
	}
	}

	pub struct BatchIterator {
	/// requests is always populated on construction and is removed in run().
	/// This is an option as run() needs to consume it, but the Drop impl prevents us
	/// from unpacking BatchIterator.
	requests: Option<BatchIteratorRequestStream>,
	stats: Arc<BatchIteratorConnectionStats>,
	data: FormattedStream,
	truncation_counter: Option<Arc<Mutex<SchemaTruncationCounter>>>,
	}

	// Checks if a given schema is within a components budget, and if it is, updates the budget,
	// then returns true. Otherwise, if the schema is not within budget, returns false.
	fn maybe_update_budget(
	budget_map: &mut HashMap<ImmutableString, usize>,
	moniker: &str,
	bytes: usize,
	byte_limit: usize,
	) -> bool {
	if let Some(remaining_budget) = budget_map.get_mut(moniker) {
	if *remaining_budget + bytes > byte_limit {
	false
	} else {
	*remaining_budget += bytes;
	true
	}
	} else if bytes > byte_limit {
	budget_map.insert(moniker.to_string().into_boxed_str(), 0);
	false
	} else {
	budget_map.insert(moniker.to_string().into_boxed_str(), bytes);
	true
	}
	}

	impl BatchIterator {
	pub fn new<Items, D>(
	data: Items,
	requests: BatchIteratorRequestStream,
	mode: StreamMode,
	stats: Arc<BatchIteratorConnectionStats>,
	per_component_byte_limit_opt: Option<usize>,
	) -> Result<Self, AccessorError>
	where
	Items: Stream<Item = Data<D>> + Send + 'static,
	D: DiagnosticsData,
	{
	let result_stats = stats.clone();

	let mut budget_tracker: HashMap<ImmutableString, usize> = HashMap::new();

	let truncation_counter = SchemaTruncationCounter::new();
	let stream_owned_counter = truncation_counter.clone();

	let data = data.map(move \|d\| {
	let mut unlocked_counter = stream_owned_counter.lock().unwrap();
	unlocked_counter.total_schemas += 1;
	if D::has_errors(&d.metadata) {
	result_stats.add_result_error();
	}

	match JsonString::serialize(&d, D::DATA_TYPE) {
	Err(e) => {
	result_stats.add_result_error();
	Err(e)
	}
	Ok(contents) => {
	result_stats.add_result();
	match per_component_byte_limit_opt {
	Some(x) => {
	if maybe_update_budget(
	&mut budget_tracker,
	&d.moniker,
	contents.size as usize,
	x,
	) {
	Ok(contents)
	} else {
	result_stats.add_schema_truncated();
	unlocked_counter.truncated_schemas += 1;

	let new_data = d.dropped_payload_schema(
	"Schema failed to fit component budget.".to_string(),
	);

	// TODO(66085): If a payload is truncated, cache the
	// new schema so that we can reuse if other schemas from
	// the same component get dropped.
	JsonString::serialize(&new_data, D::DATA_TYPE)
	}
	}
	None => Ok(contents),
	}
	}
	}
	});

	Self::new_inner(
	new_batcher(data, stats.clone(), mode),
	requests,
	stats,
	Some(truncation_counter),
	)
	}

	pub fn new_serving_arrays<D, S>(
	data: S,
	requests: BatchIteratorRequestStream,
	mode: StreamMode,
	stats: Arc<BatchIteratorConnectionStats>,
	) -> Result<Self, AccessorError>
	where
	D: Serialize + Send + 'static,
	S: Stream<Item = D> + Send + Unpin + 'static,
	{
	let data =
	JsonPacketSerializer::new(stats.clone(), FORMATTED_CONTENT_CHUNK_SIZE_TARGET, data);
	Self::new_inner(new_batcher(data, stats.clone(), mode), requests, stats, None)
	}

	fn new_inner(
	data: FormattedStream,
	requests: BatchIteratorRequestStream,
	stats: Arc<BatchIteratorConnectionStats>,
	truncation_counter: Option<Arc<Mutex<SchemaTruncationCounter>>>,
	) -> Result<Self, AccessorError> {
	stats.open_connection();
	Ok(Self { data, requests: Some(requests), stats, truncation_counter })
	}

	pub async fn run(mut self) -> Result<(), AccessorError> {
	let (serve_inner, terminated) =
	self.requests.take().expect("request stream should be present").into_inner();
	let serve_inner_clone = serve_inner.clone();
	let channel_closed_fut =
	fasync::OnSignals::new(serve_inner_clone.channel(), zx::Signals::CHANNEL_PEER_CLOSED)
	.shared();
	let mut requests = BatchIteratorRequestStream::from_inner(serve_inner, terminated);

	while let Some(res) = requests.next().await {
	let BatchIteratorRequest::GetNext { responder } = res?;
	self.stats.add_request();
	let start_time = zx::Time::get_monotonic();
	let batch = match select(self.data.next(), channel_closed_fut.clone()).await {
	// if we get None back, treat that as a terminal batch with an empty vec
	Either::Left((batch_option, _)) => batch_option.unwrap_or_default(),
	// if the client closes the channel, stop waiting and terminate.
	Either::Right(_) => break,
	};

	// turn errors into epitaphs -- we drop intermediate items if there was an error midway
	let batch = batch.into_iter().collect::<Result<Vec<_>, _>>()?;

	// increment counters
	self.stats.add_response();
	if batch.is_empty() {
	if let Some(truncation_count) = &self.truncation_counter {
	let unlocked_count = truncation_count.lock().unwrap();
	if unlocked_count.total_schemas > 0 {
	self.stats.global_stats().record_percent_truncated_schemas(
	((unlocked_count.truncated_schemas as f32
	/ unlocked_count.total_schemas as f32)
	* 100.0)
	.round() as u64,
	);
	}
	}
	self.stats.add_terminal();
	}
	self.stats.global_stats().record_batch_duration(zx::Time::get_monotonic() - start_time);

	let mut response = Ok(batch);
	responder.send(&mut response)?;
	}
	Ok(())
	}
	}

	impl Drop for BatchIterator {
	fn drop(&mut self) {
	self.stats.close_connection();
	}
	}

	pub struct PerformanceConfig {
	pub batch_timeout_sec: i64,
	pub aggregated_content_limit_bytes: Option<u64>,
	}

	impl TryFrom<&StreamParameters> for PerformanceConfig {
	type Error = AccessorError;
	fn try_from(params: &StreamParameters) -> Result<PerformanceConfig, Self::Error> {
	let batch_timeout_sec_opt = match params {
	// If only nested batch retrieval timeout is definitely not set,
	// use the optional outer field.
	StreamParameters {
	batch_retrieval_timeout_seconds,
	performance_configuration: None,
	..
	}
	\| StreamParameters {
	batch_retrieval_timeout_seconds,
	performance_configuration:
	Some(PerformanceConfiguration { batch_retrieval_timeout_seconds: None, .. }),
	..
	} => batch_retrieval_timeout_seconds,
	// If the outer field is definitely not set, and the inner field might be,
	// use the inner field.
	StreamParameters {
	batch_retrieval_timeout_seconds: None,
	performance_configuration:
	Some(PerformanceConfiguration { batch_retrieval_timeout_seconds, .. }),
	..
	} => batch_retrieval_timeout_seconds,
	// Both the inner and outer fields are set, which is an error.
	_ => return Err(AccessorError::DuplicateBatchTimeout),
	};

	let aggregated_content_limit_bytes = match params {
	StreamParameters {
	performance_configuration:
	Some(PerformanceConfiguration { max_aggregate_content_size_bytes, .. }),
	..
	} => *max_aggregate_content_size_bytes,
	_ => None,
	};

	Ok(PerformanceConfig {
	batch_timeout_sec: batch_timeout_sec_opt
	.unwrap_or(constants::PER_COMPONENT_ASYNC_TIMEOUT_SECONDS),
	aggregated_content_limit_bytes,
	})
	}
	}

	#[cfg(test)]
	mod tests {
	use super::*;
	use crate::{pipeline::Pipeline, repository::DataRepo};
	use assert_matches::assert_matches;
	use fidl_fuchsia_diagnostics::{ArchiveAccessorMarker, BatchIteratorMarker};
	use fuchsia_inspect::Node;
	use fuchsia_zircon_status as zx_status;
	use futures::channel::mpsc;
	use parking_lot::RwLock;

	#[fuchsia::test]
	async fn logs_only_accept_basic_component_selectors() {
	let (accessor, stream) =
	fidl::endpoints::create_proxy_and_stream::<ArchiveAccessorMarker>().unwrap();
	let (snd, _rcv) = mpsc::unbounded();
	fasync::Task::spawn(async move {
	let pipeline = Arc::new(RwLock::new(Pipeline::for_test(None, DataRepo::default())));
	let accessor =
	ArchiveAccessor::new(pipeline, Arc::new(AccessorStats::new(Node::default())));
	accessor.spawn_server(stream, snd);
	})
	.detach();

	// A selector of the form `component:node/path:property` is rejected.
	let (batch_iterator, server_end) =
	fidl::endpoints::create_proxy::<BatchIteratorMarker>().unwrap();
	assert!(accessor
	.r#stream_diagnostics(
	StreamParameters {
	data_type: Some(DataType::Logs),
	stream_mode: Some(StreamMode::SnapshotThenSubscribe),
	format: Some(Format::Json),
	client_selector_configuration: Some(ClientSelectorConfiguration::Selectors(
	vec![SelectorArgument::RawSelector("foo:root/bar:baz".to_string()),]
	)),
	..StreamParameters::EMPTY
	},
	server_end
	)
	.is_ok());
	assert_matches!(
	batch_iterator.get_next().await,
	Err(fidl::Error::ClientChannelClosed { status: zx_status::Status::INVALID_ARGS, .. })
	);

	// A selector of the form `component:root` is accepted.
	let (batch_iterator, server_end) =
	fidl::endpoints::create_proxy::<BatchIteratorMarker>().unwrap();
	assert!(accessor
	.r#stream_diagnostics(
	StreamParameters {
	data_type: Some(DataType::Logs),
	stream_mode: Some(StreamMode::Snapshot),
	format: Some(Format::Json),
	client_selector_configuration: Some(ClientSelectorConfiguration::Selectors(
	vec![SelectorArgument::RawSelector("foo:root".to_string()),]
	)),
	..StreamParameters::EMPTY
	},
	server_end
	)
	.is_ok());

	assert!(batch_iterator.get_next().await.is_ok());
	}

	#[fuchsia::test]
	async fn accessor_skips_invalid_selectors() {
	let (accessor, stream) =
	fidl::endpoints::create_proxy_and_stream::<ArchiveAccessorMarker>().unwrap();
	let (snd, _rcv) = mpsc::unbounded();
	fasync::Task::spawn(async move {
	let pipeline = Arc::new(RwLock::new(Pipeline::for_test(None, DataRepo::default())));
	let accessor =
	ArchiveAccessor::new(pipeline, Arc::new(AccessorStats::new(Node::default())));
	accessor.spawn_server(stream, snd);
	})
	.detach();

	// A selector of the form `component:node/path:property` is rejected.
	let (batch_iterator, server_end) =
	fidl::endpoints::create_proxy::<BatchIteratorMarker>().unwrap();

	assert!(accessor
	.r#stream_diagnostics(
	StreamParameters {
	data_type: Some(DataType::Inspect),
	stream_mode: Some(StreamMode::Snapshot),
	format: Some(Format::Json),
	client_selector_configuration: Some(ClientSelectorConfiguration::Selectors(
	vec![
	SelectorArgument::RawSelector("invalid".to_string()),
	SelectorArgument::RawSelector("valid:root".to_string()),
	]
	)),
	..StreamParameters::EMPTY
	},
	server_end
	)
	.is_ok());

	// The batch iterator proxy should remain valid and providing responses regardless of the
	// invalid selectors that were given.
	assert!(batch_iterator.get_next().await.is_ok());
	}

	#[fuchsia::test]
	fn batch_iterator_terminates_on_client_disconnect() {
	let mut executor = fasync::TestExecutor::new().expect("create executor");
	let (batch_iterator_proxy, stream) =
	fidl::endpoints::create_proxy_and_stream::<BatchIteratorMarker>().unwrap();
	// Create a batch iterator that uses a hung stream to serve logs.
	let batch_iterator = BatchIterator::new(
	futures::stream::pending::<diagnostics_data::Data<diagnostics_data::Logs>>(),
	stream,
	StreamMode::Subscribe,
	Arc::new(AccessorStats::new(Node::default()).new_inspect_batch_iterator()),
	None,
	)
	.expect("create batch iterator");

	let mut batch_iterator_fut = batch_iterator.run().boxed();
	assert!(executor.run_until_stalled(&mut batch_iterator_fut).is_pending());

	// After sending a request, the request should be unfulfilled.
	let mut iterator_request_fut = batch_iterator_proxy.get_next();
	assert!(executor.run_until_stalled(&mut iterator_request_fut).is_pending());
	assert!(executor.run_until_stalled(&mut batch_iterator_fut).is_pending());
	assert!(executor.run_until_stalled(&mut iterator_request_fut).is_pending());

	// After closing the client end of the channel, the server should terminate and release
	// resources.
	drop(iterator_request_fut);
	drop(batch_iterator_proxy);
	assert_matches!(
	executor.run_until_stalled(&mut batch_iterator_fut),
	core::task::Poll::Ready(Ok(()))
	);
	}
	}