LatestProcessedVersion step (#18)

* add version tracker

* move files around

* make it work with events processor

Author: rtso

rust/sdk-examples/src/common_steps/latest_processed_version_tracker.rs

@@ -0,0 +1,181 @@
+use crate::{
+    config::indexer_processor_config::DbConfig,
+    db::models::processor_status::ProcessorStatus,
+    schema::processor_status,
+    utils::database::{execute_with_better_error, new_db_pool, ArcDbPool},
+};
+use ahash::AHashMap;
+use anyhow::{Context, Result};
+use aptos_indexer_processor_sdk::utils::time::parse_timestamp;
+use aptos_indexer_processor_sdk::{
+    steps::{pollable_async_step::PollableAsyncRunType, PollableAsyncStep},
+    traits::{NamedStep, Processable},
+    types::transaction_context::TransactionContext,
+};
+use async_trait::async_trait;
+use diesel::{upsert::excluded, ExpressionMethods};
+
+const UPDATE_PROCESSOR_STATUS_SECS: u64 = 1;
+
+pub struct LatestVersionProcessedTracker<T>
+where
+    Self: Sized + Send + 'static,
+    T: Send + 'static,
+{
+    conn_pool: ArcDbPool,
+    tracker_name: String,
+    // Next version to process that we expect.
+    next_version: u64,
+    // Last successful batch of sequentially processed transactions. Includes metadata to write to storage.
+    last_success_batch: Option<TransactionContext<T>>,
+    // Tracks all the versions that have been processed out of order.
+    seen_versions: AHashMap<u64, TransactionContext<T>>,
+}
+
+impl<T> LatestVersionProcessedTracker<T>
+where
+    Self: Sized + Send + 'static,
+    T: Send + 'static,
+{
+    pub async fn new(
+        db_config: DbConfig,
+        starting_version: u64,
+        tracker_name: String,
+    ) -> Result<Self> {
+        let conn_pool = new_db_pool(
+            &db_config.postgres_connection_string,
+            Some(db_config.db_pool_size),
+        )
+        .await
+        .context("Failed to create connection pool")?;
+        Ok(Self {
+            conn_pool,
+            tracker_name,
+            next_version: starting_version,
+            last_success_batch: None,
+            seen_versions: AHashMap::new(),
+        })
+    }
+
+    fn update_last_success_batch(&mut self, current_batch: TransactionContext<T>) {
+        let mut new_prev_batch = current_batch;
+        // While there are batches in seen_versions that are in order, update the new_prev_batch to the next batch.
+        while let Some(next_version) = self.seen_versions.remove(&(new_prev_batch.end_version + 1))
+        {
+            new_prev_batch = next_version;
+        }
+        self.next_version = new_prev_batch.end_version + 1;
+        self.last_success_batch = Some(new_prev_batch);
+    }
+}
+
+#[async_trait]
+impl<T> Processable for LatestVersionProcessedTracker<T>
+where
+    Self: Sized + Send + 'static,
+    T: Send + 'static,
+{
+    type Input = T;
+    type Output = T;
+    type RunType = PollableAsyncRunType;
+
+    async fn process(
+        &mut self,
+        current_batch: TransactionContext<T>,
+    ) -> Option<TransactionContext<T>> {
+        // If there's a gap in the next_version and current_version, save the current_version to seen_versions for
+        // later processing.
+        if self.next_version != current_batch.start_version {
+            tracing::debug!(
+                next_version = self.next_version,
+                step = self.name(),
+                "Gap detected starting from version: {}",
+                current_batch.start_version
+            );
+            self.seen_versions.insert(
+                current_batch.start_version,
+                TransactionContext {
+                    data: vec![], // No data is needed for tracking. This is to avoid clone.
+                    start_version: current_batch.start_version,
+                    end_version: current_batch.end_version,
+                    start_transaction_timestamp: current_batch.start_transaction_timestamp.clone(),
+                    end_transaction_timestamp: current_batch.end_transaction_timestamp.clone(),
+                    total_size_in_bytes: current_batch.total_size_in_bytes,
+                },
+            );
+        } else {
+            tracing::debug!("No gap detected");
+            // If the current_batch is the next expected version, update the last success batch
+            self.update_last_success_batch(TransactionContext {
+                data: vec![], // No data is needed for tracking. This is to avoid clone.
+                start_version: current_batch.start_version,
+                end_version: current_batch.end_version,
+                start_transaction_timestamp: current_batch.start_transaction_timestamp.clone(),
+                end_transaction_timestamp: current_batch.end_transaction_timestamp.clone(),
+                total_size_in_bytes: current_batch.total_size_in_bytes,
+            });
+        }
+        // Pass through
+        Some(current_batch)
+    }
+}
+
+#[async_trait]
+impl<T: Send + 'static> PollableAsyncStep for LatestVersionProcessedTracker<T>
+where
+    Self: Sized + Send + 'static,
+    T: Send + 'static,
+{
+    fn poll_interval(&self) -> std::time::Duration {
+        std::time::Duration::from_secs(UPDATE_PROCESSOR_STATUS_SECS)
+    }
+
+    async fn poll(&mut self) -> Option<Vec<TransactionContext<T>>> {
+        // TODO: Add metrics for gap count
+        // Update the processor status
+        if let Some(last_success_batch) = self.last_success_batch.as_ref() {
+            let end_timestamp = last_success_batch
+                .end_transaction_timestamp
+                .as_ref()
+                .map(|t| parse_timestamp(t, last_success_batch.end_version as i64))
+                .map(|t| t.naive_utc());
+            let status = ProcessorStatus {
+                processor: self.tracker_name.clone(),
+                last_success_version: last_success_batch.end_version as i64,
+                last_transaction_timestamp: end_timestamp,
+            };
+            execute_with_better_error(
+                self.conn_pool.clone(),
+                diesel::insert_into(processor_status::table)
+                    .values(&status)
+                    .on_conflict(processor_status::processor)
+                    .do_update()
+                    .set((
+                        processor_status::last_success_version
+                            .eq(excluded(processor_status::last_success_version)),
+                        processor_status::last_updated.eq(excluded(processor_status::last_updated)),
+                        processor_status::last_transaction_timestamp
+                            .eq(excluded(processor_status::last_transaction_timestamp)),
+                    )),
+                Some(" WHERE processor_status.last_success_version <= EXCLUDED.last_success_version "),
+            )
+            .await
+            .expect("Failed to update processor status");
+        }
+        // Nothing should be returned
+        None
+    }
+}
+
+impl<T> NamedStep for LatestVersionProcessedTracker<T>
+where
+    Self: Sized + Send + 'static,
+    T: Send + 'static,
+{
+    fn name(&self) -> String {
+        format!(
+            "LatestVersionProcessedTracker: {}",
+            std::any::type_name::<T>()
+        )
+    }
+}

rust/sdk-examples/src/common_steps/mod.rs

@@ -0,0 +1 @@
+pub mod latest_processed_version_tracker;

rust/sdk-examples/src/config/indexer_processor_config.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::processor_config::ProcessorConfig;
-use crate::events_processor::processor::EventsProcessor;
+use crate::processors::events::events_processor::EventsProcessor;
 use anyhow::Result;
 use aptos_indexer_transaction_stream::TransactionStreamConfig;
 use sdk_server_framework::RunnableConfig;

rust/sdk-examples/src/db/mod.rs

@@ -0,0 +1 @@
+pub mod models;

rust/sdk-examples/src/events_processor/events_models.rs renamed to rust/sdk-examples/src/db/models/events_models.rs

@@ -0,0 +1,2 @@
+pub mod events_models;
+pub mod processor_status;

rust/sdk-examples/src/db/models/mod.rs

@@ -0,0 +1,40 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+#![allow(clippy::extra_unused_lifetimes)]
+
+use crate::{schema::processor_status, utils::database::DbPoolConnection};
+use diesel::{AsChangeset, ExpressionMethods, Insertable, OptionalExtension, QueryDsl, Queryable};
+use diesel_async::RunQueryDsl;
+
+#[derive(AsChangeset, Debug, Insertable)]
+#[diesel(table_name = processor_status)]
+/// Only tracking the latest version successfully processed
+pub struct ProcessorStatus {
+    pub processor: String,
+    pub last_success_version: i64,
+    pub last_transaction_timestamp: Option<chrono::NaiveDateTime>,
+}
+
+#[derive(AsChangeset, Debug, Queryable)]
+#[diesel(table_name = processor_status)]
+/// Only tracking the latest version successfully processed
+pub struct ProcessorStatusQuery {
+    pub processor: String,
+    pub last_success_version: i64,
+    pub last_updated: chrono::NaiveDateTime,
+    pub last_transaction_timestamp: Option<chrono::NaiveDateTime>,
+}
+
+impl ProcessorStatusQuery {
+    pub async fn get_by_processor(
+        processor_name: &str,
+        conn: &mut DbPoolConnection<'_>,
+    ) -> diesel::QueryResult<Option<Self>> {
+        processor_status::table
+            .filter(processor_status::processor.eq(processor_name))
+            .first::<Self>(conn)
+            .await
+            .optional()
+    }
+}

rust/sdk-examples/src/db/models/processor_status.rs

@@ -1,4 +1,6 @@
+pub mod common_steps;
 pub mod config;
-pub mod events_processor;
+pub mod db;
+pub mod processors;
 pub mod schema;
 pub mod utils;

rust/sdk-examples/src/lib.rs

@@ -1,4 +1,4 @@
-use super::events_models::EventModel;
+use crate::db::models::events_models::EventModel;
 use aptos_indexer_processor_sdk::{
     steps::{async_step::AsyncRunType, AsyncStep},
     traits::{NamedStep, Processable},

rust/sdk-examples/src/events_processor/events_extractor.rs renamed to rust/sdk-examples/src/processors/events/events_extractor.rs

@@ -1,5 +1,8 @@
 use super::{events_extractor::EventsExtractor, events_storer::EventsStorer};
-use crate::config::indexer_processor_config::DbConfig;
+use crate::{
+    common_steps::latest_processed_version_tracker::LatestVersionProcessedTracker,
+    config::indexer_processor_config::DbConfig,
+};
 use anyhow::Result;
 use aptos_indexer_processor_sdk::{
     builder::ProcessorBuilder,
@@ -26,21 +29,29 @@ impl EventsProcessor {
     pub async fn run_processor(self) -> Result<()> {
         let (_input_sender, input_receiver) = instrumented_bounded_channel("input", 1);
 
-        let transaction_stream = TransactionStreamStep::new(self.transaction_stream_config).await?;
+        let transaction_stream =
+            TransactionStreamStep::new(self.transaction_stream_config.clone()).await?;
         let transaction_stream_with_input = RunnableStepWithInputReceiver::new(
             input_receiver,
             transaction_stream.into_runnable_step(),
         );
         let events_extractor = EventsExtractor {};
         let events_storer = EventsStorer::new(self.db_config.clone()).await?;
         let timed_buffer = TimedBuffer::new(Duration::from_secs(1));
+        let version_tracker = LatestVersionProcessedTracker::new(
+            self.db_config,
+            self.transaction_stream_config.starting_version,
+            "events_processor".to_string(),
+        )
+        .await?;
 
         let (_, buffer_receiver) = ProcessorBuilder::new_with_runnable_input_receiver_first_step(
             transaction_stream_with_input,
         )
         .connect_to(events_extractor.into_runnable_step(), 10)
         .connect_to(timed_buffer.into_runnable_step(), 10)
         .connect_to(events_storer.into_runnable_step(), 10)
+        .connect_to(version_tracker.into_runnable_step(), 10)
         .end_and_return_output_receiver(10);
 
         loop {