khadhroony-bot3/ks-pipeline/src/account_state_normalization.rs

// file: ks-pipeline/src/account_state_normalization.rs
// version: 1

//! Generic Solana account-observation to Core account-state normalization pipeline.

use sha2::Digest; // rust-rules: trait-import

/// Stable processor name used by generic account-state normalization.
pub const ACCOUNT_STATE_NORMALIZER_NAME: &str = "account_observation_to_core";
/// Current generic account-state normalizer implementation version.
pub const ACCOUNT_STATE_NORMALIZER_VERSION: &str = "1";

/// Bounded generic account-state normalization request.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AccountStateNormalizationRequest {
    /// Source observation selection.
    pub selection: ks_store::AccountStateSelectionFilter,
    /// Forces replacement even when the same processor version and input hash already succeeded.
    pub force_replay: bool,
}

impl crate::AccountStateNormalizationRequest {
    /// Builds a request bound to the active generic account-state normalizer identity.
    pub fn new(
        account_keys: std::vec::Vec<std::string::String>,
        min_slot: std::option::Option<u64>,
        max_slot: std::option::Option<u64>,
        force_replay: bool,
        limit: u32,
    ) -> ks_core::Result<Self> {
        let selection_result = ks_store::AccountStateSelectionFilter::new(
            crate::ACCOUNT_STATE_NORMALIZER_NAME,
            crate::ACCOUNT_STATE_NORMALIZER_VERSION,
            account_keys,
            min_slot,
            max_slot,
            force_replay,
            limit,
        );
        let selection = match selection_result {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        return std::result::Result::Ok(Self { selection, force_replay });
    }
}

/// Summary returned by one generic account-state normalization campaign.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct AccountStateNormalizationSummary {
    /// Complete observations selected from N1.
    pub selected: u64,
    /// Core account states inserted or replaced in N2.
    pub normalized: u64,
    /// Inputs skipped because the same processor version and hash are already current.
    pub skipped: u64,
    /// Inputs that failed normalization or persistence.
    pub failed: u64,
}

