From c6f78d5f27a34ef5b03bfc6d0aaee52cf459c2d1 Mon Sep 17 00:00:00 2001
From: Niko Matsakis <niko@alum.mit.edu>
Date: Thu, 11 Oct 2018 16:48:20 -0400
Subject: [PATCH] have queries freeze the revision

---
 src/derived.rs |  9 ++++++---
 src/runtime.rs | 19 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/derived.rs b/src/derived.rs
index e827f71b..f61858b7 100644
--- a/src/derived.rs
+++ b/src/derived.rs
@@ -171,7 +171,11 @@ where
         key: &Q::Key,
         descriptor: &DB::QueryDescriptor,
     ) -> Result<StampedValue<Q::Value>, CycleDetected> {
-        let revision_now = db.salsa_runtime().current_revision();
+        let runtime = db.salsa_runtime();
+
+        let _read_lock = runtime.freeze_revision();
+
+        let revision_now = runtime.current_revision();
 
         debug!(
             "{:?}({:?}): invoked at {:?}",
@@ -239,7 +243,6 @@ where
 
         // Query was not previously executed, or value is potentially
         // stale, or value is absent. Let's execute!
-        let runtime = db.salsa_runtime();
         let (mut stamped_value, inputs) = runtime.execute_query_implementation(descriptor, || {
             debug!("{:?}({:?}): executing query", Q::default(), key);
 
@@ -254,7 +257,7 @@ where
         // not mutate the "inputs" to the query system. Sanity check
         // that assumption here, at least to the best of our ability.
         assert_eq!(
-            db.salsa_runtime().current_revision(),
+            runtime.current_revision(),
             revision_now,
             "revision altered during query execution",
         );
diff --git a/src/runtime.rs b/src/runtime.rs
index 34a0ed1f..ad08afbf 100644
--- a/src/runtime.rs
+++ b/src/runtime.rs
@@ -2,7 +2,7 @@ use crate::Database;
 use crate::Query;
 use crate::QueryFunction;
 use log::debug;
-use parking_lot::{RwLock, RwLockUpgradableReadGuard};
+use parking_lot::{RwLock, RwLockReadGuard, RwLockUpgradableReadGuard};
 use rustc_hash::FxHasher;
 use std::cell::RefCell;
 use std::fmt::Write;
@@ -69,6 +69,23 @@ where
         self.increment_revision();
     }
 
+    /// Indicates that a derived query has begun to execute; if this is the
+    /// first derived query on this thread, then acquires a read-lock on the
+    /// runtime to prevent us from moving to a new revision until that query
+    /// completes.
+    ///
+    /// (However, if other threads invoke `increment_revision`, then
+    /// the current revision may be considered cancelled, which can be
+    /// observed through `is_current_revision_canceled`.)
+    pub(crate) fn freeze_revision(&self) -> Option<RwLockReadGuard<'_, ()>> {
+        let local_state = self.local_state.borrow();
+        if local_state.query_stack.is_empty() {
+            Some(self.shared_state.revision_lock.read())
+        } else {
+            None
+        }
+    }
+
     /// Read current value of the revision counter.
     pub(crate) fn current_revision(&self) -> Revision {
         Revision {