objectstore.rs

//! Objectstore service for uploading attachments.
use std::array::TryFromSliceError;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;

use bytes::Bytes;
use futures::StreamExt;
use objectstore_client::{Client, ExpirationPolicy, PutBuilder, Session, Usecase};
use relay_config::ObjectstoreServiceConfig;
use relay_system::{
    Addr, AsyncResponse, FromMessage, Interface, LoadShed, NoResponse, Sender, SimpleService,
};
use sentry_protos::snuba::v1::TraceItem;
use uuid::Uuid;

use crate::constants::DEFAULT_ATTACHMENT_RETENTION;
use crate::envelope::ItemType;
use crate::managed::{
    Counted, Managed, ManagedEnvelope, ManagedResult, OutcomeError, Quantities, Rejected,
};
use crate::processing::utils::store::item_id_to_uuid;
use crate::services::outcome::DiscardReason;
use crate::services::store::{Store, StoreAttachment, StoreEnvelope, StoreTraceItem};
use crate::services::upload;
use crate::statsd::{RelayCounters, RelayTimers};

use super::outcome::Outcome;

/// Messages that the objectstore service can handle.
pub enum Objectstore {
    Envelope(StoreEnvelope),
    TraceAttachment(Managed<StoreTraceAttachment>),
    EventAttachment(Managed<StoreAttachment>),
    Stream {
        message: upload::Stream,
        sender: Sender<Result<ObjectstoreKey, Error>>,
    },
}

impl Objectstore {
    fn ty(&self) -> &str {
        match self {
            Objectstore::Envelope(_) => "envelope",
            Objectstore::TraceAttachment(_) => "attachment_v2",
            Objectstore::EventAttachment(_) => "attachment",
            Objectstore::Stream { .. } => "stream",
        }
    }

    fn attachment_count(&self) -> usize {
        match self {
            Self::Envelope(StoreEnvelope { envelope }) => envelope
                .envelope()
                .items()
                .filter(|item| *item.ty() == ItemType::Attachment)
                .count(),
            Self::TraceAttachment(_) => 1,
            Self::EventAttachment(_) => 1,
            Self::Stream { .. } => 1,
        }
    }
}

impl Interface for Objectstore {}

impl FromMessage<StoreEnvelope> for Objectstore {
    type Response = NoResponse;

    fn from_message(message: StoreEnvelope, _sender: ()) -> Self {
        Self::Envelope(message)
    }
}

impl FromMessage<Managed<StoreTraceAttachment>> for Objectstore {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreTraceAttachment>, _sender: ()) -> Self {
        Self::TraceAttachment(message)
    }
}

impl FromMessage<Managed<StoreAttachment>> for Objectstore {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreAttachment>, _sender: ()) -> Self {
        Self::EventAttachment(message)
    }
}

impl FromMessage<upload::Stream> for Objectstore {
    type Response = AsyncResponse<Result<ObjectstoreKey, Error>>;

    fn from_message(
        message: upload::Stream,
        sender: Sender<Result<ObjectstoreKey, Error>>,
    ) -> Self {
        Self::Stream { message, sender }
    }
}

/// An attachment that is ready for upload / EAP storage.
pub struct StoreTraceAttachment {
    /// The body to be uploaded to objectstore.
    pub body: Bytes,
    /// The trace item to be published via Kafka.
    pub trace_item: TraceItem,
}

impl Counted for StoreTraceAttachment {
    fn quantities(&self) -> Quantities {
        self.trace_item.quantities()
    }
}

/// Errors that can occur when trying to upload an attachment.
#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("load shed")]
    LoadShed,
    #[error("upload failed: {0}")]
    UploadFailed(#[from] objectstore_client::Error),
    #[error("UUID conversion failed: {0}")]
    Uuid(#[from] TryFromSliceError),
}

impl Error {
    fn as_str(&self) -> &'static str {
        match self {
            Error::LoadShed => "load-shed",
            Error::UploadFailed(_) => "upload_failed",
            Error::Uuid(_) => "uuid",
        }
    }
}

impl OutcomeError for Error {
    type Error = Self;

    fn consume(self) -> (Option<Outcome>, Self::Error) {
        (Some(Outcome::Invalid(DiscardReason::Internal)), self)
    }
}

/// The objectstore key that identifies a successful upload.
#[derive(Debug, PartialEq)]
pub struct ObjectstoreKey(String);

impl ObjectstoreKey {
    /// Returns the underlying [`String`].
    pub fn into_inner(self) -> String {
        self.0
    }
}

impl fmt::Display for ObjectstoreKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

/// The objectstore service that uploads attachments.
///
/// Accepts upload requests and maintains a list of concurrent uploads.
#[derive(Clone)]
pub struct ObjectstoreService {
    inner: Arc<ObjectstoreServiceInner>,
}

