Difficulty creating a mock 'sink' for testing purposes in Rust - rust

Question

Context
I have a function that will use rayon to paralellize some work. Each thread should write its output to a separate consumer. In practice, these consumers will usually write to (separate) files on disc. For testing, I want them to append to vectors, which I can then make assertions on.
Problem
In fact I hit an issue even with a serial version of this code.
The function I want to test is func_under_test. I want get_sink to create and hand out a new 'sink' (a Vec<i32>) each time it is called, which will be inside func_under_test. But I also want to keep a reference to these vectors in my main function, so I can make assertions on them. I've tried writing this:
fn func_under_test<GetSink, F>(
mut get_sink: GetSink
)
where
GetSink: FnMut(usize) -> F,
F: FnMut(i32) -> (),
{
let idxs : Vec<usize> = (0..2).collect();
idxs.iter().for_each(|&i| {
let mut sink = get_sink(i);
sink(0i32);
sink(1i32);
sink(2i32);
});
}
fn main() {
let mut sinks : Vec<&Vec<i32>> = vec![];
let get_sink = |i: usize| {
let mut sink : Vec<i32> = vec![];
sinks.push(&sink);
move |x : i32| sink.push(x)
};
func_under_test(get_sink);
assert_eq!(**sinks.get(0).unwrap(), vec![0i32, 1i32, 2i32]);
}
But I get this compilation error:
error[E0597]: `sink` does not live long enough
--> src/main.rs:23:20
|
20 | let mut sinks : Vec<&Vec<i32>> = vec![];
| --------- lifetime `'1` appears in the type of `sinks`
...
23 | sinks.push(&sink);
| -----------^^^^^-
| | |
| | borrowed value does not live long enough
| argument requires that `sink` is borrowed for `'1`
24 | move |x : i32| sink.push(x)
25 | };
| - `sink` dropped here while still borrowed
error[E0505]: cannot move out of `sink` because it is borrowed
--> src/main.rs:24:9
|
20 | let mut sinks : Vec<&Vec<i32>> = vec![];
| --------- lifetime `'1` appears in the type of `sinks`
...
23 | sinks.push(&sink);
| -----------------
| | |
| | borrow of `sink` occurs here
| argument requires that `sink` is borrowed for `'1`
24 | move |x : i32| sink.push(x)
| ^^^^^^^^^^^^^^ ---- move occurs due to use in closure
| |
| move out of `sink` occurs here
Am I correct in thinking that I need to use something like an Rc in this context, since I need to have two references to each Vec<i32>?
Since only one of those references needs to be mutable, do I need to use Rc<Cell<Vec<i32>>> in one case, and just Rc<Vec<i32>> in the other?

Answer 1

You could use an Arc<Mutex<i32>> to share mutable access to the vectors. (Rc is not thread-safe.)
A tidier solution would be typed_arena::Arena, which would allow you to write essentially the code you tried, but with the roles swapped: the Vec<i32>s are always owned by the arena (thus it outlives func_under_tests). Then, after all of the references to the arena are gone, you can convert it into a vector of its elements.
use typed_arena::Arena;
fn main() {
let sinks: Arena<Vec<i32>> = Arena::new();
let get_sink = |_i: usize| {
let sink_ref = sinks.alloc(Vec::new());
move |x| sink_ref.push(x)
};
func_under_test(get_sink);
assert_eq!(sinks.into_vec()[0], vec![0i32, 1i32, 2i32]);
}

Answer 2

I think I've figured out the answer to this. The thread doing the testing needs to have access to the inner vectors, and so does the thread that is populating this. So they need to be wrapped in an Arc and a Mutex. But the outer vector does not, since it is only borrowed by the get_sink closure.
So sinks can have type Vec<Arc<Mutex<Vec<i32>>>.
Anyway, I can write this and it does compile and behave and expected:
use rayon::prelude::*;
use std::sync::{Arc, Mutex};
fn func_under_test<G, S>(
get_sink: G)
where
G: Fn(&usize) -> S + Sync,
S: FnMut(i32) -> (),
{
let idxs : Vec<usize> = vec![0, 1, 2];
idxs.par_iter().for_each(|i| {
let mut sink : S = get_sink(i);
for j in 0..3 {
sink(j + 3*(*i) as i32);
}
});
}
fn main() {
let sinks : Vec<Arc<Mutex<Vec<i32>>>> = vec![
Arc::new(Mutex::new(vec![])),
Arc::new(Mutex::new(vec![])),
Arc::new(Mutex::new(vec![])),
];
let get_sink = |i: &usize| {
let m : Arc<Mutex<Vec<i32>>> = sinks.get(*i).unwrap().clone();
move |x: i32| {
m.lock().unwrap().push(x);
}
};
func_under_test(get_sink);
assert_eq!(*sinks.get(0).unwrap().lock().unwrap(), vec![0, 1, 2]);
assert_eq!(*sinks.get(1).unwrap().lock().unwrap(), vec![3, 4, 5]);
assert_eq!(*sinks.get(2).unwrap().lock().unwrap(), vec![6, 7, 8]);
}
It also seems like G implements the correct traits here: it is Sync, since it needs to be shared between threads. But is is not Send, since it does not need to send information between threads.