use crate::{

buffers::Acker,

config::{log_schema, DataType, GenerateConfig, SinkConfig, SinkContext, SinkDescription},

event::Event,

internal_events::PulsarEncodeEventFailed,

sinks::util::encoding::{EncodingConfig, EncodingConfiguration},

};

use futures::{future::BoxFuture, ready, stream::FuturesUnordered, FutureExt, Sink, Stream};

use pulsar::{

message::proto, producer::SendFuture, proto::CommandSendReceipt, Authentication,

Error as PulsarError, Producer, Pulsar, TokioExecutor,

};

use serde::{Deserialize, Serialize};

use snafu::{ResultExt, Snafu};

use std::{

collections::HashSet,

pin::Pin,

task::{Context, Poll},

};

#[derive(Debug, Snafu)]

enum BuildError {

#[snafu(display("creating pulsar producer failed: {}", source))]

CreatePulsarSink { source: PulsarError },

}

#[derive(Serialize, Deserialize, Clone, Debug)]

pub struct PulsarSinkConfig {

// Deprecated name

#[serde(alias = "address")]

endpoint: String,

topic: String,

encoding: EncodingConfig<Encoding>,

auth: Option<AuthConfig>,

}

#[derive(Serialize, Deserialize, Clone, Debug)]

pub struct AuthConfig {

name: String, // "token"

token: String, // <jwt token>

}

#[derive(Clone, Copy, Debug, Derivative, Deserialize, Serialize, Eq, PartialEq)]

#[serde(rename_all = "snake_case")]

pub enum Encoding {

Text,

Json,

Avro,

}

type PulsarProducer = Producer<TokioExecutor>;

type BoxedPulsarProducer = Box<PulsarProducer>;

enum PulsarSinkState {

None,

Ready(BoxedPulsarProducer),

Sending(BoxFuture<'static, (BoxedPulsarProducer, Result<SendFuture, PulsarError>)>),

}

struct PulsarSink {

encoding: EncodingConfig<Encoding>,

avro_schema: Option<avro_rs::Schema>,

state: PulsarSinkState,

in_flight:

FuturesUnordered<BoxFuture<'static, (usize, Result<CommandSendReceipt, PulsarError>)>>,

acker: Acker,

seq_head: usize,

seq_tail: usize,

pending_acks: HashSet<usize>,

}

inventory::submit! {

SinkDescription::new::<PulsarSinkConfig>("pulsar")

}

impl GenerateConfig for PulsarSinkConfig {

fn generate_config() -> toml::Value {

toml::Value::try_from(Self {

endpoint: "pulsar://127.0.0.1:6650".to_string(),

topic: "topic-1234".to_string(),

encoding: Encoding::Text.into(),

auth: None,

})

.unwrap()

}

}

#[async_trait::async_trait]

#[typetag::serde(name = "pulsar")]

impl SinkConfig for PulsarSinkConfig {

async fn build(

&self,

cx: SinkContext,

) -> crate::Result<(super::VectorSink, super::Healthcheck)> {

let producer = self

.create_pulsar_producer()

.await

.context(CreatePulsarSink)?;

let sink = PulsarSink::new(producer, self.encoding.clone(), cx.acker())?;

let producer = self

.create_pulsar_producer()

.await

.context(CreatePulsarSink)?;

let healthcheck = healthcheck(producer).boxed();

Ok((super::VectorSink::Sink(Box::new(sink)), healthcheck))

}

fn input_type(&self) -> DataType {

DataType::Log

}

fn sink_type(&self) -> &'static str {

"pulsar"

}

}

impl PulsarSinkConfig {

async fn create_pulsar_producer(&self) -> Result<PulsarProducer, PulsarError> {

let mut builder = Pulsar::builder(&self.endpoint, TokioExecutor);

if let Some(auth) = &self.auth {

builder = builder.with_auth(Authentication {

name: auth.name.clone(),

data: auth.token.as_bytes().to_vec(),

});

}

if let Some(avro_schema) = &self.encoding.schema() {

let pulsar = builder.build().await?;

pulsar

.producer()

.with_options(pulsar::producer::ProducerOptions {

schema: Some(proto::Schema {

schema_data: avro_schema.to_string().into_bytes(),

type_: proto::schema::Type::Avro as i32,

..Default::default()

}),

..Default::default()

})

.with_topic(&self.topic)

.build()

.await

} else {

let pulsar = builder.build().await?;

pulsar.producer().with_topic(&self.topic).build().await

}

}

}

