fix snapshotting not working for sources, add example of stateful source

Author: Damion Werner

malstrom-core/Cargo.toml

@@ -37,3 +37,7 @@ slatedb = ["dep:slatedb", "dep:object_store", "dep:tokio-stream"]
 [[example]]
 name = "slatedb_backend"
 required-features = ["slatedb"]
+
+[[example]]
+name = "slatedb_backend_failing"
+required-features = ["slatedb"]

malstrom-core/examples/slatedb_backend_failing.rs

@@ -0,0 +1,106 @@
+//! Using SlateDB as a persistence backend
+use malstrom::keyed::partitioners::rendezvous_select;
+use malstrom::operators::*;
+use malstrom::sinks::{StatelessSink, StdOutSink};
+use malstrom::sources::{StatefulSource, StatefulSourceImpl, StatefulSourcePartition};
+use malstrom::{
+    runtime::SingleThreadRuntime,
+    snapshot::SlateDbBackend,
+    worker::StreamProvider,
+};
+use object_store::{local::LocalFileSystem, path::Path};
+use std::sync::Arc;
+use std::thread::sleep;
+use std::time::{Duration, Instant};
+
+fn main() {
+    let filesystem = LocalFileSystem::new();
+    let persistence = SlateDbBackend::new(Arc::new(filesystem), Path::from("/tmp")).unwrap();
+
+    loop {
+        let job = SingleThreadRuntime::builder()
+            .persistence(persistence.clone())
+            .snapshots(Duration::from_secs(1))
+            .build(build_dataflow);
+        let thread = std::thread::spawn(move || {
+            job.execute().unwrap()
+        });
+        match thread.join() {
+            Ok(_) => return,
+            Err(_) => {
+                println!("Restarting worker");
+                continue;
+            },
+        }
+    }
+}
+
+fn build_dataflow(provider: &mut dyn StreamProvider) {
+    let start_time = Instant::now();
+    let fail_interval = Duration::from_secs(10);
+    provider
+        .new_stream()
+        .source(
+            "iter-source",
+            StatefulSource::new(StatefulNumberSource(0)),
+        )
+        .key_distribute("key-by-value", |x| x.value & 1 == 1, rendezvous_select)
+        .stateful_map("sum", |_key, value, state: i32| {
+            let state = state + value;
+            (state, Some(state))
+        })
+        .inspect("expensive-operation", |_msg, _ctx| {
+            // we need this to not overflow the sum before "crashing"
+            sleep(Duration::from_millis(100))
+        })
+        .inspect("fail-random", move |_msg, _ctx| {
+            if Instant::now().duration_since(start_time) > fail_interval {
+                panic!("Oh no!")
+            }
+        })
+        .sink("stdout", StatelessSink::new(StdOutSink));
+}
+
+struct StatefulNumberSource(i32);
+
+impl StatefulSourceImpl<i32, i32> for StatefulNumberSource {
+    type Part = ();
+    type PartitionState = i32;
+    type SourcePartition = Self;
+
+    fn list_parts(&self) -> Vec<Self::Part> {
+        vec![()]
+    }
+
+    fn build_part(
+        &mut self,
+        _part: &Self::Part,
+        part_state: Option<Self::PartitionState>,
+    ) -> Self::SourcePartition {
+        println!("Build with {part_state:?}");
+        Self(part_state.unwrap_or_default())
+    }
+}
+
+impl StatefulSourcePartition<i32, i32> for StatefulNumberSource {
+    type PartitionState = i32;
+
+    fn poll(&mut self) -> Option<(i32, i32)> {
+        let out = Some((self.0, self.0));
+        self.0 += 1;
+        out
+    }
+
+    fn is_finished(&mut self) -> bool {
+        false
+    }
+
+    fn snapshot(&self) -> Self::PartitionState {
+        println!("SNAPSHOTTING SOURCE");
+        self.0
+    }
+
+    fn collect(self) -> Self::PartitionState {
+        self.0
+    }
+}

malstrom-core/src/sources/stateful.rs

