Reflect user types in retries and responses

This change adds mechanisms to 1) inform users when a (possibly
buffered) message has been published, and 2) retry operations, both now
in the user's message type. Previously retries were only possible in an
opaque pubsub grpc type, and the publish listening wasn't available.

Author: rnarubin

Cargo.toml

@@ -23,16 +23,16 @@ maintenance = { status = "actively-developed" }
 default = []
 
 # Backends
-google = ["ya-gcp", "tracing"]
-mock = ["async-channel"]
+google = ["ya-gcp", "tracing", "parking_lot"]
+mock = ["async-channel", "parking_lot"]
 
 # Validators
 json-schema = ["valico", "serde_json", "serde"]
 protobuf = ["prost"]
 
 [[example]]
 name = "googlepubsub"
-required-features = ["google", "json-schema"]
+required-features = ["google", "protobuf"]
 
 [dependencies]
 async-trait = { version = "0.1" }
@@ -47,17 +47,21 @@ uuid = { version = "^0.8", features = ["v4"], default-features = false }
 async-channel = { version = "1.6", optional = true }
 serde = { version = "^1.0", optional = true, default-features = false }
 serde_json = { version = "^1", features = ["std"], optional = true, default-features = false }
-ya-gcp = { version = "0.6", features = ["pubsub"], optional = true }
+parking_lot = { version = "0.11", optional = true }
 prost = { version = "0.8", optional = true, features = ["std"], default-features = false }
 tracing = { version = "0.1.26", optional = true }
 valico = { version = "^3.2", optional = true, default-features = false }
+ya-gcp = { version = "0.6.3", features = ["pubsub"], optional = true }
 
 [dev-dependencies]
 async-channel = { version = "1.6" }
+futures-channel = "0.3.17"
+parking_lot = { version = "0.11" }
 prost = { version = "0.8", features = ["std", "prost-derive"] }
 tokio = { version = "1", features = ["macros", "rt"] }
+tonic = "0.5"
 serde = { version = "1", features = ["derive"] }
-ya-gcp = { version = "0.6", features = ["pubsub", "emulators"] }
+ya-gcp = { version = "0.6.3", features = ["pubsub", "emulators"] }
 structopt = "0.3"
 
 [package.metadata.docs.rs]

examples/googlepubsub.rs

@@ -1,33 +1,41 @@
-use futures_util::{SinkExt, StreamExt};
+//! An example of ingesting messages from a PubSub subscription, applying a
+//! transformation, then submitting those transformations to another PubSub topic.
+
+use futures_util::{SinkExt, StreamExt, TryFutureExt};
 use hedwig::{
     googlepubsub::{
-        AuthFlow, ClientBuilder, ClientBuilderConfig, PubSubConfig, ServiceAccountAuth,
-        StreamSubscriptionConfig, SubscriptionConfig, SubscriptionName, TopicConfig, TopicName,
+        AcknowledgeToken, AuthFlow, ClientBuilder, ClientBuilderConfig, PubSubConfig,
+        PubSubMessage, PublishError, ServiceAccountAuth, StreamSubscriptionConfig,
+        SubscriptionConfig, SubscriptionName, TopicConfig, TopicName,
     },
     validators, Consumer, DecodableMessage, EncodableMessage, Headers, Publisher,
 };
-use std::time::SystemTime;
+use std::{error::Error as StdError, time::SystemTime};
 use structopt::StructOpt;
 
-#[derive(Clone, PartialEq, Eq, prost::Message)]
+const USER_CREATED_TOPIC: &str = "user.created";
+const USER_UPDATED_TOPIC: &str = "user.updated";
+
+/// The input data, representing some user being created with the given name
+#[derive(PartialEq, Eq, prost::Message)]
 struct UserCreatedMessage {
     #[prost(string, tag = "1")]
-    payload: String,
+    name: String,
 }
 