async fn healthcheck(producer: PulsarProducer) -> crate::Result<()> {

producer.check_connection().await.map_err(Into::into)

}

impl PulsarSink {

fn new(

producer: PulsarProducer,

encoding: EncodingConfig<Encoding>,

acker: Acker,

) -> crate::Result<Self> {

let schema = match &encoding.codec() {

Encoding::Avro => {

if let Some(schema) = &encoding.schema() {

avro_rs::Schema::parse_str(schema).ok()

} else {

return Err(

"Avro requires a schema, specify a schema file with `encoding.schema`."

.into(),

);

}

}

_ => None,

};

Ok(Self {

encoding,

avro_schema: schema,

state: PulsarSinkState::Ready(Box::new(producer)),

in_flight: FuturesUnordered::new(),

acker,

seq_head: 0,

seq_tail: 0,

pending_acks: HashSet::new(),

})

}

fn poll_in_flight_prepare(&mut self, cx: &mut Context<'_>) -> Poll<()> {

if let PulsarSinkState::Sending(fut) = &mut self.state {

let (producer, result) = ready!(fut.as_mut().poll(cx));

let seqno = self.seq_head;

self.seq_head += 1;

self.state = PulsarSinkState::Ready(producer);

self.in_flight.push(Box::pin(async move {

let result = match result {

Ok(fut) => fut.await,

Err(error) => Err(error),

};

(seqno, result)

}));

}

Poll::Ready(())

}

}

impl Sink<Event> for PulsarSink {

type Error = ();

fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {

ready!(self.poll_in_flight_prepare(cx));

Poll::Ready(Ok(()))

}

fn start_send(mut self: Pin<&mut Self>, item: Event) -> Result<(), Self::Error> {

assert!(

matches!(self.state, PulsarSinkState::Ready(_)),

"Expected `poll_ready` to be called first."

);

let message = encode_event(item, &self.encoding, &self.avro_schema).map_err(|e| {

emit!(PulsarEncodeEventFailed {

error: &*e.to_string()

})

})?;

let mut producer = match std::mem::replace(&mut self.state, PulsarSinkState::None) {

PulsarSinkState::Ready(producer) => producer,

_ => unreachable!(),

};

let _ = std::mem::replace(

&mut self.state,

PulsarSinkState::Sending(Box::pin(async move {

let result = producer.send(message).await;

(producer, result)

})),

);

Ok(())

}

fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {

ready!(self.poll_in_flight_prepare(cx));

let this = Pin::into_inner(self);

while !this.in_flight.is_empty() {

match ready!(Pin::new(&mut this.in_flight).poll_next(cx)) {

Some((seqno, Ok(result))) => {

trace!(

message = "Pulsar sink produced message.",

message_id = ?result.message_id,

producer_id = %result.producer_id,

sequence_id = %result.sequence_id,

);

this.pending_acks.insert(seqno);

let mut num_to_ack = 0;

while this.pending_acks.remove(&this.seq_tail) {

num_to_ack += 1;

this.seq_tail += 1

}

this.acker.ack(num_to_ack);

}

Some((_, Err(error))) => {

error!(message = "Pulsar sink generated an error.", %error);

return Poll::Ready(Err(()));

}

None => break,

}

}

Poll::Ready(Ok(()))

}

fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {

self.poll_flush(cx)

}

}

fn encode_event(

mut item: Event,

encoding: &EncodingConfig<Encoding>,

avro_schema: &Option<avro_rs::Schema>,

) -> crate::Result<Vec<u8>> {

encoding.apply_rules(&mut item);

let log = item.into_log();

Ok(match encoding.codec() {

Encoding::Json => serde_json::to_vec(&log)?,

Encoding::Text => log

.get(log_schema().message_key())

.map(|v| v.as_bytes().to_vec())

.unwrap_or_default(),

Encoding::Avro => {

let value = avro_rs::to_value(log)?;

let resolved_value =

avro_rs::types::Value::resolve(value, avro_schema.as_ref().unwrap())?;

avro_rs::to_avro_datum(

&avro_schema

.as_ref()

.expect("Avro encoding selected but no schema found. Please report this."),

resolved_value,

)?

}

})

}

