kinetic/sinks/aws/

s3.rs

//! AWS S3 Sink.

use arrow_select::concat::concat_batches;
use aws_common::{client::create_s3_client, config::S3SinkConfig};
use aws_sdk_s3::Client as S3Client;
use kinetic_buffers::BufferReceiver;
use kinetic_core::{EventBatch, ShutdownSignal};
use metrics::{Label, counter};
use time::OffsetDateTime;
use tokio::time::{Duration, interval};
use tracing::{debug, error, info};

use kinetic_core::encode::Encoder;
use std::sync::Arc;

use kinetic_common::register;
use kinetic_common::telemetry::EventDuration;
use std::time::Instant;

pub struct S3SinkTask {
    config: S3SinkConfig,
    component_id: String,
    receiver: BufferReceiver,
    shutdown: ShutdownSignal,
    encoder: Arc<dyn Encoder>,
    labels: Arc<[Label]>,
    event_duration: EventDuration,
}

use crate::topology::healthcheck::Healthcheck;
use async_trait::async_trait;

#[async_trait]
impl Healthcheck for S3SinkTask {
    async fn check(&self) -> anyhow::Result<()> {
        let aws = self.config.auth.as_ref().cloned().unwrap_or_default();
        let client = create_s3_client(&aws).await;
        client
            .head_bucket()
            .bucket(&self.config.bucket)
            .send()
            .await
            .map_err(|e| {
                anyhow::anyhow!("Failed to head bucket '{}': {}", self.config.bucket, e)
            })?;
        Ok(())
    }
}

impl S3SinkTask {
    pub fn new(
        config: S3SinkConfig,
        component_id: String,
        receiver: BufferReceiver,
        shutdown: ShutdownSignal,
        encoder: Arc<dyn Encoder>,
    ) -> Self {
        let labels: Arc<[Label]> = Arc::new([
            Label::new("component_id", component_id.clone()),
            Label::new("component_type", "sink"),
            Label::new("component_kind", "aws_s3"),
        ]);
        let event_duration = register!(EventDuration::new(component_id.clone(), "sink"));

        Self {
            config,
            component_id,
            receiver,
            shutdown,
            encoder,
            labels,
            event_duration,
        }
    }

    pub async fn run(mut self) {
        info!(
            "Starting S3 sink '{}' to bucket: {}",
            self.component_id, self.config.bucket
        );

        let aws = self.config.auth.as_ref().cloned().unwrap_or_default();
        let client = create_s3_client(&aws).await;

        let mut current_batch: Vec<EventBatch> = Vec::new();
        let mut current_rows = 0;

        let timeout = self.config.batch.timeout_secs.unwrap_or(30);
        let mut flush_interval = interval(Duration::from_secs(timeout));
        let max_size = self.config.batch.max_size.unwrap_or(1000);

        let mut shutdown = self.shutdown.clone();

        loop {
            tokio::select! {
                reason = shutdown.recv() => {
                    info!("S3 sink '{}' received shutdown signal: {:?}", self.component_id, reason);
                    if !current_batch.is_empty() {
                        // Attempt one last flush. If it fails, we might lose data but we are shutting down.
                        if let Err(e) = self.flush_batch(&client, &mut current_batch, &mut current_rows).await {
                            error!("Final flush failed in S3 sink {}: {}", self.component_id, e);
                        }
                    }
                    break;
                }
                _ = flush_interval.tick() => {
                    if !current_batch.is_empty() {
                        // We don't block the main loop for scheduled flushes, just try it.
                        // If it fails, the next batch or next tick will try again.
                        if let Err(e) = self.flush_batch(&client, &mut current_batch, &mut current_rows).await {
                            error!(
                                "Scheduled flush failed in S3 sink {}: {}",
                                self.component_id, e
                            );
                        }
                    }
                }
                batch = self.receiver.recv() => {
                    let Some(batch) = batch else {
                        info!("S3 sink '{}' receiver closed", self.component_id);
                        if !current_batch.is_empty()
                        {
                            let start = Instant::now();
                            if let Err(e) = self
                                .flush_batch(&client, &mut current_batch, &mut current_rows)
                                .await
                            {
                                error!(
                                    "Final flush after receiver closed failed in S3 sink {}: {}",
                                    self.component_id, e
                                );
                            }
                            self.event_duration.emit(start.elapsed());
                        }
                        break;
                    };
                    let start = Instant::now();

                    counter!("component_received_events_total", self.labels.iter()).increment(batch.num_rows() as u64);
                    counter!("component_received_event_bytes_total", self.labels.iter()).increment(batch.estimated_size() as u64);

                    // 1. Check for schema change. If changed, we MUST flush existing data first.
                    if !current_batch.is_empty() && batch.payload.schema() != current_batch[0].payload.schema() {
                        debug!("Schema change detected in S3 sink {}, flushing current buffer before accepting new schema", self.component_id);
                        self.flush_with_retry(&client, &mut current_batch, &mut current_rows, &mut shutdown).await;
                        // If we returned from flush_with_retry and it's still not empty, it means we are shutting down.
                        if !current_batch.is_empty() {
                            break;
                        }
                        flush_interval.reset();
                    }

                    // 2. Add to current batch
                    current_rows += batch.num_rows();
                    current_batch.push(batch);

                    // 3. If over size, flush with retry to apply backpressure
                    if current_rows >= max_size {
                        self.flush_with_retry(&client, &mut current_batch, &mut current_rows, &mut shutdown).await;
                        flush_interval.reset();
                    }
                    self.event_duration.emit(start.elapsed());
                }
            }
        }

        info!("S3 sink task {} shutting down", self.component_id);
    }

