relay_server/services/

store.rs

//! This module contains the service that forwards events and attachments to the Sentry store.
//! The service uses Kafka topics to forward data to Sentry

use std::borrow::Cow;
use std::collections::BTreeMap;
use std::error::Error;
use std::pin::Pin;
use std::sync::Arc;
use std::task;

use bytes::Bytes;
use chrono::{DateTime, SecondsFormat, Utc};
use prost::Message as _;
use relay_base_schema::events::EventType;
use sentry_protos::snuba::v1::{TraceItem, TraceItemType};
use serde::Serialize;
use uuid::Uuid;

use relay_base_schema::data_category::DataCategory;
use relay_base_schema::organization::OrganizationId;
use relay_base_schema::project::ProjectId;
use relay_common::time::UnixTimestamp;
use relay_config::Config;
use relay_event_schema::protocol::{Event, EventId, SpanV2, datetime_to_timestamp};
use relay_kafka::{ClientError, KafkaClient, KafkaTopic, Message, SerializationOutput};
use relay_metrics::{
    Bucket, BucketView, BucketViewValue, BucketsView, ByNamespace, GaugeValue, MetricName,
    MetricNamespace, SetView,
};
use relay_protocol::{Annotated, FiniteF64, SerializableAnnotated};
use relay_quotas::Scoping;
use relay_statsd::metric;
use relay_system::{FromMessage, Interface, NoResponse, Service};
use relay_threading::AsyncPool;

use crate::envelope::{AttachmentPlaceholder, AttachmentType, ContentType, Item};
use crate::managed::{Counted, Managed, OutcomeError, Quantities, Rejected};
use crate::metrics::{ArrayEncoding, BucketEncoder, MetricOutcomes};

use crate::service::ServiceError;
use crate::services::global_config::GlobalConfigHandle;
use crate::services::objectstore::ObjectstoreKey;
use crate::services::outcome::{self, DiscardReason, Outcome, OutcomeId};
use crate::services::upload::{Final, SignedLocation};
use crate::statsd::{RelayCounters, RelayGauges, RelayTimers};
use crate::utils;

mod sessions;

/// Fallback name used for attachment items without a `filename` header.
const UNNAMED_ATTACHMENT: &str = "Unnamed Attachment";

#[derive(Debug, thiserror::Error)]
pub enum StoreError {
    #[error("failed to send the message to kafka: {0}")]
    SendFailed(#[from] ClientError),
    #[error("failed to encode data: {0}")]
    EncodingFailed(std::io::Error),
    #[error("failed to serialize data: {0}")]
    Serialize(#[from] serde_json::Error),
    #[error("failed to store event because event id was missing")]
    NoEventId,
    #[error("invalid attachment reference")]
    InvalidAttachmentRef,
}

impl OutcomeError for StoreError {
    type Error = Self;

    fn consume(self) -> (Option<Outcome>, Self::Error) {
        let outcome = match self {
            StoreError::SendFailed(_)
            | StoreError::EncodingFailed(_)
            | StoreError::Serialize(_)
            | StoreError::NoEventId => Some(Outcome::Invalid(DiscardReason::Internal)),
            StoreError::InvalidAttachmentRef => {
                Some(Outcome::Invalid(DiscardReason::InvalidAttachmentRef))
            }
        };
        (outcome, self)
    }
}

struct Producer {
    client: KafkaClient,
}

impl Producer {
    pub fn create(config: &Config) -> anyhow::Result<Self> {
        let mut client_builder = KafkaClient::builder();

        for topic in KafkaTopic::iter() {
            let kafka_configs = config.kafka_configs(*topic)?;
            client_builder = client_builder
                .add_kafka_topic_config(*topic, &kafka_configs, config.kafka_validate_topics())
                .map_err(|e| ServiceError::Kafka(e.to_string()))?;
        }

        Ok(Self {
            client: client_builder.build(),
        })
    }
}

/// Publishes an [`Event`] and its attachments to Kafka.
#[derive(Debug)]
pub struct StoreEvent {
    /// The data category the [`Self::event`] is counted in.
    pub event_category: DataCategory,
    /// The event to be stored.
    pub event: Annotated<Event>,
    /// A list of attachments associated with the event.
    pub attachments: Vec<Item>,
    /// A list of user reports associated with the event.
    pub user_reports: Vec<Item>,
    /// Event retention in days.
    pub retention_days: u16,
}

impl Counted for StoreEvent {
    fn quantities(&self) -> Quantities {
        let mut quantities = smallvec::smallvec![(self.event_category, 1)];
        quantities.extend(self.attachments.quantities());
        quantities.extend(self.user_reports.quantities());
        quantities
    }
}

/// Publishes a list of [`Bucket`]s to the Sentry core application through Kafka topics.
#[derive(Clone, Debug)]
pub struct StoreMetrics {
    pub buckets: Vec<Bucket>,
    pub scoping: Scoping,
    pub retention: u16,
}

/// Publishes a log item to the Sentry core application through Kafka.
#[derive(Debug)]
pub struct StoreTraceItem {
    /// The final trace item which will be produced to Kafka.
    pub trace_item: TraceItem,
}

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

/// Publishes a span item to the Sentry core application through Kafka.
#[derive(Debug)]
pub struct StoreSpanV2 {
    /// Routing key to assign a Kafka partition.
    pub routing_key: Option<Uuid>,
    /// Default retention of the span.
    pub retention_days: u16,
    /// Downsampled retention of the span.
    pub downsampled_retention_days: u16,
    /// Optional event id, if the span was extracted from a transaction.
    pub event_id: Option<EventId>,
    /// The final Sentry compatible span item.
    pub item: SpanV2,
    // Whether to run issue detection on the transaction or on the segment span.
    // (only true for segment spans created from a transaction).
    pub performance_issues_spans: bool,
}

impl Counted for StoreSpanV2 {
    fn quantities(&self) -> Quantities {
        smallvec::smallvec![(DataCategory::SpanIndexed, 1)]
    }
}

/// Publishes a singular profile chunk to Kafka.
#[derive(Debug)]
pub struct StoreProfileChunk {
    /// Default retention of the span.
    pub retention_days: u16,
    /// The serialized profile chunk payload.
    pub payload: Bytes,
    /// Additional attachments associated with this profile chunk.
    pub attachments: Vec<ProfileAttachment>,
    /// Outcome quantities associated with this profile.
    ///
    /// Quantities are different for backend and ui profile chunks.
    pub quantities: Quantities,
}

/// Optional raw binary blob associated with [`profile chunk`](StoreProfileChunk).
#[derive(Debug)]
pub struct ProfileAttachment {
    /// Name of the attachment,
    pub name: String,
    /// Content type of the attachment.
    pub content_type: ContentType,
    /// Objectstore id of the attachment.
    ///
    /// Using this id the attachment can be retrieved again.
    pub stored_id: ObjectstoreKey,
}

impl Counted for StoreProfileChunk {
    fn quantities(&self) -> Quantities {
        self.quantities.clone()
    }
}

/// A replay to be stored to Kafka.
#[derive(Debug)]
pub struct StoreReplay {
    /// The event ID.
    pub event_id: EventId,
    /// Number of days to retain.
    pub retention_days: u16,
    /// The recording payload (rrweb data).
    pub recording: Bytes,
    /// Optional replay event payload (JSON).
    pub event: Option<Bytes>,
    /// Optional replay video.
    pub video: Option<Bytes>,
    /// Outcome quantities associated with this replay.
    ///
    /// Quantities are different for web and native replays.
    pub quantities: Quantities,
}

impl Counted for StoreReplay {
    fn quantities(&self) -> Quantities {
        self.quantities.clone()
    }
}

/// An attachment to be stored to Kafka.
#[derive(Debug)]
pub struct StoreAttachment {
    /// The event ID.
    pub event_id: EventId,
    /// That attachment item.
    pub attachment: Item,
    /// Outcome quantities associated with this attachment.
    pub quantities: Quantities,
    /// Data retention in days for this attachment.
    pub retention: u16,
}

impl Counted for StoreAttachment {
    fn quantities(&self) -> Quantities {
        self.quantities.clone()
    }
}

/// A user report to be stored to Kafka.
#[derive(Debug)]
pub struct StoreUserReport {
    /// The event ID.
    pub event_id: EventId,
    /// The user report.
    pub report: Item,
}

impl Counted for StoreUserReport {
    fn quantities(&self) -> Quantities {
        smallvec::smallvec![(DataCategory::UserReportV2, 1)]
    }
}

/// A profile to be stored to Kafka.
#[derive(Debug)]
pub struct StoreProfile {
    /// Number of days to retain.
    pub retention_days: u16,
    /// The profile.
    pub profile: Item,
    /// Outcome quantities associated with this profile.
    pub quantities: Quantities,
}

impl Counted for StoreProfile {
    fn quantities(&self) -> Quantities {
        self.quantities.clone()
    }
}

/// A monitor check-in to be stored to Kafka.
#[derive(Debug)]
pub struct StoreCheckIn {
    /// The serialized check-in.
    pub check_in: Item,
    /// The SDK client which produced the check-in.
    pub sdk: Option<String>,
    /// Check-in retention in days.
    pub retention_days: u16,
}

impl Counted for StoreCheckIn {
    fn quantities(&self) -> Quantities {
        self.check_in.quantities()
    }
}

/// The asynchronous thread pool used for scheduling storing tasks in the envelope store.
pub type StoreServicePool = AsyncPool<StoreTask>;

/// Service interface for storing items in Kafka.
#[derive(Debug)]
pub enum Store {
    /// An [`Event`] and its associated items.
    Event(Managed<Box<StoreEvent>>),
    /// Aggregated generic metrics.
    Metrics(StoreMetrics),
    /// A singular [`TraceItem`].
    TraceItem(Managed<StoreTraceItem>),
    /// A singular Span.
    Span(Managed<Box<StoreSpanV2>>),
    /// A singular profile chunk.
    ProfileChunk(Managed<StoreProfileChunk>),
    /// A singular replay.
    Replay(Managed<StoreReplay>),
    /// A singular attachment.
    Attachment(Managed<StoreAttachment>),
    /// A singular user report.
    UserReport(Managed<StoreUserReport>),
    /// A single profile.
    Profile(Managed<StoreProfile>),
    /// A singular monitor check-in.
    CheckIn(Managed<StoreCheckIn>),
}

impl Store {
    /// Returns the name of the message variant.
    fn variant(&self) -> &'static str {
        match self {
            Store::Event(_) => "event",
            Store::Metrics(_) => "metrics",
            Store::TraceItem(_) => "trace_item",
            Store::Span(_) => "span",
            Store::ProfileChunk(_) => "profile_chunk",
            Store::Replay(_) => "replay",
            Store::Attachment(_) => "attachment",
            Store::UserReport(_) => "user_report",
            Store::Profile(_) => "profile",
            Store::CheckIn(_) => "check_in",
        }
    }
}

impl Interface for Store {}

impl FromMessage<Managed<Box<StoreEvent>>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<Box<StoreEvent>>, _: ()) -> Self {
        Self::Event(message)
    }
}

