Timeout (#48)

Timeout (#48)

Makefile.toml

@@ -46,6 +46,7 @@ dependencies = ["release"]
 script = '''
     mkdir -p ${DESTDIR}/usr/local/bin
     sudo cp target/release/kinesis-tailr ${DESTDIR}/usr/local/bin/
+    echo Installed to ${DESTDIR}/usr/local/bin/kinesis-tailr
 '''
 
 [tasks.install]

README.md

@@ -38,6 +38,7 @@ The [release page](https://github.com/grumlimited/kinesis-tailr/releases) provid
         --from-datetime <FROM_DATETIME>      Start datetime position to tail from. ISO 8601 format
         --to-datetime <TO_DATETIME>          End datetime position to tail up to. ISO 8601 format
         --max-messages <MAX_MESSAGES>        Maximum number of messages to retrieve
+        --timeout <TIMEOUT>                  Exit if no messages received after <timeout> seconds
         --max-attempts <MAX_ATTEMPTS>        Maximum number of aws sdk retries. Increase if you are seeing throttling errors [default: 3]
         --no-color                           Disable color output
         --print-delimiter                    Print a delimiter between each payload

src/cli_helpers.rs

@@ -35,6 +35,10 @@ pub struct Opt {
     #[structopt(long)]
     pub max_messages: Option<u32>,
 
+    /// Exit if no messages received after <timeout> seconds.
+    #[structopt(long)]
+    pub timeout: Option<u16>,
+
     /// Maximum number of aws sdk retries. Increase if you are seeing throttling errors.
     #[structopt(long)]
     #[clap(default_value_t = 3)]

src/kinesis.rs

@@ -200,6 +200,7 @@ where
                     .send(TickerMessage::CountUpdate(ShardCountUpdate {
                         shard_id: self.get_config().shard_id.clone(),
                         millis_behind,
+                        nb_records,
                     }))
                     .await
                     .expect("Could not send TickerUpdate to tx_ticker_updates");
@@ -236,8 +237,7 @@ where
 
             tx_shard_iterator_progress
                 .send(shard_iterator_progress)
-                .await
-                .unwrap();
+                .await?;
         } else {
             debug!(
                 "{} records in batch for shard-id {} and {} records before {}",

src/kinesis/helpers.rs

@@ -3,11 +3,10 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use anyhow::Result;
-use aws_sdk_kinesis::error::SdkError;
 use aws_sdk_kinesis::operation::get_shard_iterator::{
     GetShardIteratorError, GetShardIteratorOutput,
 };
-use aws_sdk_kinesis::operation::list_shards::{ListShardsError, ListShardsOutput};
+use aws_sdk_kinesis::operation::list_shards::ListShardsOutput;
 use chrono::Utc;
 use log::{debug, info};
 use tokio::sync::mpsc::Sender;
@@ -180,15 +179,7 @@ pub async fn get_shards(client: &AwsKinesisClient, stream: &str) -> io::Result<V
 
             Ok(shards)
         }
-        Err(e) => {
-            let message = match e.downcast_ref::<SdkError<ListShardsError>>() {
-                Some(SdkError::ServiceError(inner)) => inner.err().to_string(),
-                Some(other) => other.to_string(),
-                _ => e.to_string(),
-            };
-
-            Err(io::Error::new(io::ErrorKind::Other, message))
-        }
+        Err(e) => Err(io::Error::new(io::ErrorKind::Other, format!("{:?}", e))),
     }
 }
 

src/kinesis/models.rs

@@ -7,7 +7,7 @@ use anyhow::Result;
 use async_trait::async_trait;
 use aws_sdk_kinesis::operation::get_shard_iterator::GetShardIteratorOutput;
 use aws_sdk_kinesis::primitives::DateTime;
-use chrono::Utc;
+use chrono::{Duration, Utc};
 use std::fmt::Debug;
 use std::sync::Arc;
 use thiserror::Error;
@@ -32,6 +32,8 @@ pub enum ShardProcessorADT {
 pub enum ProcessError {
     #[error("The stream panicked: {0}")]
     PanicError(String),
+    #[error("The stream timed out after {0}.")]
+    Timeout(Duration),
 }
 
 #[derive(Debug, Clone, PartialEq)]

src/kinesis/tests.rs

@@ -130,7 +130,8 @@ async fn produced_record_is_processed() {
         ticker_update,
         TickerMessage::CountUpdate(ShardCountUpdate {
             shard_id: "shardId-000000000000".to_string(),
-            millis_behind: 1000
+            millis_behind: 1000,
+            nb_records: 1
         })
     );
 
@@ -176,7 +177,8 @@ async fn beyond_to_timestamp_is_received() {
         ticker_update,
         TickerMessage::CountUpdate(ShardCountUpdate {
             shard_id: "shardId-000000000000".to_string(),
-            millis_behind: 1000
+            millis_behind: 1000,
+            nb_records: 1
         })
     );