/// Executes one bounded replayable N1 account-observation to N2 Core account-state campaign.
pub async fn execute_account_state_normalization<S>(
    store: &S,
    request: &crate::AccountStateNormalizationRequest,
) -> ks_core::Result<crate::AccountStateNormalizationSummary>
where
    S: ks_store::AccountStateStore + Sync,
{
    if request.selection.processor_name != crate::ACCOUNT_STATE_NORMALIZER_NAME
        || request.selection.processor_version != crate::ACCOUNT_STATE_NORMALIZER_VERSION
    {
        return std::result::Result::Err(ks_core::Error::invalid_state(
            "account-state selection processor identity does not match the active normalizer",
        ));
    }
    let mut effective_selection = request.selection.clone();
    effective_selection.include_current = request.force_replay;
    let observations_result =
        store.list_account_observations_for_normalization(&effective_selection).await;
    let observations = match observations_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let mut summary = crate::AccountStateNormalizationSummary {
        selected: match u64::try_from(observations.len()) {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => {
                return std::result::Result::Err(ks_core::Error::invalid_state(format!(
                    "account-state selection length conversion failed: {error}"
                )));
            },
        },
        ..crate::AccountStateNormalizationSummary::default()
    };
    tracing::info!(target: crate::TRACING_TARGET, action = "execute_account_state_normalization", selected = summary.selected, force_replay = request.force_replay, processor_name = crate::ACCOUNT_STATE_NORMALIZER_NAME, processor_version = crate::ACCOUNT_STATE_NORMALIZER_VERSION, "account-state normalization campaign started");
    for observation in &observations {
        let identity_result = ledger_identity_from_observation(observation);
        let identity = match identity_result {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        if !request.force_replay {
            let current_result = store.is_account_state_current(&identity).await;
            match current_result {
                std::result::Result::Ok(true) => {
                    summary.skipped += 1;
                    continue;
                },
                std::result::Result::Ok(false) => {},
                std::result::Result::Err(error) => return std::result::Result::Err(error),
            }
        }
        let state_result = core_state_from_observation(observation);
        let state = match state_result {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => {
                let failure_result =
                    persist_normalization_failure(store, observation.id, identity, &error).await;
                if let std::result::Result::Err(persistence_error) = failure_result {
                    return std::result::Result::Err(persistence_error);
                }
                summary.failed += 1;
                continue;
            },
        };
        let bundle =
            ks_store::AccountStatePersistenceBundle { ledger_identity: identity.clone(), state };
        let persistence_result = store.persist_account_state(&bundle, request.force_replay).await;
        match persistence_result {
            std::result::Result::Ok(outcome) => {
                if outcome.inserted_count > 0 || outcome.updated_count > 0 {
                    summary.normalized += 1;
                } else {
                    summary.skipped += 1;
                }
            },
            std::result::Result::Err(error) => {
                let failure_result =
                    persist_normalization_failure(store, observation.id, identity, &error).await;
                if let std::result::Result::Err(persistence_error) = failure_result {
                    return std::result::Result::Err(persistence_error);
                }
                summary.failed += 1;
            },
        }
    }
    tracing::info!(target: crate::TRACING_TARGET, action = "execute_account_state_normalization", selected = summary.selected, normalized = summary.normalized, skipped = summary.skipped, failed = summary.failed, force_replay = request.force_replay, processor_name = crate::ACCOUNT_STATE_NORMALIZER_NAME, processor_version = crate::ACCOUNT_STATE_NORMALIZER_VERSION, "account-state normalization campaign completed");
    return std::result::Result::Ok(summary);
}

fn ledger_identity_from_observation(
    observation: &ks_store::AccountObservationRow,
) -> ks_core::Result<ks_store::ProcessingLedgerIdentity> {
    let hash = canonical_account_observation_hash(observation);
    return ks_store::ProcessingLedgerIdentity::new(
        ks_store::ACCOUNT_STATE_NORMALIZATION_STAGE,
        crate::ACCOUNT_STATE_NORMALIZER_NAME,
        crate::ACCOUNT_STATE_NORMALIZER_VERSION,
        observation.observation_key.clone(),
        hash,
    );
}

fn canonical_account_observation_hash(
    observation: &ks_store::AccountObservationRow,
) -> std::string::String {
    let owner = match observation.owner.as_deref() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => "",
    };
    let lamports = match observation.lamports.as_deref() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => "",
    };
    let rent_epoch = match observation.rent_epoch.as_deref() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => "",
    };
    let data_hash = match observation.data_hash.as_deref() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => "",
    };
    let space = match observation.space {
        std::option::Option::Some(value) => value,
        std::option::Option::None => 0_i64,
    };
    let executable = match observation.executable {
        std::option::Option::Some(value) => value,
        std::option::Option::None => false,
    };
    let mut hasher = sha2::Sha256::new();
    for value in [
        observation.observation_key.as_str(),
        observation.account_key.as_str(),
        owner,
        lamports,
        rent_epoch,
        data_hash,
    ] {
        hasher.update(value.as_bytes());
        hasher.update([0_u8]);
    }
    hasher.update(observation.context_slot.to_le_bytes());
    hasher.update(space.to_le_bytes());
    hasher.update([u8::from(executable)]);
    let digest = hasher.finalize();
    let mut output = std::string::String::with_capacity(64);
    for byte in digest {
        output.push_str(format!("{byte:02x}").as_str());
    }
    return output;
}

fn core_state_from_observation(
    observation: &ks_store::AccountObservationRow,
) -> ks_core::Result<ks_store::CoreAccountStateInsert> {
    let slot_result = u64::try_from(observation.context_slot);
    let slot = match slot_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => {
            return std::result::Result::Err(ks_core::Error::invalid_state(format!(
                "account observation context slot is negative: {error}"
            )));
        },
    };
    let owner = match &observation.owner {
        std::option::Option::Some(value) => value.clone(),
        std::option::Option::None => {
            return std::result::Result::Err(ks_core::Error::invalid_state(
                "account observation owner is missing",
            ));
        },
    };
    let lamports_result = parse_u64_text(observation.lamports.as_deref(), "lamports");
    let lamports = match lamports_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let executable = match observation.executable {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(ks_core::Error::invalid_state(
                "account observation executable flag is missing",
            ));
        },
    };
    let rent_epoch_result = parse_u64_text(observation.rent_epoch.as_deref(), "rent epoch");
    let rent_epoch = match rent_epoch_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let space_sql = match observation.space {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(ks_core::Error::invalid_state(
                "account observation space is missing",
            ));
        },
    };
    let space_result = u64::try_from(space_sql);
    let space = match space_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => {
            return std::result::Result::Err(ks_core::Error::invalid_state(format!(
                "account observation space is negative: {error}"
            )));
        },
    };
    let data_base64 = match &observation.data_base64 {
        std::option::Option::Some(value) => value.clone(),
        std::option::Option::None => {
            return std::result::Result::Err(ks_core::Error::invalid_state(
                "account observation data is missing",
            ));
        },
    };
    let data_hash = match &observation.data_hash {
        std::option::Option::Some(value) => value.clone(),
        std::option::Option::None => {
            return std::result::Result::Err(ks_core::Error::invalid_state(
                "account observation data hash is missing",
            ));
        },
    };
    let state = ks_store::CoreAccountStateInsert {
        source_observation_id: observation.id,
        account_key: observation.account_key.clone(),
        slot,
        owner,
        lamports,
        executable,
        rent_epoch,
        space,
        data_base64,
        data_hash,
    };
    let validation_result = state.validate();
    if let std::result::Result::Err(error) = validation_result {
        return std::result::Result::Err(error);
    }
    return std::result::Result::Ok(state);
}