impl FromMessage<StoreMetrics> for Store {
    type Response = NoResponse;

    fn from_message(message: StoreMetrics, _: ()) -> Self {
        Self::Metrics(message)
    }
}

impl FromMessage<Managed<StoreTraceItem>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreTraceItem>, _: ()) -> Self {
        Self::TraceItem(message)
    }
}

impl FromMessage<Managed<Box<StoreSpanV2>>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<Box<StoreSpanV2>>, _: ()) -> Self {
        Self::Span(message)
    }
}

impl FromMessage<Managed<StoreProfileChunk>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreProfileChunk>, _: ()) -> Self {
        Self::ProfileChunk(message)
    }
}

impl FromMessage<Managed<StoreReplay>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreReplay>, _: ()) -> Self {
        Self::Replay(message)
    }
}

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

    fn from_message(message: Managed<StoreAttachment>, _: ()) -> Self {
        Self::Attachment(message)
    }
}

impl FromMessage<Managed<StoreUserReport>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreUserReport>, _: ()) -> Self {
        Self::UserReport(message)
    }
}

impl FromMessage<Managed<StoreProfile>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreProfile>, _: ()) -> Self {
        Self::Profile(message)
    }
}

impl FromMessage<Managed<StoreCheckIn>> for Store {
    type Response = NoResponse;

    fn from_message(message: Managed<StoreCheckIn>, _: ()) -> Self {
        Self::CheckIn(message)
    }
}

/// Service implementing the [`Store`] interface.
pub struct StoreService {
    pool: StoreServicePool,
    config: Arc<Config>,
    global_config: GlobalConfigHandle,
    metric_outcomes: MetricOutcomes,
    producer: Producer,
}

impl StoreService {
    pub fn create(
        pool: StoreServicePool,
        config: Arc<Config>,
        global_config: GlobalConfigHandle,
        metric_outcomes: MetricOutcomes,
    ) -> anyhow::Result<Self> {
        let producer = Producer::create(&config)?;
        Ok(Self {
            pool,
            config,
            global_config,
            metric_outcomes,
            producer,
        })
    }

    fn handle_message(&self, message: Store) {
        let ty = message.variant();
        relay_statsd::metric!(timer(RelayTimers::StoreServiceDuration), message = ty, {
            let result = match message {
                Store::Event(message) => self.handle_store_event(message),
                Store::Metrics(message) => {
                    self.handle_store_metrics(message);
                    Ok(())
                }
                Store::TraceItem(message) => self.handle_store_trace_item(message),
                Store::Span(message) => self.handle_store_span(message),
                Store::ProfileChunk(message) => self.handle_store_profile_chunk(message),
                Store::Replay(message) => self.handle_store_replay(message),
                Store::Attachment(message) => self.handle_store_attachment(message),
                Store::UserReport(message) => self.handle_user_report(message),
                Store::Profile(message) => self.handle_profile(message),
                Store::CheckIn(message) => self.handle_check_in(message),
            };
            if let Err(error) = result {
                relay_log::error!(
                    error = &error as &dyn Error,
                    tags.message = ty,
                    "failed to store message"
                );
            }
        })
    }

