feat(subscriber): Reduce retention period to fit in max message size

If the initial update message would be too big for tokio-console's grpc
decoder, reduce the retention period and try again.

Currently the default retention period is 1 hour. That can easily grow
to more than the max grpc decode message size (4 MiB), at which point
tokio-console won't connect.

There's really no minimum safe duration for retention. It depends on how
busy the app is and on how much trace data runtime and tokio log. Here
we repeatedly divide the retention period in half until it fits in the
message.

Author: grahamking

Cargo.lock

@@ -46,6 +46,7 @@ hdrhistogram = { version = "7.3.0", default-features = false, features = ["seria
 # feature to also enable `tracing-subscriber`'s parking_lot feature flag.
 parking_lot_crate = { package = "parking_lot", version = "0.12", optional = true }
 humantime = "2.1.0"
+prost = "0.12"
 prost-types = "0.12.0"
 
 # Required for recording:

console-subscriber/Cargo.toml

@@ -1,26 +1,34 @@
-use super::{Command, Event, Shared, Watch};
-use crate::{
-    stats::{self, Unsent},
-    ToProto, WatchRequest,
-};
-use console_api as proto;
-use proto::resources::resource;
-use tokio::sync::{mpsc, Notify};
-
 use std::{
     sync::{
         atomic::{AtomicBool, Ordering::*},
         Arc,
     },
     time::{Duration, Instant},
 };
+
+use console_api as proto;
+use prost::Message;
+use proto::resources::resource;
+use tokio::sync::{mpsc, Notify};
 use tracing_core::{span::Id, Metadata};
 
+use super::{Command, Event, Shared, Watch};
+use crate::{
+    stats::{self, Unsent},
+    ToProto, WatchRequest,
+};
+
 mod id_data;
 mod shrink;
 use self::id_data::{IdData, Include};
 use self::shrink::{ShrinkMap, ShrinkVec};
 
+/// Should match tonic's (private) codec::DEFAULT_MAX_RECV_MESSAGE_SIZE
+const MAX_MESSAGE_SIZE: usize = 4 * 1024 * 1024;
+
+/// The smallest amount we will shrink retention to in an attempt to fit data in MAX_MESSAGE_SIZE
+const MIN_RETENTION: Duration = Duration::from_secs(1);
+
 /// Aggregates instrumentation traces and prepares state for the instrument
 /// server.
 ///
@@ -274,24 +282,48 @@ impl Aggregator {
     /// Add the task subscription to the watchers after sending the first update
     fn add_instrument_subscription(&mut self, subscription: Watch<proto::instrument::Update>) {
         tracing::debug!("new instrument subscription");
-
-        let task_update = Some(self.task_update(Include::All));
-        let resource_update = Some(self.resource_update(Include::All));
-        let async_op_update = Some(self.async_op_update(Include::All));
         let now = Instant::now();
 
-        let update = &proto::instrument::Update {
-            task_update,
-            resource_update,
-            async_op_update,
-            now: Some(self.base_time.to_timestamp(now)),
-            new_metadata: Some(proto::RegisterMetadata {
-                metadata: (*self.all_metadata).clone(),
-            }),
+        let update = loop {
+            let update = proto::instrument::Update {
+                task_update: Some(self.task_update(Include::All)),
+                resource_update: Some(self.resource_update(Include::All)),
+                async_op_update: Some(self.async_op_update(Include::All)),
+                now: Some(self.base_time.to_timestamp(now)),
+                new_metadata: Some(proto::RegisterMetadata {
+                    metadata: (*self.all_metadata).clone(),
+                }),
+            };
+            let el = update.encoded_len();
+            if el < MAX_MESSAGE_SIZE {
+                // normal case
+                break Some(update);
+            }
+            // If the grpc message is bigger than tokio-console will accept throw away the oldest
+            // inactive data and try again
+            self.retention /= 2;
+            self.cleanup_closed();
+            tracing::debug!(
+                retention = ?self.retention,
+                message_size = el,
+                max_message_size = MAX_MESSAGE_SIZE,
+                "Message too big, reduced retention",
+            );
+
+            if self.retention <= MIN_RETENTION {
+                self.retention = MIN_RETENTION;
+                break None;
+            }
         };
+        if update.is_none() {
+            tracing::error!(min_retention = ?MIN_RETENTION, "Message too big. Start with smaller retention.");
+            // User will only get updates
+            self.watchers.push(subscription);
+            return;
+        }
 
         // Send the initial state --- if this fails, the subscription is already dead
-        if subscription.update(update) {
+        if subscription.update(&update.unwrap()) {
             self.watchers.push(subscription)
         }
     }