fn parse_u64_text(value: std::option::Option<&str>, label: &str) -> ks_core::Result<u64> {
    let text = match value {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(ks_core::Error::invalid_state(format!(
                "account observation {label} is missing"
            )));
        },
    };
    let parse_result = text.parse::<u64>();
    return match parse_result {
        std::result::Result::Ok(value) => std::result::Result::Ok(value),
        std::result::Result::Err(error) => std::result::Result::Err(ks_core::Error::invalid_state(
            format!("account observation {label} is invalid: {error}"),
        )),
    };
}

async fn persist_normalization_failure<S>(
    store: &S,
    source_observation_id: i64,
    identity: ks_store::ProcessingLedgerIdentity,
    error: &ks_core::Error,
) -> ks_core::Result<()>
where
    S: ks_store::AccountStateStore + Sync,
{
    let message = bounded_error_message(error.to_string().as_str());
    let failure = ks_store::AccountStateNormalizationFailure {
        source_observation_id,
        ledger_identity: identity,
        error_code: "account_state_normalization_failed".to_string(),
        error_message: message,
    };
    let persistence_result = store.mark_account_state_failed(&failure).await;
    return match persistence_result {
        std::result::Result::Ok(_outcome) => std::result::Result::Ok(()),
        std::result::Result::Err(persistence_error) => std::result::Result::Err(persistence_error),
    };
}

fn bounded_error_message(value: &str) -> std::string::String {
    const MAX_ERROR_MESSAGE_CHARS: usize = 512;
    return value.chars().take(MAX_ERROR_MESSAGE_CHARS).collect();
}

#[cfg(test)]
mod tests {
    #[test]
    fn request_constructor_binds_active_processor_identity() {
        let request_result = crate::AccountStateNormalizationRequest::new(
            std::vec::Vec::new(),
            std::option::Option::None,
            std::option::Option::None,
            true,
            10,
        );
        let request = match request_result {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => panic!("unexpected request error: {error}"),
        };
        assert_eq!(request.selection.processor_name, crate::ACCOUNT_STATE_NORMALIZER_NAME);
        assert_eq!(request.selection.processor_version, crate::ACCOUNT_STATE_NORMALIZER_VERSION);
        assert!(request.selection.include_current);
        assert!(request.force_replay);
    }

    #[test]
    fn request_constructor_excludes_current_inputs_without_force_replay() {
        let request_result = crate::AccountStateNormalizationRequest::new(
            std::vec::Vec::new(),
            std::option::Option::None,
            std::option::Option::None,
            false,
            10,
        );
        let request = match request_result {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => panic!("unexpected request error: {error}"),
        };
        assert!(!request.selection.include_current);
        assert!(!request.force_replay);
    }

    #[test]
    fn unsigned_decimal_parser_accepts_full_u64_range() {
        assert_eq!(
            super::parse_u64_text(std::option::Option::Some("18446744073709551615"), "lamports"),
            std::result::Result::Ok(u64::MAX)
        );
    }

    #[test]
    fn unsigned_decimal_parser_rejects_negative_values() {
        assert!(super::parse_u64_text(std::option::Option::Some("-1"), "lamports").is_err());
    }

    #[test]
    fn error_messages_are_bounded() {
        assert_eq!(super::bounded_error_message("x".repeat(600).as_str()).len(), 512);
    }
}