    fn handle_store_event(
        &self,
        message: Managed<Box<StoreEvent>>,
    ) -> Result<(), Rejected<StoreError>> {
        let received_at = message.received_at();
        let scoping = message.scoping();
        let remote_addr = message.remote_addr().map(|ip| ip.to_string());

        message.try_accept(|m| self.do_store_event(*m, scoping, received_at, remote_addr))
    }

    fn do_store_event(
        &self,
        store: StoreEvent,
        scoping: Scoping,
        received_at: DateTime<Utc>,
        remote_addr: Option<String>,
    ) -> Result<(), StoreError> {
        let event_id = store.event.value().and_then(|e| e.id.value()).copied();
        let event_id = event_id.ok_or(StoreError::NoEventId)?;

        let event_type = store.event.value().and_then(|e| e.ty.value());
        let send_individual_attachments = matches!(
            event_type,
            Some(&EventType::Transaction) | Some(&EventType::UserReportV2)
        );

        let mut attachments = Vec::new();
        for attachment in store.attachments {
            // Note: Technically when an error occurs we may have already produced some items to Kafka,
            // but since we don't keep track of that properly and just error out here, outcomes will
            // also report items which were already produced to Kafka as dropped.
            //
            // We specifically accept this here, since this is an extreme edge case with minimal practical
            // consequences.
            if let Some(attachment) = self.produce_attachment(
                event_id,
                scoping.project_id,
                scoping.organization_id,
                &attachment,
                send_individual_attachments,
                store.retention_days,
            )? {
                attachments.push(attachment);
            }
        }

        for user_report in &store.user_reports {
            self.produce_user_report(
                event_id,
                scoping.project_id,
                scoping.organization_id,
                received_at,
                user_report,
            )?;
        }

        let event_topic = if event_type == Some(&EventType::Transaction) {
            KafkaTopic::Transactions
        } else if event_type == Some(&EventType::UserReportV2) {
            KafkaTopic::Feedback
        } else if !attachments.is_empty() || !store.user_reports.is_empty() {
            KafkaTopic::Attachments
        } else {
            KafkaTopic::Events
        };

        let payload = store.event.to_json()?.into_bytes().into();
        self.produce(
            event_topic,
            KafkaMessage::Event(EventKafkaMessage {
                payload,
                start_time: safe_timestamp(received_at),
                event_id,
                project_id: scoping.project_id,
                org_id: scoping.organization_id,
                remote_addr,
                attachments,
            }),
        )
    }

    fn handle_store_metrics(&self, message: StoreMetrics) {
        let StoreMetrics {
            buckets,
            scoping,
            retention,
        } = message;

        let batch_size = self.config.metrics_max_batch_size_bytes();
        let mut error = None;

        let global_config = self.global_config.current().unwrap_or_default();
        let mut encoder = BucketEncoder::new(&global_config);

        let emit_sessions_to_eap = utils::is_rolled_out(
            scoping.organization_id.value(),
            global_config.options.sessions_eap_rollout_rate,
        )
        .is_keep();

        let now = UnixTimestamp::now();
        let mut delay_stats = ByNamespace::<(u64, u64, u64)>::default();

        for mut bucket in buckets {
            let namespace = encoder.prepare(&mut bucket);

            if let Some(received_at) = bucket.metadata.received_at {
                let delay = now.as_secs().saturating_sub(received_at.as_secs());
                let (total, count, max) = delay_stats.get_mut(namespace);
                *total += delay;
                *count += 1;
                *max = (*max).max(delay);
            }

            // Create a local bucket view to avoid splitting buckets unnecessarily. Since we produce
            // each bucket separately, we only need to split buckets that exceed the size, but not
            // batches.
            for view in BucketsView::new(std::slice::from_ref(&bucket))
                .by_size(batch_size)
                .flatten()
            {
                let message =
                    self.create_metric_message(&scoping, &mut encoder, namespace, &view, retention);

                let result =
                    message.and_then(|message| self.send_metric_message(namespace, message));

                let outcome = match result {
                    Ok(()) => Outcome::Accepted,
                    Err(e) => {
                        error.get_or_insert(e);
                        Outcome::Invalid(DiscardReason::Internal)
                    }
                };

                self.metric_outcomes.track(scoping, &[view], outcome);
            }

            if emit_sessions_to_eap
                && let Some(trace_item) = sessions::to_trace_item(scoping, bucket, retention)
            {
                let message = KafkaMessage::for_item(scoping, trace_item);
                let res = self.produce(KafkaTopic::Items, message);
                if let Err(error) = res {
                    relay_log::error!(
                        error = &error as &dyn std::error::Error,
                        "failed to produce session metrics to EAP"
                    )
                }
            }
        }

        if let Some(error) = error {
            relay_log::error!(
                error = &error as &dyn std::error::Error,
                "failed to produce metric buckets: {error}"
            );
        }

        for (namespace, (total, count, max)) in delay_stats {
            if count == 0 {
                continue;
            }
            metric!(
                counter(RelayCounters::MetricDelaySum) += total,
                namespace = namespace.as_str()
            );
            metric!(
                counter(RelayCounters::MetricDelayCount) += count,
                namespace = namespace.as_str()
            );
            metric!(
                gauge(RelayGauges::MetricDelayMax) = max,
                namespace = namespace.as_str()
            );
        }
    }

    fn handle_store_trace_item(
        &self,
        message: Managed<StoreTraceItem>,
    ) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();

        message.try_accept(|item| {
            let message = KafkaMessage::for_item(scoping, item.trace_item);
            self.produce(KafkaTopic::Items, message)
        })?;