    async fn flush_with_retry(
        &self,
        client: &S3Client,
        batches: &mut Vec<EventBatch>,
        rows: &mut usize,
        shutdown: &mut ShutdownSignal,
    ) {
        let mut backoff = Duration::from_secs(1);
        while !batches.is_empty() {
            match self.flush_batch(client, batches, rows).await {
                Ok(_) => break,
                Err(e) => {
                    error!(
                        "S3 flush failed for component {}: {}. Retrying in {:?}...",
                        self.component_id, e, backoff
                    );
                    tokio::select! {
                        _ = shutdown.recv() => {
                            error!("Shutdown received during S3 flush retry for {}. Data in current buffer may be lost.", self.component_id);
                            return;
                        }
                        _ = tokio::time::sleep(backoff) => {
                            backoff = (backoff * 2).min(Duration::from_secs(60));
                        }
                    }
                }
            }
        }
    }

    async fn flush_batch(
        &self,
        client: &S3Client,
        batches: &mut Vec<EventBatch>,
        rows: &mut usize,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if batches.is_empty() {
            return Ok(());
        }

        debug!("S3 sink {} flushing {} rows", self.component_id, rows);

        // Prepare the merged batch
        let record_batches: Vec<_> = batches.iter().map(|b| b.payload.clone()).collect();
        let schema = record_batches[0].schema();
        let merged_batch = concat_batches(&schema, record_batches.iter())?;

        // Wrap the merged batch in a temporary EventBatch for encoding
        let metadata = batches[0].metadata.clone();
        let event_batch = EventBatch::new(merged_batch, metadata)?;

        let buf = self.encoder.encode(&event_batch)?;

        let now = OffsetDateTime::now_utc();
        let extension = if self.config.encoding.as_deref() == Some("parquet") {
            "parquet"
        } else {
            "json"
        };
        let key = format!(
            "{}{}.{}",
            self.config.key_prefix,
            now.unix_timestamp_nanos(),
            extension
        );

        let body = aws_sdk_s3::primitives::ByteStream::from(buf.to_vec());
        let buf_len = buf.len();

        // The actual S3 upload
        match client
            .put_object()
            .bucket(&self.config.bucket)
            .key(&key)
            .body(body)
            .send()
            .await
        {
            Ok(_) => {
                counter!("component_sent_network_bytes_total", self.labels.iter())
                    .increment(buf_len as u64);
                counter!("component_sent_events_total", self.labels.iter()).increment(*rows as u64);
                // For simplicity, we use the estimated size of the EventBatch as sent_event_bytes
                counter!("component_sent_event_bytes_total", self.labels.iter())
                    .increment(event_batch.estimated_size() as u64);
            }
            Err(e) => {
                counter!("component_errors_total", self.labels.iter()).increment(1);
                return Err(e.into());
            }
        }

        info!(
            "Successfully uploaded {} rows to s3://{}/{}",
            rows, self.config.bucket, key
        );

        // Success! Now clear and acknowledge all batches
        for mut batch in batches.drain(..) {
            if let Some(token) = batch.ack_token.take() {
                token.ack();
            }
        }

        *rows = 0;
        Ok(())
    }
}