Fix deadlocks (#34)

Author: osipovartem

crates/catalog/src/lib.rs

@@ -1,6 +1,9 @@
 use error::Result;
+use futures::TryFutureExt;
+use futures::executor::block_on;
 use snafu::ResultExt;
-use tokio::runtime::Builder;
+use tokio::runtime::{Builder, Handle, RuntimeFlavor};
+use tokio::task;
 
 #[allow(clippy::module_inception)]
 pub mod catalog;
@@ -34,6 +37,23 @@ where
     .unwrap_or_else(|_| error::ThreadPanickedWhileExecutingFutureSnafu.fail()?)
 }
 
+fn block_on_without_deadlock<F>(future: F) -> F::Output
+where
+    F: Future + Send + 'static,
+    F::Output: Send + 'static,
+{
+    match Handle::try_current() {
+        Ok(handle) => match handle.runtime_flavor() {
+            RuntimeFlavor::CurrentThread => block_on(
+                task::spawn_blocking(|| block_on(future))
+                    .unwrap_or_else(|err| std::panic::resume_unwind(err.into_panic())),
+            ),
+            _ => task::block_in_place(|| handle.block_on(future)),
+        },
+        Err(_) => block_on(future),
+    }
+}
+
 pub mod test_utils {
     use datafusion::arrow::array::{ArrayRef, RecordBatch};
     use datafusion::arrow::compute::{

crates/catalog/src/schema.rs

@@ -1,18 +1,16 @@
-use crate::df_error;
 use crate::snowflake_table::CaseInsensitiveTable;
 use crate::table::{CachingTable, IcebergTableBuilder};
+use crate::{block_on_without_deadlock, df_error};
 use async_trait::async_trait;
 use dashmap::DashMap;
 use datafusion::catalog::{SchemaProvider, TableProvider};
 use datafusion_common::DataFusionError;
 use datafusion_expr::TableType;
 use datafusion_iceberg::DataFusionTable;
-use futures::executor::block_on;
 use iceberg_rust::catalog::Catalog;
 use iceberg_rust::catalog::tabular::Tabular as IcebergTabular;
 use iceberg_rust_spec::identifier::Identifier;
 use snafu::ResultExt;
-use snafu::futures::TryFutureExt;
 use std::any::Any;
 use std::sync::Arc;
 
@@ -89,11 +87,15 @@ impl SchemaProvider for CachingSchema {
             && let Some(iceberg_builder) = table.as_any().downcast_ref::<IcebergTableBuilder>()
             && table.table_type() != TableType::View
         {
-            let ident = Identifier::new(std::slice::from_ref(&self.name), &name);
-            block_on(async move {
-                let mut builder = iceberg_builder.builder.clone();
+            let catalog = Arc::clone(catalog);
+            let mut builder = iceberg_builder.builder.clone();
+            let namespace = vec![self.name.clone()];
+            let table_name = name.clone();
+
+            block_on_without_deadlock(async move {
+                let ident = Identifier::new(&namespace, &table_name);
                 let iceberg_table = builder
-                    .build(ident.namespace(), catalog.clone())
+                    .build(ident.namespace(), catalog)
                     .await
                     .context(df_error::IcebergSnafu)?;
                 let tabular = IcebergTabular::Table(iceberg_table);
@@ -121,8 +123,17 @@ impl SchemaProvider for CachingSchema {
         if let Some((_, caching_table)) = table {
             if caching_table.table_type() != TableType::View {
                 if let Some(catalog) = &self.iceberg_catalog {
-                    let ident = Identifier::new(std::slice::from_ref(&self.name), name);
-                    block_on(catalog.drop_table(&ident).context(df_error::IcebergSnafu))?;
+                    let catalog = Arc::clone(catalog);
+                    let namespace = vec![self.name.clone()];
+                    let table_name = name.to_string();
+
+                    block_on_without_deadlock(async move {
+                        let ident = Identifier::new(&namespace, &table_name);
+                        catalog
+                            .drop_table(&ident)
+                            .await
+                            .context(df_error::IcebergSnafu)
+                    })?;
                 } else {
                     return self.schema.deregister_table(name);
                 }