        Ok(())
    }

    fn handle_store_span(
        &self,
        message: Managed<Box<StoreSpanV2>>,
    ) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();
        let received_at = message.received_at();

        let meta = SpanMeta {
            organization_id: scoping.organization_id,
            project_id: scoping.project_id,
            key_id: scoping.key_id,
            event_id: message.event_id,
            retention_days: message.retention_days,
            downsampled_retention_days: message.downsampled_retention_days,
            received: datetime_to_timestamp(received_at),
            performance_issues_spans: message.performance_issues_spans,
        };

        message.try_accept(|span| {
            let item = Annotated::new(span.item);
            let message = KafkaMessage::SpanV2 {
                routing_key: span.routing_key,
                headers: BTreeMap::from([(
                    "project_id".to_owned(),
                    scoping.project_id.to_string(),
                )]),
                message: SpanKafkaMessage {
                    meta,
                    span: SerializableAnnotated(&item),
                },
                org_id: scoping.organization_id,
            };

            self.produce(KafkaTopic::Spans, message)
        })?;

        relay_statsd::metric!(
            counter(RelayCounters::SpanV2Produced) += 1,
            via = "processing"
        );

        Ok(())
    }

    fn handle_store_profile_chunk(
        &self,
        message: Managed<StoreProfileChunk>,
    ) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();
        let received_at = message.received_at();

        message.try_accept(|message| {
            let message = ProfileChunkKafkaMessage {
                organization_id: scoping.organization_id,
                project_id: scoping.project_id,
                received: safe_timestamp(received_at),
                retention_days: message.retention_days,
                headers: BTreeMap::from([(
                    "project_id".to_owned(),
                    scoping.project_id.to_string(),
                )]),
                payload: message.payload,
                attachments: message
                    .attachments
                    .into_iter()
                    .map(|attachment| ProfileChunkKafkaAttachment {
                        name: attachment.name,
                        content_type: attachment.content_type.as_str(),
                        stored_id: attachment.stored_id.into_inner(),
                    })
                    .collect(),
            };

            self.produce(KafkaTopic::Profiles, KafkaMessage::ProfileChunk(message))
        })
    }

    fn handle_store_replay(
        &self,
        message: Managed<StoreReplay>,
    ) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();
        let received_at = message.received_at();

        message.try_accept(|replay| {
            let kafka_msg =
                KafkaMessage::ReplayRecordingNotChunked(ReplayRecordingNotChunkedKafkaMessage {
                    replay_id: replay.event_id,
                    key_id: scoping.key_id,
                    org_id: scoping.organization_id,
                    project_id: scoping.project_id,
                    received: safe_timestamp(received_at),
                    retention_days: replay.retention_days,
                    payload: &replay.recording,
                    replay_event: replay.event.as_deref(),
                    replay_video: replay.video.as_deref(),
                    // Hardcoded to `true` to indicate to the consumer that it should always publish the
                    // replay_event as relay no longer does it.
                    relay_snuba_publish_disabled: true,
                });
            self.produce(KafkaTopic::ReplayRecordings, kafka_msg)
        })
    }

    fn handle_store_attachment(
        &self,
        message: Managed<StoreAttachment>,
    ) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();
        message.try_accept(|attachment| {
            let result = self.produce_attachment(
                attachment.event_id,
                scoping.project_id,
                scoping.organization_id,
                &attachment.attachment,
                // Hardcoded to `true` since standalone attachments are 'individual attachments'.
                true,
                attachment.retention,
            );
            // Since we are sending an 'individual attachment' this function should never return a
            // `ChunkedAttachment`.
            debug_assert!(!matches!(result, Ok(Some(_))));
            result.map(|_| ())
        })
    }

    fn handle_user_report(
        &self,
        message: Managed<StoreUserReport>,
    ) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();
        let received_at = message.received_at();

        message.try_accept(|report| {
            let kafka_msg = KafkaMessage::UserReport(UserReportKafkaMessage {
                project_id: scoping.project_id,
                event_id: report.event_id,
                start_time: safe_timestamp(received_at),
                payload: report.report.payload(),
                org_id: scoping.organization_id,
            });
            self.produce(KafkaTopic::Attachments, kafka_msg)
        })
    }

    fn handle_profile(&self, message: Managed<StoreProfile>) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();
        let received_at = message.received_at();

        message.try_accept(|profile| {
            self.produce_profile(
                scoping.organization_id,
                scoping.project_id,
                scoping.key_id,
                received_at,
                profile.retention_days,
                &profile.profile,
            )
        })
    }

    fn handle_check_in(&self, message: Managed<StoreCheckIn>) -> Result<(), Rejected<StoreError>> {
        let scoping = message.scoping();
        let received_at = message.received_at();

        message.try_accept(|check_in| {
            let message = KafkaMessage::CheckIn(CheckInKafkaMessage {
                message_type: CheckInMessageType::CheckIn,
                project_id: scoping.project_id,
                org_id: scoping.organization_id,
                retention_days: check_in.retention_days,
                start_time: safe_timestamp(received_at),
                sdk: check_in.sdk,
                payload: check_in.check_in.payload(),
                routing_key_hint: check_in.check_in.routing_hint(),
            });

            self.produce(KafkaTopic::Monitors, message)
        })
    }

    fn create_metric_message<'a>(
        &self,
        scoping: &Scoping,
        encoder: &'a mut BucketEncoder,
        namespace: MetricNamespace,
        view: &BucketView<'a>,
        retention_days: u16,
    ) -> Result<MetricKafkaMessage<'a>, StoreError> {
        let value = match view.value() {
            BucketViewValue::Counter(c) => MetricValue::Counter(c),
            BucketViewValue::Distribution(data) => MetricValue::Distribution(
                encoder
                    .encode_distribution(namespace, data)
                    .map_err(StoreError::EncodingFailed)?,
            ),
            BucketViewValue::Set(data) => MetricValue::Set(
                encoder
                    .encode_set(namespace, data)
                    .map_err(StoreError::EncodingFailed)?,
            ),
            BucketViewValue::Gauge(g) => MetricValue::Gauge(g),
        };

        Ok(MetricKafkaMessage {
            org_id: scoping.organization_id,
            project_id: scoping.project_id,
            key_id: scoping.key_id,
            name: view.name(),
            value,
            timestamp: view.timestamp(),
            tags: view.tags(),
            retention_days,
            received_at: view.metadata().received_at,
        })
    }

    fn produce(
        &self,
        topic: KafkaTopic,
        // Takes message by value to ensure it is not being produced twice.
        message: KafkaMessage,
    ) -> Result<(), StoreError> {
        relay_log::trace!(
            "Sending kafka message of type {} to {topic:?}",
            message.variant()
        );

        let topic_name = self.producer.client.send_message(topic, &message)?;

        match &message {
            KafkaMessage::Metric {
                message: metric, ..
            } => {
                metric!(
                    counter(RelayCounters::ProcessingMessageProduced) += 1,
                    event_type = message.variant(),
                    topic = topic_name,
                    metric_type = metric.value.variant(),
                    metric_encoding = metric.value.encoding().unwrap_or(""),
                );
            }
            KafkaMessage::ReplayRecordingNotChunked(replay) => {
                let has_video = replay.replay_video.is_some();

                metric!(
                    counter(RelayCounters::ProcessingMessageProduced) += 1,
                    event_type = message.variant(),
                    topic = topic_name,
                    has_video = bool_to_str(has_video),
                );
            }
            message => {
                metric!(
                    counter(RelayCounters::ProcessingMessageProduced) += 1,
                    event_type = message.variant(),
                    topic = topic_name,
                );
            }
        }

        Ok(())
    }

    fn chunked_attachment_from_placeholder(
        &self,
        item: &Item,
        retention_days: u16,
    ) -> Result<ChunkedAttachment, StoreError> {
        debug_assert!(
            item.stored_key().is_none(),
            "AttachmentRef should not have been uploaded to objectstore"
        );

        let payload = item.payload();
        let placeholder: AttachmentPlaceholder<'_> =
            serde_json::from_slice(&payload).map_err(|_| StoreError::InvalidAttachmentRef)?;
        let location = SignedLocation::<Final>::try_from_str(placeholder.location)
            .ok_or(StoreError::InvalidAttachmentRef)?
            .verify(Utc::now(), &self.config)
            .map_err(|_| StoreError::InvalidAttachmentRef)?;

        let store_key = location.key;

        Ok(ChunkedAttachment {
            id: Uuid::new_v4().to_string(),
            name: item.filename().unwrap_or(UNNAMED_ATTACHMENT).to_owned(),
            rate_limited: item.rate_limited(),
            content_type: placeholder.content_type,
            attachment_type: item.attachment_type().unwrap_or_default(),
            size: item.attachment_body_size(),
            retention_days,
            payload: AttachmentPayload::Stored(store_key),
        })
    }

    fn chunked_attachment_from_attachment(
        &self,
        event_id: EventId,
        project_id: ProjectId,
        org_id: OrganizationId,
        item: &Item,
        send_individual_attachments: bool,
        retention_days: u16,
    ) -> Result<ChunkedAttachment, StoreError> {
        let id = Uuid::new_v4().to_string();

        let payload = item.payload();
        let size = item.len();
        let max_chunk_size = self.config.attachment_chunk_size();

        let payload = if size == 0 {
            AttachmentPayload::Chunked(0)
        } else if let Some(stored_key) = item.stored_key() {
            AttachmentPayload::Stored(stored_key.into())
        } else if send_individual_attachments && size < max_chunk_size {
            // When sending individual attachments, and we have a single chunk, we want to send the
            // `data` inline in the `attachment` message.
            // This avoids a needless roundtrip through the attachments cache on the Sentry side.
            AttachmentPayload::Inline(payload)
        } else {
            let mut chunk_index = 0;
            let mut offset = 0;
            // This skips chunks for empty attachments. The consumer does not require chunks for
            // empty attachments. `chunks` will be `0` in this case.
            while offset < size {
                let chunk_size = std::cmp::min(max_chunk_size, size - offset);
                let chunk_message = AttachmentChunkKafkaMessage {
                    payload: payload.slice(offset..offset + chunk_size),
                    event_id,
                    project_id,
                    id: id.clone(),
                    chunk_index,
                    org_id,
                };

                self.produce(
                    KafkaTopic::Attachments,
                    KafkaMessage::AttachmentChunk(chunk_message),
                )?;
                offset += chunk_size;
                chunk_index += 1;
            }

            // The chunk_index is incremented after every loop iteration. After we exit the loop, it
            // is one larger than the last chunk, so it is equal to the number of chunks.
            AttachmentPayload::Chunked(chunk_index)
        };

        Ok(ChunkedAttachment {
            id,
            name: match item.filename() {
                Some(name) => name.to_owned(),
                None => UNNAMED_ATTACHMENT.to_owned(),
            },
            rate_limited: item.rate_limited(),
            content_type: item.raw_content_type().map(|s| s.to_ascii_lowercase()),
            attachment_type: item.attachment_type().unwrap_or_default(),
            size,
            retention_days,
            payload,
        })
    }

    /// Produces Kafka messages for the content and metadata of an attachment item.
    ///
    /// The `send_individual_attachments` controls whether the metadata of an attachment
    /// is produced directly as an individual `attachment` message, or returned from this function
    /// to be later sent as part of an `event` message.
    ///
    /// Attachment contents are chunked and sent as multiple `attachment_chunk` messages,
    /// unless the `send_individual_attachments` flag is set, and the content is small enough
    /// to fit inside a message.
    /// In that case, no `attachment_chunk` is produced, but the content is sent as part
    /// of the `attachment` message instead.
    fn produce_attachment(
        &self,
        event_id: EventId,
        project_id: ProjectId,
        org_id: OrganizationId,
        item: &Item,
        send_individual_attachments: bool,
        retention_days: u16,
    ) -> Result<Option<ChunkedAttachment>, StoreError> {
        let attachment = if item.is_attachment_ref() {
            self.chunked_attachment_from_placeholder(item, retention_days)
        } else {
            self.chunked_attachment_from_attachment(
                event_id,
                project_id,
                org_id,
                item,
                send_individual_attachments,
                retention_days,
            )
        }?;

        if send_individual_attachments {
            let message = KafkaMessage::Attachment(AttachmentKafkaMessage {
                event_id,
                project_id,
                attachment,
                org_id,
            });
            self.produce(KafkaTopic::Attachments, message)?;
            Ok(None)
        } else {
            Ok(Some(attachment))
        }
    }

    fn produce_user_report(
        &self,
        event_id: EventId,
        project_id: ProjectId,
        org_id: OrganizationId,
        received_at: DateTime<Utc>,
        item: &Item,
    ) -> Result<(), StoreError> {
        let message = KafkaMessage::UserReport(UserReportKafkaMessage {
            project_id,
            event_id,
            start_time: safe_timestamp(received_at),
            payload: item.payload(),
            org_id,
        });

        self.produce(KafkaTopic::Attachments, message)
    }

    fn send_metric_message(
        &self,
        namespace: MetricNamespace,
        message: MetricKafkaMessage,
    ) -> Result<(), StoreError> {
        let topic = match namespace {
            MetricNamespace::Sessions => KafkaTopic::MetricsSessions,
            MetricNamespace::Outcomes => {
                return self.send_metric_based_outcome(message);
            }
            MetricNamespace::Unsupported => {
                relay_log::error!(
                    metric_message.name = message.name.as_ref(),
                    "store service dropping unknown metric usecase"
                );
                return Ok(());
            }
            _ => KafkaTopic::MetricsGeneric,
        };

        let headers = BTreeMap::from([("namespace".to_owned(), namespace.to_string())]);
        self.produce(topic, KafkaMessage::Metric { headers, message })?;
        Ok(())
    }

    fn send_metric_based_outcome(&self, message: MetricKafkaMessage) -> Result<(), StoreError> {
        let Some(outcome) = outcome::metric::to_outcome_id(message.name) else {
            relay_log::error!(
                mri = message.name.as_ref(),
                "invalid outcome metric, cannot infer outcome id from metric name"
            );
            return Ok(());
        };
        let quantity = match message.value {
            MetricValue::Counter(c) => c.to_f64() as _,
            v => {
                relay_log::error!(
                    mri = message.name.as_ref(),
                    "invalid outcome metric, expected a counter got '{}'",
                    v.variant()
                );
                return Ok(());
            }
        };

        let outcome = OutcomeMessage {
            timestamp: message
                .timestamp
                .as_datetime()
                .unwrap_or_else(Utc::now)
                .to_rfc3339_opts(SecondsFormat::Micros, true),
            org_id: Some(message.org_id).filter(|id| id.value() != 0),
            project_id: message.project_id,
            key_id: message.key_id,
            outcome,
            reason: message.tags.get("reason").map(|s| s.as_str()),
            event_id: message.tags.get("event_id").map(|s| s.as_str()),
            remote_addr: message.tags.get("remote_addr").map(|s| s.as_str()),
            source: message.tags.get("source").map(|s| s.as_str()),
            category: message.tags.get("category").and_then(|s| s.parse().ok()),
            quantity: Some(quantity),
        };

        let topic = match outcome.outcome.is_billing() {
            true => KafkaTopic::OutcomesBilling,
            false => KafkaTopic::Outcomes,
        };

        self.produce(topic, KafkaMessage::Outcome(outcome))
    }

    fn produce_profile(
        &self,
        organization_id: OrganizationId,
        project_id: ProjectId,
        key_id: Option<u64>,
        received_at: DateTime<Utc>,
        retention_days: u16,
        item: &Item,
    ) -> Result<(), StoreError> {
        let message = ProfileKafkaMessage {
            organization_id,
            project_id,
            key_id,
            received: safe_timestamp(received_at),
            retention_days,
            headers: BTreeMap::from([
                (
                    "sampled".to_owned(),
                    if item.sampled() { "true" } else { "false" }.to_owned(),
                ),
                ("project_id".to_owned(), project_id.to_string()),
            ]),
            payload: item.payload(),
        };
        self.produce(KafkaTopic::Profiles, KafkaMessage::Profile(message))?;
        Ok(())
    }
}