#[cfg(test)]

mod tests {

use super::*;

use std::collections::HashMap;

#[test]

fn generate_config() {

crate::test_util::test_generate_config::<PulsarSinkConfig>();

}

#[test]

fn pulsar_event_json() {

let msg = "hello_world".to_owned();

let mut evt = Event::from(msg.clone());

evt.as_mut_log().insert("key", "value");

let result = encode_event(evt, &EncodingConfig::from(Encoding::Json), &None).unwrap();

let map: HashMap<String, String> = serde_json::from_slice(&result[..]).unwrap();

assert_eq!(msg, map[&log_schema().message_key().to_string()]);

}

#[test]

fn pulsar_event_text() {

let msg = "hello_world".to_owned();

let evt = Event::from(msg.clone());

let event = encode_event(evt, &EncodingConfig::from(Encoding::Text), &None).unwrap();

assert_eq!(&event[..], msg.as_bytes());

}

#[test]

fn pulsar_event_avro() {

let raw_schema = r#"

{

"type": "record",

"name": "Log",

"fields": [

{"name": "message","type": ["null","string"]}

]

}

"#;

let msg = "hello_world".to_owned();

let mut evt = Event::from(msg);

evt.as_mut_log().insert("key", "value");

let mut encoding = EncodingConfig::from(Encoding::Avro);

encoding.schema = Some(raw_schema.to_string());

let schema = avro_rs::Schema::parse_str(&raw_schema).unwrap();

let result = encode_event(evt.clone(), &encoding, &Some(schema.clone())).unwrap();

let value = avro_rs::to_value(evt.into_log()).unwrap();

let resolved_value = avro_rs::types::Value::resolve(value, &schema).unwrap();

let must_be = avro_rs::to_avro_datum(&schema, resolved_value).unwrap();

assert_eq!(result, must_be);

}

#[test]

fn pulsar_encode_event() {

let msg = "hello_world";

let mut evt = Event::from(msg);

evt.as_mut_log().insert("key", "value");

let event = encode_event(

evt,

&EncodingConfig {

codec: Encoding::Json,

schema: None,

only_fields: None,

except_fields: Some(vec!["key".into()]),

timestamp_format: None,

},

&None,

)

.unwrap();

let map: HashMap<String, String> = serde_json::from_slice(&event[..]).unwrap();

assert!(!map.contains_key("key"));

}

}

#[cfg(feature = "pulsar-integration-tests")]

#[cfg(test)]

mod integration_tests {

use super::*;

use crate::test_util::{random_lines_with_stream, random_string, trace_init};

use futures::StreamExt;

use pulsar::SubType;

#[tokio::test]

async fn pulsar_happy() {

trace_init();

let num_events = 1_000;

let (_input, events) = random_lines_with_stream(100, num_events);

let topic = format!("test-{}", random_string(10));

let cnf = PulsarSinkConfig {

endpoint: "pulsar://127.0.0.1:6650".to_owned(),

topic: topic.clone(),

encoding: Encoding::Text.into(),

auth: None,

};

let pulsar = Pulsar::<TokioExecutor>::builder(&cnf.endpoint, TokioExecutor)

.build()

.await

.unwrap();

let mut consumer = pulsar

.consumer()

.with_topic(&topic)

.with_consumer_name("VectorTestConsumer")

.with_subscription_type(SubType::Shared)

.with_subscription("VectorTestSub")

.build::<String>()

.await

.unwrap();

let (acker, ack_counter) = Acker::new_for_testing();

let producer = cnf.create_pulsar_producer().await.unwrap();

let sink = PulsarSink::new(producer, cnf.encoding, acker).unwrap();

events.map(Ok).forward(sink).await.unwrap();

assert_eq!(

ack_counter.load(std::sync::atomic::Ordering::Relaxed),

num_events

);

for _ in 0..num_events {

let msg = match consumer.next().await.unwrap() {

Ok(msg) => msg,

Err(error) => panic!("{:?}", error),

};

consumer.ack(&msg).await.unwrap();

}

}

}