dag: do not resolve id or names remotely inside a nonblocking context

Summary:
Otherwise it might panic (ex. calling into tokio without entering a tokio
runtime). This can happen in, for example, `Debug::fmt(&IdStaticSet, ...)`.

Reviewed By: sfilipco

Differential Revision: D27581487

fbshipit-source-id: feec53e088706adcc6710afcf24fa70598f886cf

eden/scm/lib/dag/src/namedag.rs

@@ -39,6 +39,7 @@ use crate::ops::PrefixLookup;
 use crate::ops::ToIdSet;
 use crate::ops::TryClone;
 use crate::protocol;
+use crate::protocol::is_remote_protocol_disabled;
 use crate::protocol::AncestorPath;
 use crate::protocol::Process;
 use crate::protocol::RemoteIdConvertProtocol;
@@ -56,6 +57,7 @@ use parking_lot::Mutex;
 use parking_lot::RwLock;
 use std::collections::{HashMap, HashSet};
 use std::fmt;
+use std::io;
 use std::ops::Deref;
 use std::sync::Arc;
 
@@ -520,6 +522,13 @@ where
         if names.is_empty() {
             return Ok(Vec::new());
         }
+        if is_remote_protocol_disabled() {
+            return Err(io::Error::new(
+                io::ErrorKind::WouldBlock,
+                "resolving vertexes remotely disabled",
+            )
+            .into());
+        }
         if names.len() < 30 {
             tracing::debug!("resolve names {:?} remotely", &names);
         } else {
@@ -550,6 +559,13 @@ where
         if ids.is_empty() {
             return Ok(Vec::new());
         }
+        if is_remote_protocol_disabled() {
+            return Err(io::Error::new(
+                io::ErrorKind::WouldBlock,
+                "resolving ids remotely disabled",
+            )
+            .into());
+        }
         if ids.len() < 30 {
             tracing::debug!("resolve ids {:?} remotely", &ids);
         } else {

eden/scm/lib/dag/src/nameset/id_lazy.rs

@@ -11,6 +11,7 @@ use super::BoxVertexStream;
 use super::{AsyncNameSetQuery, Hints};
 use crate::ops::DagAlgorithm;
 use crate::ops::IdConvert;
+use crate::protocol::disable_remote_protocol;
 use crate::Group;
 use crate::Id;
 use crate::IdSet;
@@ -154,7 +155,9 @@ impl fmt::Debug for IdLazySet {
         f.debug_list()
             .entries(inner.visited.iter().take(limit).map(|&id| DebugId {
                 id,
-                name: non_blocking_result(self.map.vertex_name(id)).ok(),
+                name: disable_remote_protocol(|| {
+                    non_blocking_result(self.map.vertex_name(id)).ok()
+                }),
             }))
             .finish()?;
         let remaining = inner.visited.len().max(limit) - limit;

eden/scm/lib/dag/src/nameset/id_static.rs

@@ -10,6 +10,7 @@ use super::BoxVertexStream;
 use super::{AsyncNameSetQuery, Hints};
 use crate::ops::DagAlgorithm;
 use crate::ops::IdConvert;
+use crate::protocol::disable_remote_protocol;
 use crate::Group;
 use crate::IdSpan;
 use crate::Result;
@@ -136,8 +137,12 @@ impl fmt::Debug for IdStaticSet {
         f.debug_list()
             .entries(spans.iter().take(limit).map(|span| DebugSpan {
                 span: *span,
-                low_name: non_blocking_result(self.map.vertex_name(span.low)).ok(),
-                high_name: non_blocking_result(self.map.vertex_name(span.high)).ok(),
+                low_name: disable_remote_protocol(|| {
+                    non_blocking_result(self.map.vertex_name(span.low)).ok()
+                }),
+                high_name: disable_remote_protocol(|| {
+                    non_blocking_result(self.map.vertex_name(span.high)).ok()
+                }),
             }))
             .finish()?;
         match spans.len().max(limit) - limit {

eden/scm/lib/dag/src/protocol.rs

@@ -33,7 +33,9 @@ use futures::stream::StreamExt;
 use futures::stream::TryStreamExt;
 use futures::TryFutureExt;
 use serde::{Deserialize, Serialize};
+use std::cell::RefCell;
 use std::fmt;
+use std::thread_local;
 
 // Request and Response structures -------------------------------------------
 
@@ -413,3 +415,27 @@ impl<'a, DagStore: IdDagStore> Process<ResponseIdNamePair, ()>
         Ok(())
     }
 }
+
+// Disable remote protocol temporarily ---------------------------------------
+// This can be useful for Debug::fmt to disable remote fetching which might
+// panic (ex. calling tokio without tokio runtime) when executing futures
+// via nonblocking.
+
+thread_local! {
+    static NON_BLOCKING_DEPTH: RefCell<usize> = RefCell::new(0);
+}
+
+/// Check if the current future is running inside a "non-blocking" block.
+pub(crate) fn disable_remote_protocol<F, R>(f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    NON_BLOCKING_DEPTH.with(|v| *v.borrow_mut() += 1);
+    let result = f();
+    NON_BLOCKING_DEPTH.with(|v| *v.borrow_mut() -= 1);
+    result
+}
+
+pub(crate) fn is_remote_protocol_disabled() -> bool {
+    NON_BLOCKING_DEPTH.with(|v| *v.borrow() != 0)
+}