impl ObjectstoreService {
    pub fn new(
        config: &ObjectstoreServiceConfig,
        store: Option<Addr<Store>>,
    ) -> anyhow::Result<Option<Self>> {
        let Some(store) = store else { return Ok(None) };
        let ObjectstoreServiceConfig {
            objectstore_url,
            max_concurrent_requests: _,
            max_backlog: _,
            timeout,
        } = config;
        let Some(objectstore_url) = objectstore_url else {
            return Ok(None);
        };

        let objectstore_client = Client::builder(objectstore_url)
            .configure_reqwest(|builder| builder.timeout(Duration::from_secs(*timeout)))
            .build()?;
        let event_attachments = Usecase::new("attachments")
            .with_expiration_policy(ExpirationPolicy::TimeToLive(DEFAULT_ATTACHMENT_RETENTION));
        let trace_attachments = Usecase::new("trace_attachments")
            .with_expiration_policy(ExpirationPolicy::TimeToLive(DEFAULT_ATTACHMENT_RETENTION));

        let inner = ObjectstoreServiceInner {
            store,
            objectstore_client,
            event_attachments,
            trace_attachments,
        };

        Ok(Some(Self {
            inner: Arc::new(inner),
        }))
    }
}

impl SimpleService for ObjectstoreService {
    type Interface = Objectstore;

    async fn handle_message(&self, message: Self::Interface) {
        self.inner.handle_message(message).await
    }
}

impl LoadShed<Objectstore> for ObjectstoreService {
    fn handle_loadshed(&self, message: Objectstore) {
        relay_statsd::metric!(
            counter(RelayCounters::AttachmentUpload) += message.attachment_count() as u64,
            result = "load_shed",
            type = message.ty(),
        );
        match message {
            Objectstore::Envelope(envelope) => {
                // Event attachments can still go the old route.

                self.inner.store.send(envelope);
            }
            Objectstore::EventAttachment(message) => {
                // Event attachments can still go the old route.
                self.inner.store.send(message);
            }
            Objectstore::TraceAttachment(managed) => {
                let _ = managed.reject_err(Error::LoadShed);
            }
            Objectstore::Stream { message: _, sender } => {
                sender.send(Err(Error::LoadShed));
            }
        }
    }
}

struct ObjectstoreServiceInner {
    store: Addr<Store>,

    objectstore_client: Client,
    event_attachments: Usecase,
    trace_attachments: Usecase,
}

impl ObjectstoreServiceInner {
    async fn handle_message(&self, message: Objectstore) {
        match message {
            Objectstore::Envelope(StoreEnvelope { envelope }) => {
                self.handle_envelope(envelope).await;
            }
            Objectstore::TraceAttachment(attachment) => {
                self.handle_trace_attachment(attachment).await
            }
            Objectstore::EventAttachment(attachment) => {
                self.handle_event_attachment(attachment).await
            }
            Objectstore::Stream {
                message: managed,
                sender,
            } => self.handle_stream(managed, sender).await,
        }
    }

    /// Uploads all attachments belonging to the given envelope.
    ///
    /// This mutates the attachment items in-place, setting their `stored_key` field to the key
    /// in objectstore.
    async fn handle_envelope(&self, mut envelope: ManagedEnvelope) -> () {
        let scoping = envelope.scoping();
        let session = self
            .event_attachments
            .for_project(scoping.organization_id.value(), scoping.project_id.value())
            .session(&self.objectstore_client);

        let attachments = envelope
            .envelope_mut()
            .items_mut()
            .filter(|item| *item.ty() == ItemType::Attachment);

        match session {
            Err(error) => {
                relay_log::error!(error = &error as &dyn std::error::Error, "session error");
                relay_statsd::metric!(
                    counter(RelayCounters::AttachmentUpload) += attachments.count() as u64,
                    result = error.to_string().as_str(),
                    type = "envelope",
                );
            }
            Ok(session) => {
                for attachment in attachments {
                    // we are not storing zero-size attachments in objectstore
                    if attachment.is_empty() {
                        continue;
                    }
                    let result = self
                        .upload_bytes("envelope", &session, attachment.payload(), None)
                        .await;

                    relay_statsd::metric!(
                        counter(RelayCounters::AttachmentUpload) += 1,
                        result = match &result {
                            Ok(_) => "success",
                            Err(e) => e.as_str(),
                        },
                        type = "envelope",
                    );
                    if let Ok(stored_key) = result {
                        attachment.set_stored_key(stored_key.into_inner());
                    }
                }
            }
        }

        // last but not least, forward the envelope to the store endpoint
        self.store.send(StoreEnvelope { envelope });
    }

