lite-rpc/blockstore/src/block_stores/postgres/postgres_block_store_writer.rs

397 lines
14 KiB
Rust

use std::time::{Duration, Instant};
use crate::block_stores::postgres::{LITERPC_QUERY_ROLE, LITERPC_ROLE};
use anyhow::{bail, Context, Result};
use itertools::Itertools;
use log::{debug, info, trace, warn};
use solana_lite_rpc_core::structures::epoch::EpochRef;
use solana_lite_rpc_core::structures::{epoch::EpochCache, produced_block::ProducedBlock};
use solana_sdk::commitment_config::CommitmentLevel;
use solana_sdk::slot_history::Slot;
use tokio_postgres::error::SqlState;
use super::postgres_block::*;
use super::postgres_config::*;
use super::postgres_epoch::*;
use super::postgres_session::*;
use super::postgres_transaction::*;
const PARALLEL_WRITE_SESSIONS: usize = 4;
const MIN_WRITE_CHUNK_SIZE: usize = 500;
#[derive(Clone)]
pub struct PostgresBlockStore {
session_cache: PostgresSessionCache,
// use this session only for the write path!
write_sessions: Vec<PostgresWriteSession>,
epoch_schedule: EpochCache,
}
impl PostgresBlockStore {
pub async fn new(epoch_schedule: EpochCache, pg_session_config: PostgresSessionConfig) -> Self {
let session_cache = PostgresSessionCache::new(pg_session_config.clone())
.await
.unwrap();
let mut write_sessions = Vec::new();
for _i in 0..PARALLEL_WRITE_SESSIONS {
write_sessions.push(
PostgresWriteSession::new(pg_session_config.clone())
.await
.unwrap(),
);
}
assert!(
!write_sessions.is_empty(),
"must have at least one write session"
);
Self::check_write_role(&session_cache).await;
Self {
session_cache,
write_sessions,
epoch_schedule,
}
}
async fn check_write_role(session_cache: &PostgresSessionCache) {
let role = LITERPC_ROLE;
let statement = format!("SELECT 1 FROM pg_roles WHERE rolname='{role}'");
let count = session_cache
.get_session()
.await
.expect("must get session")
.execute(&statement, &[])
.await
.expect("must execute query to check for role");
if count == 0 {
panic!(
"Missing mandatory postgres write/ownership role '{}' for Lite RPC - see permissions.sql",
role
);
} else {
info!(
"Self check - found postgres write role/ownership '{}'",
role
);
}
}
// return true if schema was actually created
async fn start_new_epoch_if_necessary(&self, epoch: EpochRef) -> Result<bool> {
// create schema for new epoch
let schema_name = PostgresEpoch::build_schema_name(epoch);
let session = self.get_session().await;
let statement = PostgresEpoch::build_create_schema_statement(epoch);
// note: requires GRANT CREATE ON DATABASE xyz
let result_create_schema = session.execute_multiple(&statement).await;
if let Err(err) = result_create_schema {
if err
.code()
.map(|sqlstate| sqlstate == &SqlState::DUPLICATE_SCHEMA)
.unwrap_or(false)
{
// TODO: do we want to allow this; continuing with existing epoch schema might lead to inconsistent data in blocks and transactions table
info!(
"Schema {} for epoch {} already exists - data will be appended",
schema_name, epoch
);
return Ok(false);
} else {
return Err(err).context("create schema for new epoch");
}
}
// set permissions for new schema
let statement = build_assign_permissions_statements(epoch);
session
.execute_multiple(&statement)
.await
.context("Set postgres permissions for new schema")?;
// Create blocks table
let statement = PostgresBlock::build_create_table_statement(epoch);
session
.execute_multiple(&statement)
.await
.context("create blocks table for new epoch")?;
// create transaction table
let statement = PostgresTransaction::build_create_table_statement(epoch);
session
.execute_multiple(&statement)
.await
.context("create transaction table for new epoch")?;
// add foreign key constraint between transactions and blocks
let statement = PostgresTransaction::build_foreign_key_statement(epoch);
session
.execute_multiple(&statement)
.await
.context("create foreign key constraint between transactions and blocks")?;
info!("Start new epoch in postgres schema {}", schema_name);
Ok(true)
}
async fn get_session(&self) -> PostgresSession {
self.session_cache
.get_session()
.await
.expect("should get new postgres session")
}
// optimistically try to progress commitment level for a block that is already stored
pub async fn progress_block_commitment_level(&self, block: &ProducedBlock) -> Result<()> {
// ATM we only support updating confirmed block to finalized
if block.commitment_config.commitment == CommitmentLevel::Finalized {
debug!(
"Checking block {} if we can progress it to finalized ...",
block.slot
);
// TODO model commitment levels in new table
}
Ok(())
}
pub async fn save_block(&self, block: &ProducedBlock) -> Result<()> {
self.progress_block_commitment_level(block).await?;
// let PostgresData { current_epoch, .. } = { *self.postgres_data.read().await };
trace!(
"Saving block {}@{} to postgres storage...",
block.slot,
block.commitment_config.commitment
);
let slot = block.slot;
let transactions = block
.transactions
.iter()
.map(|x| PostgresTransaction::new(x, slot))
.collect_vec();
let postgres_block = PostgresBlock::from(block);
let epoch = self.epoch_schedule.get_epoch_at_slot(slot);
let write_session_single = self.write_sessions[0].get_write_session().await;
let started_block = Instant::now();
let inserted = postgres_block
.save(&write_session_single, epoch.into())
.await?;
if !inserted {
debug!("Block {} already exists - skip update", slot);
return Ok(());
}
let elapsed_block_insert = started_block.elapsed();
let started_txs = Instant::now();
let mut queries_fut = Vec::new();
let chunk_size =
div_ceil(transactions.len(), self.write_sessions.len()).max(MIN_WRITE_CHUNK_SIZE);
let chunks = transactions.chunks(chunk_size).collect_vec();
assert!(
chunks.len() <= self.write_sessions.len(),
"cannot have more chunks than session"
);
for (i, chunk) in chunks.iter().enumerate() {
let session = self.write_sessions[i].get_write_session().await.clone();
let future =
PostgresTransaction::save_transactions_from_block(session, epoch.into(), chunk);
queries_fut.push(future);
}
let all_results: Vec<Result<()>> = futures_util::future::join_all(queries_fut).await;
for result in all_results {
result.expect("Save query must succeed");
}
let elapsed_txs_insert = started_txs.elapsed();
info!(
"Saving block {}@{} to postgres took {:.2}ms for block and {:.2}ms for {} transactions ({}x{} chunks)",
slot, block.commitment_config.commitment,
elapsed_block_insert.as_secs_f64() * 1000.0,
elapsed_txs_insert.as_secs_f64() * 1000.0,
transactions.len(),
chunks.len(),
chunk_size,
);
Ok(())
}
// ATM we focus on blocks as this table gets INSERTS and does deduplication checks (i.e. heavy reads on index pk_block_slot)
pub async fn optimize_blocks_table(&self, slot: Slot) -> Result<()> {
let started = Instant::now();
let epoch: EpochRef = self.epoch_schedule.get_epoch_at_slot(slot).into();
let random_session = slot as usize % self.write_sessions.len();
let write_session_single = self.write_sessions[random_session]
.get_write_session()
.await;
let statement = format!(
r#"
ANALYZE (SKIP_LOCKED) {schema}.blocks;
"#,
schema = PostgresEpoch::build_schema_name(epoch),
);
tokio::spawn(async move {
write_session_single
.execute_multiple(&statement)
.await
.unwrap();
let elapsed = started.elapsed();
debug!(
"Postgres analyze of blocks table took {:.2}ms",
elapsed.as_secs_f64() * 1000.0
);
if elapsed > Duration::from_millis(500) {
warn!(
"Very slow postgres ANALYZE on slot {} - took {:.2}ms",
slot,
elapsed.as_secs_f64() * 1000.0
);
}
});
Ok(())
}
// create current + next epoch
// true if anything was created; false if a NOOP
pub async fn prepare_epoch_schema(&self, slot: Slot) -> anyhow::Result<bool> {
let epoch = self.epoch_schedule.get_epoch_at_slot(slot);
let current_epoch = epoch.into();
let created_current = self.start_new_epoch_if_necessary(current_epoch).await?;
let next_epoch = current_epoch.get_next_epoch();
let created_next = self.start_new_epoch_if_necessary(next_epoch).await?;
Ok(created_current || created_next)
}
// used for testing only ATM
pub async fn drop_epoch_schema(&self, epoch: EpochRef) -> anyhow::Result<()> {
// create schema for new epoch
let schema_name = PostgresEpoch::build_schema_name(epoch);
let session = self.get_session().await;
let statement = PostgresEpoch::build_drop_schema_statement(epoch);
let result_drop_schema = session.execute_multiple(&statement).await;
match result_drop_schema {
Ok(_) => {
warn!("Dropped schema {}", schema_name);
Ok(())
}
Err(_err) => {
bail!("Error dropping schema {}", schema_name)
}
}
}
}
fn build_assign_permissions_statements(epoch: EpochRef) -> String {
let schema = PostgresEpoch::build_schema_name(epoch);
format!(
r#"
GRANT USAGE ON SCHEMA {schema} TO {role};
GRANT ALL ON ALL TABLES IN SCHEMA {schema} TO {role};
ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT ALL ON TABLES TO {role};
GRANT USAGE ON SCHEMA {schema} TO {query_role};
ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT SELECT ON TABLES TO {query_role};
"#,
role = LITERPC_ROLE,
query_role = LITERPC_QUERY_ROLE,
)
}
fn div_ceil(a: usize, b: usize) -> usize {
(a.saturating_add(b).saturating_sub(1)).saturating_div(b)
}
#[cfg(test)]
mod tests {
use super::*;
use solana_lite_rpc_core::structures::produced_block::TransactionInfo;
use solana_sdk::commitment_config::CommitmentConfig;
use solana_sdk::signature::Signature;
use std::str::FromStr;
#[tokio::test]
#[ignore]
async fn postgres_write_session() {
let write_session = PostgresWriteSession::new_from_env().await.unwrap();
let row_role = write_session
.get_write_session()
.await
.query_one("SELECT current_role", &[])
.await
.unwrap();
info!("row: {:?}", row_role);
}
#[tokio::test]
#[ignore]
async fn test_save_block() {
tracing_subscriber::fmt::init();
let pg_session_config = PostgresSessionConfig {
pg_config: "host=localhost dbname=literpc3 user=literpc_app password=litelitesecret"
.to_string(),
ssl: None,
};
let _postgres_session_cache = PostgresSessionCache::new(pg_session_config.clone())
.await
.unwrap();
let epoch_cache = EpochCache::new_for_tests();
let postgres_block_store =
PostgresBlockStore::new(epoch_cache.clone(), pg_session_config.clone()).await;
postgres_block_store
.save_block(&create_test_block())
.await
.unwrap();
}
fn create_test_block() -> ProducedBlock {
let sig1 = Signature::from_str("5VBroA4MxsbZdZmaSEb618WRRwhWYW9weKhh3md1asGRx7nXDVFLua9c98voeiWdBE7A9isEoLL7buKyaVRSK1pV").unwrap();
let sig2 = Signature::from_str("3d9x3rkVQEoza37MLJqXyadeTbEJGUB6unywK4pjeRLJc16wPsgw3dxPryRWw3UaLcRyuxEp1AXKGECvroYxAEf2").unwrap();
ProducedBlock {
block_height: 42,
blockhash: "blockhash".to_string(),
previous_blockhash: "previous_blockhash".to_string(),
parent_slot: 666,
slot: 223555999,
transactions: vec![create_test_tx(sig1), create_test_tx(sig2)],
// TODO double if this is unix millis or seconds
block_time: 1699260872000,
commitment_config: CommitmentConfig::finalized(),
leader_id: None,
rewards: None,
}
}
fn create_test_tx(signature: Signature) -> TransactionInfo {
TransactionInfo {
signature: signature.to_string(),
is_vote: false,
err: None,
cu_requested: Some(40000),
prioritization_fees: Some(5000),
cu_consumed: Some(32000),
recent_blockhash: "recent_blockhash".to_string(),
message: "some message".to_string(),
writable_accounts: vec![],
readable_accounts: vec![],
}
}
}