Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Indexer-Grpc-V2] Add ConnectionManager. #15750

Merged
merged 1 commit into from
Jan 19, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,16 @@ rust-version = { workspace = true }
[dependencies]
anyhow = { workspace = true }
aptos-indexer-grpc-server-framework = { workspace = true }
aptos-indexer-grpc-utils = { workspace = true }
aptos-protos = { workspace = true }
async-trait = { workspace = true }
clap = { workspace = true }
dashmap = { workspace = true }
rand = { workspace = true }
serde = { workspace = true }
tokio = { workspace = true }
tonic = { workspace = true }
tracing = { workspace = true }

[target.'cfg(unix)'.dependencies]
jemallocator = { workspace = true }
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ use anyhow::Result;
use aptos_indexer_grpc_server_framework::RunnableConfig;
use serde::{Deserialize, Serialize};

pub(crate) const MAX_MESSAGE_SIZE: usize = 256 * (1 << 20);

#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct IndexerGrpcDataServiceConfig {}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,291 @@
// Copyright (c) Aptos Foundation
// SPDX-License-Identifier: Apache-2.0

use crate::config::MAX_MESSAGE_SIZE;
use aptos_indexer_grpc_utils::{system_time_to_proto, timestamp_now_proto};
use aptos_protos::indexer::v1::{
grpc_manager_client::GrpcManagerClient, service_info::Info, ActiveStream, HeartbeatRequest,
HistoricalDataServiceInfo, LiveDataServiceInfo, ServiceInfo, StreamInfo, StreamProgress,
StreamProgressSampleProto,
};
use dashmap::DashMap;
use rand::prelude::*;
use std::{
collections::VecDeque,
sync::atomic::{AtomicU64, Ordering},
time::{Duration, SystemTime},
};
use tonic::{codec::CompressionEncoding, transport::channel::Channel};
use tracing::{info, warn};

pub static MAX_HEARTBEAT_RETRIES: usize = 3;

static OLD_PROGRESS_SAMPLING_RATE: Duration = Duration::from_secs(60);
static RECENT_PROGRESS_SAMPLING_RATE: Duration = Duration::from_secs(1);
static MAX_RECENT_SAMPLES_TO_KEEP: usize = 60;
static MAX_OLD_SAMPLES_TO_KEEP: usize = 60;

#[derive(Default)]
struct StreamProgressSamples {
old_samples: VecDeque<StreamProgressSample>,
recent_samples: VecDeque<StreamProgressSample>,
}

impl StreamProgressSamples {
fn new() -> Self {
Default::default()
}

fn to_proto(&self) -> Vec<StreamProgressSampleProto> {
self.old_samples
.iter()
.chain(self.recent_samples.iter())
.map(|sample| StreamProgressSampleProto {
timestamp: Some(system_time_to_proto(sample.timestamp)),
version: sample.version,
size_bytes: sample.size_bytes,
})
.collect()
}

fn maybe_add_sample(&mut self, version: u64, size_bytes: u64) {
Copy link
Contributor

@larry-aptos larry-aptos Jan 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does ringbuffer make this simpler since the sample size is fixed?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not really, since I'm keeping two VecDeque. (the push doesn't return the oldest element when it's full in ringbuffer crate)

let now = SystemTime::now();
let sample = StreamProgressSample {
timestamp: now,
version,
size_bytes,
};

if Self::accept_sample(&self.recent_samples, &sample, RECENT_PROGRESS_SAMPLING_RATE) {
self.recent_samples.push_back(sample);
if self.recent_samples.len() > MAX_RECENT_SAMPLES_TO_KEEP {
let sample = self.recent_samples.pop_front().unwrap();
if Self::accept_sample(&self.old_samples, &sample, OLD_PROGRESS_SAMPLING_RATE) {
self.old_samples.push_back(sample);
if self.old_samples.len() > MAX_OLD_SAMPLES_TO_KEEP {
self.old_samples.pop_front();
}
}
}
}
}

fn accept_sample(
samples: &VecDeque<StreamProgressSample>,
sample: &StreamProgressSample,
sampling_rate: Duration,
) -> bool {
if let Some(last_sample) = samples.back() {
if let Ok(delta) = sample.timestamp.duration_since(last_sample.timestamp) {
if delta >= sampling_rate {
return true;
}
}
} else {
return true;
};

false
}
}

struct StreamProgressSample {
timestamp: SystemTime,
version: u64,
size_bytes: u64,
}

pub(crate) struct ConnectionManager {
Copy link
Contributor

@larry-aptos larry-aptos Jan 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: some documentation comments here. the ConnectionManager is critical in the stack based on my understanding.

also i see that String is often used here as the service address. let's either use url/socketaddr or typedef it to avoid it mixed with regular strings.

chain_id: u64,
grpc_manager_connections: DashMap<String, GrpcManagerClient<Channel>>,
self_advertised_address: String,
known_latest_version: AtomicU64,
active_streams: DashMap<String, (ActiveStream, StreamProgressSamples)>,
is_live_data_service: bool,
}

