Clean up segment path by using Rc<TempDir> (risc0#1421)

At this time, the zkVM fails to clean up segment information. Fix this
by reference-counting each segment file if it uses a tempdir. We will
let the user manage the lifetime of the segment if they designate a
specific path to store segments. Some of this is based off of Austin's
work. In this implementation, I used `Arc` because `SegmentRef`
implements the `send` trait.

---------

Co-authored-by: Austin Abell <[email protected]>
Co-authored-by: Frank Laub <[email protected]>

risc0/zkvm/Cargo.toml

@@ -104,6 +104,7 @@ client = [
   "dep:prost",
   "dep:prost-build",
   "dep:protobuf-src",
+  "dep:tempfile",
   "std",
 ]
 cuda = [

risc0/zkvm/src/host/api/server.rs

@@ -43,7 +43,6 @@ pub struct Server {
 #[derive(Clone, Serialize, Deserialize)]
 struct EmptySegmentRef;
 
-#[typetag::serde]
 impl SegmentRef for EmptySegmentRef {
     fn resolve(&self) -> Result<Segment> {
         Err(anyhow!("Segment resolution not supported"))

risc0/zkvm/src/host/client/env.rs

@@ -21,13 +21,15 @@ use std::{
     mem,
     path::{Path, PathBuf},
     rc::Rc,
+    sync::Arc,
 };
 
 use anyhow::Result;
 use bytemuck::Pod;
 use bytes::Bytes;
 use risc0_zkvm_platform::{self, fileno};
 use serde::Serialize;
+use tempfile::TempDir;
 
 use crate::serde::to_vec;
 use crate::{
@@ -69,6 +71,13 @@ pub(crate) struct Assumptions {
     pub(crate) accessed: Vec<Assumption>,
 }
 
+#[allow(dead_code)]
+#[derive(Clone)]
+pub enum SegmentPath {
+    TempDir(Arc<TempDir>),
+    Path(PathBuf),
+}
+
 /// The [crate::Executor] is configured from this object.
 ///
 /// The executor environment holds configuration details that inform how the
@@ -84,7 +93,7 @@ pub struct ExecutorEnv<'a> {
     pub(crate) input: Vec<u8>,
     pub(crate) trace: Vec<Rc<RefCell<dyn TraceCallback + 'a>>>,
     pub(crate) assumptions: Rc<RefCell<Assumptions>>,
-    pub(crate) segment_path: Option<PathBuf>,
+    pub(crate) segment_path: Option<SegmentPath>,
     pub(crate) pprof_out: Option<PathBuf>,
 }
 
@@ -352,7 +361,7 @@ impl<'a> ExecutorEnvBuilder<'a> {
 
     /// Set the path where segments will be stored.
     pub fn segment_path<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
-        self.inner.segment_path = Some(path.as_ref().to_path_buf());
+        self.inner.segment_path = Some(SegmentPath::Path(path.as_ref().to_path_buf()));
         self
     }
 

risc0/zkvm/src/host/server/exec/executor.rs

@@ -14,7 +14,7 @@
 
 //! This module implements the Executor.
 
-use std::{cell::RefCell, fmt::Debug, io::Write, mem, rc::Rc};
+use std::{cell::RefCell, fmt::Debug, io::Write, mem, rc::Rc, sync::Arc};
 
 use addr2line::{
     fallible_iterator::FallibleIterator,
@@ -52,7 +52,7 @@ use super::{monitor::MemoryMonitor, profiler::Profiler, syscall::SyscallTable};
 use crate::{
     align_up,
     host::{
-        client::exec::TraceEvent,
+        client::{env::SegmentPath, exec::TraceEvent},
         receipt::Assumption,
         server::opcode::{MajorType, OpCode},
     },
@@ -238,7 +238,7 @@ impl<'a> ExecutorImpl<'a> {
     /// of the execution.
     pub fn run(&mut self) -> Result<Session> {
         if self.env.segment_path.is_none() {
-            self.env.segment_path = Some(tempdir()?.into_path());
+            self.env.segment_path = Some(SegmentPath::TempDir(Arc::new(tempdir()?)));
         }
 
         let path = self.env.segment_path.clone().unwrap();

risc0/zkvm/src/host/server/session.rs

@@ -27,8 +27,9 @@ use risc0_zkvm_platform::WORD_SIZE;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    host::server::exec::executor::SyscallRecord, sha::Digest, Assumption, Assumptions, ExitCode,
-    Journal, Output, ReceiptClaim,
+    host::{client::env::SegmentPath, server::exec::executor::SyscallRecord},
+    sha::Digest,
+    Assumption, Assumptions, ExitCode, Journal, Output, ReceiptClaim,
 };
 
 #[derive(Clone, Default, Serialize, Deserialize, Debug)]
@@ -43,7 +44,6 @@ pub struct PageFaults {
 /// initial memory image (which includes the starting PC) and proceeds until
 /// either a sys_halt or a sys_pause syscall is encountered. This record is
 /// stored as a vector of [Segment]s.
-#[derive(Serialize, Deserialize)]
 pub struct Session {
     /// The constituent [Segment]s of the Session. The final [Segment] will have
     /// an [ExitCode] of [Halted](ExitCode::Halted), [Paused](ExitCode::Paused),
@@ -64,7 +64,6 @@ pub struct Session {
     pub assumptions: Vec<Assumption>,
 
     /// The hooks to be called during the proving phase.
-    #[serde(skip)]
     pub hooks: Vec<Box<dyn SessionEvents>>,
 
     /// The number of user cycles without any overhead for continuations or po2
@@ -87,7 +86,6 @@ pub struct Session {
 /// This allows implementors to determine the best way to represent this in an
 /// pluggable manner. See the [SimpleSegmentRef] for a very basic
 /// implmentation.
-#[typetag::serde(tag = "type")]
 pub trait SegmentRef: Send {
     /// Resolve this reference into an actual [Segment].
     fn resolve(&self) -> Result<Segment>;
@@ -300,7 +298,6 @@ pub struct SimpleSegmentRef {
     segment: Segment,
 }
 
-#[typetag::serde]
 impl SegmentRef for SimpleSegmentRef {
     fn resolve(&self) -> Result<Segment> {
         Ok(self.segment.clone())
@@ -322,12 +319,11 @@ impl SimpleSegmentRef {
 /// There is an example of using [FileSegmentRef] in our [EVM example][1]
 ///
 /// [1]: https://github.com/risc0/risc0/blob/main/examples/zkevm-demo/src/main.rs
-#[derive(Clone, Serialize, Deserialize)]
 pub struct FileSegmentRef {
     path: PathBuf,
+    _dir: SegmentPath,
 }
 
-#[typetag::serde]
 impl SegmentRef for FileSegmentRef {
     fn resolve(&self) -> Result<Segment> {
         let contents = fs::read(&self.path)?;
@@ -336,13 +332,25 @@ impl SegmentRef for FileSegmentRef {
     }
 }
 
+impl SegmentPath {
+    pub(crate) fn path(&self) -> &Path {
+        match self {
+            Self::TempDir(dir) => dir.path(),
+            Self::Path(path) => path.as_path(),
+        }
+    }
+}
+
 impl FileSegmentRef {
     /// Construct a [FileSegmentRef]
     ///
     /// This builds a FileSegmentRef that stores `segment` in a file at `path`.
-    pub fn new(segment: &Segment, path: &Path) -> Result<Self> {
-        let path = path.join(format!("{}.bincode", segment.index));
+    pub fn new(segment: &Segment, dir: &SegmentPath) -> Result<Self> {
+        let path = dir.path().join(format!("{}.bincode", segment.index));
         fs::write(&path, bincode::serialize(&segment)?)?;
-        Ok(Self { path })
+        Ok(Self {
+            path,
+            _dir: dir.clone(),
+        })
     }
 }