[Serverless Mini Agent] Run in Azure Spring Apps (#547)

* add azure spring app environment type for serverless mini agent * add config.statsd_port to mini agent info endpoint * update mini agent trace endpoint status code and response body to work with java tracer * use different environment variable to identify azure spring apps * only update http response for success responses to traces endpoint * address lint errors * updates comment and formatting * update trace-mini-agent description to include Azure Spring Apps * fix formatting
DataDog · Jul 31, 2024 · 6eaa738 · 6eaa738
1 parent f148c87
commit 6eaa738
Show file tree

Hide file tree

Showing 7 changed files with 36 additions and 4 deletions.
diff --git a/trace-mini-agent/Cargo.toml b/trace-mini-agent/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "datadog-trace-mini-agent"
-description = "A subset of the trace agent that is shipped alongside tracers in a few serverless use cases (Google Cloud Functions and Azure Functions)"
+description = "A subset of the trace agent that is shipped alongside tracers in a few serverless use cases (Google Cloud Functions, Azure Functions, and Azure Spring Apps)"
 edition.workspace = true
 version.workspace = true
 rust-version.workspace = true

diff --git a/trace-mini-agent/src/env_verifier.rs b/trace-mini-agent/src/env_verifier.rs
@@ -140,6 +140,9 @@ impl EnvVerifier for ServerlessEnvVerifier {
                     .verify_gcp_environment_or_exit(verify_env_timeout)
                     .await;
             }
+            trace_utils::EnvironmentType::AzureSpringApp => {
+                trace_utils::MiniAgentMetadata::default()
+            }
             trace_utils::EnvironmentType::LambdaFunction => {
                 trace_utils::MiniAgentMetadata::default()
             }

diff --git a/trace-mini-agent/src/http_utils.rs b/trace-mini-agent/src/http_utils.rs
@@ -31,6 +31,26 @@ pub fn log_and_create_http_response(
     Response::builder().status(status).body(Body::from(body))
 }
 
+/// Does two things:
+/// 1. Logs the given message
+/// 2. Returns the rate_by_service map to use to set the sampling priority in the body of JSON
+///    response with the given status code.
+///
+/// Response body format:
+/// {
+///     "rate_by_service": {
+///         "service:,env:":1
+///     }
+/// }
+pub fn log_and_create_traces_success_http_response(
+    message: &str,
+    status: StatusCode,
+) -> http::Result<Response<Body>> {
+    info!("{message}");
+    let body = json!({"rate_by_service":{"service:,env:":1}}).to_string();
+    Response::builder().status(status).body(Body::from(body))
+}
+
 /// Takes a request's header map, and verifies that the "content-length" header is present, valid,
 /// and less than the given max_content_length.
 ///

diff --git a/trace-mini-agent/src/mini_agent.rs b/trace-mini-agent/src/mini_agent.rs
@@ -192,6 +192,9 @@ impl MiniAgent {
                     INFO_ENDPOINT_PATH
                 ],
                 "client_drop_p0s": true,
+                "config": {
+                    "statsd_port": MINI_AGENT_PORT
+                }
             }
         );
         Response::builder()

diff --git a/trace-mini-agent/src/trace_processor.rs b/trace-mini-agent/src/trace_processor.rs
@@ -17,7 +17,7 @@ use datadog_trace_utils::tracer_payload::TraceEncoding;
 
 use crate::{
     config::Config,
-    http_utils::{self, log_and_create_http_response},
+    http_utils::{self, log_and_create_http_response, log_and_create_traces_success_http_response},
 };
 
 #[async_trait]
@@ -105,9 +105,9 @@ impl TraceProcessor for ServerlessTraceProcessor {
         // send trace payload to our trace flusher
         match tx.send(send_data).await {
             Ok(_) => {
-                return log_and_create_http_response(
+                return log_and_create_traces_success_http_response(
                     "Successfully buffered traces to be flushed.",
-                    StatusCode::ACCEPTED,
+                    StatusCode::OK,
                 );
             }
             Err(err) => {

diff --git a/trace-utils/src/config_utils.rs b/trace-utils/src/config_utils.rs
@@ -25,6 +25,10 @@ pub fn read_cloud_env() -> Option<(String, trace_utils::EnvironmentType)> {
         // Set by Azure Functions
         return Some((res, trace_utils::EnvironmentType::AzureFunction));
     }
+    if let Ok(res) = env::var("ASCSVCRT_SPRING__APPLICATION__NAME") {
+        // Set by Azure Spring Apps
+        return Some((res, trace_utils::EnvironmentType::AzureSpringApp));
+    }
     None
 }
 

diff --git a/trace-utils/src/trace_utils.rs b/trace-utils/src/trace_utils.rs
@@ -429,6 +429,7 @@ pub fn set_serverless_root_span_tags(
     let origin_tag = match env_type {
         EnvironmentType::CloudFunction => "cloudfunction",
         EnvironmentType::AzureFunction => "azurefunction",
+        EnvironmentType::AzureSpringApp => "azurespringapp",
         EnvironmentType::LambdaFunction => "lambda", // historical reasons
     };
     span.meta
@@ -451,6 +452,7 @@ fn update_tracer_top_level(span: &mut pb::Span) {
 pub enum EnvironmentType {
     CloudFunction,
     AzureFunction,
+    AzureSpringApp,
     LambdaFunction,
 }