impl Service for StoreService {
    type Interface = Store;

    async fn run(self, mut rx: relay_system::Receiver<Self::Interface>) {
        let this = Arc::new(self);

        relay_log::info!("store forwarder started");

        while let Some(message) = rx.recv().await {
            let task = StoreTask {
                service: Arc::clone(&this),
                message: Some(message),
            };
            this.pool.spawn_async(task).await;
        }

        relay_log::info!("store forwarder stopped");
    }
}

/// Task executed by [`StoreService`].
pub struct StoreTask {
    service: Arc<StoreService>,
    message: Option<Store>,
}

impl Future for StoreTask {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, _: &mut task::Context<'_>) -> task::Poll<Self::Output> {
        let message = self
            .message
            .take()
            .expect("StoreTask polled after completion");
        let () = relay_log::with_scope(|_| {}, || self.service.handle_message(message));
        task::Poll::Ready(())
    }
}

/// This signifies how the attachment payload is being transfered.
#[derive(Debug, Serialize)]
enum AttachmentPayload {
    /// The payload has been split into multiple chunks.
    ///
    /// The individual chunks are being sent as separate [`AttachmentChunkKafkaMessage`] messages.
    /// If the payload `size == 0`, the number of chunks will also be `0`.
    #[serde(rename = "chunks")]
    Chunked(usize),

