iota_data_ingestion/workers/
kv_store.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    borrow::Borrow,
7    collections::{HashMap, HashSet, VecDeque},
8    iter::repeat,
9    sync::Arc,
10    time::Duration,
11};
12
13use anyhow::Result;
14use async_trait::async_trait;
15use aws_config::{BehaviorVersion, timeout::TimeoutConfig};
16use aws_sdk_dynamodb::{
17    Client,
18    config::{Credentials, Region},
19    primitives::Blob,
20    types::{AttributeValue, PutRequest, WriteRequest},
21};
22use backoff::{ExponentialBackoff, backoff::Backoff};
23use bytes::Bytes;
24use iota_config::object_storage_config::ObjectStoreConfig;
25use iota_data_ingestion_core::Worker;
26use iota_storage::http_key_value_store::{ItemType, TaggedKey};
27use iota_types::{full_checkpoint_content::CheckpointData, storage::ObjectKey};
28use object_store::{DynObjectStore, ObjectStoreExt, path::Path};
29use serde::{Deserialize, Serialize};
30use tracing::{error, info, warn};
31
32use crate::common;
33
34#[derive(Serialize, Deserialize, Clone, Debug)]
35#[serde(rename_all = "kebab-case")]
36pub struct KVStoreTaskConfig {
37    pub object_store_config: ObjectStoreConfig,
38    pub dynamodb_config: DynamoDBConfig,
39}
40
41#[derive(Serialize, Deserialize, Clone, Debug)]
42#[serde(rename_all = "kebab-case")]
43pub struct DynamoDBConfig {
44    pub aws_access_key_id: String,
45    pub aws_secret_access_key: String,
46    pub aws_region: String,
47    pub aws_endpoint: Option<String>,
48    pub table_name: String,
49}
50
51#[derive(Clone)]
52pub struct KVStoreWorker {
53    dynamo_client: Client,
54    remote_store: Arc<DynObjectStore>,
55    table_name: String,
56}
57
58impl KVStoreWorker {
59    pub async fn new(config: KVStoreTaskConfig) -> anyhow::Result<Self> {
60        let credentials = Credentials::new(
61            &config.dynamodb_config.aws_access_key_id,
62            &config.dynamodb_config.aws_secret_access_key,
63            None,
64            None,
65            "dynamodb",
66        );
67        let timeout_config = TimeoutConfig::builder()
68            .operation_timeout(Duration::from_secs(3))
69            .operation_attempt_timeout(Duration::from_secs(10))
70            .connect_timeout(Duration::from_secs(3))
71            .build();
72
73        let http_client = aws_smithy_http_client::Builder::new()
74            .tls_provider(aws_smithy_http_client::tls::Provider::Rustls(
75                aws_smithy_http_client::tls::rustls_provider::CryptoMode::Ring,
76            ))
77            .build_https();
78        let mut aws_config_loader = aws_config::defaults(BehaviorVersion::latest())
79            .http_client(http_client)
80            .credentials_provider(credentials)
81            .region(Region::new(config.dynamodb_config.aws_region))
82            .timeout_config(timeout_config);
83
84        if let Some(url) = config.dynamodb_config.aws_endpoint {
85            aws_config_loader = aws_config_loader.endpoint_url(url);
86        }
87        let aws_config = aws_config_loader.load().await;
88        Ok(Self {
89            dynamo_client: Client::new(&aws_config),
90            remote_store: config.object_store_config.make()?,
91            table_name: config.dynamodb_config.table_name,
92        })
93    }
94
95    async fn multi_set<V: Serialize>(
96        &self,
97        item_type: ItemType,
98        values: impl IntoIterator<Item = (Vec<u8>, V)> + std::marker::Send,
99    ) -> anyhow::Result<()> {
100        let mut items = vec![];
101        let mut seen = HashSet::new();
102        for (digest, value) in values {
103            if seen.contains(&digest) {
104                continue;
105            }
106            seen.insert(digest.clone());
107            let item = WriteRequest::builder()
108                .set_put_request(Some(
109                    PutRequest::builder()
110                        .item("digest", AttributeValue::B(Blob::new(digest)))
111                        .item("type", AttributeValue::S(item_type.to_string()))
112                        .item(
113                            "bcs",
114                            AttributeValue::B(Blob::new(bcs::to_bytes(value.borrow())?)),
115                        )
116                        .build()?,
117                ))
118                .build();
119            items.push(item);
120        }
121        if items.is_empty() {
122            return Ok(());
123        }
124        let mut backoff = ExponentialBackoff::default();
125        let mut queue: VecDeque<Vec<_>> = items.chunks(25).map(|ck| ck.to_vec()).collect();
126        while let Some(chunk) = queue.pop_front() {
127            let response = self
128                .dynamo_client
129                .batch_write_item()
130                .set_request_items(Some(HashMap::from([(
131                    self.table_name.clone(),
132                    chunk.to_vec(),
133                )])))
134                .send()
135                .await
136                .inspect_err(|sdk_err| {
137                    error!(
138                        "{:?}",
139                        sdk_err.as_service_error().map(|e| e.meta().to_string())
140                    )
141                })?;
142            if let Some(response) = response.unprocessed_items {
143                if let Some(unprocessed) = response.into_iter().next() {
144                    if !unprocessed.1.is_empty() {
145                        if queue.is_empty() {
146                            if let Some(duration) = backoff.next_backoff() {
147                                tokio::time::sleep(duration).await;
148                            }
149                        }
150                        queue.push_back(unprocessed.1);
151                    }
152                }
153            }
154        }
155        Ok(())
156    }
157
158    /// Uploads checkpoint contents to storage, with automatic fallback from
159    /// DynamoDB to S3.
160    ///
161    /// This function attempts to store checkpoint contents in DynamoDB first.
162    /// If that fails (typically due to size constraints), it automatically
163    /// falls back to uploading the contents to S3 instead.
164    async fn upload_checkpoint_contents<V: Serialize + std::marker::Send>(
165        &self,
166        key: Vec<u8>,
167        value: V,
168    ) -> anyhow::Result<()> {
169        let bcs_bytes = bcs::to_bytes(value.borrow())?;
170
171        let attributes = HashMap::from([
172            (
173                "digest".to_string(),
174                AttributeValue::B(Blob::new(key.clone())),
175            ),
176            (
177                "type".to_string(),
178                AttributeValue::S(ItemType::CheckpointContents.to_string()),
179            ),
180            (
181                "bcs".to_string(),
182                AttributeValue::B(Blob::new(bcs_bytes.clone())),
183            ),
184        ]);
185
186        let res = self
187            .dynamo_client
188            .put_item()
189            .table_name(&self.table_name)
190            .set_item(Some(attributes))
191            .send()
192            .await
193            .inspect_err(|err| warn!("dynamodb error: {err}"));
194
195        if res.is_err() {
196            info!("attempt to store checkpoint contents on S3");
197            let location = Path::from(base64_url::encode(&key));
198            self.remote_store
199                .put(&location, Bytes::from(bcs_bytes).into())
200                .await?;
201        }
202
203        Ok(())
204    }
205}
206
207#[async_trait]
208impl Worker for KVStoreWorker {
209    type Message = ();
210    type Error = anyhow::Error;
211
212    async fn process_checkpoint(
213        &self,
214        checkpoint: Arc<CheckpointData>,
215    ) -> Result<Self::Message, Self::Error> {
216        let mut transactions = vec![];
217        let mut effects = vec![];
218        let mut events = vec![];
219        let mut objects = vec![];
220        let mut transactions_to_checkpoint = vec![];
221        let checkpoint_number = checkpoint.checkpoint_summary.sequence_number;
222
223        for transaction in &checkpoint.transactions {
224            let transaction_digest = transaction.transaction.digest().into_bytes().to_vec();
225            effects.push((transaction_digest.clone(), transaction.effects.clone()));
226            transactions_to_checkpoint.push((transaction_digest.clone(), checkpoint_number));
227            transactions.push((transaction_digest.clone(), transaction.transaction.clone()));
228
229            if let Some(tx_events) = &transaction.events {
230                events.push((tx_events.digest().into_bytes().to_vec(), tx_events));
231            }
232            for object in &transaction.output_objects {
233                let object_key = ObjectKey(object.id(), object.version());
234                objects.push((bcs::to_bytes(&object_key)?, object));
235            }
236        }
237        self.multi_set(ItemType::Transaction, transactions).await?;
238        self.multi_set(ItemType::TransactionEffects, effects)
239            .await?;
240        self.multi_set(ItemType::EventTransactionDigest, events)
241            .await?;
242        self.multi_set(ItemType::Object, objects).await?;
243        self.multi_set(
244            ItemType::TransactionToCheckpoint,
245            transactions_to_checkpoint,
246        )
247        .await?;
248
249        let serialized_checkpoint_number =
250            bcs::to_bytes(&TaggedKey::CheckpointSequenceNumber(checkpoint_number))?;
251        let checkpoint_summary = &checkpoint.checkpoint_summary;
252
253        self.upload_checkpoint_contents(
254            serialized_checkpoint_number.clone(),
255            checkpoint.checkpoint_contents.clone(),
256        )
257        .await?;
258
259        self.multi_set(
260            ItemType::CheckpointSummary,
261            [
262                serialized_checkpoint_number,
263                checkpoint_summary.digest().into_bytes().to_vec(),
264            ]
265            .into_iter()
266            .zip(repeat(checkpoint_summary)),
267        )
268        .await?;
269
270        if checkpoint_number.is_multiple_of(common::PROGRESS_LOG_INTERVAL) {
271            info!("stored checkpoint {checkpoint_number} in the KV store");
272        }
273
274        Ok(())
275    }
276}
iota_data_ingestion/workers/kv_store.rs

iota_data_ingestion/workers/
kv_store.rs
