Skip to content

perf(storage-cassandra): restore copy writes plain INSERTs, not the LWT path #527

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LeeroyHannigan wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(storage-cassandra): restore copy writes plain INSERTs, not the LWT path #527

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 57 additions & 52 deletions crates/storage-cassandra/src/backup_engine.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,8 +73,8 @@ use extenddb_storage::backup_definition::{
};
use extenddb_storage::error::StorageError;
use extenddb_storage::{BackupEngine, BackupRestoreSource, MetadataEngine, RestoreRequest};
use futures::StreamExt;
use futures::future::BoxFuture;
use futures::{StreamExt, TryStreamExt};

use crate::CassandraEngine;
use crate::create_table::RestoreTargetColumns;
Expand All @@ -87,8 +87,13 @@ use crate::delete_table::DeleteTableCaller;
const BACKUP_ITEMS_PER_BUCKET: i64 = 64;
const BACKUP_SCAN_PAGE_SIZE: i64 = 1_000;

/// Items of one bucket a restore copy writes concurrently.
const RESTORE_WRITE_CONCURRENCY: usize = 8;
/// Items a restore copy writes concurrently. Each write is a plain INSERT
/// per row (`data::restore_copy`), so this bounds in-flight statements per
/// index-less item; an item with indexes adds one statement per index row.
const RESTORE_WRITE_CONCURRENCY: usize = 64;

/// Backup payload buckets a restore copy reads ahead of its writes.
const RESTORE_READ_AHEAD_BUCKETS: usize = 2;