@@ -177,7 +177,7 @@ where
     ) -> Self {
         let comm_clients =
             ctx.create_all_communication_clients::<PartitionFinished<Builder::Part>>();
-        Self {
+        let mut this = Self {
             partitions: IndexMap::new(),
             part_builder,
             all_partitions,
@@ -186,7 +186,14 @@ where
             // times should not be an issue
             max_t: Some(T::MAX),
             _phantom: PhantomData,
+        };
+
+        if let Some(state) = ctx.load_state::<IndexMap<Builder::Part, Builder::PartitionState>>() {
+            for (k, v) in state.into_iter() {
+                this.add_partition(k, Some(v));
+            }
         }
+        this
     }
 
     fn add_partition(&mut self, part: Builder::Part, part_state: Option<Builder::PartitionState>) {
@@ -264,7 +271,7 @@ where
         _output: &mut Output<Builder::Part, VO, TO>,
         ctx: &mut OperatorContext,
     ) {
-        let state: Vec<_> = self
+        let state: IndexMap<Builder::Part, Builder::PartitionState> = self
             .partitions
             .iter()
             .map(|(k, v)| (k.clone(), v.snapshot()))
@@ -333,3 +340,129 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{sync::Mutex, time::Duration};
+
+    use crate::{
+        operators::*,
+        runtime::SingleThreadRuntime,
+        sinks::{StatelessSink, VecSink},
+        sources::{StatefulSource, StatefulSourceImpl, StatefulSourcePartition},
+        testing::CapturingPersistenceBackend,
+        worker::StreamProvider,
+    };
+
+    struct MockSource(i32);
+    struct MockSourcePartition {
+        max: i32,
+        next: i32,
+        was_snapshotted: Mutex<bool>,
+    }
+
+    impl StatefulSourceImpl<i32, i32> for MockSource {
+        type Part = ();
+
+        type PartitionState = i32;
+
+        type SourcePartition = MockSourcePartition;
+
+        fn list_parts(&self) -> Vec<Self::Part> {
+            vec![()]
+        }
+
+        fn build_part(
+            &mut self,
+            _part: &Self::Part,
+            part_state: Option<Self::PartitionState>,
+        ) -> Self::SourcePartition {
+            MockSourcePartition {
+                max: self.0,
+                next: part_state.unwrap_or_default(),
+                was_snapshotted: Mutex::new(false),
+            }
+        }
+    }
+
+    impl StatefulSourcePartition<i32, i32> for MockSourcePartition {
+        type PartitionState = i32;
+
+        fn poll(&mut self) -> Option<(i32, i32)> {
+            if self.next > self.max {
+                None
+            } else {
+                let out = (self.next, self.next);
+                self.next += 1;
+                Some(out)
+            }
+        }
+
+        fn is_finished(&mut self) -> bool {
+            // only terminate after we have made a snapshot
+            self.next > self.max && *self.was_snapshotted.lock().unwrap()
+        }
+
+        fn snapshot(&self) -> Self::PartitionState {
+            *self.was_snapshotted.lock().unwrap() = true;
+            self.next
+        }
+
+        fn collect(self) -> Self::PartitionState {
+            self.next
+        }
+    }
+
+    /// Check that state gets loaded from persistence backend
+    /// on initial start
+    #[test]
+    fn test_state_is_loaded_from_persistence() {
+        let persistence = CapturingPersistenceBackend::default();
+
+        let first_sink = VecSink::new();
+        let first_collected = first_sink.clone();
+
+        // execute once, this will finish as soon as a snapshot was taken
+        let rt = SingleThreadRuntime::builder()
+            .snapshots(Duration::from_millis(50))
+            .persistence(persistence.clone())
+            .build(move |provider: &mut dyn StreamProvider| {
+                provider
+                    .new_stream()
+                    .source("mock-source", StatefulSource::new(MockSource(10)))
+                    .sink("vec-sink", StatelessSink::new(first_sink));
+            });
+        rt.execute().unwrap();
+        let result: Vec<_> = first_collected
+            .drain_vec(..)
+            .iter()
+            .map(|x| x.value)
+            .collect();
+        let expected: Vec<_> = (0..=10).collect();
+        assert_eq!(result, expected);
+
+        // execute again, only numbers 11-15 should have been counted since we started from the
+        // state which had already counted to 10
+        let second_sink = VecSink::new();
+        let second_collected = second_sink.clone();
+
+        // execute again
+        let rt = SingleThreadRuntime::builder()
+            .snapshots(Duration::from_millis(50))
+            .persistence(persistence)
+            .build(move |provider: &mut dyn StreamProvider| {
+                provider
+                    .new_stream()
+                    .source("mock-source", StatefulSource::new(MockSource(15)))
+                    .sink("vec-sink", StatelessSink::new(second_sink));
+            });
+        rt.execute().unwrap();
+        let result: Vec<_> = second_collected
+            .drain_vec(..)
+            .iter()
+            .map(|x| x.value)
+            .collect();
+        let expected: Vec<_> = (11..=15).collect();
+        assert_eq!(result, expected);
+    }
+}

malstrom-core/src/testing/mod.rs

@@ -1,3 +1,4 @@
+use std::sync::Arc;
 use std::{collections::HashMap, rc::Rc, sync::Mutex};
 
 use crate::keyed::distributed::{Acquire, Collect, Interrogate};
@@ -37,13 +38,13 @@ where
 /// If you have a clone of this backend you can retrieve the state using
 /// the corresponding operator_id
 pub struct CapturingPersistenceBackend {
-    capture: Rc<Mutex<HashMap<OperatorId, Vec<u8>>>>,
+    capture: Arc<Mutex<HashMap<OperatorId, Vec<u8>>>>,
 }
 impl PersistenceBackend for CapturingPersistenceBackend {
     type Client = CapturingPersistenceBackend;
 
     fn last_commited(&self) -> Option<SnapshotVersion> {
-        None
+        Some(SnapshotVersion::default())
     }
 
     fn for_version(
@@ -55,8 +56,7 @@ impl PersistenceBackend for CapturingPersistenceBackend {
     }
 
     fn commit_version(&self, _snapshot_version: &crate::snapshot::SnapshotVersion) {
-        // TODO
-        todo!()
+        // nothing happening here
     }
 }