impl ConnectionManager {
pub(crate) async fn new(
chain_id: u64,
grpc_manager_addresses: Vec<String>,
self_advertised_address: String,
is_live_data_service: bool,
) -> Self {
let grpc_manager_connections = DashMap::new();
grpc_manager_addresses.into_iter().for_each(|address| {
grpc_manager_connections
.insert(address.clone(), Self::create_client_from_address(&address));
});
let res = Self {
chain_id,
grpc_manager_connections,
self_advertised_address,
known_latest_version: AtomicU64::new(0),
active_streams: DashMap::new(),
is_live_data_service,
};

// Keep fetching latest version until it is available.
while res.known_latest_version.load(Ordering::SeqCst) == 0 {
for entry in res.grpc_manager_connections.iter() {
let address = entry.key();
if let Err(e) = res.heartbeat(address).await {
warn!("Error during heartbeat: {e}.");
}
}
tokio::time::sleep(Duration::from_secs(1)).await;
}

res
}

pub(crate) async fn start(&self) {
loop {
for entry in self.grpc_manager_connections.iter() {
let address = entry.key();
let mut retries = 0;
loop {
let result = self.heartbeat(address).await;
if result.is_ok() {
break;
}
retries += 1;
if retries > MAX_HEARTBEAT_RETRIES {
warn!("Failed to send heartbeat to GrpcManager at {address}, last error: {result:?}.");
}
}
continue;
}

tokio::time::sleep(Duration::from_secs(1)).await;
}
}

pub(crate) fn get_grpc_manager_client_for_request(&self) -> GrpcManagerClient<Channel> {
let mut rng = thread_rng();
self.grpc_manager_connections
.iter()
.choose(&mut rng)
.map(|kv| kv.value().clone())
.unwrap()
}

pub(crate) fn known_latest_version(&self) -> u64 {
self.known_latest_version.load(Ordering::SeqCst)
}

pub(crate) fn update_known_latest_version(&self, version: u64) {
self.known_latest_version
.fetch_max(version, Ordering::SeqCst);
}

pub(crate) fn insert_active_stream(
&self,
id: &str,
start_version: u64,
end_version: Option<u64>,
) {
self.active_streams.insert(
id.to_owned(),
(
ActiveStream {
id: id.to_owned(),
start_time: Some(timestamp_now_proto()),
start_version,
end_version,
progress: None,
},
StreamProgressSamples::new(),
),
);
}

pub(crate) fn remove_active_stream(&self, id: &String) {
self.active_streams.remove(id);
}

pub(crate) fn update_stream_progress(&self, id: &str, version: u64, size_bytes: u64) {
self.active_streams
.get_mut(id)
.unwrap()
.1
.maybe_add_sample(version, size_bytes);
}

pub(crate) fn get_active_streams(&self) -> Vec<ActiveStream> {
self.active_streams
.iter()
.map(|entry| {
let (active_stream, samples) = entry.value();
let mut active_stream = active_stream.clone();
active_stream.progress = Some(StreamProgress {
samples: samples.to_proto(),
});
active_stream
})
.collect()
}

async fn heartbeat(&self, address: &str) -> Result<(), tonic::Status> {
info!("Sending heartbeat to GrpcManager {address}.");
let timestamp = Some(timestamp_now_proto());
let known_latest_version = Some(self.known_latest_version());
let stream_info = Some(StreamInfo {
active_streams: self.get_active_streams(),
});

let info = if self.is_live_data_service {
Some(Info::LiveDataServiceInfo(LiveDataServiceInfo {
chain_id: self.chain_id,
timestamp,
known_latest_version,
stream_info,
// TODO(grao): Populate min_servable_version.
min_servable_version: None,
}))
} else {
Some(Info::HistoricalDataServiceInfo(HistoricalDataServiceInfo {
chain_id: self.chain_id,
timestamp,
known_latest_version,
stream_info,
}))
};
let service_info = ServiceInfo {
address: Some(self.self_advertised_address.clone()),
info,
};
let request = HeartbeatRequest {
service_info: Some(service_info),
};
let response = self
.grpc_manager_connections
.get(address)
// TODO(grao): Consider to not use unwrap here.
.unwrap()
.clone()
.heartbeat(request)
.await?
.into_inner();
if let Some(known_latest_version) = response.known_latest_version {
info!("Received known_latest_version ({known_latest_version}) from GrpcManager {address}.");
self.update_known_latest_version(known_latest_version);
} else {
warn!("HeartbeatResponse doesn't contain known_latest_version, GrpcManager address: {address}");
}

Ok(())
}

fn create_client_from_address(address: &str) -> GrpcManagerClient<Channel> {
info!("Creating GrpcManagerClient for {address}.");
let channel = Channel::from_shared(address.to_string())
.expect("Bad address.")
.connect_lazy();
GrpcManagerClient::new(channel)
.send_compressed(CompressionEncoding::Zstd)
.accept_compressed(CompressionEncoding::Zstd)
.max_decoding_message_size(MAX_MESSAGE_SIZE)
.max_encoding_message_size(MAX_MESSAGE_SIZE)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,5 @@
// SPDX-License-Identifier: Apache-2.0

pub mod config;
#[allow(dead_code)]
mod connection_manager;
2 changes: 1 addition & 1 deletion protos/proto/aptos/indexer/v1/grpc.proto
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ message StreamProgress {
}

message ActiveStream {
optional string id = 1;
string id = 1;
optional aptos.util.timestamp.Timestamp start_time = 2;
uint64 start_version = 3;
optional uint64 end_version = 4;
Expand Down
Loading
Loading