    /// The payload is inlined here directly, and thus into the [`ChunkedAttachment`].
    #[serde(rename = "data")]
    Inline(Bytes),

    /// The attachment has already been stored into the objectstore, with the given Id.
    #[serde(rename = "stored_id")]
    Stored(String),
}

/// Common attributes for both standalone attachments and processing-relevant attachments.
#[derive(Debug, Serialize)]
struct ChunkedAttachment {
    /// The attachment ID within the event.
    ///
    /// The triple `(project_id, event_id, id)` identifies an attachment uniquely.
    id: String,

    /// File name of the attachment file.
    name: String,

    /// Whether this attachment was rate limited and should be removed after processing.
    ///
    /// By default, rate limited attachments are immediately removed from Envelopes. For processing,
    /// native crash reports still need to be retained. These attachments are marked with the
    /// `rate_limited` header, which signals to the processing pipeline that the attachment should
    /// not be persisted after processing.
    rate_limited: bool,

    /// Content type of the attachment payload.
    #[serde(skip_serializing_if = "Option::is_none")]
    content_type: Option<String>,

    /// The Sentry-internal attachment type used in the processing pipeline.
    #[serde(serialize_with = "serialize_attachment_type")]
    attachment_type: AttachmentType,

    /// The size of the attachment in bytes.
    size: usize,

    /// The retention in days for this attachment.
    retention_days: u16,

    /// The attachment payload, chunked, inlined, or already stored.
    #[serde(flatten)]
    payload: AttachmentPayload,
}

/// A hack to make rmp-serde behave more like serde-json when serializing enums.
///
/// Cannot serialize bytes.
///
/// See <https://github.com/3Hren/msgpack-rust/pull/214>
fn serialize_attachment_type<S, T>(t: &T, serializer: S) -> Result<S::Ok, S::Error>
where
    S: serde::Serializer,
    T: serde::Serialize,
{
    serde_json::to_value(t)
        .map_err(|e| serde::ser::Error::custom(e.to_string()))?
        .serialize(serializer)
}

/// Container payload for event messages.
#[derive(Debug, Serialize)]
struct EventKafkaMessage {
    /// Raw event payload.
    payload: Bytes,
    /// Time at which the event was received by Relay.
    start_time: u64,
    /// The event id.
    event_id: EventId,
    /// The project id for the current event.
    project_id: ProjectId,
    /// The client ip address.
    remote_addr: Option<String>,
    /// Attachments that are potentially relevant for processing.
    attachments: Vec<ChunkedAttachment>,

    /// Used for [`KafkaMessage::key`]
    #[serde(skip)]
    org_id: OrganizationId,
}

/// Container payload for chunks of attachments.
#[derive(Debug, Serialize)]
struct AttachmentChunkKafkaMessage {
    /// Chunk payload of the attachment.
    payload: Bytes,
    /// The event id.
    event_id: EventId,
    /// The project id for the current event.
    project_id: ProjectId,
    /// The attachment ID within the event.
    ///
    /// The triple `(project_id, event_id, id)` identifies an attachment uniquely.
    id: String,
    /// Sequence number of chunk. Starts at 0 and ends at `AttachmentKafkaMessage.num_chunks - 1`.
    chunk_index: usize,

    /// Used for [`KafkaMessage::key`]
    #[serde(skip)]
    org_id: OrganizationId,
}

/// A "standalone" attachment.
///
/// Still belongs to an event but can be sent independently (like UserReport) and is not
/// considered in processing.
#[derive(Debug, Serialize)]
struct AttachmentKafkaMessage {
    /// The event id.
    event_id: EventId,
    /// The project id for the current event.
    project_id: ProjectId,
    /// The attachment.
    attachment: ChunkedAttachment,

    /// Used for [`KafkaMessage::key`]
    #[serde(skip)]
    org_id: OrganizationId,
}

#[derive(Debug, Serialize)]
struct ReplayRecordingNotChunkedKafkaMessage<'a> {
    replay_id: EventId,
    key_id: Option<u64>,
    org_id: OrganizationId,
    project_id: ProjectId,
    received: u64,
    retention_days: u16,
    #[serde(with = "serde_bytes")]
    payload: &'a [u8],
    #[serde(with = "serde_bytes")]
    replay_event: Option<&'a [u8]>,
    #[serde(with = "serde_bytes")]
    replay_video: Option<&'a [u8]>,
    relay_snuba_publish_disabled: bool,
}

