From c736c712e53d1e352a8c182507656e22413661bc Mon Sep 17 00:00:00 2001
From: Pat Hickey <phickey@fastly.com>
Date: Fri, 19 Jan 2024 11:07:38 -0800
Subject: [PATCH] wasmtime: add EngineWeak which has ::upgrade, created by
 Engine::weak (#7797)

* wasmtime: add EngineWeak which has ::upgrade, created by Engine::weak

Engine is, internally, just an Arc<EngineInner>, so this is trivial to implement -
EngineWeak is a Weak<EngineInner>.

This behavior is desirable because `Engine::increment_epoch` typically
happens in a worker thread, which in turn requires additional machinery
to discard the `Engine` once it is no longer needed. If instead the
worker thread holds an `EngineWeak`, it can stop ticking when all
consumers of the `Engine` have dropped it. This has been documented as a
suggestion in `increment_epoch`.

For an example of additional machinery which is simplified by this change, see https://github.com/fastly/Viceroy/blob/25edee0700ec0b20b1b56db0a3a8d6f090397b3a/lib/src/execute.rs#L108-L116)

Co-authored-by: John Van Enk <vanenkj@gmail.com>

* add a test

---------

Co-authored-by: John Van Enk <vanenkj@gmail.com>
---
 crates/wasmtime/src/engine.rs | 39 +++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/crates/wasmtime/src/engine.rs b/crates/wasmtime/src/engine.rs
index bc754e0c01..1d91ef69ab 100644
--- a/crates/wasmtime/src/engine.rs
+++ b/crates/wasmtime/src/engine.rs
@@ -180,6 +180,11 @@ impl Engine {
     /// for an introduction to epoch-based interruption and pointers
     /// to the other relevant methods.
     ///
+    /// When performing `increment_epoch` in a separate thread, consider using
+    /// [`Engine::weak`] to hold an [`EngineWeak`] and performing
+    /// [`EngineWeak::upgrade`] on each tick, so that the epoch ticking thread
+    /// does not keep an [`Engine`] alive longer than any of its consumers.
+    ///
     /// ## Signal Safety
     ///
     /// This method is signal-safe: it does not make any syscalls, and
@@ -654,6 +659,13 @@ impl Engine {
     pub fn detect_precompiled_file(&self, path: impl AsRef<Path>) -> Result<Option<Precompiled>> {
         serialization::detect_precompiled_file(path)
     }
+
+    /// Take a weak reference to this engine.
+    pub fn weak(&self) -> EngineWeak {
+        EngineWeak {
+            inner: Arc::downgrade(&self.inner),
+        }
+    }
 }
 
 impl Default for Engine {
@@ -671,6 +683,20 @@ pub enum Precompiled {
     Component,
 }
 
+/// A weak reference to an [`Engine`].
+#[derive(Clone)]
+pub struct EngineWeak {
+    inner: std::sync::Weak<EngineInner>,
+}
+
+impl EngineWeak {
+    /// Upgrade this weak reference into an [`Engine`]. Returns `None` if
+    /// strong references (the [`Engine`] type itself) no longer exist.
+    pub fn upgrade(&self) -> Option<Engine> {
+        std::sync::Weak::upgrade(&self.inner).map(|inner| Engine { inner })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::{
@@ -787,4 +813,17 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn engine_weak_upgrades() {
+        let engine = Engine::default();
+        let weak = engine.weak();
+        weak.upgrade()
+            .expect("engine is still alive, so weak reference can upgrade");
+        drop(engine);
+        assert!(
+            weak.upgrade().is_none(),
+            "engine was dropped, so weak reference cannot upgrade"
+        );
+    }
 }