-impl<'a> EncodableMessage for UserCreatedMessage {
+impl EncodableMessage for UserCreatedMessage {
     type Error = validators::ProstValidatorError;
     type Validator = validators::ProstValidator;
     fn topic(&self) -> hedwig::Topic {
-        "user.created".into()
+        USER_CREATED_TOPIC.into()
     }
-    fn encode(self, validator: &Self::Validator) -> Result<hedwig::ValidatedMessage, Self::Error> {
+    fn encode(&self, validator: &Self::Validator) -> Result<hedwig::ValidatedMessage, Self::Error> {
         Ok(validator.validate(
             uuid::Uuid::new_v4(),
             SystemTime::now(),
             "user.created/1.0",
             Headers::new(),
-            &self,
+            self,
         )?)
     }
 }
@@ -42,6 +50,46 @@ impl DecodableMessage for UserCreatedMessage {
     }
 }
 
+/// The output data, where the given user has now been assigned an ID and some metadata
+#[derive(PartialEq, Eq, prost::Message)]
+struct UserUpdatedMessage {
+    #[prost(string, tag = "1")]
+    name: String,
+
+    #[prost(int64, tag = "2")]
+    id: i64,
+
+    #[prost(string, tag = "3")]
+    metadata: String,
+}
+
+#[derive(Debug)]
+struct TransformedMessage {
+    // keep the input's ack token to ack when the output is successfully published, or nack on
+    // failure
+    input_token: AcknowledgeToken,
+    output: UserUpdatedMessage,
+}
+
+impl EncodableMessage for TransformedMessage {
+    type Error = validators::ProstValidatorError;
+    type Validator = validators::ProstValidator;
+
+    fn topic(&self) -> hedwig::Topic {
+        USER_UPDATED_TOPIC.into()
+    }
+
+    fn encode(&self, validator: &Self::Validator) -> Result<hedwig::ValidatedMessage, Self::Error> {
+        Ok(validator.validate(
+            uuid::Uuid::new_v4(),
+            SystemTime::now(),
+            "user.updated/1.0",
+            Headers::new(),
+            &self.output,
+        )?)
+    }
+}
+
 #[derive(Debug, StructOpt)]
 struct Args {
     /// The name of the pubsub project
@@ -50,7 +98,7 @@ struct Args {
 }
 
 #[tokio::main(flavor = "current_thread")]
-async fn main() -> Result<(), Box<dyn std::error::Error + 'static>> {
+async fn main() -> Result<(), Box<dyn StdError>> {
     let args = Args::from_args();
 
     println!("Building PubSub clients");
@@ -61,49 +109,59 @@ async fn main() -> Result<(), Box<dyn std::error::Error + 'static>> {
     )
     .await?;
 
-    let topic_name = TopicName::new("pubsub_example_topic");
-    let subscription_name = SubscriptionName::new("pubsub_example_subscription");
+    let input_topic_name = TopicName::new(USER_CREATED_TOPIC);
+    let subscription_name = SubscriptionName::new("user-metadata-updaters");
+
+    let output_topic_name = TopicName::new(USER_UPDATED_TOPIC);
+    const APP_NAME: &str = "user-metadata-updater";
 
     let mut publisher_client = builder
-        .build_publisher(&args.project_name, "myapp_publisher")
-        .await?;
-    let mut consumer_client = builder
-        .build_consumer(&args.project_name, "myapp_consumer")
+        .build_publisher(&args.project_name, APP_NAME)
         .await?;
+    let mut consumer_client = builder.build_consumer(&args.project_name, APP_NAME).await?;
 
-    println!("Creating topic {:?}", &topic_name);
+    for topic_name in [&input_topic_name, &output_topic_name] {
+        println!("Creating topic {:?}", topic_name);
 
-    publisher_client
-        .create_topic(TopicConfig {
-            name: topic_name.clone(),
-            ..TopicConfig::default()
-        })
-        .await?;
+        publisher_client
+            .create_topic(TopicConfig {
+                name: topic_name.clone(),
+                ..TopicConfig::default()
+            })
+            .await?;
+    }
 
     println!("Creating subscription {:?}", &subscription_name);
 
     consumer_client
         .create_subscription(SubscriptionConfig {
-            topic: topic_name.clone(),
+            topic: input_topic_name.clone(),
             name: subscription_name.clone(),
             ..SubscriptionConfig::default()
         })
         .await?;
 
-    println!("Publishing message to topic");
+    println!(
+        "Synthesizing input messages for topic {:?}",
+        &input_topic_name
+    );
 
-    let validator = hedwig::validators::ProstValidator::new();
-    let mut publisher = publisher_client.publisher().publish_sink(validator);
+    {
+        let validator = validators::ProstValidator::new();
+        let mut input_sink =
+            Publisher::<UserCreatedMessage>::publish_sink(publisher_client.publisher(), validator);
 
-    for i in 1..=10 {
-        let message = UserCreatedMessage {
-            payload: format!("this is message #{}", i),
-        };
+        for i in 1..=10 {
+            let message = UserCreatedMessage {
+                name: format!("Example Name #{}", i),
+            };
 
-        publisher.send(message).await?;
+            input_sink.feed(message).await?;
+        }
+        input_sink.flush().await?;
     }
 
-    println!("Reading back published message");
+    println!("Ingesting input messages, applying transformations, and publishing to destination");
 
     let mut read_stream = consumer_client
         .stream_subscription(
@@ -114,35 +172,70 @@ async fn main() -> Result<(), Box<dyn std::error::Error + 'static>> {
             hedwig::validators::prost::ExactSchemaMatcher::new("user.created/1.0"),
         ));
 
+    let mut output_sink = Publisher::<TransformedMessage, _>::publish_sink_with_responses(
+        publisher_client.publisher(),
+        validators::ProstValidator::new(),
+        futures_util::sink::unfold((), |_, message: TransformedMessage| async move {
+            // if the output is successfully sent, ack the input to mark it as processed
+            message.input_token.ack().await
+        }),
+    );
+
     for i in 1..=10 {
-        let message = read_stream
+        let PubSubMessage { ack_token, message } = read_stream
             .next()
             .await
-            .ok_or("unexpected end of stream")??
-            .ack()
-            .await?;
+            .expect("stream should have 10 elements")?;
 
-        assert_eq!(
-            message,
-            UserCreatedMessage {
-                payload: format!("this is message #{}", i)
-            }
-        );
+        assert_eq!(&message.name, &format!("Example Name #{}", i));
+
+        let transformed = TransformedMessage {
+            output: UserUpdatedMessage {
+                name: message.name,
+                id: random_id(),
+                metadata: "some metadata".into(),
+            },
+            input_token: ack_token,
+        };
+
+        output_sink
+            .feed(transformed)
+            .or_else(|publish_error| async move {
+                // if publishing fails, nack the failed messages to allow later retries
+                Err(match publish_error {
+                    PublishError::Publish { cause, messages } => {
+                        for failed_transform in messages {
+                            failed_transform.input_token.nack().await?;
+                        }
+                        Box::<dyn StdError>::from(cause)
+                    }
+                    err => Box::<dyn StdError>::from(err),
+                })
+            })
+            .await?
     }
+    output_sink.flush().await?;
 
-    println!("All messages matched!");
+    println!("All messages matched and published successfully!");
 
     println!("Deleting subscription {:?}", &subscription_name);
 
     consumer_client
         .delete_subscription(subscription_name)
         .await?;
 
-    println!("Deleting topic {:?}", &topic_name);
+    for topic_name in [input_topic_name, output_topic_name] {
+        println!("Deleting topic {:?}", &topic_name);
 
-    publisher_client.delete_topic(topic_name).await?;
+        publisher_client.delete_topic(topic_name).await?;
+    }
 
     println!("Done");
 
     Ok(())
 }
+
+fn random_id() -> i64 {
+    4 // chosen by fair dice roll.
+      // guaranteed to be random.
+}

src/backends/googlepubsub/mod.rs

@@ -7,8 +7,9 @@ use std::{borrow::Cow, fmt::Display};
 pub use ya_gcp::{
     grpc::StatusCodeSet,
     pubsub::{
-        AcknowledgeError, BuildError, Error as PubSubError, MakeConnection, ModifyAcknowledgeError,
-        PubSubConfig, SinkError, StreamSubscriptionConfig, Uri, DEFAULT_RETRY_CODES,
+        AcknowledgeError, AcknowledgeToken, BuildError, Error as PubSubError, MakeConnection,
+        ModifyAcknowledgeError, PubSubConfig, SinkError, StreamSubscriptionConfig, Uri,
+        DEFAULT_RETRY_CODES,
     },
     retry_policy, AuthFlow, ClientBuilderConfig, Connect, CreateBuilderError, DefaultConnector,
     ServiceAccountAuth,