use super::{ShardWriter, SinkPoolConfig, SinkProbeFn};
pub trait SinkBundle {
type Writer: ShardWriter;
fn into_parts(self) -> SinkParts<Self::Writer>;
}
#[non_exhaustive]
pub struct SinkParts<W: ShardWriter> {
pub writer: W,
pub shard_endpoints: Vec<Vec<W::Endpoint>>,
pub pool: SinkPoolConfig,
pub component_type: String,
pub replica_labels: Vec<Vec<String>>,
pub probe: Option<SinkProbeFn>,
}
impl<W: ShardWriter> SinkParts<W> {
pub fn new(writer: W, shard_endpoints: Vec<Vec<W::Endpoint>>, pool: SinkPoolConfig) -> Self {
SinkParts {
writer,
shard_endpoints,
pool,
component_type: "custom".to_string(),
replica_labels: Vec::new(),
probe: None,
}
}
#[must_use]
pub fn with_component_type(mut self, component_type: impl Into<String>) -> Self {
self.component_type = component_type.into();
self
}
#[must_use]
pub fn with_replica_labels(mut self, labels: Vec<Vec<String>>) -> Self {
self.replica_labels = labels;
self
}
#[must_use]
pub fn with_probe(mut self, probe: SinkProbeFn) -> Self {
self.probe = Some(probe);
self
}
pub fn effective_replica_labels(&self) -> Vec<Vec<String>> {
if self.replica_labels.is_empty() {
self.shard_endpoints
.iter()
.enumerate()
.map(|(shard, replicas)| {
(0..replicas.len())
.map(|replica| format!("{}-{shard}-{replica}", self.component_type))
.collect()
})
.collect()
} else {
self.replica_labels.clone()
}
}
}
impl<W: ShardWriter> std::fmt::Debug for SinkParts<W> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SinkParts")
.field("component_type", &self.component_type)
.field(
"shards",
&self
.shard_endpoints
.iter()
.map(Vec::len)
.collect::<Vec<_>>(),
)
.field("pool", &self.pool)
.field("probe", &self.probe.is_some())
.finish_non_exhaustive()
}
}
impl<W: ShardWriter> SinkBundle for SinkParts<W> {
type Writer = W;
fn into_parts(self) -> SinkParts<W> {
self
}
}
#[cfg(all(test, not(loom)))]
mod tests {
use super::*;
use crate::error::SinkError;
use crate::sink::SealedBatch;
struct NullWriter;
impl ShardWriter for NullWriter {
type Endpoint = ();
async fn write_batch(&self, (): &(), _batch: &SealedBatch) -> Result<(), SinkError> {
Ok(())
}
}
#[test]
fn default_replica_labels_follow_topology_shape() {
let parts = SinkParts::new(
NullWriter,
vec![vec![(), ()], vec![()]],
SinkPoolConfig::default(),
)
.with_component_type("stdout");
assert_eq!(
parts.effective_replica_labels(),
vec![
vec!["stdout-0-0".to_string(), "stdout-0-1".to_string()],
vec!["stdout-1-0".to_string()]
]
);
}
#[test]
fn explicit_replica_labels_win() {
let parts = SinkParts::new(NullWriter, vec![vec![()]], SinkPoolConfig::default())
.with_replica_labels(vec![vec!["primary".to_string()]]);
assert_eq!(
parts.effective_replica_labels(),
vec![vec!["primary".to_string()]]
);
}
#[test]
fn sink_parts_round_trips_through_the_trait() {
let parts = SinkParts::new(NullWriter, vec![vec![()]], SinkPoolConfig::default())
.with_component_type("capture");
let parts = SinkBundle::into_parts(parts);
assert_eq!(parts.component_type, "capture");
assert_eq!(parts.shard_endpoints.len(), 1);
}
}