linera_storage_service/
client.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    mem,
6    sync::{
7        atomic::{AtomicBool, Ordering},
8        Arc,
9    },
10};
11
12use async_lock::{Semaphore, SemaphoreGuard};
13use futures::future::join_all;
14use linera_base::{ensure, util::future::FutureSyncExt as _};
15#[cfg(with_metrics)]
16use linera_views::metering::MeteredDatabase;
17#[cfg(with_testing)]
18use linera_views::store::TestKeyValueDatabase;
19use linera_views::{
20    batch::{Batch, WriteOperation},
21    lru_caching::LruCachingDatabase,
22    store::{KeyValueDatabase, ReadableKeyValueStore, WithError, WritableKeyValueStore},
23};
24use serde::de::DeserializeOwned;
25use tonic::transport::{Channel, Endpoint};
26
27#[cfg(with_testing)]
28use crate::common::storage_service_test_endpoint;
29use crate::{
30    common::{
31        KeyPrefix, StorageServiceStoreError, StorageServiceStoreInternalConfig, MAX_PAYLOAD_SIZE,
32    },
33    key_value_store::{
34        statement::Operation, storage_service_client::StorageServiceClient, KeyValue,
35        KeyValueAppend, ReplyContainsKey, ReplyContainsKeys, ReplyExistsNamespace,
36        ReplyFindKeyValuesByPrefix, ReplyFindKeysByPrefix, ReplyListAll, ReplyListRootKeys,
37        ReplyReadMultiValues, ReplyReadValue, ReplySpecificChunk, RequestContainsKey,
38        RequestContainsKeys, RequestCreateNamespace, RequestDeleteNamespace,
39        RequestExistsNamespace, RequestFindKeyValuesByPrefix, RequestFindKeysByPrefix,
40        RequestListRootKeys, RequestReadMultiValues, RequestReadValue, RequestSpecificChunk,
41        RequestWriteBatchExtended, Statement,
42    },
43};
44
45// The maximum key size is set to 1M rather arbitrarily.
46const MAX_KEY_SIZE: usize = 1000000;
47
48// The shared store client.
49// * Interior mutability is required for the client because
50// accessing requires mutability while the KeyValueStore
51// does not allow it.
52// * The semaphore works as other stores.
53//
54// The encoding of namespaces is done by taking their
55// serialization. This works because the set of serialization
56// of strings is prefix free.
57// The data is stored in the following way.
58// * A `key` in a `namespace` is stored as
59//   [`KeyPrefix::Key`] + namespace + key
60// * An additional key with empty value is stored at
61//   [`KeyPrefix::Namespace`] + namespace
62//   is stored to indicate the existence of a namespace.
63// * A key with empty value is stored at
64//   [`KeyPrefix::RootKey`] + namespace + root_key
65//   to indicate the existence of a root key.
66/// A database handle connected to a remote storage service.
67#[derive(Clone)]
68pub struct StorageServiceDatabaseInternal {
69    channel: Channel,
70    semaphore: Option<Arc<Semaphore>>,
71    namespace: Vec<u8>,
72}
73
74/// A store within a namespace and root key, backed by a remote storage service.
75#[derive(Clone)]
76pub struct StorageServiceStoreInternal {
77    channel: Channel,
78    semaphore: Option<Arc<Semaphore>>,
79    prefix_len: usize,
80    start_key: Vec<u8>,
81    root_key_written: Arc<AtomicBool>,
82}
83
84impl WithError for StorageServiceDatabaseInternal {
85    type Error = StorageServiceStoreError;
86}
87
88impl WithError for StorageServiceStoreInternal {
89    type Error = StorageServiceStoreError;
90}
91
92impl ReadableKeyValueStore for StorageServiceStoreInternal {
93    const MAX_KEY_SIZE: usize = MAX_KEY_SIZE;
94
95    fn root_key(&self) -> Result<Vec<u8>, StorageServiceStoreError> {
96        let root_key = bcs::from_bytes(&self.start_key[self.prefix_len..])?;
97        Ok(root_key)
98    }
99
100    async fn read_value_bytes(
101        &self,
102        key: &[u8],
103    ) -> Result<Option<Vec<u8>>, StorageServiceStoreError> {
104        ensure!(
105            key.len() <= MAX_KEY_SIZE,
106            StorageServiceStoreError::KeyTooLong
107        );
108        let mut full_key = self.start_key.clone();
109        full_key.extend(key);
110        let query = RequestReadValue { key: full_key };
111        let request = tonic::Request::new(query);
112        let channel = self.channel.clone();
113        let mut client = StorageServiceClient::new(channel);
114        let _guard = self.acquire().await;
115        let response = client.process_read_value(request).make_sync().await?;
116        let response = response.into_inner();
117        let ReplyReadValue {
118            value,
119            message_index,
120            num_chunks,
121        } = response;
122        if num_chunks == 0 {
123            Ok(value)
124        } else {
125            self.read_entries(message_index, num_chunks).await
126        }
127    }
128
129    async fn contains_key(&self, key: &[u8]) -> Result<bool, StorageServiceStoreError> {
130        ensure!(
131            key.len() <= MAX_KEY_SIZE,
132            StorageServiceStoreError::KeyTooLong
133        );
134        let mut full_key = self.start_key.clone();
135        full_key.extend(key);
136        let query = RequestContainsKey { key: full_key };
137        let request = tonic::Request::new(query);
138        let channel = self.channel.clone();
139        let mut client = StorageServiceClient::new(channel);
140        let _guard = self.acquire().await;
141        let response = client.process_contains_key(request).make_sync().await?;
142        let response = response.into_inner();
143        let ReplyContainsKey { test } = response;
144        Ok(test)
145    }
146
147    async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, StorageServiceStoreError> {
148        let mut full_keys = Vec::new();
149        for key in keys {
150            ensure!(
151                key.len() <= MAX_KEY_SIZE,
152                StorageServiceStoreError::KeyTooLong
153            );
154            let mut full_key = self.start_key.clone();
155            full_key.extend(key);
156            full_keys.push(full_key);
157        }
158        let query = RequestContainsKeys { keys: full_keys };
159        let request = tonic::Request::new(query);
160        let channel = self.channel.clone();
161        let mut client = StorageServiceClient::new(channel);
162        let _guard = self.acquire().await;
163        let response = client.process_contains_keys(request).make_sync().await?;
164        let response = response.into_inner();
165        let ReplyContainsKeys { tests } = response;
166        Ok(tests)
167    }
168
169    async fn read_multi_values_bytes(
170        &self,
171        keys: &[Vec<u8>],
172    ) -> Result<Vec<Option<Vec<u8>>>, StorageServiceStoreError> {
173        let mut full_keys = Vec::new();
174        for key in keys {
175            ensure!(
176                key.len() <= MAX_KEY_SIZE,
177                StorageServiceStoreError::KeyTooLong
178            );
179            let mut full_key = self.start_key.clone();
180            full_key.extend(key);
181            full_keys.push(full_key);
182        }
183        let query = RequestReadMultiValues { keys: full_keys };
184        let request = tonic::Request::new(query);
185        let channel = self.channel.clone();
186        let mut client = StorageServiceClient::new(channel);
187        let _guard = self.acquire().await;
188        let response = client
189            .process_read_multi_values(request)
190            .make_sync()
191            .await?;
192        let response = response.into_inner();
193        let ReplyReadMultiValues {
194            values,
195            message_index,
196            num_chunks,
197        } = response;
198        if num_chunks == 0 {
199            let values = values.into_iter().map(|x| x.value).collect::<Vec<_>>();
200            Ok(values)
201        } else {
202            self.read_entries(message_index, num_chunks).await
203        }
204    }
205
206    async fn find_keys_by_prefix(
207        &self,
208        key_prefix: &[u8],
209    ) -> Result<Vec<Vec<u8>>, StorageServiceStoreError> {
210        ensure!(
211            key_prefix.len() <= MAX_KEY_SIZE,
212            StorageServiceStoreError::KeyTooLong
213        );
214        let mut full_key_prefix = self.start_key.clone();
215        full_key_prefix.extend(key_prefix);
216        let query = RequestFindKeysByPrefix {
217            key_prefix: full_key_prefix,
218        };
219        let request = tonic::Request::new(query);
220        let channel = self.channel.clone();
221        let mut client = StorageServiceClient::new(channel);
222        let _guard = self.acquire().await;
223        let response = client
224            .process_find_keys_by_prefix(request)
225            .make_sync()
226            .await?;
227        let response = response.into_inner();
228        let ReplyFindKeysByPrefix {
229            keys,
230            message_index,
231            num_chunks,
232        } = response;
233        if num_chunks == 0 {
234            Ok(keys)
235        } else {
236            self.read_entries(message_index, num_chunks).await
237        }
238    }
239
240    async fn find_key_values_by_prefix(
241        &self,
242        key_prefix: &[u8],
243    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageServiceStoreError> {
244        ensure!(
245            key_prefix.len() <= MAX_KEY_SIZE,
246            StorageServiceStoreError::KeyTooLong
247        );
248        let mut full_key_prefix = self.start_key.clone();
249        full_key_prefix.extend(key_prefix);
250        let query = RequestFindKeyValuesByPrefix {
251            key_prefix: full_key_prefix,
252        };
253        let request = tonic::Request::new(query);
254        let channel = self.channel.clone();
255        let mut client = StorageServiceClient::new(channel);
256        let _guard = self.acquire().await;
257        let response = client
258            .process_find_key_values_by_prefix(request)
259            .make_sync()
260            .await?;
261        let response = response.into_inner();
262        let ReplyFindKeyValuesByPrefix {
263            key_values,
264            message_index,
265            num_chunks,
266        } = response;
267        if num_chunks == 0 {
268            let key_values = key_values
269                .into_iter()
270                .map(|x| (x.key, x.value))
271                .collect::<Vec<_>>();
272            Ok(key_values)
273        } else {
274            self.read_entries(message_index, num_chunks).await
275        }
276    }
277}
278
279impl WritableKeyValueStore for StorageServiceStoreInternal {
280    const MAX_VALUE_SIZE: usize = usize::MAX;
281
282    async fn write_batch(&self, batch: Batch) -> Result<(), StorageServiceStoreError> {
283        if batch.operations.is_empty() {
284            return Ok(());
285        }
286        let mut statements = Vec::new();
287        let mut chunk_size = 0;
288
289        if !self.root_key_written.fetch_or(true, Ordering::SeqCst) {
290            let mut full_key = self.start_key.clone();
291            full_key[0] = KeyPrefix::RootKey as u8;
292            let operation = Operation::Put(KeyValue {
293                key: full_key,
294                value: vec![],
295            });
296            let statement = Statement {
297                operation: Some(operation),
298            };
299            statements.push(statement);
300            chunk_size += self.start_key.len();
301        }
302
303        let bcs_root_key_len = self.start_key.len() - self.prefix_len;
304        for operation in batch.operations {
305            let (key_len, value_len) = match &operation {
306                WriteOperation::Delete { key } => (key.len(), 0),
307                WriteOperation::Put { key, value } => (key.len(), value.len()),
308                WriteOperation::DeletePrefix { key_prefix } => (key_prefix.len(), 0),
309            };
310            let operation_size = key_len + value_len + bcs_root_key_len;
311            ensure!(
312                key_len <= MAX_KEY_SIZE,
313                StorageServiceStoreError::KeyTooLong
314            );
315            if operation_size + chunk_size < MAX_PAYLOAD_SIZE {
316                let statement = self.get_statement(operation);
317                statements.push(statement);
318                chunk_size += operation_size;
319            } else {
320                self.submit_statements(mem::take(&mut statements)).await?;
321                chunk_size = 0;
322                if operation_size > MAX_PAYLOAD_SIZE {
323                    // One single operation is especially big. So split it in chunks.
324                    let WriteOperation::Put { key, value } = operation else {
325                        // Only the put can go over the limit
326                        unreachable!();
327                    };
328                    let mut full_key = self.start_key.clone();
329                    full_key.extend(key);
330                    let value_chunks = value
331                        .chunks(MAX_PAYLOAD_SIZE)
332                        .map(|x| x.to_vec())
333                        .collect::<Vec<_>>();
334                    let num_chunks = value_chunks.len();
335                    for (i_chunk, value) in value_chunks.into_iter().enumerate() {
336                        let last = i_chunk + 1 == num_chunks;
337                        let operation = Operation::Append(KeyValueAppend {
338                            key: full_key.clone(),
339                            value,
340                            last,
341                        });
342                        statements = vec![Statement {
343                            operation: Some(operation),
344                        }];
345                        self.submit_statements(mem::take(&mut statements)).await?;
346                    }
347                } else {
348                    // The operation is small enough, it is just that we have many so we need to split.
349                    let statement = self.get_statement(operation);
350                    statements.push(statement);
351                    chunk_size = operation_size;
352                }
353            }
354        }
355        self.submit_statements(mem::take(&mut statements)).await
356    }
357
358    async fn clear_journal(&self) -> Result<(), StorageServiceStoreError> {
359        Ok(())
360    }
361}
362
363impl StorageServiceStoreInternal {
364    /// Obtains the semaphore lock on the database if needed.
365    async fn acquire(&self) -> Option<SemaphoreGuard<'_>> {
366        match &self.semaphore {
367            None => None,
368            Some(count) => Some(count.acquire().await),
369        }
370    }
371
372    async fn submit_statements(
373        &self,
374        statements: Vec<Statement>,
375    ) -> Result<(), StorageServiceStoreError> {
376        if !statements.is_empty() {
377            let query = RequestWriteBatchExtended { statements };
378            let request = tonic::Request::new(query);
379            let channel = self.channel.clone();
380            let mut client = StorageServiceClient::new(channel);
381            let _guard = self.acquire().await;
382            let _response = client
383                .process_write_batch_extended(request)
384                .make_sync()
385                .await?;
386        }
387        Ok(())
388    }
389
390    fn get_statement(&self, operation: WriteOperation) -> Statement {
391        let operation = match operation {
392            WriteOperation::Delete { key } => {
393                let mut full_key = self.start_key.clone();
394                full_key.extend(key);
395                Operation::Delete(full_key)
396            }
397            WriteOperation::Put { key, value } => {
398                let mut full_key = self.start_key.clone();
399                full_key.extend(key);
400                Operation::Put(KeyValue {
401                    key: full_key,
402                    value,
403                })
404            }
405            WriteOperation::DeletePrefix { key_prefix } => {
406                let mut full_key_prefix = self.start_key.clone();
407                full_key_prefix.extend(key_prefix);
408                Operation::DeletePrefix(full_key_prefix)
409            }
410        };
411        Statement {
412            operation: Some(operation),
413        }
414    }
415
416    async fn read_single_entry(
417        &self,
418        message_index: i64,
419        index: i32,
420    ) -> Result<Vec<u8>, StorageServiceStoreError> {
421        let channel = self.channel.clone();
422        let query = RequestSpecificChunk {
423            message_index,
424            index,
425        };
426        let request = tonic::Request::new(query);
427        let mut client = StorageServiceClient::new(channel);
428        let response = client.process_specific_chunk(request).make_sync().await?;
429        let response = response.into_inner();
430        let ReplySpecificChunk { chunk } = response;
431        Ok(chunk)
432    }
433
434    async fn read_entries<S: DeserializeOwned>(
435        &self,
436        message_index: i64,
437        num_chunks: i32,
438    ) -> Result<S, StorageServiceStoreError> {
439        let mut handles = Vec::new();
440        for index in 0..num_chunks {
441            let handle = self.read_single_entry(message_index, index);
442            handles.push(handle);
443        }
444        let mut value = Vec::new();
445        for chunk in join_all(handles).await {
446            let chunk = chunk?;
447            value.extend(chunk);
448        }
449        Ok(bcs::from_bytes(&value)?)
450    }
451}
452
453impl KeyValueDatabase for StorageServiceDatabaseInternal {
454    type Config = StorageServiceStoreInternalConfig;
455    type Store = StorageServiceStoreInternal;
456
457    fn get_name() -> String {
458        "service store".to_string()
459    }
460
461    async fn connect(
462        config: &Self::Config,
463        namespace: &str,
464    ) -> Result<Self, StorageServiceStoreError> {
465        let semaphore = config
466            .max_concurrent_queries
467            .map(|n| Arc::new(Semaphore::new(n)));
468        let namespace = bcs::to_bytes(namespace)?;
469        let endpoint = config.http_address();
470        let endpoint = Endpoint::from_shared(endpoint)?;
471        let channel = endpoint.connect_lazy();
472        Ok(Self {
473            channel,
474            semaphore,
475            namespace,
476        })
477    }
478
479    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, StorageServiceStoreError> {
480        let channel = self.channel.clone();
481        let semaphore = self.semaphore.clone();
482        let mut start_key = vec![KeyPrefix::Key as u8];
483        start_key.extend(&self.namespace);
484        start_key.extend(bcs::to_bytes(root_key)?);
485        let prefix_len = self.namespace.len() + 1;
486        Ok(StorageServiceStoreInternal {
487            channel,
488            semaphore,
489            prefix_len,
490            start_key,
491            root_key_written: Arc::new(AtomicBool::new(false)),
492        })
493    }
494
495    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
496        self.open_shared(root_key)
497    }
498
499    async fn list_all(config: &Self::Config) -> Result<Vec<String>, StorageServiceStoreError> {
500        let endpoint = config.http_address();
501        let endpoint = Endpoint::from_shared(endpoint)?;
502        let mut client = StorageServiceClient::connect(endpoint).make_sync().await?;
503        let response = client.process_list_all(()).make_sync().await?;
504        let response = response.into_inner();
505        let ReplyListAll { namespaces } = response;
506        let namespaces = namespaces
507            .into_iter()
508            .map(|x| bcs::from_bytes(&x))
509            .collect::<Result<_, _>>()?;
510        Ok(namespaces)
511    }
512
513    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, StorageServiceStoreError> {
514        let query = RequestListRootKeys {
515            namespace: self.namespace.clone(),
516        };
517        let request = tonic::Request::new(query);
518        let mut client = StorageServiceClient::new(self.channel.clone());
519        let response = client.process_list_root_keys(request).make_sync().await?;
520        let response = response.into_inner();
521        let ReplyListRootKeys { root_keys } = response;
522        Ok(root_keys)
523    }
524
525    async fn delete_all(config: &Self::Config) -> Result<(), StorageServiceStoreError> {
526        let endpoint = config.http_address();
527        let endpoint = Endpoint::from_shared(endpoint)?;
528        let mut client = StorageServiceClient::connect(endpoint).make_sync().await?;
529        let _response = client.process_delete_all(()).make_sync().await?;
530        Ok(())
531    }
532
533    async fn exists(
534        config: &Self::Config,
535        namespace: &str,
536    ) -> Result<bool, StorageServiceStoreError> {
537        let namespace = bcs::to_bytes(namespace)?;
538        let query = RequestExistsNamespace { namespace };
539        let request = tonic::Request::new(query);
540        let endpoint = config.http_address();
541        let endpoint = Endpoint::from_shared(endpoint)?;
542        let mut client = StorageServiceClient::connect(endpoint).make_sync().await?;
543        let response = client.process_exists_namespace(request).make_sync().await?;
544        let response = response.into_inner();
545        let ReplyExistsNamespace { exists } = response;
546        Ok(exists)
547    }
548
549    async fn create(
550        config: &Self::Config,
551        namespace: &str,
552    ) -> Result<(), StorageServiceStoreError> {
553        if StorageServiceDatabaseInternal::exists(config, namespace).await? {
554            return Err(StorageServiceStoreError::StoreAlreadyExists);
555        }
556        let namespace = bcs::to_bytes(namespace)?;
557        let query = RequestCreateNamespace { namespace };
558        let request = tonic::Request::new(query);
559        let endpoint = config.http_address();
560        let endpoint = Endpoint::from_shared(endpoint)?;
561        let mut client = StorageServiceClient::connect(endpoint).make_sync().await?;
562        let _response = client.process_create_namespace(request).make_sync().await?;
563        Ok(())
564    }
565
566    async fn delete(
567        config: &Self::Config,
568        namespace: &str,
569    ) -> Result<(), StorageServiceStoreError> {
570        let namespace = bcs::to_bytes(namespace)?;
571        let query = RequestDeleteNamespace { namespace };
572        let request = tonic::Request::new(query);
573        let endpoint = config.http_address();
574        let endpoint = Endpoint::from_shared(endpoint)?;
575        let mut client = StorageServiceClient::connect(endpoint).make_sync().await?;
576        let _response = client.process_delete_namespace(request).make_sync().await?;
577        Ok(())
578    }
579}
580
581#[cfg(with_testing)]
582impl TestKeyValueDatabase for StorageServiceDatabaseInternal {
583    async fn new_test_config() -> Result<StorageServiceStoreInternalConfig, StorageServiceStoreError>
584    {
585        let endpoint = storage_service_test_endpoint()?;
586        service_config_from_endpoint(&endpoint)
587    }
588}
589
590/// Creates a `StorageServiceStoreConfig` from an endpoint.
591pub fn service_config_from_endpoint(
592    endpoint: &str,
593) -> Result<StorageServiceStoreInternalConfig, StorageServiceStoreError> {
594    Ok(StorageServiceStoreInternalConfig {
595        endpoint: endpoint.to_string(),
596        max_concurrent_queries: None,
597    })
598}
599
600/// Checks that endpoint is truly absent.
601pub async fn storage_service_check_absence(
602    endpoint: &str,
603) -> Result<bool, StorageServiceStoreError> {
604    let endpoint = Endpoint::from_shared(endpoint.to_string())?;
605    let result = StorageServiceClient::connect(endpoint).await;
606    Ok(result.is_err())
607}
608
609/// Checks whether an endpoint is valid or not.
610pub async fn storage_service_check_validity(
611    endpoint: &str,
612) -> Result<(), StorageServiceStoreError> {
613    let config = service_config_from_endpoint(endpoint).unwrap();
614    let namespace = "namespace";
615    let database = StorageServiceDatabaseInternal::connect(&config, namespace).await?;
616    let store = database.open_shared(&[])?;
617    let _value = store.read_value_bytes(&[42]).await?;
618    Ok(())
619}
620
621/// The service database client with metrics
622#[cfg(with_metrics)]
623pub type StorageServiceDatabase =
624    MeteredDatabase<LruCachingDatabase<MeteredDatabase<StorageServiceDatabaseInternal>>>;
625
626/// The service database client without metrics
627#[cfg(not(with_metrics))]
628pub type StorageServiceDatabase = LruCachingDatabase<StorageServiceDatabaseInternal>;
linera_storage_service/client.rs

linera_storage_service/
client.rs
