Skip to content

perf(init): issue telemetry subscribe earlier, overlapping service construction #1281

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
duncanista wants to merge 1 commit into
base: jordan.gonzalez/cold-start-instrumentation/feature
Choose a base branch
from
Draft

perf(init): issue telemetry subscribe earlier, overlapping service construction #1281

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 103 additions & 66 deletions bottlecap/src/bin/bottlecap/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,79 +373,116 @@ async fn extension_loop_active(
log_init_checkpoint(start_time, "dogstatsd_started");

let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(config)));
// Lifecycle Invocation Processor
let (invocation_processor_handle, invocation_processor_service) =
InvocationProcessorService::new(
Arc::clone(&tags_provider),
Arc::clone(config),
Arc::clone(&aws_config),
metrics_aggregator_handle.clone(),
Arc::clone(&propagator),
durable_context_tx,
);
tokio::spawn(async move {
invocation_processor_service.run().await;
});

// AppSec processor (if enabled)
let appsec_processor = match AppSecProcessor::new(config) {
Ok(p) => Some(Arc::new(TokioMutex::new(p))),
Err(AppSecFeatureDisabled) => None,
Err(e) => {
error!(
"AAP | error creating App & API Protection processor, the feature will be disabled: {e}"
);
None
}
};

let (
trace_agent_channel,
trace_flusher,
trace_processor,
stats_flusher,
proxy_flusher,
trace_agent_shutdown_token,
stats_concentrator,
trace_aggregator_handle,
) = start_trace_agent(
config,
&api_key_factory,
&tags_provider,
invocation_processor_handle.clone(),
appsec_processor.clone(),
&shared_client,
);
log_init_checkpoint(start_time, "trace_agent_started");

let api_runtime_proxy_shutdown_signal = start_api_runtime_proxy(
config,
Arc::clone(&aws_config),
&invocation_processor_handle,
appsec_processor.as_ref(),
Arc::clone(&propagator),
);

let lifecycle_listener =
LifecycleListener::new(invocation_processor_handle.clone(), Arc::clone(&propagator));
let lifecycle_listener_shutdown_token = lifecycle_listener.get_shutdown_token();
// TODO(astuyve): deprioritize this task after the first request
tokio::spawn(async move {
if let Err(e) = lifecycle_listener.start().await {
error!("Error starting lifecycle listener: {e:?}");
}
});

let telemetry_listener_cancel_token = setup_telemetry_client(
// The Lambda Extensions API ends the INIT phase at the first `/next` call, so the
// serialized critical path before that call directly inflates cold-start time. The
// telemetry subscription only depends on `logs_agent_channel` (already available),
// and its `subscribe` HTTP round-trip is independent of constructing the trace agent,
// AppSec, API-runtime proxy, and lifecycle listener below. We therefore issue the
// subscribe and run that remaining (synchronous) service construction concurrently
// via `tokio::join!`: the subscribe future is polled first so its network round-trip
// is in flight while construction proceeds, instead of being serialized behind it.
//
// Correctness: `setup_telemetry_client` binds the telemetry listener socket
// (synchronously, before `subscribe` returns — see `TelemetryListener::start`) prior
// to subscribing, so the listener is already accepting connections when the Telemetry
// API begins delivering events. No early `platform.initStart`/`initReport` or logs are
// dropped. The construction branch only spawns background tasks and returns handles;
// it performs no I/O that the subscribe depends on, so the two are order-independent.
let telemetry_setup = setup_telemetry_client(
client,
&r.extension_id,
&aws_config.runtime_api,
logs_agent_channel,
event_bus_tx.clone(),
config.ext.serverless_logs_enabled,
aws_config.is_managed_instance_mode(),
)
.await?;
);

