Skip to content

Commit c5ee4fc

Browse files
jchrostek-ddjchrostek-dd
authored
remove invocation lock (#894)
## Task https://datadoghq.atlassian.net/browse/SVLS-7465?atlOrigin=eyJpIjoiYzI2ZWQ3NmU1M2U4NGFiMGI1ZjE3MWY5MTM2MzlkN2MiLCJwIjoiaiJ9 ## Overview Remove InvocationProcessor lock and replace with an aggregator service.
1 parent 50fca62 commit c5ee4fc

7 files changed

Lines changed: 721 additions & 135 deletions

File tree

‎bottlecap/src/bin/bottlecap/main.rs‎

Lines changed: 94 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ use bottlecap::{
4040
fips::{log_fips_status, prepare_client_provider},
4141
lifecycle::{
4242
flush_control::{FlushControl, FlushDecision},
43-
invocation::processor::Processor as InvocationProcessor,
43+
invocation::processor_service::{InvocationProcessorHandle, InvocationProcessorService},
4444
listener::Listener as LifecycleListener,
4545
},
4646
logger,
@@ -403,13 +403,18 @@ async fn extension_loop_active(
403403

404404
let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(config)));
405405
// Lifecycle Invocation Processor
406-
let invocation_processor = Arc::new(TokioMutex::new(InvocationProcessor::new(
407-
Arc::clone(&tags_provider),
408-
Arc::clone(config),
409-
Arc::clone(&aws_config),
410-
metrics_aggregator_handle.clone(),
411-
Arc::clone(&propagator),
412-
)));
406+
let (invocation_processor_handle, invocation_processor_service) =
407+
InvocationProcessorService::new(
408+
Arc::clone(&tags_provider),
409+
Arc::clone(config),
410+
Arc::clone(&aws_config),
411+
metrics_aggregator_handle.clone(),
412+
Arc::clone(&propagator),
413+
);
414+
tokio::spawn(async move {
415+
invocation_processor_service.run().await;
416+
});
417+
413418
// AppSec processor (if enabled)
414419
let appsec_processor = match AppSecProcessor::new(config) {
415420
Ok(p) => Some(Arc::new(TokioMutex::new(p))),
@@ -434,20 +439,20 @@ async fn extension_loop_active(
434439
config,
435440
&api_key_factory,
436441
&tags_provider,
437-
Arc::clone(&invocation_processor),
442+
invocation_processor_handle.clone(),
438443
appsec_processor.clone(),
439444
);
440445

441446
let api_runtime_proxy_shutdown_signal = start_api_runtime_proxy(
442447
config,
443448
Arc::clone(&aws_config),
444-
&invocation_processor,
449+
&invocation_processor_handle,
445450
appsec_processor.as_ref(),
446451
Arc::clone(&propagator),
447452
);
448453

449454
let lifecycle_listener =
450-
LifecycleListener::new(Arc::clone(&invocation_processor), Arc::clone(&propagator));
455+
LifecycleListener::new(invocation_processor_handle.clone(), Arc::clone(&propagator));
451456
let lifecycle_listener_shutdown_token = lifecycle_listener.get_shutdown_token();
452457
// TODO(astuyve): deprioritize this task after the first request
453458
tokio::spawn(async move {
@@ -488,7 +493,7 @@ async fn extension_loop_active(
488493
extension::next_event(client, &aws_config.runtime_api, &r.extension_id).await;
489494
// first invoke we must call next
490495
let mut pending_flush_handles = PendingFlushHandles::new();
491-
handle_next_invocation(next_lambda_response, invocation_processor.clone()).await;
496+
handle_next_invocation(next_lambda_response, &invocation_processor_handle).await;
492497
loop {
493498
let maybe_shutdown_event;
494499

@@ -503,7 +508,7 @@ async fn extension_loop_active(
503508
tokio::select! {
504509
biased;
505510
Some(event) = event_bus.rx.recv() => {
506-
if let Some(telemetry_event) = handle_event_bus_event(event, invocation_processor.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await {
511+
if let Some(telemetry_event) = handle_event_bus_event(event, invocation_processor_handle.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await {
507512
if let TelemetryRecord::PlatformRuntimeDone{ .. } = telemetry_event.record {
508513
break 'flush_end;
509514
}
@@ -541,7 +546,7 @@ async fn extension_loop_active(
541546
let next_response =
542547
extension::next_event(client, &aws_config.runtime_api, &r.extension_id).await;
543548
maybe_shutdown_event =
544-
handle_next_invocation(next_response, invocation_processor.clone()).await;
549+
handle_next_invocation(next_response, &invocation_processor_handle).await;
545550
}
546551
FlushDecision::Continuous | FlushDecision::Periodic | FlushDecision::Dont => {
547552
match current_flush_decision {
@@ -625,12 +630,12 @@ async fn extension_loop_active(
625630
tokio::select! {
626631
biased;
627632
next_response = &mut next_lambda_response => {
628-
maybe_shutdown_event = handle_next_invocation(next_response, invocation_processor.clone()).await;
633+
maybe_shutdown_event = handle_next_invocation(next_response, &invocation_processor_handle).await;
629634
// Need to break here to re-call next
630635
break 'next_invocation;
631636
}
632637
Some(event) = event_bus.rx.recv() => {
633-
handle_event_bus_event(event, invocation_processor.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await;
638+
handle_event_bus_event(event, invocation_processor_handle.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await;
634639
}
635640
_ = race_flush_interval.tick() => {
636641
if flush_control.flush_strategy == FlushStrategy::Default {
@@ -685,7 +690,7 @@ async fn extension_loop_active(
685690
loop {
686691
match event_bus.rx.try_recv() {
687692
Ok(event) => {
688-
handle_event_bus_event(event, invocation_processor.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await;
693+
handle_event_bus_event(event, invocation_processor_handle.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await;
689694
},
690695
Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => break 'shutdown,
691696
// Empty signals there are still outstanding senders
@@ -695,7 +700,7 @@ async fn extension_loop_active(
695700
}
696701
}
697702
} else {
698-
handle_event_bus_event(event, invocation_processor.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await;
703+
handle_event_bus_event(event, invocation_processor_handle.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await;
699704
}
700705
}
701706
// Add timeout to prevent hanging indefinitely
@@ -769,9 +774,10 @@ async fn blocking_flush_all(
769774
race_flush_interval.reset();
770775
}
771776

777+
#[allow(clippy::too_many_arguments)]
772778
async fn handle_event_bus_event(
773779
event: Event,
774-
invocation_processor: Arc<TokioMutex<InvocationProcessor>>,
780+
invocation_processor_handle: InvocationProcessorHandle,
775781
appsec_processor: Option<Arc<TokioMutex<AppSecProcessor>>>,
776782
tags_provider: Arc<TagProvider>,
777783
trace_processor: Arc<trace_processor::ServerlessTraceProcessor>,
@@ -780,35 +786,47 @@ async fn handle_event_bus_event(
780786
) -> Option<TelemetryEvent> {
781787
match event {
782788
Event::OutOfMemory(event_timestamp) => {
783-
let mut p = invocation_processor.lock().await;
784-
p.on_out_of_memory_error(event_timestamp);
785-
drop(p);
789+
if let Err(e) = invocation_processor_handle
790+
.on_out_of_memory_error(event_timestamp)
791+
.await
792+
{
793+
error!("Failed to send out of memory error to processor: {}", e);
794+
}
786795
}
787796
Event::Telemetry(event) => {
788797
debug!("Telemetry event received: {:?}", event);
789798
match event.record {
790799
TelemetryRecord::PlatformInitStart { .. } => {
791-
let mut p = invocation_processor.lock().await;
792-
p.on_platform_init_start(event.time);
793-
drop(p);
800+
if let Err(e) = invocation_processor_handle
801+
.on_platform_init_start(event.time)
802+
.await
803+
{
804+
error!("Failed to send platform init start to processor: {}", e);
805+
}
794806
}
795807
TelemetryRecord::PlatformInitReport {
796808
metrics,
797809
initialization_type,
798810
..
799811
} => {
800-
let mut p = invocation_processor.lock().await;
801-
p.on_platform_init_report(
802-
initialization_type,
803-
metrics.duration_ms,
804-
event.time.timestamp(),
805-
);
806-
drop(p);
812+
if let Err(e) = invocation_processor_handle
813+
.on_platform_init_report(
814+
initialization_type,
815+
metrics.duration_ms,
816+
event.time.timestamp(),
817+
)
818+
.await
819+
{
820+
error!("Failed to send platform init report to processor: {}", e);
821+
}
807822
}
808823
TelemetryRecord::PlatformStart { request_id, .. } => {
809-
let mut p = invocation_processor.lock().await;
810-
p.on_platform_start(request_id, event.time);
811-
drop(p);
824+
if let Err(e) = invocation_processor_handle
825+
.on_platform_start(request_id, event.time)
826+
.await
827+
{
828+
error!("Failed to send platform start to processor: {}", e);
829+
}
812830
}
813831
TelemetryRecord::PlatformRuntimeDone {
814832
ref request_id,
@@ -817,35 +835,40 @@ async fn handle_event_bus_event(
817835
ref error_type,
818836
..
819837
} => {
820-
let mut p = invocation_processor.lock().await;
821-
p.on_platform_runtime_done(
822-
request_id,
823-
metrics,
824-
status,
825-
error_type.clone(),
826-
tags_provider.clone(),
827-
Arc::new(SendingTraceProcessor {
828-
appsec: appsec_processor.clone(),
829-
processor: trace_processor.clone(),
830-
trace_tx: trace_agent_channel.clone(),
831-
stats_generator: Arc::new(StatsGenerator::new(
832-
stats_concentrator.clone(),
833-
)),
834-
}),
835-
event.time.timestamp(),
836-
)
837-
.await;
838-
drop(p);
838+
if let Err(e) = invocation_processor_handle
839+
.on_platform_runtime_done(
840+
request_id.clone(),
841+
metrics,
842+
status,
843+
error_type.clone(),
844+
tags_provider.clone(),
845+
Arc::new(SendingTraceProcessor {
846+
appsec: appsec_processor.clone(),
847+
processor: trace_processor.clone(),
848+
trace_tx: trace_agent_channel.clone(),
849+
stats_generator: Arc::new(StatsGenerator::new(
850+
stats_concentrator.clone(),
851+
)),
852+
}),
853+
event.time.timestamp(),
854+
)
855+
.await
856+
{
857+
error!("Failed to send platform runtime done to processor: {}", e);
858+
}
839859
return Some(event);
840860
}
841861
TelemetryRecord::PlatformReport {
842862
ref request_id,
843863
metrics,
844864
..
845865
} => {
846-
let mut p = invocation_processor.lock().await;
847-
p.on_platform_report(request_id, metrics, event.time.timestamp());
848-
drop(p);
866+
if let Err(e) = invocation_processor_handle
867+
.on_platform_report(request_id.clone(), metrics, event.time.timestamp())
868+
.await
869+
{
870+
error!("Failed to send platform report to processor: {}", e);
871+
}
849872
return Some(event);
850873
}
851874
_ => {
@@ -861,7 +884,7 @@ async fn handle_event_bus_event(
861884

862885
async fn handle_next_invocation(
863886
next_response: Result<NextEventResponse, ExtensionError>,
864-
invocation_processor: Arc<TokioMutex<InvocationProcessor>>,
887+
invocation_processor_handle: &InvocationProcessorHandle,
865888
) -> NextEventResponse {
866889
match next_response {
867890
Ok(NextEventResponse::Invoke {
@@ -875,16 +898,20 @@ async fn handle_next_invocation(
875898
deadline_ms,
876899
invoked_function_arn.clone()
877900
);
878-
let mut p = invocation_processor.lock().await;
879-
p.on_invoke_event(request_id.into());
880-
drop(p);
901+
if let Err(e) = invocation_processor_handle
902+
.on_invoke_event(request_id.into())
903+
.await
904+
{
905+
error!("Failed to send invoke event to processor: {}", e);
906+
}
881907
}
882908
Ok(NextEventResponse::Shutdown {
883909
ref shutdown_reason,
884910
deadline_ms,
885911
}) => {
886-
let mut p = invocation_processor.lock().await;
887-
p.on_shutdown_event();
912+
if let Err(e) = invocation_processor_handle.on_shutdown_event().await {
913+
error!("Failed to send shutdown event to processor: {}", e);
914+
}
888915
println!("Exiting: {shutdown_reason}, deadline: {deadline_ms}");
889916
}
890917
Err(ref err) => {
@@ -952,7 +979,7 @@ fn start_trace_agent(
952979
config: &Arc<Config>,
953980
api_key_factory: &Arc<ApiKeyFactory>,
954981
tags_provider: &Arc<TagProvider>,
955-
invocation_processor: Arc<TokioMutex<InvocationProcessor>>,
982+
invocation_processor_handle: InvocationProcessorHandle,
956983
appsec_processor: Option<Arc<TokioMutex<AppSecProcessor>>>,
957984
) -> (
958985
Sender<SendDataBuilderInfo>,
@@ -1015,7 +1042,7 @@ fn start_trace_agent(
10151042
stats_aggregator,
10161043
stats_processor,
10171044
proxy_aggregator,
1018-
invocation_processor,
1045+
invocation_processor_handle,
10191046
appsec_processor,
10201047
Arc::clone(tags_provider),
10211048
stats_concentrator_handle.clone(),
@@ -1227,7 +1254,7 @@ fn start_otlp_agent(
12271254
fn start_api_runtime_proxy(
12281255
config: &Arc<Config>,
12291256
aws_config: Arc<AwsConfig>,
1230-
invocation_processor: &Arc<TokioMutex<InvocationProcessor>>,
1257+
invocation_processor_handle: &InvocationProcessorHandle,
12311258
appsec_processor: Option<&Arc<TokioMutex<AppSecProcessor>>>,
12321259
propagator: Arc<DatadogCompositePropagator>,
12331260
) -> Option<CancellationToken> {
@@ -1236,11 +1263,10 @@ fn start_api_runtime_proxy(
12361263
return None;
12371264
}
12381265

1239-
let invocation_processor = invocation_processor.clone();
12401266
let appsec_processor = appsec_processor.map(Arc::clone);
12411267
interceptor::start(
12421268
aws_config,
1243-
invocation_processor,
1269+
invocation_processor_handle.clone(),
12441270
appsec_processor,
12451271
propagator,
12461272
)

‎bottlecap/src/lifecycle/invocation/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use tracing::debug;
99

1010
pub mod context;
1111
pub mod processor;
12+
pub mod processor_service;
1213
pub mod span_inferrer;
1314
pub mod triggers;
1415

0 commit comments

Comments
 (0)