/// How often a running restore copy refreshes its target's heartbeat.
pub const RESTORE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
Expand Down Expand Up @@ -972,64 +977,64 @@ impl CassandraEngine {
key_schema: job.key_schema.clone(),
base_key_schema: job.key_schema.clone(),
attribute_definitions: job.attribute_definitions.clone(),
// The put path reads the target's indexes from the catalog by
// table_id; these routing fields are not consulted by it.
// The restore write path reads the target's indexes from the
// catalog by table_id; these routing fields are not consulted by it.
has_lsi: false,
global_secondary_indexes: Vec::new(),
local_secondary_indexes: Vec::new(),
stream_specification: None,
vector_indexes: Vec::new(),
};
let account_keyspace = self.account_keyspace(&job.account_id);
let query = format!(
"SELECT item_data FROM {account_keyspace}.backup_items WHERE backup_arn = ? AND bucket = ?"
);
let writer = self
.restore_writer(&key_info)
.await
.map_err(CopyError::Failed)?;
let read_bucket = |bucket: i64| {
let account_keyspace = &account_keyspace;
async move {
if heartbeat.is_lost() {
return Err(CopyError::OwnershipLost);
}
self.read_backup_bucket(account_keyspace, &job.backup_arn, bucket)
.await
.map_err(CopyError::Failed)
}
};
// Buckets are read ahead of the writes so the copy is never waiting on
// a payload read with nothing in flight; the write stage runs
// RESTORE_WRITE_CONCURRENCY items at a time across bucket boundaries.
let mut restored_count = 0_i64;
let mut restored_size = 0_i64;
for bucket in 0..bucket_count(job.item_count) {
if heartbeat.is_lost() {
return Err(CopyError::OwnershipLost);
}
#[allow(clippy::cast_possible_truncation)]
let bucket_i32 = bucket as i32;
let rows = self
.session
.query_with_values(
&query,
cdrs_tokio::query_values!(job.backup_arn.as_str(), bucket_i32),
)
.await
.map_err(|e| {
CopyError::Failed(StorageError::Internal(format!("Read backup payload: {e}")))
})?
.response_body()
.map_err(|e| {
CopyError::Failed(StorageError::Internal(format!("Parse backup payload: {e}")))
})?
.into_rows()
.unwrap_or_default();
let mut items = Vec::with_capacity(rows.len());
for row in rows {
let item_data: String = row.get_r_by_name("item_data").map_err(|e| {
CopyError::Failed(StorageError::Internal(format!("Parse backup item: {e}")))
})?;
#[allow(clippy::cast_possible_wrap)]
let item_len = item_data.len() as i64;
restored_size = restored_size.saturating_add(item_len);
let item = serde_json::from_str(&item_data).map_err(|e| {
CopyError::Failed(StorageError::Internal(format!(
"Deserialize backup item: {e}"
)))
})?;
items.push(item);
}
let mut writes = futures::stream::iter(items)
.map(|item| self.put_restored_item(&key_info, item))
.buffer_unordered(RESTORE_WRITE_CONCURRENCY);
while let Some(result) = writes.next().await {
result.map_err(CopyError::Failed)?;
restored_count += 1;
}
let mut writes = futures::stream::iter(0..bucket_count(job.item_count))
.map(read_bucket)
.buffered(RESTORE_READ_AHEAD_BUCKETS)
.map_ok(|items| futures::stream::iter(items.into_iter().map(Ok)))
.try_flatten()
.map_ok(|(item, item_len)| {
let writer = &writer;
let heartbeat = &heartbeat;
async move {
// The read stage checks ownership per bucket; with two
// buckets read ahead and 64 writes in flight, up to three
// buckets of items would otherwise still be written into
// a target the sweep has claimed.
if heartbeat.is_lost() {
return Err(CopyError::OwnershipLost);
}
self.write_restored_item(writer, &item)
.await
.map_err(CopyError::Failed)?;
Ok::<usize, CopyError>(item_len)
}
})
.try_buffer_unordered(RESTORE_WRITE_CONCURRENCY);
while let Some(result) = writes.next().await {
let item_len = result?;
#[allow(clippy::cast_possible_wrap)]
let item_len = item_len as i64;
restored_size = restored_size.saturating_add(item_len);
restored_count += 1;
}
if restored_count != job.item_count {
return Err(CopyError::Failed(StorageError::Internal(format!(
Expand Down
33 changes: 29 additions & 4 deletions crates/storage-cassandra/src/data/index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -506,6 +506,34 @@ pub(crate) fn insert_index_row_multi(
idx_sks: &[(&str, ScalarAttributeType)],
base_sks: &[(&str, ScalarAttributeType)],
) -> Result<(), StorageError> {
let (cql, qv) = index_row_insert(
account_keyspace,
idx_table,
item,
projected,
index_ks,
base_ks,
idx_sks,
base_sks,
)?;
let old = std::mem::replace(batch, BatchQueryBuilder::new());
*batch = old.add_query(cql, qv);
Ok(())
}

/// The parameterized INSERT of `item`'s row in the index table `idx_table`,
/// as a statement and its bound values.
#[allow(clippy::too_many_arguments)]
pub(crate) fn index_row_insert(
account_keyspace: &str,
idx_table: &str,
item: &Item,
projected: &Item,
index_ks: &[KeySchemaElement],
base_ks: &[KeySchemaElement],
idx_sks: &[(&str, ScalarAttributeType)],
base_sks: &[(&str, ScalarAttributeType)],
) -> Result<(String, cdrs_tokio::query::QueryValues), StorageError> {
let idx_pk_text = composite_pk_to_text(item, index_ks)?;
let base_pk_text = composite_pk_to_text(item, base_ks)?;

Expand Down Expand Up @@ -555,10 +583,7 @@ pub(crate) fn insert_index_row_multi(
cols.join(", "),
placeholders
);
let qv = cdrs_tokio::query::QueryValues::SimpleValues(values);
let old = std::mem::replace(batch, BatchQueryBuilder::new());
*batch = old.add_query(cql, qv);
Ok(())
Ok((cql, cdrs_tokio::query::QueryValues::SimpleValues(values)))
}

/// Convert a `SortKeyValue` to a `cdrs_tokio` bound `Value`.
Expand Down
1 change: 1 addition & 0 deletions crates/storage-cassandra/src/data/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ pub mod index;
mod put_get_item;
mod query;
mod query_helpers;
mod restore_copy;
mod scan;
pub mod transaction_ledger;
mod transactions;
Expand Down
79 changes: 3 additions & 76 deletions crates/storage-cassandra/src/data/put_get_item.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,6 @@ use super::{json_to_item, query_with_pk_sk};
use crate::CassandraEngine;
use crate::stream_util::stream_record_statement;

/// How a put maintains GSIs.
#[derive(Clone, Copy, PartialEq, Eq)]
enum IndexWrites {
/// Each GSI by its propagation delay: synchronous at 0, otherwise through
/// the `gsi_pending` queue.
AsConfigured,
/// Every GSI in the item's own batch, as with a 0 delay.
AllSynchronous,
}

impl CassandraEngine {
/// Implementation of `DataEngine::put_item`.
///
Expand All @@ -42,68 +32,12 @@ impl CassandraEngine {
condition: Option<&Expr>,
maps: &ExpressionMaps,
stream: Option<&StreamCapture>,
) -> Result<Option<Item>, StorageError> {
self.put_item_with_retry(
key_info,
item,
return_old,
condition,
maps,
stream,
IndexWrites::AsConfigured,
)
.await
}

/// The write a restore copy makes for one backup item: the ordinary put
/// path (same base-row write, same index maintenance helpers, same
/// projection), with every GSI written in the item's own logged batch as
/// LSIs always are, instead of through the asynchronous propagation queue.
/// A restore target becomes ACTIVE only after its last item is written,
/// and DynamoDB reports a restored table's indexes as complete at ACTIVE;
/// queueing the rows would let ACTIVE precede them.
pub(crate) async fn put_restored_item(
&self,
key_info: &TableKeyInfo,
item: Item,
) -> Result<(), StorageError> {
self.put_item_with_retry(
key_info,
item,
false,
None,
&ExpressionMaps::default(),
None,
IndexWrites::AllSynchronous,
)
.await
.map(|_| ())
}

#[allow(clippy::too_many_arguments)]
async fn put_item_with_retry(
&self,
key_info: &TableKeyInfo,
item: Item,
return_old: bool,
condition: Option<&Expr>,
maps: &ExpressionMaps,
stream: Option<&StreamCapture>,
index_writes: IndexWrites,
) -> Result<Option<Item>, StorageError> {
use super::delete_item::{TTL_CLAIM_MAX_RETRIES, ttl_claim_backoff};

for attempt in 0..=TTL_CLAIM_MAX_RETRIES {
match self
.put_item_impl_inner(
key_info,
item.clone(),
return_old,
condition,
maps,
stream,
index_writes,
)
.put_item_impl_inner(key_info, item.clone(), return_old, condition, maps, stream)
.await
{
Err(StorageError::TransactionConflict(message))
Expand All @@ -126,7 +60,6 @@ impl CassandraEngine {
unreachable!("loop returns on the final attempt")
}

#[allow(clippy::too_many_arguments)]
async fn put_item_impl_inner(
&self,
key_info: &TableKeyInfo,
Expand All @@ -135,7 +68,6 @@ impl CassandraEngine {
condition: Option<&Expr>,
maps: &ExpressionMaps,
stream: Option<&StreamCapture>,
index_writes: IndexWrites,
) -> Result<Option<Item>, StorageError> {
let data_keyspace = self.account_keyspace(&key_info.account_id);
let catalog_keyspace = self.catalog_keyspace();
Expand All @@ -149,7 +81,7 @@ impl CassandraEngine {

// Fetch indexes for GSI/LSI updates
// Both are catalog reads with no data dependency; overlap them.
let (mut indexes, ttl_config) = futures::try_join!(
let (indexes, ttl_config) = futures::try_join!(
super::index::fetch_indexes_for_table(
&key_info.table_id,
&self.session,
Expand All @@ -174,12 +106,7 @@ impl CassandraEngine {
)
.map_err(|e| StorageError::Validation(e.to_string()))?;
}
if index_writes == IndexWrites::AllSynchronous {
for index in &mut indexes {
index.propagation_delay_ms = Some(0);
}
}
let sys_delay = if indexes.is_empty() || index_writes == IndexWrites::AllSynchronous {
let sys_delay = if indexes.is_empty() {
0
} else {
self.gsi_default_delay_ms
Expand Down
Loading
Loading