malstrom 0.1.0

Malstrom is a distributed, stateful stream processing framework written in Rust.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use indexmap::IndexSet;
use tracing::info;

use super::{super::types::*, MessageRouter, NormalRouter};
use crate::{channels::operator_io::Output, keyed::distributed::Remotes, types::*};
/// A collect route which has finished its local collection cycles
/// and is now just waiting for all remotes to finish
#[derive(Debug)]
pub(crate) struct FinishedRouter {
    pub(super) version: Version,
    old_worker_set: IndexSet<WorkerId>,
    pub(super) new_worker_set: IndexSet<WorkerId>,
    trigger: RescaleMessage,
}

impl FinishedRouter {
    pub(super) fn new(
        old_worker_set: IndexSet<WorkerId>,
        new_worker_set: IndexSet<WorkerId>,
        trigger: RescaleMessage,
    ) -> Self {
        Self {
            version: trigger.get_version(),
            old_worker_set,
            new_worker_set,
            trigger,
        }
    }

    pub(super) fn route_message<K, V, T>(
        &mut self,
        key: &K,
        partitioner: WorkerPartitioner<K>,
        this_worker: WorkerId,
        sender: WorkerId,
        remotes: &Remotes<K, V, T>,
    ) -> WorkerId {
        let new_target = partitioner(key, &self.new_worker_set);
        if new_target == this_worker {
            let old_target = partitioner(key, &self.old_worker_set);
            let old_target_version = remotes.get(&old_target).map(|x| x.1.last_version);
            info!(?new_target, ?old_target, ?old_target_version, ?self.version);
            if old_target == new_target {
                new_target
            }
            // TODO: I think we might not need this condition
            else if old_target == sender {
                this_worker
            // if the old target is already at our version (or does not exist),
            // we know it has already
            // given up the state for this key (or did not have it in the first place)
            // and we can safely keep the message
            } else if remotes
                .get(&old_target)
                .map(|v| v.1.last_version == Some(self.version))
                .unwrap_or(false)
            {
                new_target
            } else {
                old_target
            }
        } else {
            new_target
        }
    }

    pub(crate) fn lifecycle<K, V, T>(
        self,
        _partitioner: WorkerPartitioner<K>,
        output: &mut Output<K, V, T>,
        remotes: &mut Remotes<K, V, T>,
    ) -> MessageRouter<K, V, T>
    where
        K: Key,
        V: MaybeData,
        T: MaybeTime,
    {
        if remotes.values().all(|(_, state)| {
            state.last_version.map(|v| v == self.version).unwrap_or(false)
                // we cannot progress if they did not acknowledge our version because
                // in that case they might still try to send us messages
                && state.last_ack_version.map(|v| v == self.version).unwrap_or(false)
        }) {
            for wid in self.old_worker_set.difference(&self.new_worker_set) {
                remotes.swap_remove(wid);
            }
            remotes.shrink_to_fit();
            let normal_router = NormalRouter::new(self.new_worker_set, self.version);
            output.send(Message::Rescale(self.trigger));
            MessageRouter::Normal(normal_router)
        } else {
            MessageRouter::Finished(self)
        }
    }
}

#[cfg(test)]
mod tests {
    use indexmap::IndexSet;

    use crate::{
        channels::operator_io::{full_broadcast, link, Input, Output},
        keyed::{distributed::Remotes, partitioners::rendezvous_select},
        types::Message,
    };

    use super::{FinishedRouter, RescaleMessage};

    /// It should emit the original rescale message upon returning to the normal state
    #[test]
    fn emit_original_rescale() {
        let router = FinishedRouter::new(
            IndexSet::from([0]),
            IndexSet::from([0, 1]),
            RescaleMessage::new(IndexSet::from([1]), 0),
        );
        let mut output: Output<usize, usize, usize> = Output::new_unlinked(full_broadcast);
        let mut input = Input::new_unlinked();
        link(&mut output, &mut input);
        let mut remotes = Remotes::new();
        router.lifecycle(rendezvous_select, &mut output, &mut remotes);
        let msg = input.recv().unwrap();
        assert!(matches!(msg, Message::Rescale(_)))
    }
}