dynamo_runtime/pipeline/nodes/sinks/
segment.rs1use super::*;
17use crate::Error;
18
19impl<Req: PipelineIO, Resp: PipelineIO> SegmentSink<Req, Resp> {
20 pub fn new() -> Arc<Self> {
21 Arc::new(Self::default())
22 }
23
24 pub fn attach(&self, engine: ServiceEngine<Req, Resp>) -> Result<(), PipelineError> {
25 self.engine
26 .set(engine)
27 .map_err(|_| PipelineError::EdgeAlreadySet)
28 }
29}
30
31impl<Req: PipelineIO, Resp: PipelineIO> Default for SegmentSink<Req, Resp> {
32 fn default() -> Self {
33 Self {
34 engine: OnceLock::new(),
35 inner: SinkEdge::default(),
36 }
37 }
38}
39
40#[async_trait]
41impl<Req: PipelineIO + Sync, Resp: PipelineIO> Sink<Req> for SegmentSink<Req, Resp> {
42 async fn on_data(&self, data: Req, _: Token) -> Result<(), Error> {
43 let stream = self
44 .engine
45 .get()
46 .ok_or(PipelineError::NoNetworkEdge)?
47 .generate(data)
48 .await?;
49 self.on_next(stream, Token).await
50 }
51}
52
53#[async_trait]
54impl<Req: PipelineIO, Resp: PipelineIO> Source<Resp> for SegmentSink<Req, Resp> {
55 async fn on_next(&self, data: Resp, _: Token) -> Result<(), Error> {
56 self.inner.on_next(data, Token).await
57 }
58
59 fn set_edge(&self, edge: Edge<Resp>, _: Token) -> Result<(), PipelineError> {
60 self.inner.set_edge(edge, Token)
61 }
62}