From 3a27069a15e4859e7aa3cd177cf932b7b411f14e Mon Sep 17 00:00:00 2001
From: Mikhail Zabaluev <mikhail.zabaluev@movementlabs.xyz>
Date: Wed, 16 Oct 2024 15:52:09 +0300
Subject: [PATCH] fix(tracing): work around another shutdown problem

---
 Cargo.lock                    |  1 +
 util/tracing/Cargo.toml       |  1 +
 util/tracing/src/telemetry.rs | 22 ++++++++++++++++++----
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2cf8a1a69..b047f80a0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -9680,6 +9680,7 @@ dependencies = [
  "opentelemetry-otlp",
  "opentelemetry-semantic-conventions",
  "opentelemetry_sdk",
+ "tokio",
  "tracing",
  "tracing-opentelemetry",
  "tracing-subscriber 0.3.18",
diff --git a/util/tracing/Cargo.toml b/util/tracing/Cargo.toml
index 6c9fad7f8..eceb5d5a7 100644
--- a/util/tracing/Cargo.toml
+++ b/util/tracing/Cargo.toml
@@ -16,6 +16,7 @@ opentelemetry = { workspace = true }
 opentelemetry_sdk = { workspace = true, features = ["rt-tokio"] }
 opentelemetry-otlp = { workspace = true }
 opentelemetry-semantic-conventions = { workspace = true }
+tokio = { workspace = true }
 tracing = { workspace = true }
 tracing-opentelemetry = { workspace = true }
 tracing-subscriber = { workspace = true }
diff --git a/util/tracing/src/telemetry.rs b/util/tracing/src/telemetry.rs
index a34e64a44..8a3286700 100644
--- a/util/tracing/src/telemetry.rs
+++ b/util/tracing/src/telemetry.rs
@@ -11,8 +11,9 @@ use crate::Config;
 use opentelemetry::{trace::TracerProvider as _, KeyValue};
 use opentelemetry_otlp::WithExportConfig as _;
 use opentelemetry_sdk::trace::{Config as TraceConfig, TracerProvider};
-use opentelemetry_sdk::{runtime, Resource};
+use opentelemetry_sdk::{runtime::Tokio as TokioRuntimeSelector, Resource};
 use opentelemetry_semantic_conventions::resource::{SERVICE_NAME, SERVICE_VERSION};
+use tokio::runtime;
 use tracing::{error, Level, Subscriber};
 use tracing_opentelemetry::OpenTelemetryLayer;
 use tracing_subscriber::filter;
@@ -28,12 +29,25 @@ pub struct ScopeGuard(Option<TracerProvider>);
 
 impl Drop for ScopeGuard {
 	fn drop(&mut self) {
-		if let Some(tracer_provider) = &self.0 {
-			// Make sure all batched traces are exported.
+		fn shutdown_provider(tracer_provider: &TracerProvider) {
 			if let Err(e) = tracer_provider.shutdown() {
 				error!("OpenTelemetry tracer provider shutdown failed: {e}");
 			}
 		}
+
+		if let Some(tracer_provider) = &self.0 {
+			// Make sure all batched traces are exported.
+			if let Ok(handle) = runtime::Handle::try_current() {
+				// Can't call shutdown in async context due to
+				// https://github.com/open-telemetry/opentelemetry-rust/issues/2047#issuecomment-2416480148
+				let tracer_provider = tracer_provider.clone();
+				handle.spawn_blocking(move || {
+					shutdown_provider(&tracer_provider);
+				});
+			} else {
+				shutdown_provider(tracer_provider);
+			}
+		}
 	}
 }
 
@@ -61,7 +75,7 @@ where
 				KeyValue::new(SERVICE_NAME, service_name),
 				KeyValue::new(SERVICE_VERSION, service_version),
 			])))
-			.install_batch(runtime::Tokio)?;
+			.install_batch(TokioRuntimeSelector)?;
 		let layer = OpenTelemetryLayer::new(provider.tracer("movement"))
 			.with_filter(filter::Targets::new().with_target("movement_telemetry", Level::INFO));
 		(Some(provider), Some(layer))