/// User report for an event wrapped up in a message ready for consumption in Kafka.
///
/// Is always independent of an event and can be sent as part of any envelope.
#[derive(Debug, Serialize)]
struct UserReportKafkaMessage {
    /// The project id for the current event.
    project_id: ProjectId,
    start_time: u64,
    payload: Bytes,

    /// Used for [`KafkaMessage::key`]
    #[serde(skip)]
    event_id: EventId,
    /// Used for [`KafkaMessage::key`]
    #[serde(skip)]
    org_id: OrganizationId,
}

#[derive(Clone, Debug, Serialize)]
struct MetricKafkaMessage<'a> {
    org_id: OrganizationId,
    project_id: ProjectId,
    #[serde(skip)]
    key_id: Option<u64>,
    name: &'a MetricName,
    #[serde(flatten)]
    value: MetricValue<'a>,
    timestamp: UnixTimestamp,
    tags: &'a BTreeMap<String, String>,
    retention_days: u16,
    #[serde(skip_serializing_if = "Option::is_none")]
    received_at: Option<UnixTimestamp>,
}

#[derive(Clone, Debug, Serialize)]
#[serde(tag = "type", content = "value")]
enum MetricValue<'a> {
    #[serde(rename = "c")]
    Counter(FiniteF64),
    #[serde(rename = "d")]
    Distribution(ArrayEncoding<'a, &'a [FiniteF64]>),
    #[serde(rename = "s")]
    Set(ArrayEncoding<'a, SetView<'a>>),
    #[serde(rename = "g")]
    Gauge(GaugeValue),
}

impl MetricValue<'_> {
    fn variant(&self) -> &'static str {
        match self {
            Self::Counter(_) => "counter",
            Self::Distribution(_) => "distribution",
            Self::Set(_) => "set",
            Self::Gauge(_) => "gauge",
        }
    }

    fn encoding(&self) -> Option<&'static str> {
        match self {
            Self::Distribution(ae) => Some(ae.name()),
            Self::Set(ae) => Some(ae.name()),
            _ => None,
        }
    }
}

/// Raw representation of an outcome for Kafka.
#[derive(Debug, Serialize, Clone)]
pub struct OutcomeMessage<'a> {
    /// The timespan of the event outcome.
    timestamp: String,
    /// Organization id.
    #[serde(skip_serializing_if = "Option::is_none")]
    org_id: Option<OrganizationId>,
    /// Project id.
    project_id: ProjectId,
    /// The DSN project key id.
    #[serde(skip_serializing_if = "Option::is_none")]
    key_id: Option<u64>,
    /// The outcome.
    outcome: OutcomeId,
    /// Reason for the outcome.
    #[serde(skip_serializing_if = "Option::is_none")]
    reason: Option<&'a str>,
    /// The event id.
    #[serde(skip_serializing_if = "Option::is_none")]
    event_id: Option<&'a str>,
    /// The client ip address.
    #[serde(skip_serializing_if = "Option::is_none")]
    remote_addr: Option<&'a str>,
    /// The source of the outcome (which Relay sent it)
    #[serde(skip_serializing_if = "Option::is_none")]
    source: Option<&'a str>,
    /// The event's data category.
    #[serde(skip_serializing_if = "Option::is_none")]
    category: Option<u8>,
    /// The number of events or total attachment size in bytes.
    #[serde(skip_serializing_if = "Option::is_none")]
    quantity: Option<u64>,
}

#[derive(Clone, Debug, Serialize)]
struct ProfileKafkaMessage {
    organization_id: OrganizationId,
    project_id: ProjectId,
    key_id: Option<u64>,
    received: u64,
    retention_days: u16,
    #[serde(skip)]
    headers: BTreeMap<String, String>,
    payload: Bytes,
}

/// Used to discriminate cron monitor ingestion messages.
///
/// There are two types of messages that end up in the ingest-monitors kafka topic, "check_in" (the
/// ones produced here in relay) and "clock_pulse" messages, which are produced externally and are
/// intended to ensure the clock continues to run even when ingestion volume drops.
#[allow(dead_code)]
#[derive(Debug, Serialize)]
#[serde(rename_all = "snake_case")]
enum CheckInMessageType {
    ClockPulse,
    CheckIn,
}

#[derive(Debug, Serialize)]
struct CheckInKafkaMessage {
    /// Used by the consumer to discrinminate the message.
    message_type: CheckInMessageType,
    /// Raw event payload.
    payload: Bytes,
    /// Time at which the event was received by Relay.
    start_time: u64,
    /// The SDK client which produced the event.
    sdk: Option<String>,
    /// The project id for the current event.
    project_id: ProjectId,
    /// Number of days to retain.
    retention_days: u16,

    /// Used for [`KafkaMessage::key`]
    #[serde(skip)]
    routing_key_hint: Option<Uuid>,
    /// Used for [`KafkaMessage::key`]
    #[serde(skip)]
    org_id: OrganizationId,
}

#[derive(Debug, Serialize)]
struct SpanKafkaMessage<'a> {
    #[serde(flatten)]
    meta: SpanMeta,
    #[serde(flatten)]
    span: SerializableAnnotated<'a, SpanV2>,
}

#[derive(Debug, Serialize)]
struct SpanMeta {
    organization_id: OrganizationId,
    project_id: ProjectId,
    // Required for the buffer to emit outcomes scoped to the DSN.
    #[serde(skip_serializing_if = "Option::is_none")]
    key_id: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    event_id: Option<EventId>,
    /// Time at which the event was received by Relay. Not to be confused with `start_timestamp_ms`.
    received: f64,
    /// Number of days until these data should be deleted.
    retention_days: u16,
    /// Number of days until the downsampled version of this data should be deleted.
    downsampled_retention_days: u16,
    /// Whether the segment span should be used for issue detection instead of the transaction.
    #[serde(rename = "_performance_issues_spans", skip_serializing_if = "is_false")]
    performance_issues_spans: bool,
}

fn is_false(val: &bool) -> bool {
    !val
}

#[derive(Clone, Debug, Serialize)]
struct ProfileChunkKafkaMessage {
    organization_id: OrganizationId,
    project_id: ProjectId,
    received: u64,
    retention_days: u16,
    #[serde(skip)]
    headers: BTreeMap<String, String>,
    payload: Bytes,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    attachments: Vec<ProfileChunkKafkaAttachment>,
}

#[derive(Clone, Debug, Serialize)]
struct ProfileChunkKafkaAttachment {
    name: String,
    content_type: &'static str,
    stored_id: String,
}

