use std::sync::Arc;
use dynomite::embed::Datastore;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::sync::mpsc;
use crate::mapreduce::job::{Inputs, KeyDatum, MapReduceJob};
use crate::mapreduce::phase::Phase;
use crate::mapreduce::registry::PhaseRegistry;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum MrError {
#[error("unknown {kind} function: {name}")]
UnknownFunction {
kind: &'static str,
name: String,
},
#[error("phase {phase} ({kind}) failed: {message}")]
PhaseFailed {
phase: u32,
kind: &'static str,
message: String,
},
#[error("unsupported MapReduce inputs: {0}")]
UnsupportedInputs(&'static str),
#[error("wasm-module phases are not enabled on this executor")]
WasmNotImplemented,
#[error("wasm module not found: {0}")]
WasmModuleNotFound(String),
#[error("wasm phase exceeded execution time / fuel limit")]
WasmExecutionTimeout,
#[error("wasm phase exceeded memory limit")]
WasmMemoryLimit,
#[error("wasm runtime error: {0}")]
WasmRuntime(String),
#[error("wasm phase encoding error: {0}")]
WasmEncoding(String),
#[error("link phases require a datastore-backed job")]
LinkNotImplemented,
#[error("internal pipeline error: {0}")]
Pipeline(String),
#[error("json: {0}")]
Json(String),
}
pub trait WasmHook: Send + Sync {
fn apply_phase(
&self,
module_id: &str,
fn_name: &str,
inputs: &[Value],
) -> Result<Vec<Value>, MrError>;
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct PhaseOutput {
pub phase: u32,
pub value: Value,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct PhaseBatch {
pub phase: u32,
pub data: Vec<Value>,
}
pub async fn run_job(
job: MapReduceJob,
registry: Arc<PhaseRegistry>,
) -> Result<Vec<PhaseOutput>, MrError> {
run_job_full(job, registry, None, None).await
}
#[must_use]
pub fn run_job_streaming(
job: MapReduceJob,
registry: Arc<PhaseRegistry>,
) -> mpsc::Receiver<Result<PhaseBatch, MrError>> {
run_job_streaming_full(job, registry, None, None)
}
#[must_use]
pub fn run_job_streaming_with_wasm(
job: MapReduceJob,
registry: Arc<PhaseRegistry>,
wasm: Option<Arc<dyn WasmHook>>,
) -> mpsc::Receiver<Result<PhaseBatch, MrError>> {
run_job_streaming_full(job, registry, wasm, None)
}
#[must_use]
pub fn run_job_streaming_full(
job: MapReduceJob,
registry: Arc<PhaseRegistry>,
wasm: Option<Arc<dyn WasmHook>>,
datastore: Option<Arc<dyn Datastore>>,
) -> mpsc::Receiver<Result<PhaseBatch, MrError>> {
let (tx, rx) = mpsc::channel::<Result<PhaseBatch, MrError>>(4);
tokio::spawn(async move {
let result = stream_job_inner(job, registry, wasm, datastore, tx.clone()).await;
if let Err(e) = result {
let _ = tx.send(Err(e)).await;
}
});
rx
}
async fn stream_job_inner(
job: MapReduceJob,
registry: Arc<PhaseRegistry>,
wasm: Option<Arc<dyn WasmHook>>,
datastore: Option<Arc<dyn Datastore>>,
tx: mpsc::Sender<Result<PhaseBatch, MrError>>,
) -> Result<(), MrError> {
let initial = resolve_inputs(&job.inputs, datastore.as_ref()).await?;
materialised_initial_inputs_must_be_iterable(&initial);
if job.phases.is_empty() {
if !initial.is_empty() {
let batch = PhaseBatch {
phase: 0,
data: initial,
};
let _ = tx.send(Ok(batch)).await;
}
return Ok(());
}
let n_phases = job.phases.len();
let mut current: Vec<Value> = initial;
for (idx, phase) in job.phases.iter().enumerate() {
let phase_idx = u32::try_from(idx)
.map_err(|_| MrError::Pipeline("phase index exceeds u32 range".into()))?;
let is_last = idx + 1 == n_phases;
let outputs = run_phase(
phase_idx,
phase,
current,
®istry,
wasm.as_ref(),
datastore.as_ref(),
)
.await?;
if (phase.keep() || is_last) && !outputs.is_empty() {
let batch = PhaseBatch {
phase: phase_idx,
data: outputs.clone(),
};
if tx.send(Ok(batch)).await.is_err() {
return Ok(());
}
}
current = outputs;
}
Ok(())
}
pub async fn run_job_with_wasm(
job: MapReduceJob,
registry: Arc<PhaseRegistry>,
wasm: Option<Arc<dyn WasmHook>>,
) -> Result<Vec<PhaseOutput>, MrError> {
run_job_full(job, registry, wasm, None).await
}
pub async fn run_job_full(
job: MapReduceJob,
registry: Arc<PhaseRegistry>,
wasm: Option<Arc<dyn WasmHook>>,
datastore: Option<Arc<dyn Datastore>>,
) -> Result<Vec<PhaseOutput>, MrError> {
let initial = resolve_inputs(&job.inputs, datastore.as_ref()).await?;
materialised_initial_inputs_must_be_iterable(&initial);
if job.phases.is_empty() {
let out: Vec<PhaseOutput> = initial
.into_iter()
.map(|v| PhaseOutput { phase: 0, value: v })
.collect();
return Ok(out);
}
let n_phases = job.phases.len();
let mut current: Vec<Value> = initial;
let mut captured: Vec<PhaseOutput> = Vec::new();
for (idx, phase) in job.phases.iter().enumerate() {
let phase_idx = u32::try_from(idx)
.map_err(|_| MrError::Pipeline("phase index exceeds u32 range".into()))?;
let is_last = idx + 1 == n_phases;
let outputs = run_phase(
phase_idx,
phase,
current,
®istry,
wasm.as_ref(),
datastore.as_ref(),
)
.await?;
if phase.keep() || is_last {
for v in &outputs {
captured.push(PhaseOutput {
phase: phase_idx,
value: v.clone(),
});
}
}
current = outputs;
}
Ok(captured)
}
async fn run_phase(
phase_idx: u32,
phase: &Phase,
inputs: Vec<Value>,
registry: &Arc<PhaseRegistry>,
wasm: Option<&Arc<dyn WasmHook>>,
datastore: Option<&Arc<dyn Datastore>>,
) -> Result<Vec<Value>, MrError> {
let (tx_in, rx_in) = mpsc::channel::<Value>(64);
let (tx_out, mut rx_out) = mpsc::channel::<Value>(64);
tokio::spawn(async move {
for v in inputs {
if tx_in.send(v).await.is_err() {
return;
}
}
});
let phase_clone = phase.clone();
let registry_clone = Arc::clone(registry);
let wasm_clone = wasm.cloned();
let datastore_clone = datastore.cloned();
let phase_join = tokio::spawn(async move {
run_phase_task(
phase_idx,
phase_clone,
rx_in,
tx_out,
registry_clone,
wasm_clone,
datastore_clone,
)
.await
});
let mut outputs = Vec::new();
while let Some(v) = rx_out.recv().await {
outputs.push(v);
}
match phase_join.await {
Ok(Ok(())) => Ok(outputs),
Ok(Err(e)) => Err(e),
Err(e) => Err(MrError::Pipeline(format!("join: {e}"))),
}
}
async fn run_phase_task(
phase_idx: u32,
phase: Phase,
mut rx: mpsc::Receiver<Value>,
tx: mpsc::Sender<Value>,
registry: Arc<PhaseRegistry>,
wasm: Option<Arc<dyn WasmHook>>,
datastore: Option<Arc<dyn Datastore>>,
) -> Result<(), MrError> {
match phase {
Phase::Map { fn_name, arg, .. } => {
let f = registry.map_fn(&fn_name).ok_or(MrError::UnknownFunction {
kind: "map",
name: fn_name.clone(),
})?;
let f = f.clone();
while let Some(v) = rx.recv().await {
let outs = (f)(&v, arg.as_ref()).map_err(|e| MrError::PhaseFailed {
phase: phase_idx,
kind: "map",
message: e.to_string(),
})?;
for o in outs {
if tx.send(o).await.is_err() {
return Err(MrError::Pipeline(
"downstream phase dropped its inbound channel".into(),
));
}
}
}
Ok(())
}
Phase::Reduce { fn_name, arg, .. } => {
let f = registry
.reduce_fn(&fn_name)
.ok_or(MrError::UnknownFunction {
kind: "reduce",
name: fn_name.clone(),
})?;
let f = f.clone();
let mut buf: Vec<Value> = Vec::new();
while let Some(v) = rx.recv().await {
buf.push(v);
}
let outs = (f)(&buf, arg.as_ref()).map_err(|e| MrError::PhaseFailed {
phase: phase_idx,
kind: "reduce",
message: e.to_string(),
})?;
for o in outs {
if tx.send(o).await.is_err() {
return Err(MrError::Pipeline(
"downstream phase dropped its inbound channel".into(),
));
}
}
Ok(())
}
Phase::Link { bucket, tag, .. } => {
run_link_phase(
phase_idx,
bucket.as_deref(),
tag.as_deref(),
&mut rx,
&tx,
datastore,
)
.await
}
Phase::WasmModule {
module_id, fn_name, ..
} => {
let hook = wasm.ok_or(MrError::WasmNotImplemented)?;
let mut buf: Vec<Value> = Vec::new();
while let Some(v) = rx.recv().await {
buf.push(v);
}
let mid = module_id.clone();
let fname = fn_name.clone();
let outs = tokio::task::spawn_blocking(move || hook.apply_phase(&mid, &fname, &buf))
.await
.map_err(|e| MrError::Pipeline(format!("wasm join: {e}")))?
.map_err(|e| match e {
MrError::WasmModuleNotFound(_)
| MrError::WasmExecutionTimeout
| MrError::WasmMemoryLimit
| MrError::WasmRuntime(_)
| MrError::WasmEncoding(_)
| MrError::WasmNotImplemented => e,
other => MrError::PhaseFailed {
phase: phase_idx,
kind: "wasm",
message: other.to_string(),
},
})?;
for o in outs {
if tx.send(o).await.is_err() {
return Err(MrError::Pipeline(
"downstream phase dropped its inbound channel".into(),
));
}
}
Ok(())
}
}
}
async fn run_link_phase(
phase_idx: u32,
bucket: Option<&str>,
tag: Option<&str>,
rx: &mut mpsc::Receiver<Value>,
tx: &mpsc::Sender<Value>,
datastore: Option<Arc<dyn Datastore>>,
) -> Result<(), MrError> {
let store = datastore.ok_or(MrError::LinkNotImplemented)?;
while let Some(v) = rx.recv().await {
let (in_bucket, in_key) = link_input_target(&v).ok_or_else(|| MrError::PhaseFailed {
phase: phase_idx,
kind: "link",
message: "link-phase input is missing bucket/key".into(),
})?;
let stored = store
.riak_get(in_bucket.as_bytes(), in_key.as_bytes())
.await
.map_err(|e| MrError::PhaseFailed {
phase: phase_idx,
kind: "link",
message: format!("riak get {in_bucket}/{in_key}: {e}"),
})?;
let Some(stored) = stored else {
continue;
};
let obj =
crate::proto::http::object::HttpObject::from_storage_bytes(&stored).map_err(|e| {
MrError::PhaseFailed {
phase: phase_idx,
kind: "link",
message: format!("decode {in_bucket}/{in_key}: {e}"),
}
})?;
for link in &obj.links {
let bucket_ok = bucket.is_none_or(|b| b == link.bucket);
let tag_ok = tag.is_none_or(|t| t == link.tag);
if bucket_ok && tag_ok {
let out = KeyDatum::pair(link.bucket.clone(), link.key.clone()).to_value();
if tx.send(out).await.is_err() {
return Err(MrError::Pipeline(
"downstream phase dropped its inbound channel".into(),
));
}
}
}
}
Ok(())
}
fn materialised_initial_inputs_must_be_iterable<T>(_: &[T]) {}
fn link_input_target(v: &Value) -> Option<(String, String)> {
let bucket = v.get("bucket")?.as_str()?.to_string();
let key = v.get("key")?.as_str()?.to_string();
Some((bucket, key))
}
async fn resolve_inputs(
inputs: &Inputs,
datastore: Option<&Arc<dyn Datastore>>,
) -> Result<Vec<Value>, MrError> {
if let Some(items) = inputs.items() {
return Ok(items.into_iter().map(|kd| kd.to_value()).collect());
}
let Inputs::Bucket(bucket) = inputs else {
return Err(MrError::UnsupportedInputs("unrunnable input spec"));
};
let store = datastore.ok_or(MrError::UnsupportedInputs("bucket scan"))?;
let mut stream = store.list_keys_stream(bucket.as_bytes());
let mut out = Vec::new();
while let Some(item) = stream.next().await {
let key = item.map_err(|e| MrError::Pipeline(format!("bucket scan: {e}")))?;
let key = String::from_utf8_lossy(&key).into_owned();
out.push(KeyDatum::pair(bucket.clone(), key).to_value());
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::mapreduce::builtins::default_registry;
use crate::mapreduce::job::{Inputs, KeyDatum};
fn registry() -> Arc<PhaseRegistry> {
Arc::new(default_registry())
}
struct ScriptedStore {
body: Option<Result<Option<Vec<u8>>, ()>>,
}
impl dynomite::embed::Datastore for ScriptedStore {
fn protocol(&self) -> dynomite::embed::Protocol {
dynomite::embed::Protocol::Custom
}
fn dispatch(
&self,
req: dynomite::msg::Msg,
) -> dynomite::embed::BoxFuture<
'_,
Result<dynomite::msg::Msg, dynomite::embed::DatastoreError>,
> {
Box::pin(async move {
Ok(dynomite::msg::Msg::new(
req.id(),
dynomite::msg::MsgType::Unknown,
false,
))
})
}
fn riak_get<'a>(
&'a self,
_bucket: &'a [u8],
_key: &'a [u8],
) -> dynomite::embed::BoxFuture<'a, Result<Option<Vec<u8>>, dynomite::embed::DatastoreError>>
{
let body = self.body.clone();
Box::pin(async move {
match body {
Some(Ok(v)) => Ok(v),
Some(Err(())) | None => Err(dynomite::embed::DatastoreError::Backend(
"scripted failure".into(),
)),
}
})
}
}
fn store_with(body: Option<Result<Option<Vec<u8>>, ()>>) -> Arc<dyn Datastore> {
Arc::new(ScriptedStore { body })
}
#[tokio::test]
async fn empty_phase_list_is_identity() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::with_value("b", "k", serde_json::json!(1))]),
phases: vec![],
timeout_ms: None,
};
let out = run_job(job, registry()).await.expect("ok");
assert_eq!(out.len(), 1);
assert_eq!(out[0].phase, 0);
assert_eq!(out[0].value["bucket"], "b");
}
#[tokio::test]
async fn map_then_reduce_pipeline() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![
KeyDatum::with_value("b", "k1", serde_json::json!(2)),
KeyDatum::with_value("b", "k2", serde_json::json!(3)),
KeyDatum::with_value("b", "k3", serde_json::json!(4)),
]),
phases: vec![
Phase::Map {
fn_name: "map_object_value".into(),
arg: None,
keep: false,
},
Phase::Reduce {
fn_name: "reduce_sum".into(),
arg: None,
keep: true,
},
],
timeout_ms: None,
};
let out = run_job(job, registry()).await.expect("ok");
assert_eq!(out.len(), 1);
assert_eq!(out[0].phase, 1);
assert_eq!(out[0].value, serde_json::json!(9));
}
#[tokio::test]
async fn keep_intermediate_phase_outputs_are_captured() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![
KeyDatum::with_value("b", "k1", serde_json::json!(5)),
KeyDatum::with_value("b", "k2", serde_json::json!(7)),
]),
phases: vec![
Phase::Map {
fn_name: "map_object_value".into(),
arg: None,
keep: true,
},
Phase::Reduce {
fn_name: "reduce_sum".into(),
arg: None,
keep: true,
},
],
timeout_ms: None,
};
let out = run_job(job, registry()).await.expect("ok");
assert_eq!(out.len(), 3);
assert_eq!(out[0].phase, 0);
assert_eq!(out[1].phase, 0);
assert_eq!(out[2].phase, 1);
assert_eq!(out[2].value, serde_json::json!(12));
}
#[tokio::test]
async fn unknown_map_function_is_typed_error() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::with_value("b", "k", serde_json::json!(1))]),
phases: vec![Phase::Map {
fn_name: "no_such_function".into(),
arg: None,
keep: false,
}],
timeout_ms: None,
};
let err = run_job(job, registry()).await.expect_err("error");
assert!(matches!(err, MrError::UnknownFunction { kind: "map", .. }));
}
#[tokio::test]
async fn unknown_reduce_function_is_typed_error() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![]),
phases: vec![Phase::Reduce {
fn_name: "no_such_reduce".into(),
arg: None,
keep: false,
}],
timeout_ms: None,
};
let err = run_job(job, registry()).await.expect_err("error");
assert!(matches!(
err,
MrError::UnknownFunction { kind: "reduce", .. }
));
}
#[tokio::test]
async fn bucket_inputs_without_datastore_are_unsupported() {
let job = MapReduceJob {
inputs: Inputs::Bucket("b".into()),
phases: vec![],
timeout_ms: None,
};
let err = run_job(job, registry()).await.expect_err("error");
assert!(matches!(err, MrError::UnsupportedInputs(_)));
}
#[tokio::test]
async fn bucket_inputs_enumerate_keys_through_datastore() {
use dynomite::embed::hooks::MemoryDatastore;
let ds = MemoryDatastore::new();
ds.insert(b"users", b"alice");
ds.insert(b"users", b"bob");
ds.insert(b"users", b"carol");
ds.insert(b"orders", b"o1");
let ds: Arc<dyn Datastore> = Arc::new(ds);
let job = MapReduceJob {
inputs: Inputs::Bucket("users".into()),
phases: vec![],
timeout_ms: None,
};
let out = run_job_full(job, registry(), None, Some(ds))
.await
.expect("ok");
assert_eq!(out.len(), 3);
let mut keys: Vec<String> = out
.iter()
.map(|o| o.value["key"].as_str().unwrap().to_string())
.collect();
keys.sort();
assert_eq!(keys, vec!["alice", "bob", "carol"]);
for o in &out {
assert_eq!(o.value["bucket"], "users");
assert!(o.value["value"].is_null());
}
}
#[tokio::test]
async fn bucket_inputs_feed_map_reduce_pipeline() {
use dynomite::embed::hooks::MemoryDatastore;
let ds = MemoryDatastore::new();
for i in 0..5u32 {
ds.insert(b"nums", format!("k{i}").as_bytes());
}
let ds: Arc<dyn Datastore> = Arc::new(ds);
let job = MapReduceJob {
inputs: Inputs::Bucket("nums".into()),
phases: vec![
Phase::Map {
fn_name: "map_identity".into(),
arg: None,
keep: false,
},
Phase::Reduce {
fn_name: "reduce_count".into(),
arg: None,
keep: true,
},
],
timeout_ms: None,
};
let out = run_job_full(job, registry(), None, Some(ds))
.await
.expect("ok");
assert_eq!(out.len(), 1);
assert_eq!(out[0].value, serde_json::json!(5));
}
#[tokio::test]
async fn bucket_inputs_over_empty_bucket_yield_empty_result() {
use dynomite::embed::hooks::MemoryDatastore;
let ds: Arc<dyn Datastore> = Arc::new(MemoryDatastore::new());
let job = MapReduceJob {
inputs: Inputs::Bucket("nonexistent".into()),
phases: vec![Phase::Reduce {
fn_name: "reduce_sum".into(),
arg: None,
keep: true,
}],
timeout_ms: None,
};
let out = run_job_full(job, registry(), None, Some(ds))
.await
.expect("ok, not error");
assert_eq!(
out,
vec![PhaseOutput {
phase: 0,
value: serde_json::json!(0)
}]
);
}
#[tokio::test]
async fn bucket_inputs_stream_through_streaming_entry_point() {
use dynomite::embed::hooks::MemoryDatastore;
let ds = MemoryDatastore::new();
ds.insert(b"b", b"k1");
ds.insert(b"b", b"k2");
let ds: Arc<dyn Datastore> = Arc::new(ds);
let job = MapReduceJob {
inputs: Inputs::Bucket("b".into()),
phases: vec![],
timeout_ms: None,
};
let rx = run_job_streaming_full(job, registry(), None, Some(ds));
let items = drain_stream(rx).await;
assert_eq!(items.len(), 1);
let b = items[0].as_ref().expect("ok");
assert_eq!(b.phase, 0);
assert_eq!(b.data.len(), 2);
}
#[tokio::test]
async fn wasm_phase_returns_typed_error() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::with_value("b", "k", serde_json::json!(1))]),
phases: vec![Phase::WasmModule {
module_id: "m".into(),
fn_name: "f".into(),
arg: None,
keep: false,
}],
timeout_ms: None,
};
let err = run_job(job, registry()).await.expect_err("error");
assert!(matches!(err, MrError::WasmNotImplemented));
}
#[tokio::test]
async fn link_phase_without_datastore_returns_typed_error() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::pair("b", "k")]),
phases: vec![Phase::Link {
bucket: None,
tag: None,
keep: false,
}],
timeout_ms: None,
};
let err = run_job(job, registry()).await.expect_err("error");
assert!(matches!(err, MrError::LinkNotImplemented));
}
#[tokio::test]
async fn determinism_under_one_hundred_inputs_three_phases() {
let mut data = Vec::new();
for i in 0..100u32 {
data.push(KeyDatum::with_value(
"b",
format!("k{i}"),
serde_json::json!(i),
));
}
let phases = vec![
Phase::Map {
fn_name: "map_object_value".into(),
arg: None,
keep: false,
},
Phase::Reduce {
fn_name: "reduce_sort".into(),
arg: None,
keep: false,
},
Phase::Reduce {
fn_name: "reduce_count".into(),
arg: None,
keep: true,
},
];
let job = MapReduceJob {
inputs: Inputs::KeyData(data.clone()),
phases: phases.clone(),
timeout_ms: None,
};
let job2 = MapReduceJob {
inputs: Inputs::KeyData(data),
phases,
timeout_ms: None,
};
let r1 = run_job(job, registry()).await.expect("run1");
let r2 = run_job(job2, registry()).await.expect("run2");
assert_eq!(r1, r2);
assert_eq!(r1.len(), 1);
assert_eq!(r1[0].value, serde_json::json!(100));
}
#[tokio::test]
async fn pairs_inputs_are_materialised_with_null_value() {
let job = MapReduceJob {
inputs: Inputs::Pairs(vec![("b".into(), "k1".into()), ("b".into(), "k2".into())]),
phases: vec![],
timeout_ms: None,
};
let out = run_job(job, registry()).await.expect("ok");
assert_eq!(out.len(), 2);
assert!(out[0].value["value"].is_null());
assert_eq!(out[0].value["bucket"], "b");
assert_eq!(out[0].value["key"], "k1");
}
async fn drain_stream(
mut rx: mpsc::Receiver<Result<PhaseBatch, MrError>>,
) -> Vec<Result<PhaseBatch, MrError>> {
let mut out = Vec::new();
while let Some(item) = rx.recv().await {
out.push(item);
}
out
}
#[tokio::test]
async fn streaming_emits_one_batch_per_kept_phase() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![
KeyDatum::with_value("b", "k1", serde_json::json!(1)),
KeyDatum::with_value("b", "k2", serde_json::json!(2)),
KeyDatum::with_value("b", "k3", serde_json::json!(3)),
]),
phases: vec![
Phase::Map {
fn_name: "map_object_value".into(),
arg: None,
keep: true,
},
Phase::Reduce {
fn_name: "reduce_sum".into(),
arg: None,
keep: true,
},
],
timeout_ms: None,
};
let rx = run_job_streaming(job, registry());
let items = drain_stream(rx).await;
assert_eq!(items.len(), 2, "two kept phases yield two batches");
let b0 = items[0].as_ref().expect("phase 0 ok");
assert_eq!(b0.phase, 0);
assert_eq!(b0.data.len(), 3);
let b1 = items[1].as_ref().expect("phase 1 ok");
assert_eq!(b1.phase, 1);
assert_eq!(b1.data, vec![serde_json::json!(6)]);
}
#[tokio::test]
async fn streaming_skips_non_keep_intermediate_phases() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::with_value("b", "k", serde_json::json!(2))]),
phases: vec![
Phase::Map {
fn_name: "map_object_value".into(),
arg: None,
keep: false,
},
Phase::Reduce {
fn_name: "reduce_sum".into(),
arg: None,
keep: true,
},
],
timeout_ms: None,
};
let rx = run_job_streaming(job, registry());
let items = drain_stream(rx).await;
assert_eq!(items.len(), 1, "only the kept reduce emits a batch");
let b = items[0].as_ref().expect("ok");
assert_eq!(b.phase, 1);
}
#[tokio::test]
async fn streaming_surfaces_phase_error_as_terminal_item() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::pair("b", "k")]),
phases: vec![Phase::Map {
fn_name: "no_such_function".into(),
arg: None,
keep: true,
}],
timeout_ms: None,
};
let rx = run_job_streaming(job, registry());
let items = drain_stream(rx).await;
assert_eq!(items.len(), 1);
let err = items[0].as_ref().expect_err("unknown function");
assert!(matches!(err, MrError::UnknownFunction { kind: "map", .. }));
}
#[tokio::test]
async fn streaming_empty_phase_list_emits_initial_inputs() {
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::with_value("b", "k", serde_json::json!(42))]),
phases: vec![],
timeout_ms: None,
};
let rx = run_job_streaming(job, registry());
let items = drain_stream(rx).await;
assert_eq!(items.len(), 1);
let b = items[0].as_ref().expect("ok");
assert_eq!(b.phase, 0);
assert_eq!(b.data.len(), 1);
assert_eq!(b.data[0]["value"], serde_json::json!(42));
}
fn link_job() -> MapReduceJob {
MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::pair("people", "a")]),
phases: vec![Phase::Link {
bucket: None,
tag: None,
keep: true,
}],
timeout_ms: None,
}
}
#[tokio::test]
async fn link_phase_riak_get_error_is_phase_failed() {
let ds = store_with(Some(Err(())));
let err = run_job_full(link_job(), registry(), None, Some(ds))
.await
.expect_err("riak get fails");
assert!(matches!(err, MrError::PhaseFailed { kind: "link", .. }));
}
#[tokio::test]
async fn link_phase_undecodable_object_is_phase_failed() {
let ds = store_with(Some(Ok(Some(vec![0xff, 0xff, 0xff]))));
let err = run_job_full(link_job(), registry(), None, Some(ds))
.await
.expect_err("decode fails");
assert!(matches!(err, MrError::PhaseFailed { kind: "link", .. }));
}
#[tokio::test]
async fn link_phase_missing_object_yields_empty_via_scripted_store() {
let ds = store_with(Some(Ok(None)));
let out = run_job_full(link_job(), registry(), None, Some(ds))
.await
.expect("missing is not an error");
assert!(out.is_empty());
}
#[tokio::test]
async fn link_phase_input_missing_bucket_key_is_phase_failed() {
let ds = store_with(Some(Ok(None)));
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::with_value("b", "k", serde_json::json!(1))]),
phases: vec![
Phase::Map {
fn_name: "map_object_value".into(),
arg: None,
keep: false,
},
Phase::Link {
bucket: None,
tag: None,
keep: true,
},
],
timeout_ms: None,
};
let err = run_job_full(job, registry(), None, Some(ds))
.await
.expect_err("non-routing datum into link");
assert!(matches!(err, MrError::PhaseFailed { kind: "link", .. }));
}
#[tokio::test]
async fn link_phase_emits_matching_targets_via_scripted_store() {
let obj = crate::proto::http::object::HttpObject {
value: b"src".to_vec(),
content_type: None,
indexes: Vec::new(),
links: vec![
crate::proto::http::object::HttpLink {
bucket: "people".into(),
key: "b".into(),
tag: "friend".into(),
},
crate::proto::http::object::HttpLink {
bucket: "work".into(),
key: "acme".into(),
tag: "colleague".into(),
},
],
};
let ds = store_with(Some(Ok(Some(obj.to_storage_bytes()))));
let job = MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::pair("people", "a")]),
phases: vec![Phase::Link {
bucket: Some("people".into()),
tag: Some("friend".into()),
keep: true,
}],
timeout_ms: None,
};
let out = run_job_full(job, registry(), None, Some(ds))
.await
.expect("ok");
assert_eq!(out.len(), 1);
assert_eq!(out[0].value["bucket"], "people");
assert_eq!(out[0].value["key"], "b");
}
struct ScriptedWasm {
outcome: WasmOutcome,
}
enum WasmOutcome {
Ok(Vec<Value>),
WasmError,
GenericError,
}
impl WasmHook for ScriptedWasm {
fn apply_phase(
&self,
_module_id: &str,
_fn_name: &str,
_inputs: &[Value],
) -> Result<Vec<Value>, MrError> {
match &self.outcome {
WasmOutcome::Ok(v) => Ok(v.clone()),
WasmOutcome::WasmError => Err(MrError::WasmExecutionTimeout),
WasmOutcome::GenericError => Err(MrError::Json("boom".into())),
}
}
}
fn wasm_job() -> MapReduceJob {
MapReduceJob {
inputs: Inputs::KeyData(vec![KeyDatum::with_value("b", "k", serde_json::json!(1))]),
phases: vec![Phase::WasmModule {
module_id: "m".into(),
fn_name: "f".into(),
arg: None,
keep: true,
}],
timeout_ms: None,
}
}
#[tokio::test]
async fn wasm_phase_success_threads_output() {
let hook: Arc<dyn WasmHook> = Arc::new(ScriptedWasm {
outcome: WasmOutcome::Ok(vec![serde_json::json!("done")]),
});
let out = run_job_with_wasm(wasm_job(), registry(), Some(hook))
.await
.expect("ok");
assert_eq!(out.len(), 1);
assert_eq!(out[0].value, serde_json::json!("done"));
}
#[tokio::test]
async fn wasm_phase_wasm_error_passes_through_untouched() {
let hook: Arc<dyn WasmHook> = Arc::new(ScriptedWasm {
outcome: WasmOutcome::WasmError,
});
let err = run_job_with_wasm(wasm_job(), registry(), Some(hook))
.await
.expect_err("wasm timeout");
assert!(matches!(err, MrError::WasmExecutionTimeout));
}
#[tokio::test]
async fn wasm_phase_generic_error_is_wrapped_as_phase_failed() {
let hook: Arc<dyn WasmHook> = Arc::new(ScriptedWasm {
outcome: WasmOutcome::GenericError,
});
let err = run_job_with_wasm(wasm_job(), registry(), Some(hook))
.await
.expect_err("wrapped");
assert!(matches!(err, MrError::PhaseFailed { kind: "wasm", .. }));
}
#[tokio::test]
async fn streaming_with_wasm_threads_phase_output() {
let hook: Arc<dyn WasmHook> = Arc::new(ScriptedWasm {
outcome: WasmOutcome::Ok(vec![serde_json::json!("w")]),
});
let rx = run_job_streaming_with_wasm(wasm_job(), registry(), Some(hook));
let items = drain_stream(rx).await;
assert_eq!(items.len(), 1);
let b = items[0].as_ref().expect("ok");
assert_eq!(b.data, vec![serde_json::json!("w")]);
}
#[tokio::test]
async fn bucket_scan_stream_error_surfaces_as_pipeline_error() {
let ds = store_with(Some(Ok(None)));
let job = MapReduceJob {
inputs: Inputs::Bucket("b".into()),
phases: vec![],
timeout_ms: None,
};
let err = run_job_full(job, registry(), None, Some(ds))
.await
.expect_err("scan errors");
assert!(matches!(err, MrError::Pipeline(_)));
}
#[tokio::test]
async fn scripted_store_protocol_and_dispatch_are_exercised() {
let store = ScriptedStore { body: None };
assert_eq!(store.protocol(), dynomite::embed::Protocol::Custom);
let rsp = store
.dispatch(dynomite::msg::Msg::new(
7,
dynomite::msg::MsgType::Unknown,
true,
))
.await
.expect("dispatch ok");
assert_eq!(rsp.id(), 7);
}
}