let build_services = async {
// Lifecycle Invocation Processor
let (invocation_processor_handle, invocation_processor_service) =
InvocationProcessorService::new(
Arc::clone(&tags_provider),
Arc::clone(config),
Arc::clone(&aws_config),
metrics_aggregator_handle.clone(),
Arc::clone(&propagator),
durable_context_tx,
);
tokio::spawn(async move {
invocation_processor_service.run().await;
});

// AppSec processor (if enabled)
let appsec_processor = match AppSecProcessor::new(config) {
Ok(p) => Some(Arc::new(TokioMutex::new(p))),
Err(AppSecFeatureDisabled) => None,
Err(e) => {
error!(
"AAP | error creating App & API Protection processor, the feature will be disabled: {e}"
);
None
}
};

let trace_agent_bundle = start_trace_agent(
config,
&api_key_factory,
&tags_provider,
invocation_processor_handle.clone(),
appsec_processor.clone(),
&shared_client,
);
log_init_checkpoint(start_time, "trace_agent_started");

let api_runtime_proxy_shutdown_signal = start_api_runtime_proxy(
config,
Arc::clone(&aws_config),
&invocation_processor_handle,
appsec_processor.as_ref(),
Arc::clone(&propagator),
);

let lifecycle_listener =
LifecycleListener::new(invocation_processor_handle.clone(), Arc::clone(&propagator));
let lifecycle_listener_shutdown_token = lifecycle_listener.get_shutdown_token();
// TODO(astuyve): deprioritize this task after the first request
tokio::spawn(async move {
if let Err(e) = lifecycle_listener.start().await {
error!("Error starting lifecycle listener: {e:?}");
}
});

(
invocation_processor_handle,
appsec_processor,
trace_agent_bundle,
api_runtime_proxy_shutdown_signal,
lifecycle_listener_shutdown_token,
)
};

let (
telemetry_listener_cancel_token,
(
invocation_processor_handle,
appsec_processor,
(
trace_agent_channel,
trace_flusher,
trace_processor,
stats_flusher,
proxy_flusher,
trace_agent_shutdown_token,
stats_concentrator,
trace_aggregator_handle,
),
api_runtime_proxy_shutdown_signal,
lifecycle_listener_shutdown_token,
),
) = tokio::join!(telemetry_setup, build_services);
let telemetry_listener_cancel_token = telemetry_listener_cancel_token?;
log_init_checkpoint(start_time, "telemetry_subscribed");

let otlp_cancel_token = start_otlp_agent(
Expand Down Expand Up @@ -1417,7 +1454,7 @@ async fn setup_telemetry_client(
let listener = TelemetryListener::new(EXTENSION_HOST_IP, TELEMETRY_PORT, logs_tx, event_bus_tx);

let cancel_token = listener.cancel_token();
match listener.start() {
match listener.start().await {
Ok(()) => {
// Drop the listener, so event_bus_tx is closed
drop(listener);
Expand Down
20 changes: 15 additions & 5 deletions bottlecap/src/extension/telemetry/listener.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,17 +49,27 @@ impl TelemetryListener {
self.cancel_token.clone()
}

pub fn start(&self) -> Result<(), Box<dyn std::error::Error>> {
/// Binds the telemetry listener socket and starts serving in a background task.
///
/// The bind happens synchronously (before this method returns) so the socket is
/// already accepting connections by the time the caller subscribes to the Lambda
/// Telemetry API. This makes listener readiness deterministic and avoids dropping
/// early `platform.initStart`/`initReport` events delivered immediately after
/// subscription.
///
/// # Errors
///
/// Returns an error if the listener socket cannot be bound.
pub async fn start(&self) -> Result<(), Box<dyn std::error::Error>> {
let socket = SocketAddr::from((self.host, self.port));
let router = self.make_router();

let listener = TcpListener::bind(&socket).await?;
debug!("TELEMETRY API | Starting listener on {}", socket);

let cancel_token_clone = self.cancel_token();
let event_bus_tx = self.event_bus_tx.clone();
tokio::spawn(async move {
let listener = TcpListener::bind(&socket)
.await
.expect("Failed to bind socket");
debug!("TELEMETRY API | Starting listener on {}", socket);
axum::serve(listener, router)
.with_graceful_shutdown(Self::graceful_shutdown(cancel_token_clone, event_bus_tx))
.await
Expand Down
Loading