/// An enum over all possible ingest messages.
#[derive(Debug, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
#[allow(clippy::large_enum_variant)]
enum KafkaMessage<'a> {
    Event(EventKafkaMessage),
    UserReport(UserReportKafkaMessage),
    Metric {
        #[serde(skip)]
        headers: BTreeMap<String, String>,
        #[serde(flatten)]
        message: MetricKafkaMessage<'a>,
    },
    CheckIn(CheckInKafkaMessage),
    Item {
        #[serde(skip)]
        headers: BTreeMap<String, String>,
        #[serde(skip)]
        item_type: TraceItemType,
        #[serde(skip)]
        message: TraceItem,
    },
    SpanV2 {
        #[serde(skip)]
        routing_key: Option<Uuid>,
        #[serde(skip)]
        headers: BTreeMap<String, String>,
        #[serde(flatten)]
        message: SpanKafkaMessage<'a>,

        /// Used for [`KafkaMessage::key`]
        #[serde(skip)]
        org_id: OrganizationId,
    },

    Attachment(AttachmentKafkaMessage),
    AttachmentChunk(AttachmentChunkKafkaMessage),

    Profile(ProfileKafkaMessage),
    ProfileChunk(ProfileChunkKafkaMessage),

    ReplayRecordingNotChunked(ReplayRecordingNotChunkedKafkaMessage<'a>),

    Outcome(OutcomeMessage<'a>),
}

impl KafkaMessage<'_> {
    /// Creates a [`KafkaMessage`] for a [`TraceItem`].
    fn for_item(scoping: Scoping, item: TraceItem) -> KafkaMessage<'static> {
        let item_type = item.item_type();
        KafkaMessage::Item {
            headers: BTreeMap::from([
                ("project_id".to_owned(), scoping.project_id.to_string()),
                ("item_type".to_owned(), (item_type as i32).to_string()),
            ]),
            message: item,
            item_type,
        }
    }
}

impl Message for KafkaMessage<'_> {
    fn variant(&self) -> &'static str {
        match self {
            KafkaMessage::Event(_) => "event",
            KafkaMessage::UserReport(_) => "user_report",
            KafkaMessage::Metric { message, .. } => match message.name.namespace() {
                MetricNamespace::Sessions => "metric_sessions",
                MetricNamespace::Spans => "metric_spans",
                MetricNamespace::Transactions => "metric_transactions",
                MetricNamespace::Outcomes => "metric_outcomes",
                MetricNamespace::Unsupported => "metric_unsupported",
            },
            KafkaMessage::CheckIn(_) => "check_in",
            KafkaMessage::SpanV2 { .. } => "span",
            KafkaMessage::Item { item_type, .. } => item_type.as_str_name(),

            KafkaMessage::Attachment(_) => "attachment",
            KafkaMessage::AttachmentChunk(_) => "attachment_chunk",

            KafkaMessage::Profile(_) => "profile",
            KafkaMessage::ProfileChunk(_) => "profile_chunk",

            KafkaMessage::ReplayRecordingNotChunked(_) => "replay_recording_not_chunked",

            KafkaMessage::Outcome(_) => "outcome",
        }
    }

    /// Returns the partitioning key for this Kafka message determining.
    fn key(&self) -> Option<relay_kafka::Key> {
        match self {
            Self::Event(message) => Some((message.event_id.0, message.org_id)),
            Self::UserReport(message) => Some((message.event_id.0, message.org_id)),
            Self::SpanV2 {
                routing_key,
                org_id,
                ..
            } => routing_key.map(|r| (r, *org_id)),

            // Monitor check-ins use the hinted UUID passed through from the Envelope.
            //
            // XXX(epurkhiser): In the future it would be better if all KafkaMessage's would
            // recieve the routing_key_hint form their envelopes.
            Self::CheckIn(message) => message.routing_key_hint.map(|r| (r, message.org_id)),

            Self::Attachment(message) => Some((message.event_id.0, message.org_id)),
            Self::AttachmentChunk(message) => Some((message.event_id.0, message.org_id)),

            // Random partitioning
            Self::Metric { .. }
            | Self::Item { .. }
            | Self::Profile(_)
            | Self::ProfileChunk(_)
            | Self::ReplayRecordingNotChunked(_)
            | Self::Outcome(_) => None,
        }
        .filter(|(uuid, _)| !uuid.is_nil())
        .map(|(uuid, org_id)| {
            // mix key with org id for better paritioning in case of incidents where messages
            // across orgs get the same id
            let mut res = uuid.into_bytes();
            for (i, &b) in org_id.value().to_be_bytes().iter().enumerate() {
                res[i] ^= b;
            }
            u128::from_be_bytes(res)
        })
    }

    fn headers(&self) -> Option<&BTreeMap<String, String>> {
        match &self {
            KafkaMessage::Metric { headers, .. }
            | KafkaMessage::SpanV2 { headers, .. }
            | KafkaMessage::Item { headers, .. }
            | KafkaMessage::Profile(ProfileKafkaMessage { headers, .. })
            | KafkaMessage::ProfileChunk(ProfileChunkKafkaMessage { headers, .. }) => Some(headers),

            KafkaMessage::Event(_)
            | KafkaMessage::UserReport(_)
            | KafkaMessage::CheckIn(_)
            | KafkaMessage::Attachment(_)
            | KafkaMessage::AttachmentChunk(_)
            | KafkaMessage::ReplayRecordingNotChunked(_)
            | KafkaMessage::Outcome(_) => None,
        }
    }

    fn serialize(&self) -> Result<SerializationOutput<'_>, ClientError> {
        match self {
            KafkaMessage::Metric { message, .. } => serialize_as_json(message),
            KafkaMessage::SpanV2 { message, .. } => serialize_as_json(message),
            KafkaMessage::Item { message, .. } => {
                let mut payload = Vec::new();
                match message.encode(&mut payload) {
                    Ok(_) => Ok(SerializationOutput::Protobuf(Cow::Owned(payload))),
                    Err(_) => Err(ClientError::ProtobufEncodingFailed),
                }
            }
            KafkaMessage::Outcome(outcome) => serialize_as_json(outcome),
            KafkaMessage::Event(_)
            | KafkaMessage::UserReport(_)
            | KafkaMessage::CheckIn(_)
            | KafkaMessage::Attachment(_)
            | KafkaMessage::AttachmentChunk(_)
            | KafkaMessage::Profile(_)
            | KafkaMessage::ProfileChunk(_)
            | KafkaMessage::ReplayRecordingNotChunked(_) => match rmp_serde::to_vec_named(&self) {
                Ok(x) => Ok(SerializationOutput::MsgPack(Cow::Owned(x))),
                Err(err) => Err(ClientError::InvalidMsgPack(err)),
            },
        }
    }
}

fn serialize_as_json<T: serde::Serialize>(
    value: &T,
) -> Result<SerializationOutput<'_>, ClientError> {
    match serde_json::to_vec(value) {
        Ok(vec) => Ok(SerializationOutput::Json(Cow::Owned(vec))),
        Err(err) => Err(ClientError::InvalidJson(err)),
    }
}

fn bool_to_str(value: bool) -> &'static str {
    if value { "true" } else { "false" }
}

/// Returns a safe timestamp for Kafka.
///
/// Kafka expects timestamps to be in UTC and in seconds since epoch.
fn safe_timestamp(timestamp: DateTime<Utc>) -> u64 {
    let ts = timestamp.timestamp();
    if ts >= 0 {
        return ts as u64;
    }

    // We assume this call can't return < 0.
    Utc::now().timestamp() as u64
}