use std::{net::SocketAddr, sync::Arc, time::Duration};
use tokio::task::JoinHandle;
use tonic::{Status, Streaming};
use tower::BoxError;
use zakura_chain::{
block::{self, Block, Height},
parameters::Network,
serialization::BytesInDisplayOrder,
};
use zakura_state::{
spawn_init_read_only, ChainTipBlock, ChainTipChange, ChainTipSender, CheckpointVerifiedBlock,
HashOrHeight, LatestChainTip, NonFinalizedState, ReadStateService, SemanticallyVerifiedBlock,
ValidateContextError, ZakuraDb,
};
use zakura_chain::diagnostic::task::WaitForPanics;
use crate::indexer::{
indexer_client::IndexerClient, BlockAndHash, BlockRequest, Empty,
NonFinalizedStateChangeRequest,
};
const POLL_DELAY: Duration = Duration::from_secs(5);
const STREAM_MESSAGE_TIMEOUT: Duration = Duration::from_secs(10 * 60);
const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(60);
const KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(20);
const COMMIT_RETRY_DELAY: Duration = Duration::from_secs(1);
const GET_BLOCK_TIMEOUT: Duration = Duration::from_secs(30);
const SUBSCRIBE_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Debug)]
pub struct TrustedChainSync {
pub indexer_rpc_client: IndexerClient<tonic::transport::Channel>,
db: ZakuraDb,
non_finalized_state: NonFinalizedState,
chain_tip_sender: ChainTipSender,
non_finalized_state_sender: tokio::sync::watch::Sender<NonFinalizedState>,
started_sync_sender: tokio::sync::watch::Sender<bool>,
finalized_tip_updater: Option<JoinHandle<()>>,
}
async fn stop_finalized_tip_updater(
started_sync_sender: &tokio::sync::watch::Sender<bool>,
finalized_tip_updater: JoinHandle<()>,
) {
started_sync_sender.send_replace(true);
finalized_tip_updater.wait_for_panics().await;
}
fn block_height_is_finalized(finalized_tip_height: Option<Height>, block_height: Height) -> bool {
finalized_tip_height.is_some_and(|tip| block_height <= tip)
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum CommitOutcome {
Committed,
AlreadyFinalized,
}
async fn update_finalized_chain_tip(
db: ZakuraDb,
mut indexer_rpc_client: IndexerClient<tonic::transport::Channel>,
mut finalized_chain_tip_sender: ChainTipSender,
mut started_sync_receiver: tokio::sync::watch::Receiver<bool>,
) {
let mut chain_tip_change_stream = None;
loop {
if *started_sync_receiver.borrow() {
return;
}
let Some(ref mut chain_tip_change) = chain_tip_change_stream else {
chain_tip_change_stream = match tokio::time::timeout(
SUBSCRIBE_TIMEOUT,
indexer_rpc_client.chain_tip_change(Empty {}),
)
.await
{
Ok(Ok(response)) => Some(response.into_inner()),
Ok(Err(err)) => {
tracing::warn!(?err, "failed to subscribe to chain tip changes");
tokio::time::sleep(POLL_DELAY).await;
None
}
Err(_) => {
tracing::warn!("timed out subscribing to chain tip changes");
tokio::time::sleep(POLL_DELAY).await;
None
}
};
continue;
};
let message = tokio::select! {
biased;
_ = started_sync_receiver.changed() => return,
message = tokio::time::timeout(STREAM_MESSAGE_TIMEOUT, chain_tip_change.message()) => message,
};
match message {
Ok(Ok(Some(_block_hash_and_height))) => {}
Ok(Ok(None)) => {
tracing::warn!("chain_tip_change stream ended unexpectedly");
chain_tip_change_stream = None;
continue;
}
Ok(Err(err)) => {
tracing::warn!(?err, "error receiving chain tip change");
chain_tip_change_stream = None;
continue;
}
Err(_) => {
tracing::debug!("chain tip change stream timed out, re-subscribing");
chain_tip_change_stream = None;
continue;
}
}
if *started_sync_receiver.borrow() {
return;
}
if let Err(error) = db.spawn_try_catch_up_with_primary().await {
tracing::debug!(
?error,
"failed to catch up to the primary database while updating the finalized tip"
);
continue;
}
if let Some(tip_block) = finalized_chain_tip_block(&db).await {
if *started_sync_receiver.borrow() {
return;
}
finalized_chain_tip_sender.set_finalized_tip(tip_block);
}
}
}
async fn finalized_chain_tip_block(db: &ZakuraDb) -> Option<ChainTipBlock> {
let db = db.clone();
tokio::task::spawn_blocking(move || {
let (height, hash) = db.tip()?;
db.block(height.into())
.map(|block| CheckpointVerifiedBlock::with_hash(block, hash))
.map(ChainTipBlock::from)
})
.wait_for_panics()
.await
}
impl TrustedChainSync {
pub async fn spawn(
indexer_rpc_address: SocketAddr,
db: ZakuraDb,
non_finalized_state_sender: tokio::sync::watch::Sender<NonFinalizedState>,
) -> Result<(LatestChainTip, ChainTipChange, JoinHandle<()>), BoxError> {
let non_finalized_state = NonFinalizedState::new(&db.network());
let (chain_tip_sender, latest_chain_tip, chain_tip_change) =
ChainTipSender::new(None, &db.network());
let channel =
tonic::transport::Endpoint::from_shared(format!("http://{indexer_rpc_address}"))?
.keep_alive_while_idle(true)
.http2_keep_alive_interval(KEEPALIVE_INTERVAL)
.keep_alive_timeout(KEEPALIVE_TIMEOUT)
.connect()
.await?;
let indexer_rpc_client = IndexerClient::new(channel);
let finalized_chain_tip_sender = chain_tip_sender.finalized_sender();
let (started_sync_sender, started_sync_receiver) = tokio::sync::watch::channel(false);
let finalized_tip_updater_db = db.clone();
let finalized_tip_updater_client = indexer_rpc_client.clone();
let finalized_tip_updater = tokio::spawn(async move {
update_finalized_chain_tip(
finalized_tip_updater_db,
finalized_tip_updater_client,
finalized_chain_tip_sender,
started_sync_receiver,
)
.await
});
let mut syncer = Self {
indexer_rpc_client,
db,
non_finalized_state,
chain_tip_sender,
non_finalized_state_sender,
started_sync_sender,
finalized_tip_updater: Some(finalized_tip_updater),
};
let sync_task = tokio::spawn(async move {
syncer.sync().await;
});
Ok((latest_chain_tip, chain_tip_change, sync_task))
}
async fn take_over_finalized_tip_updates(&mut self) {
let Some(finalized_tip_updater) = self.finalized_tip_updater.take() else {
return;
};
stop_finalized_tip_updater(&self.started_sync_sender, finalized_tip_updater).await;
}
#[tracing::instrument(skip_all)]
async fn sync(&mut self) {
let mut non_finalized_blocks_listener = None;
let mut last_failed_commit_hash = None;
self.try_catch_up_with_primary().await;
if let Some(finalized_tip_block) = finalized_chain_tip_block(&self.db).await {
self.chain_tip_sender.set_finalized_tip(finalized_tip_block);
}
loop {
let Some(ref mut non_finalized_state_change) = non_finalized_blocks_listener else {
non_finalized_blocks_listener = match self
.subscribe_to_non_finalized_state_change()
.await
{
Ok(listener) => Some(listener),
Err(err) => {
tracing::warn!(?err, "failed to subscribe to non-finalized state changes");
tokio::time::sleep(POLL_DELAY).await;
None
}
};
continue;
};
let message = match tokio::time::timeout(
STREAM_MESSAGE_TIMEOUT,
non_finalized_state_change.message(),
)
.await
{
Ok(Ok(Some(block_and_hash))) => block_and_hash,
Ok(Ok(None)) => {
tracing::warn!("non-finalized state change stream ended unexpectedly");
non_finalized_blocks_listener = None;
continue;
}
Ok(Err(err)) => {
tracing::warn!(?err, "error receiving non-finalized state change");
non_finalized_blocks_listener = None;
continue;
}
Err(_) => {
tracing::debug!("non-finalized state change stream timed out, re-subscribing");
non_finalized_blocks_listener = None;
continue;
}
};
let Some((block, hash)) = message.decode() else {
tracing::warn!("received malformed non-finalized state change message");
non_finalized_blocks_listener = None;
continue;
};
self.take_over_finalized_tip_updates().await;
if self.non_finalized_state.any_chain_contains(&hash) {
tracing::debug!(
?hash,
"non-finalized state already contains block, skipping"
);
continue;
}
let block = SemanticallyVerifiedBlock::with_hash(Arc::new(block), hash);
match self.try_commit(block).await {
Ok(CommitOutcome::Committed) => {
last_failed_commit_hash = None;
}
Ok(CommitOutcome::AlreadyFinalized) => {
last_failed_commit_hash = None;
}
Err(error) => {
if last_failed_commit_hash != Some(hash) {
tracing::warn!(
?error,
?hash,
"failed to commit block to non-finalized state"
);
last_failed_commit_hash = Some(hash);
}
non_finalized_blocks_listener = None;
tokio::time::sleep(COMMIT_RETRY_DELAY).await;
}
};
}
}
async fn try_commit(
&mut self,
block: SemanticallyVerifiedBlock,
) -> Result<CommitOutcome, ValidateContextError> {
self.try_catch_up_with_primary().await;
if block_height_is_finalized(self.db.finalized_tip_height(), block.height) {
tracing::debug!(
height = ?block.height,
"skipping block finalized while the secondary caught up"
);
self.publish_current_state().await;
return Ok(CommitOutcome::AlreadyFinalized);
}
if self.non_finalized_state.best_chain().is_none()
&& self.db.finalized_tip_hash() != block.block.header.previous_block_hash
{
self.fill_finalized_gap(block.height).await;
}
if block_height_is_finalized(self.db.finalized_tip_height(), block.height) {
tracing::debug!(
height = ?block.height,
"skipping block finalized while bridging the finalized gap"
);
self.publish_current_state().await;
return Ok(CommitOutcome::AlreadyFinalized);
}
self.commit(block)?;
Ok(CommitOutcome::Committed)
}
fn commit(&mut self, block: SemanticallyVerifiedBlock) -> Result<(), ValidateContextError> {
if self.db.finalized_tip_hash() == block.block.header.previous_block_hash {
let _ = self.prune_finalized();
self.non_finalized_state.commit_new_chain(block, &self.db)?;
} else {
self.non_finalized_state.commit_block(block, &self.db)?;
let _ = self.prune_finalized();
}
self.update_channels();
Ok(())
}
async fn fill_finalized_gap(&mut self, target_height: Height) {
loop {
self.try_catch_up_with_primary().await;
let Some(highest) = self
.non_finalized_state
.best_tip()
.map(|(height, _hash)| height)
.max(self.db.finalized_tip_height())
else {
return;
};
let Ok(next_height) = highest.next() else {
return;
};
if next_height >= target_height {
return;
}
let (block, hash) = match self.get_block(next_height.into()).await {
Ok(block_and_hash) => block_and_hash,
Err(error) => {
tracing::warn!(
?error,
?next_height,
"failed to fetch a block while bridging the finalized gap; \
will retry on the next subscription"
);
return;
}
};
let block = SemanticallyVerifiedBlock::with_hash(Arc::new(block), hash);
if let Err(error) = self.commit(block) {
tracing::warn!(
?error,
?next_height,
"failed to commit a block while bridging the finalized gap; \
will retry on the next subscription"
);
return;
}
}
}
async fn get_block(
&self,
hash_or_height: HashOrHeight,
) -> Result<(Block, block::Hash), Status> {
let hash_or_height = match hash_or_height {
HashOrHeight::Hash(hash) => hash.bytes_in_display_order().to_vec(),
HashOrHeight::Height(height) => height.0.to_be_bytes().to_vec(),
};
let request = BlockRequest { hash_or_height };
let response = tokio::time::timeout(
GET_BLOCK_TIMEOUT,
self.indexer_rpc_client.clone().get_block(request),
)
.await
.map_err(|_| Status::deadline_exceeded("get_block request timed out"))??;
response
.into_inner()
.decode()
.ok_or_else(|| Status::internal("failed to decode block from get_block response"))
}
async fn subscribe_to_non_finalized_state_change(
&mut self,
) -> Result<Streaming<BlockAndHash>, Status> {
let request = NonFinalizedStateChangeRequest {
chain_tip_hashes: self
.non_finalized_state
.chain_iter()
.map(|c| c.non_finalized_tip_hash().bytes_in_display_order().to_vec())
.collect(),
};
tokio::time::timeout(
SUBSCRIBE_TIMEOUT,
self.indexer_rpc_client
.clone()
.non_finalized_state_change(request),
)
.await
.map_err(|_| {
Status::deadline_exceeded("non_finalized_state_change subscription timed out")
})?
.map(|a| a.into_inner())
}
async fn try_catch_up_with_primary(&mut self) {
let _ = self.db.spawn_try_catch_up_with_primary().await;
if self.prune_finalized() {
self.publish_current_state().await;
}
}
fn prune_finalized(&mut self) -> bool {
let finalized_tip_height = self.db.finalized_tip_height().unwrap_or(Height::MIN);
let mut pruned = false;
while self
.non_finalized_state
.root_height()
.is_some_and(|root_height| root_height <= finalized_tip_height)
{
tracing::trace!("finalizing block past the reorg limit");
self.non_finalized_state.finalize();
pruned = true;
}
pruned
}
async fn publish_current_state(&mut self) {
if self.non_finalized_state.best_chain().is_some() {
self.update_channels();
return;
}
let _ = self
.non_finalized_state_sender
.send(self.non_finalized_state.clone());
if let Some(finalized_tip_block) = finalized_chain_tip_block(&self.db).await {
self.chain_tip_sender
.finalized_sender()
.set_finalized_tip(finalized_tip_block);
}
}
fn update_channels(&mut self) {
let _ = self
.non_finalized_state_sender
.send(self.non_finalized_state.clone());
let best_chain = self.non_finalized_state.best_chain().expect("unexpected empty non-finalized state: must commit at least one block before updating channels");
let tip_block = best_chain
.tip_block()
.expect(
"unexpected empty chain: must commit at least one block before updating channels",
)
.clone();
self.chain_tip_sender
.set_best_non_finalized_tip(Some(tip_block.into()));
}
}
pub fn init_read_state_with_syncer(
config: zakura_state::Config,
network: &Network,
indexer_rpc_address: SocketAddr,
) -> tokio::task::JoinHandle<
Result<
(
ReadStateService,
LatestChainTip,
ChainTipChange,
tokio::task::JoinHandle<()>,
),
BoxError,
>,
> {
let network = network.clone();
tokio::spawn(async move {
if config.ephemeral {
return Err("standalone read state service cannot be used with ephemeral state".into());
}
let (read_state, db, non_finalized_state_sender) =
spawn_init_read_only(config, &network).await??;
let (latest_chain_tip, chain_tip_change, sync_task) =
TrustedChainSync::spawn(indexer_rpc_address, db, non_finalized_state_sender).await?;
Ok((read_state, latest_chain_tip, chain_tip_change, sync_task))
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn finalized_height_check_includes_the_tip() {
let finalized_tip = Height(10);
assert!(!block_height_is_finalized(None, Height(9)));
assert!(block_height_is_finalized(Some(finalized_tip), Height(9)));
assert!(block_height_is_finalized(Some(finalized_tip), Height(10)));
assert!(!block_height_is_finalized(Some(finalized_tip), Height(11)));
}
#[tokio::test]
async fn finalized_tip_handoff_waits_for_updater() {
let (started_sender, mut started_receiver) = tokio::sync::watch::channel(false);
let (observed_sender, observed_receiver) = tokio::sync::oneshot::channel();
let (release_sender, release_receiver) = tokio::sync::oneshot::channel();
let updater = tokio::spawn(async move {
started_receiver
.changed()
.await
.expect("started sender remains open during handoff");
assert!(*started_receiver.borrow());
observed_sender
.send(())
.expect("handoff test keeps the observation receiver open");
release_receiver
.await
.expect("handoff test releases the updater");
});
let handoff_sender = started_sender.clone();
let handoff = tokio::spawn(async move {
stop_finalized_tip_updater(&handoff_sender, updater).await;
});
tokio::time::timeout(Duration::from_secs(1), observed_receiver)
.await
.expect("updater should observe the handoff signal")
.expect("updater should send the handoff observation");
assert!(
!handoff.is_finished(),
"handoff must wait for the updater to finish"
);
release_sender
.send(())
.expect("updater remains alive until it is released");
tokio::time::timeout(Duration::from_secs(1), handoff)
.await
.expect("handoff should finish after the updater exits")
.expect("handoff task should not panic");
assert!(*started_sender.borrow());
}
}