    /// Uploads the attachment.
    ///
    /// This mutates the attachment item in-place, setting the `stored_key` field to the key in the
    /// objectstore.
    async fn handle_event_attachment(&self, mut attachment: Managed<StoreAttachment>) {
        // we are not storing zero-size attachments in objectstore
        if attachment.attachment.is_empty() {
            self.store.send(attachment);
            return;
        }

        let scoping = attachment.scoping();
        let session = self
            .event_attachments
            .for_project(scoping.organization_id.value(), scoping.project_id.value())
            .session(&self.objectstore_client);

        match session {
            Err(error) => {
                relay_log::error!(error = &error as &dyn std::error::Error, "session error");
                relay_statsd::metric!(
                    counter(RelayCounters::AttachmentUpload) += 1,
                    result = error.to_string().as_str(),
                    type = "attachment",
                );
            }
            Ok(session) => {
                let result = self
                    .upload_bytes(
                        "attachment",
                        &session,
                        attachment.attachment.payload(),
                        None,
                    )
                    .await;

                relay_statsd::metric!(
                    counter(RelayCounters::AttachmentUpload) += 1,
                    result = match &result {
                        Ok(_) => "success",
                        Err(e) => e.as_str(),
                    },
                    type = "attachment",
                );

                if let Ok(stored_key) = result {
                    attachment.modify(|attachment, _| {
                        attachment
                            .attachment
                            .set_stored_key(stored_key.into_inner());
                    });
                }
            }
        }

        self.store.send(attachment)
    }

    async fn handle_trace_attachment(&self, managed: Managed<StoreTraceAttachment>) {
        let result = self.do_handle_store_attachment(managed).await;

        relay_statsd::metric!(
            counter(RelayCounters::AttachmentUpload) += 1,
            result = match result {
                Ok(()) => "success",
                Err(e) => e.into_inner().as_str(),
            },
            type = "attachment_v2",
        );
    }

    async fn do_handle_store_attachment(
        &self,
        managed: Managed<StoreTraceAttachment>,
    ) -> Result<(), Rejected<Error>> {
        let scoping = managed.scoping();
        let session = self
            .trace_attachments
            .for_project(scoping.organization_id.value(), scoping.project_id.value())
            .session(&self.objectstore_client)
            .map_err(Error::UploadFailed)
            .reject(&managed)?;

        let body = Bytes::clone(&managed.body);

        // Make sure that the attachment can be converted into a trace item:
        let trace_item = managed.try_map(|attachment, _record_keeper| {
            let StoreTraceAttachment {
                trace_item,
                body: _,
            } = attachment;
            Ok::<_, Error>(StoreTraceItem { trace_item })
        })?;

        // Upload the attachment:
        if !body.is_empty() {
            relay_log::trace!("Starting attachment upload");
            let key = item_id_to_uuid(&trace_item.trace_item.item_id)
                .map_err(Error::from)
                .reject(&trace_item)?
                .as_simple()
                .to_string();

            #[cfg(debug_assertions)]
            let original_key = key.clone();

            let _stored_key = self
                .upload_bytes("attachment_v2", &session, body, Some(key))
                .await
                .reject(&trace_item)?;

            #[cfg(debug_assertions)]
            debug_assert_eq!(_stored_key.into_inner(), original_key);
        }

        // Only after successful upload forward the attachment to the store.
        self.store.send(trace_item);

        Ok(())
    }

    async fn handle_stream(
        &self,
        upload::Stream { scoping, stream }: upload::Stream,
        sender: Sender<Result<ObjectstoreKey, Error>>,
    ) {
        let session = match self
            .event_attachments
            .for_project(scoping.organization_id.value(), scoping.project_id.value())
            .session(&self.objectstore_client)
        {
            Ok(session) => session,
            Err(error) => {
                relay_statsd::metric!(
                    counter(RelayCounters::AttachmentUpload) += 1,
                    result = error.to_string().as_str(),
                    type = "stream",
                );
                sender.send(Err(Error::UploadFailed(error)));
                return;
            }
        };

        let request = session
            .put_stream(stream.boxed())
            // generate ID here to drop the hyphens and be consistent with other attachment uploads.
            .key(Uuid::now_v7().as_simple().to_string());

        let result = self.upload("stream", request).await;

        relay_statsd::metric!(
            counter(RelayCounters::AttachmentUpload) += 1,
            result = match &result {
                Ok(_) => "success",
                Err(e) => e.as_str(),
            },
            type = "stream",
        );

        sender.send(result);
    }

    async fn upload_bytes(
        &self,
        ty: &str,
        session: &Session,
        payload: Bytes,
        key: Option<String>,
    ) -> Result<ObjectstoreKey, Error> {
        let mut request = session.put(payload);
        if let Some(key) = key {
            request = request.key(key);
        }
        self.upload(ty, request).await
    }

    async fn upload(&self, ty: &str, request: PutBuilder) -> Result<ObjectstoreKey, Error> {
        self.upload_inner(ty, request).await.inspect_err(|e| {
            relay_log::error!(
                error = e as &dyn std::error::Error,
                "objectstore upload failed"
            )
        })
    }

    async fn upload_inner(&self, ty: &str, request: PutBuilder) -> Result<ObjectstoreKey, Error> {
        relay_log::trace!("Starting attachment upload");
        let response = relay_statsd::metric!(
            timer(RelayTimers::AttachmentUploadDuration),
            type = ty,
            {
                request.send()
                    .await
                    .map_err(Error::UploadFailed)?
            }
        );

        relay_log::trace!("Finished attachment upload");

        Ok(ObjectstoreKey(response.key))
    }
}