liboxia 0.0.1

Liboxia is a Rust library designed for both native Rust applications and for integration with other languages via its C Foreign Function Interface (FFI). It serves as a client SDK for Oxia, a distributed key-value store, enabling robust, asynchronous data operations.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use crate::errors::OxiaError;
use crate::errors::OxiaError::{InternalRetryable, UnexpectedStatus};
use crate::oxia::{WriteRequest, WriteResponse};
use crate::provider_manager::ProviderManager;
use crate::shard_manager::ShardManager;
use crate::write_stream::WriteStream;
use dashmap::DashMap;
use std::sync::Arc;
use tokio::sync::{mpsc, OnceCell};
use tokio::task::JoinSet;
use tonic::codegen::tokio_stream::wrappers::UnboundedReceiverStream;
use tonic::Request;

const WRITE_STREAM_HEADER_NAMESPACE: &str = "namespace";
const WRITE_STREAM_HEADER_SHARD_ID: &str = "shard-id";

pub struct WriteStreamManager {
    namespace: String,
    shard_manager: Arc<ShardManager>,
    provider_manager: Arc<ProviderManager>,

    streams: DashMap<i64, OnceCell<WriteStream>>,
}

impl WriteStreamManager {
    pub async fn write(&self, request: WriteRequest) -> Result<WriteResponse, OxiaError> {
        // todo: make request Rc
        loop {
            let result = self.write0(request.clone()).await;
            if result.is_ok() {
                return result;
            }
            if result.is_err() {
                let error = result.unwrap_err();
                match error {
                    InternalRetryable() => continue,
                    _ => return Err(error),
                }
            }
        }
    }

    async fn write0(&self, request: WriteRequest) -> Result<WriteResponse, OxiaError> {
        let shard_id = request.shard.unwrap();
        let option = self.shard_manager.get_leader(shard_id);
        if option.is_none() {
            return Err(OxiaError::ShardLeaderNotFound(shard_id));
        }
        let defer_init = || async {
            let client = self
                .provider_manager
                .get_provider(option.unwrap().service_address)
                .await?;
            let (tx, rx) = mpsc::unbounded_channel();
            let mut client_guard = client.lock().await;
            let mut write_stream_request = Request::new(UnboundedReceiverStream::new(rx));
            let write_stream_request_metadata = write_stream_request.metadata_mut();
            write_stream_request_metadata.insert(
                WRITE_STREAM_HEADER_NAMESPACE,
                self.namespace.parse().unwrap(),
            );
            write_stream_request_metadata.insert(
                WRITE_STREAM_HEADER_SHARD_ID,
                shard_id.to_string().parse().unwrap(),
            );
            // https://github.com/hyperium/hyper/issues/3737
            // We should put the first message into the stream to trigger it to actually send.
            // otherwise, write_stream will hang forever.
            let w_stream = WriteStream::new(tx);
            w_stream.send_defer(request.clone()).await?;
            let streaming = client_guard
                .write_stream(write_stream_request)
                .await?
                .into_inner();
            w_stream.listen(streaming).await;
            Ok::<WriteStream, OxiaError>(w_stream)
        };
        let mut cell = self
            .streams
            .entry(shard_id)
            .or_insert_with(|| OnceCell::new());
        let initialized = cell.initialized();
        let w_stream = cell.get_or_try_init(defer_init).await?;
        if !w_stream.is_alive().await {
            if let Some(stream) = cell.take() {
                stream.shutdown().await?;
            }
            return Err(InternalRetryable());
        }
        if initialized {
            return w_stream.send(request).await;
        }
        w_stream.get_defer_response().await.unwrap()
    }

    pub async fn shutdown(self) -> Result<(), OxiaError> {
        let streams = self.streams;

        let mut joiner = JoinSet::new();
        for (_, stream_cell) in streams.into_iter() {
            if let Some(stream) = stream_cell.into_inner() {
                joiner.spawn(stream.shutdown());
            }
        }
        while let Some(result) = joiner.join_next().await {
            result.map_err(|err| {
                UnexpectedStatus(format!(
                    "write stream nanager background task failed to join: {}",
                    err
                ))
            })??;
        }
        Ok(())
    }

    pub fn new(
        namespace: String,
        shard_manager: Arc<ShardManager>,
        provider_manager: Arc<ProviderManager>,
    ) -> Self {
        WriteStreamManager {
            namespace,
            shard_manager,
            provider_manager,
            streams: DashMap::new(),
        }
    }
}