Line data Source code
1 : use std::time::Duration;
2 :
3 : use proto::TenantTimelineId as ProtoTenantTimelineId;
4 : use proto::broker_service_client::BrokerServiceClient;
5 : use tonic::Status;
6 : use tonic::codegen::StdError;
7 : use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
8 : use utils::id::{TenantId, TenantTimelineId, TimelineId};
9 :
10 : // Code generated by protobuf.
11 : pub mod proto {
12 : // Tonic does derives as `#[derive(Clone, PartialEq, ::prost::Message)]`
13 : // we don't use these types for anything but broker data transmission,
14 : // so it's ok to ignore this one.
15 : #![allow(clippy::derive_partial_eq_without_eq)]
16 : tonic::include_proto!("storage_broker");
17 : }
18 :
19 : pub mod metrics;
20 :
21 : // Re-exports to avoid direct tonic dependency in user crates.
22 : pub use hyper::Uri;
23 : pub use tonic::{Code, Request, Streaming};
24 :
25 : pub const DEFAULT_LISTEN_ADDR: &str = "127.0.0.1:50051";
26 : pub const DEFAULT_ENDPOINT: &str = const_format::formatcp!("http://{DEFAULT_LISTEN_ADDR}");
27 :
28 : pub const DEFAULT_KEEPALIVE_INTERVAL: &str = "5000 ms";
29 : pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_millis(5000);
30 :
31 : // BrokerServiceClient charged with tonic provided Channel transport; helps to
32 : // avoid depending on tonic directly in user crates.
33 : pub type BrokerClientChannel = BrokerServiceClient<Channel>;
34 :
35 : // Create connection object configured to run TLS if schema starts with https://
36 : // and plain text otherwise. Connection is lazy, only endpoint sanity is
37 : // validated here.
38 : //
39 : // NB: this function is not async, but still must be run on a tokio runtime thread
40 : // because that's a requirement of tonic_endpoint.connect_lazy()'s Channel::new call.
41 0 : pub fn connect<U>(endpoint: U, keepalive_interval: Duration) -> anyhow::Result<BrokerClientChannel>
42 0 : where
43 0 : U: std::convert::TryInto<Uri>,
44 0 : U::Error: std::error::Error + Send + Sync + 'static,
45 0 : {
46 0 : let uri: Uri = endpoint.try_into()?;
47 0 : let mut tonic_endpoint: Endpoint = uri.into();
48 : // If schema starts with https, start encrypted connection; do plain text
49 : // otherwise.
50 0 : if let Some("https") = tonic_endpoint.uri().scheme_str() {
51 : // if there's no default provider and both ring+aws-lc-rs are enabled
52 : // this the tls settings on tonic will not work.
53 : // erroring is ok.
54 0 : rustls::crypto::ring::default_provider()
55 0 : .install_default()
56 0 : .ok();
57 0 : let tls = ClientTlsConfig::new();
58 0 : tonic_endpoint = tonic_endpoint.tls_config(tls)?;
59 0 : }
60 0 : tonic_endpoint = tonic_endpoint
61 0 : .http2_keep_alive_interval(keepalive_interval)
62 0 : .keep_alive_while_idle(true)
63 0 : .connect_timeout(DEFAULT_CONNECT_TIMEOUT);
64 0 : // keep_alive_timeout is 20s by default on both client and server side
65 0 : let channel = tonic_endpoint.connect_lazy();
66 0 : Ok(BrokerClientChannel::new(channel))
67 0 : }
68 :
69 : impl BrokerClientChannel {
70 : /// Create a new client to the given endpoint, but don't actually connect until the first request.
71 0 : pub async fn connect_lazy<D>(dst: D) -> Result<Self, tonic::transport::Error>
72 0 : where
73 0 : D: std::convert::TryInto<tonic::transport::Endpoint>,
74 0 : D::Error: Into<StdError>,
75 0 : {
76 0 : let conn = tonic::transport::Endpoint::new(dst)?.connect_lazy();
77 0 : Ok(Self::new(conn))
78 0 : }
79 : }
80 :
81 : // parse variable length bytes from protobuf
82 2 : pub fn parse_proto_ttid(proto_ttid: &ProtoTenantTimelineId) -> Result<TenantTimelineId, Status> {
83 2 : let tenant_id = TenantId::from_slice(&proto_ttid.tenant_id)
84 2 : .map_err(|e| Status::new(Code::InvalidArgument, format!("malformed tenant_id: {}", e)))?;
85 2 : let timeline_id = TimelineId::from_slice(&proto_ttid.timeline_id).map_err(|e| {
86 0 : Status::new(
87 0 : Code::InvalidArgument,
88 0 : format!("malformed timeline_id: {}", e),
89 0 : )
90 2 : })?;
91 2 : Ok(TenantTimelineId {
92 2 : tenant_id,
93 2 : timeline_id,
94 2 : })
95 2 : }
|