github-starred/komodo
2
0
Fork 1
mirror of https://github.com/moghtech/komodo.git synced 2026-03-11 17:44:19 -05:00
Code Issues 258 Packages Projects Releases 10 Wiki Activity

create and delete connections on demand

Browse Source
This commit is contained in:
mbecker20
2025-09-24 00:21:05 -07:00
parent b400add6f1
commit 9eb8b32f4a
12 changed files with 433 additions and 442 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
bin/core/src/api/execute/build.rs
View File

@@ -249,7 +249,7 @@ impl Resolve<ExecuteArgs> for RunBuild {
_ = cancel.cancelled() => {
debug!("build cancelled during clone, cleaning up builder");
update.push_error_log("build cancelled", String::from("user cancelled build during repo clone"));
cleanup_builder_instance(cleanup_data, &mut update)
cleanup_builder_instance(periphery, cleanup_data, &mut update)
.await;
info!("builder cleaned up");
return handle_early_return(update, build.id, build.name, true).await
@@ -298,7 +298,7 @@ impl Resolve<ExecuteArgs> for RunBuild {
_ = cancel.cancelled() => {
info!("build cancelled during build, cleaning up builder");
update.push_error_log("build cancelled", String::from("user cancelled build during docker build"));
cleanup_builder_instance(cleanup_data, &mut update)
cleanup_builder_instance(periphery, cleanup_data, &mut update)
.await;
return handle_early_return(update, build.id, build.name, true).await
},
@@ -344,7 +344,8 @@ impl Resolve<ExecuteArgs> for RunBuild {
// If building on temporary cloud server (AWS),
// this will terminate the server.
cleanup_builder_instance(cleanup_data, &mut update).await;
cleanup_builder_instance(periphery, cleanup_data, &mut update)
.await;
// Need to manually update the update before cache refresh,
// and before broadcast with add_update.

5
bin/core/src/api/execute/repo.rs
View File

@@ -463,7 +463,7 @@ impl Resolve<ExecuteArgs> for BuildRepo {
_ = cancel.cancelled() => {
debug!("build cancelled during clone, cleaning up builder");
update.push_error_log("build cancelled", String::from("user cancelled build during repo clone"));
cleanup_builder_instance(cleanup_data, &mut update)
cleanup_builder_instance(periphery, cleanup_data, &mut update)
.await;
info!("builder cleaned up");
return handle_builder_early_return(update, repo.id, repo.name, true).await
@@ -510,7 +510,8 @@ impl Resolve<ExecuteArgs> for BuildRepo {
// If building on temporary cloud server (AWS),
// this will terminate the server.
cleanup_builder_instance(cleanup_data, &mut update).await;
cleanup_builder_instance(periphery, cleanup_data, &mut update)
.await;
// Need to manually update the update before cache refresh,
// and before broadcast with add_update.

27
bin/core/src/api/write/build.rs
View File

@@ -7,7 +7,6 @@ use database::{
mongo_indexed::doc, mungos::mongodb::bson::oid::ObjectId,
};
use formatting::format_serror;
use komodo_client::entities::optional_string;
use komodo_client::{
api::write::*,
entities::{
@@ -31,7 +30,7 @@ use periphery_client::api::build::{
use resolver_api::Resolve;
use tokio::fs;
use crate::connection::client::spawn_client_connection;
use crate::connection::PeripheryConnectionArgs;
use crate::periphery::PeripheryClient;
use crate::{
config::core_config,
@@ -433,21 +432,15 @@ async fn get_on_host_periphery(
BuilderConfig::Url(config) => {
// TODO: Ensure connection is actually established.
// Builder id no good because it may be active for multiple connections.
let periphery =
PeripheryClient::new_with_spawned_client_connection(
ObjectId::new().to_hex(),
&config.address,
|server_id, address| async move {
spawn_client_connection(
server_id,
address,
config.private_key,
optional_string(config.public_key),
)
.await
},
)
.await?;
let periphery = PeripheryClient::new(
ObjectId::new().to_hex(),
PeripheryConnectionArgs {
address: &config.address,
private_key: &config.private_key,
expected_public_key: &config.public_key,
},
)
.await?;
// Poll for connection to be estalished
let mut err = None;
for _ in 0..10 {

214
bin/core/src/connection/client.rs
View File

@@ -1,179 +1,95 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{sync::Arc, time::Duration};
use anyhow::Context;
use anyhow::{Context, anyhow};
use axum::http::HeaderValue;
use komodo_client::entities::{optional_string, server::Server};
use periphery_client::CONNECTION_RETRY_SECONDS;
use rustls::{ClientConfig, client::danger::ServerCertVerifier};
use tokio_tungstenite::Connector;
use tracing::{info, warn};
use transport::{
auth::{ClientLoginFlow, ConnectionIdentifiers},
fix_ws_address,
websocket::tungstenite::TungsteniteWebsocket,
};
use crate::{config::core_config, state::periphery_connections};
use crate::{
connection::PeripheryConnectionArgs, periphery::ConnectionChannels,
state::periphery_connections,
};
/// Managed connections to exactly those specified by specs (ServerId -> Address)
pub async fn manage_client_connections(servers: &[Server]) {
let periphery_connections = periphery_connections();
let specs = servers
.iter()
.filter(|s| s.config.enabled)
.map(|s| {
(
&s.id,
(
&s.config.address,
&s.config.private_key,
&s.config.public_key,
),
)
})
.collect::<HashMap<_, _>>();
// Clear non specced / enabled server connections
for server_id in periphery_connections.get_keys().await {
if !specs.contains_key(&server_id) {
info!(
"Specs do not contain {server_id}, cancelling connection"
);
periphery_connections.remove(&server_id).await;
impl PeripheryConnectionArgs<'_> {
pub async fn spawn_client_connection(
self,
server_id: String,
) -> anyhow::Result<Arc<ConnectionChannels>> {
if self.address.is_empty() {
return Err(anyhow!(
"Cannot spawn client connection with empty address"
));
}
}
// Apply latest connection specs
for (server_id, (address, private_key, expected_public_key)) in
specs
{
let address = if address.is_empty() {
address.to_string()
} else {
fix_ws_address(address)
};
match (
address.is_empty(),
periphery_connections.get(server_id).await,
) {
// Periphery -> Core connections
(true, Some(existing)) if existing.address.is_none() => {
continue;
}
(true, Some(existing)) => {
existing.cancel();
continue;
}
(true, None) => continue,
// Core -> Periphery connections
(false, Some(existing))
if existing
.address
.as_ref()
.map(|a| a == &address)
.unwrap_or_default() =>
{
// Connection OK
continue;
}
// Recreate connection cases
(false, Some(_)) => {}
(false, None) => {}
};
// If reaches here, recreate the connection.
if let Err(e) = spawn_client_connection(
server_id.clone(),
address,
private_key.clone(),
optional_string(expected_public_key),
)
.await
{
warn!(
"Failed to spawn new connnection for {server_id} | {e:#}"
);
let address = fix_ws_address(self.address);
let url = ::url::Url::parse(&address)
.context("Failed to parse server address")?;
let mut host = url.host().context("url has no host")?.to_string();
if let Some(port) = url.port() {
host.push(':');
host.push_str(&port.to_string());
}
}
}
// Assumes address already wss formatted
pub async fn spawn_client_connection(
server_id: String,
address: String,
private_key: String,
expected_public_key: Option<String>,
) -> anyhow::Result<()> {
let url = ::url::Url::parse(&address)
.context("Failed to parse server address")?;
let mut host = url.host().context("url has no host")?.to_string();
if let Some(port) = url.port() {
host.push(':');
host.push_str(&port.to_string());
}
let (connection, mut write_receiver) =
periphery_connections().insert(server_id, self).await;
let (connection, mut write_receiver) = periphery_connections()
.insert(server_id, address.clone().into())
.await;
let channels = connection.channels.clone();
let config = core_config();
let private_key = if private_key.is_empty() {
config.private_key.clone()
} else {
private_key
};
let expected_public_key = expected_public_key
.or_else(|| config.periphery_public_key.clone());
tokio::spawn(async move {
loop {
let ws = tokio::select! {
_ = connection.cancel.cancelled() => {
break
}
ws = connect_websocket(&address) => ws,
};
tokio::spawn(async move {
loop {
let ws = tokio::select! {
_ = connection.cancel.cancelled() => {
break
}
ws = connect_websocket(&address) => ws,
};
let (socket, accept) = match ws {
Ok(res) => res,
Err(e) => {
connection.set_error(e).await;
tokio::time::sleep(Duration::from_secs(
CONNECTION_RETRY_SECONDS,
))
.await;
continue;
}
};
let (socket, accept) = match ws {
Ok(res) => res,
Err(e) => {
let handler = super::WebsocketHandler {
socket,
connection_identifiers: ConnectionIdentifiers {
host: host.as_bytes(),
accept: accept.as_bytes(),
query: &[],
},
write_receiver: &mut write_receiver,
connection: &connection,
};
if let Err(e) = handler.handle::<ClientLoginFlow>().await {
if connection.cancel.is_cancelled() {
break;
}
connection.set_error(e).await;
tokio::time::sleep(Duration::from_secs(
CONNECTION_RETRY_SECONDS,
))
.await;
continue;
}
};
};
}
});
let handler = super::WebsocketHandler {
socket,
connection_identifiers: ConnectionIdentifiers {
host: host.as_bytes(),
accept: accept.as_bytes(),
query: &[],
},
private_key: &private_key,
expected_public_key: expected_public_key.as_deref(),
write_receiver: &mut write_receiver,
connection: &connection,
};
if let Err(e) = handler.handle::<ClientLoginFlow>().await {
if connection.cancel.is_cancelled() {
break;
}
connection.set_error(e).await;
tokio::time::sleep(Duration::from_secs(
CONNECTION_RETRY_SECONDS,
))
.await;
continue;
};
}
});
Ok(())
Ok(channels)
}
}
pub async fn connect_websocket(

234
bin/core/src/connection/mod.rs
View File

@@ -1,16 +1,31 @@
use anyhow::anyhow;
use std::{
sync::{
Arc,
atomic::{self, AtomicBool},
},
time::Duration,
};
use anyhow::{Context, anyhow};
use bytes::Bytes;
use cache::CloneCache;
use serror::serror_into_anyhow_error;
use tokio::sync::{
RwLock,
mpsc::{Sender, error::SendError},
};
use tokio_util::sync::CancellationToken;
use transport::{
auth::{ConnectionIdentifiers, LoginFlow, PublicKeyValidator},
channel::BufferedReceiver,
bytes::id_from_transport_bytes,
channel::{BufferedReceiver, buffered_channel},
websocket::{
Websocket, WebsocketMessage, WebsocketReceiver as _,
WebsocketSender as _,
},
};
use crate::periphery::PeripheryConnection;
use crate::{config::core_config, periphery::ConnectionChannels};
pub mod client;
pub mod server;
@@ -18,11 +33,8 @@ pub mod server;
pub struct WebsocketHandler<'a, W> {
pub socket: W,
pub connection_identifiers: ConnectionIdentifiers<'a>,
pub private_key: &'a str,
pub expected_public_key: Option<&'a str>,
pub write_receiver: &'a mut BufferedReceiver<Bytes>,
pub connection: &'a PeripheryConnection,
// pub handler: &'a MessageHandler,
}
impl<W: Websocket> WebsocketHandler<'_, W> {
@@ -30,13 +42,33 @@ impl<W: Websocket> WebsocketHandler<'_, W> {
let WebsocketHandler {
mut socket,
connection_identifiers,
private_key,
expected_public_key,
write_receiver,
connection,
// handler,
} = self;
let private_key = if connection.private_key.is_empty() {
&core_config().private_key
} else {
&connection.private_key
};
let expected_public_key = if !connection
.expected_public_key
.is_empty()
{
Some(connection.expected_public_key.as_str())
} else if connection.address.is_empty() {
// Only force periphery public key for Periphery -> Core connections
Some(
core_config()
.periphery_public_key
.as_deref()
.context("Must either configure Server 'Periphery Public Key' or set KOMODO_PERIPHERY_PUBLIC_KEY")?
)
} else {
None
};
L::login(
&mut socket,
connection_identifiers,
@@ -134,3 +166,187 @@ impl PublicKeyValidator for PeripheryPublicKeyValidator<'_> {
}
}
}
#[derive(Default)]
pub struct PeripheryConnections(
CloneCache<String, Arc<PeripheryConnection>>,
);
impl PeripheryConnections {
pub async fn insert(
&self,
server_id: String,
args: PeripheryConnectionArgs<'_>,
) -> (Arc<PeripheryConnection>, BufferedReceiver<Bytes>) {
let channels = if let Some(existing_connection) =
self.0.remove(&server_id).await
{
existing_connection.cancel();
// Keep the same channels so requests
// can handle disconnects while processing.
existing_connection.channels.clone()
} else {
Default::default()
};
let (connection, receiver) =
PeripheryConnection::new(args, channels);
self.0.insert(server_id, connection.clone()).await;
(connection, receiver)
}
pub async fn get(
&self,
server_id: &String,
) -> Option<Arc<PeripheryConnection>> {
self.0.get(server_id).await
}
/// Remove and cancel connection
pub async fn remove(
&self,
server_id: &String,
) -> Option<Arc<PeripheryConnection>> {
self
.0
.remove(server_id)
.await
.inspect(|connection| connection.cancel())
}
}
/// The configurable args of a connection
#[derive(Clone, Copy)]
pub struct PeripheryConnectionArgs<'a> {
pub address: &'a str,
pub private_key: &'a str,
pub expected_public_key: &'a str,
}
impl PeripheryConnectionArgs<'_> {
pub fn matches(&self, connection: &PeripheryConnection) -> bool {
self.address == connection.address
&& self.private_key == connection.private_key
&& self.expected_public_key == connection.expected_public_key
}
}
#[derive(Debug)]
pub struct PeripheryConnection {
/// Specify outbound connection address.
/// Inbound connections have this as empty string
pub address: String,
/// The private key to use, or empty for core private key
pub private_key: String,
/// The public key to expect Periphery to have.
/// Required non-empty for inbound connection.
pub expected_public_key: String,
/// Whether Periphery is currently connected.
pub connected: AtomicBool,
/// Stores latest connection error
pub error: RwLock<Option<serror::Serror>>,
/// Cancel the connection
pub cancel: CancellationToken,
/// Send bytes to Periphery
pub sender: Sender<Bytes>,
/// Send bytes from Periphery to channel handlers.
/// Must be maintained if new connection replaces old
/// at the same server id.
pub channels: Arc<ConnectionChannels>,
}
impl PeripheryConnection {
pub fn new(
args: PeripheryConnectionArgs<'_>,
channels: Arc<ConnectionChannels>,
) -> (Arc<PeripheryConnection>, BufferedReceiver<Bytes>) {
let (sender, receiver) = buffered_channel();
(
PeripheryConnection {
address: args.address.to_string(),
private_key: args.private_key.to_string(),
expected_public_key: args.expected_public_key.to_string(),
sender,
channels,
connected: AtomicBool::new(false),
error: RwLock::new(None),
cancel: CancellationToken::new(),
}
.into(),
receiver,
)
}
pub async fn handle_incoming_bytes(&self, bytes: Bytes) {
let id = match id_from_transport_bytes(&bytes) {
Ok(res) => res,
Err(e) => {
// TODO: handle better
warn!("Failed to read id | {e:#}");
return;
}
};
let Some(channel) = self.channels.get(&id).await else {
// TODO: handle better
warn!("Failed to send response | No response channel found");
return;
};
if let Err(e) = channel.send(bytes).await {
// TODO: handle better
warn!("Failed to send response | Channel failure | {e:#}");
}
}
pub async fn send(
&self,
value: Bytes,
) -> Result<(), SendError<Bytes>> {
self.sender.send(value).await
}
pub fn set_connected(&self, connected: bool) {
self.connected.store(connected, atomic::Ordering::Relaxed);
}
pub fn connected(&self) -> bool {
self.connected.load(atomic::Ordering::Relaxed)
}
/// Polls connected 3 times (500ms in between) before bailing.
pub async fn bail_if_not_connected(&self) -> anyhow::Result<()> {
const POLL_TIMES: usize = 3;
for i in 0..POLL_TIMES {
if self.connected() {
return Ok(());
}
if i < POLL_TIMES - 1 {
tokio::time::sleep(Duration::from_millis(500)).await;
}
}
if let Some(e) = self.error().await {
Err(serror_into_anyhow_error(e))
} else {
Err(anyhow!("Server is not currently connected"))
}
}
pub async fn error(&self) -> Option<serror::Serror> {
self.error.read().await.clone()
}
pub async fn set_error(&self, e: anyhow::Error) {
let mut error = self.error.write().await;
*error = Some(e.into());
}
pub async fn clear_error(&self) {
let mut error = self.error.write().await;
*error = None;
}
pub fn cancel(&self) {
self.cancel.cancel();
}
}

30
bin/core/src/connection/server.rs
View File

@@ -1,4 +1,4 @@
use anyhow::{Context, anyhow};
use anyhow::anyhow;
use axum::{
extract::{Query, WebSocketUpgrade},
http::{HeaderMap, StatusCode},
@@ -12,7 +12,9 @@ use transport::{
websocket::axum::AxumWebsocket,
};
use crate::{config::core_config, state::periphery_connections};
use crate::{
connection::PeripheryConnectionArgs, state::periphery_connections,
};
pub async fn handler(
Query(PeripheryConnectionQuery { server: _server }): Query<
@@ -53,17 +55,15 @@ pub async fn handler(
);
}
let expected_public_key = if server.config.public_key.is_empty() {
core_config()
.periphery_public_key
.clone()
.context("Must either configure Server 'Periphery Public Key' or set KOMODO_PERIPHERY_PUBLIC_KEY")?
} else {
server.config.public_key
};
let (connection, mut write_receiver) = periphery_connections()
.insert(server.id.clone(), None)
.insert(
server.id.clone(),
PeripheryConnectionArgs {
address: "",
private_key: &server.config.private_key,
expected_public_key: &server.config.public_key,
},
)
.await;
Ok(ws.on_upgrade(|socket| async move {
@@ -71,12 +71,6 @@ pub async fn handler(
let handler = super::WebsocketHandler {
socket: AxumWebsocket(socket),
connection_identifiers: identifiers.build(query.as_bytes()),
private_key: if server.config.private_key.is_empty() {
&core_config().private_key
} else {
&server.config.private_key
},
expected_public_key: Some(&expected_public_key),
write_receiver: &mut write_receiver,
connection: &connection,
};

66
bin/core/src/helpers/builder.rs
View File

@@ -6,7 +6,7 @@ use formatting::muted;
use komodo_client::entities::{
Version,
builder::{AwsBuilderConfig, Builder, BuilderConfig},
komodo_timestamp, optional_string,
komodo_timestamp,
server::Server,
update::{Log, Update},
};
@@ -14,10 +14,16 @@ use periphery_client::api::{self, GetVersionResponse};
use crate::{
cloud::{
BuildCleanupData,
aws::ec2::{
launch_ec2_instance, terminate_ec2_instance_with_retry, Ec2Instance
}, BuildCleanupData
}, connection::client::spawn_client_connection, helpers::update::update_update, periphery::PeripheryClient, resource
Ec2Instance, launch_ec2_instance,
terminate_ec2_instance_with_retry,
},
},
connection::PeripheryConnectionArgs,
helpers::update::update_update,
periphery::PeripheryClient,
resource,
};
use super::periphery_client;
@@ -42,21 +48,15 @@ pub async fn get_builder_periphery(
}
// TODO: Dont use builder id, or will be problems
// with simultaneous spawned builders.
let periphery =
PeripheryClient::new_with_spawned_client_connection(
ObjectId::new().to_hex(),
&config.address,
|server_id, address| async move {
spawn_client_connection(
server_id,
address,
config.private_key,
optional_string(config.public_key),
)
.await
},
)
.await?;
let periphery = PeripheryClient::new(
ObjectId::new().to_hex(),
PeripheryConnectionArgs {
address: &config.address,
private_key: &config.private_key,
expected_public_key: &config.public_key,
},
)
.await?;
periphery
.health_check()
.await
@@ -109,21 +109,15 @@ async fn get_aws_builder(
// TODO: Handle ad-hoc (non server) periphery connections. These don't have ids.
let periphery_address =
format!("{protocol}://{ip}:{}", config.port);
let periphery =
PeripheryClient::new_with_spawned_client_connection(
ObjectId::new().to_hex(),
&periphery_address,
|server_id, address| async move {
spawn_client_connection(
server_id,
address,
config.private_key,
optional_string(config.public_key),
)
.await
},
)
.await?;
let periphery = PeripheryClient::new(
ObjectId::new().to_hex(),
PeripheryConnectionArgs {
address: &periphery_address,
private_key: &config.private_key,
expected_public_key: &config.public_key,
},
)
.await?;
let start_connect_ts = komodo_timestamp();
let mut res = Ok(GetVersionResponse {
@@ -175,11 +169,13 @@ async fn get_aws_builder(
)
}
#[instrument(skip(update))]
#[instrument(skip(periphery, update))]
pub async fn cleanup_builder_instance(
periphery: PeripheryClient,
cleanup_data: BuildCleanupData,
update: &mut Update,
) {
periphery.cleanup().await;
match cleanup_data {
BuildCleanupData::Server => {
// Nothing to clean up

11
bin/core/src/helpers/mod.rs
View File

@@ -17,6 +17,7 @@ use komodo_client::entities::{
};
use rand::Rng;
use crate::connection::PeripheryConnectionArgs;
use crate::{
config::core_config, periphery::PeripheryClient, state::db_client,
};
@@ -191,7 +192,15 @@ pub async fn periphery_client(
if !server.config.enabled {
return Err(anyhow!("server not enabled"));
}
PeripheryClient::new(server.id.clone()).await
PeripheryClient::new(
server.id.clone(),
PeripheryConnectionArgs {
address: &server.config.address,
private_key: &server.config.private_key,
expected_public_key: &server.config.public_key,
},
)
.await
}
#[instrument]

2
bin/core/src/monitor/mod.rs
View File

@@ -113,8 +113,6 @@ async fn refresh_server_cache(ts: i64) {
return;
}
};
crate::connection::client::manage_client_connections(&servers)
.await;
let futures = servers.into_iter().map(|server| async move {
update_cache_for_server(&server, false).await;
});

260
bin/core/src/periphery/mod.rs
View File

@@ -1,10 +1,4 @@
use std::{
sync::{
Arc,
atomic::{self, AtomicBool},
},
time::Duration,
};
use std::{sync::Arc, time::Duration};
use anyhow::{Context, anyhow};
use bytes::Bytes;
@@ -13,24 +7,19 @@ use periphery_client::api;
use resolver_api::HasResponse;
use serde::{Serialize, de::DeserializeOwned};
use serde_json::json;
use serror::{deserialize_error_bytes, serror_into_anyhow_error};
use tokio::sync::{
RwLock,
mpsc::{self, Sender, error::SendError},
};
use tokio_util::sync::CancellationToken;
use serror::deserialize_error_bytes;
use tokio::sync::mpsc::{self, Sender};
use tracing::warn;
use transport::{
MessageState,
bytes::{
from_transport_bytes, id_from_transport_bytes, to_transport_bytes,
},
channel::{BufferedReceiver, buffered_channel},
fix_ws_address,
bytes::{from_transport_bytes, to_transport_bytes},
};
use uuid::Uuid;
use crate::state::periphery_connections;
use crate::{
connection::{PeripheryConnection, PeripheryConnectionArgs},
state::periphery_connections,
};
pub mod terminal;
@@ -44,31 +33,56 @@ pub struct PeripheryClient {
impl PeripheryClient {
pub async fn new(
server_id: String,
args: PeripheryConnectionArgs<'_>,
) -> anyhow::Result<PeripheryClient> {
Ok(PeripheryClient {
channels: periphery_connections()
.get(&server_id)
.await
.context("Periphery not connected")?
.channels
.clone(),
server_id,
})
let connections = periphery_connections();
// Spawn client side connection if one doesn't exist.
let Some(connection) = connections.get(&server_id).await else {
if args.address.is_empty() {
return Err(anyhow!("Server {server_id} is not connected"));
}
let channels =
args.spawn_client_connection(server_id.clone()).await?;
return Ok(PeripheryClient {
server_id,
channels,
});
};
// Ensure the connection args are unchanged.
if args.matches(&connection) {
return Ok(PeripheryClient {
server_id,
channels: connection.channels.clone(),
});
}
// The args have changed.
if args.address.is_empty() {
// Remove this connection, wait and see if client reconnects
connections.remove(&server_id).await;
tokio::time::sleep(Duration::from_secs(500)).await;
let connection =
connections.get(&server_id).await.with_context(|| {
format!("Server {server_id} is not connected")
})?;
Ok(PeripheryClient {
server_id,
channels: connection.channels.clone(),
})
} else {
let channels =
args.spawn_client_connection(server_id.clone()).await?;
Ok(PeripheryClient {
server_id,
channels,
})
}
}
pub async fn new_with_spawned_client_connection<
F: Future<Output = anyhow::Result<()>>,
>(
server_id: String,
address: &str,
// (Server id, address)
spawn: impl FnOnce(String, String) -> F,
) -> anyhow::Result<PeripheryClient> {
if address.is_empty() {
return Err(anyhow!("Server address cannot be empty"));
}
spawn(server_id.clone(), fix_ws_address(address)).await?;
PeripheryClient::new(server_id).await
pub async fn cleanup(self) -> Option<Arc<PeripheryConnection>> {
periphery_connections().remove(&self.server_id).await
}
#[tracing::instrument(level = "debug", skip(self))]
@@ -180,167 +194,3 @@ impl PeripheryClient {
}
}
}
#[derive(Default)]
pub struct PeripheryConnections(
CloneCache<String, Arc<PeripheryConnection>>,
);
impl PeripheryConnections {
pub async fn insert(
&self,
server_id: String,
address: Option<String>,
) -> (Arc<PeripheryConnection>, BufferedReceiver<Bytes>) {
let channels = if let Some(existing_connection) =
self.0.remove(&server_id).await
{
existing_connection.cancel();
// Keep the same channels so requests
// can handle disconnects while processing.
existing_connection.channels.clone()
} else {
Default::default()
};
let (connection, receiver) =
PeripheryConnection::new(address, channels);
self.0.insert(server_id, connection.clone()).await;
(connection, receiver)
}
pub async fn get(
&self,
server_id: &String,
) -> Option<Arc<PeripheryConnection>> {
self.0.get(server_id).await
}
/// Remove and cancel connection
pub async fn remove(
&self,
server_id: &String,
) -> Option<Arc<PeripheryConnection>> {
self
.0
.remove(server_id)
.await
.inspect(|connection| connection.cancel())
}
pub async fn get_keys(&self) -> Vec<String> {
self.0.get_keys().await
}
}
#[derive(Debug)]
pub struct PeripheryConnection {
/// Specify outbound connection address.
/// Inbound connections have this as None
pub address: Option<String>,
/// Whether Periphery is currently connected.
pub connected: AtomicBool,
/// Stores latest connection error
pub error: RwLock<Option<serror::Serror>>,
/// Cancel the connection
pub cancel: CancellationToken,
/// Send bytes to Periphery
pub sender: Sender<Bytes>,
/// Send bytes from Periphery to channel handlers.
/// Must be maintained if new connection replaces old
/// at the same server id.
pub channels: Arc<ConnectionChannels>,
}
impl PeripheryConnection {
pub fn new(
address: Option<String>,
channels: Arc<ConnectionChannels>,
) -> (Arc<PeripheryConnection>, BufferedReceiver<Bytes>) {
let (sender, receiver) = buffered_channel();
(
PeripheryConnection {
address,
sender,
channels,
connected: AtomicBool::new(false),
error: RwLock::new(None),
cancel: CancellationToken::new(),
}
.into(),
receiver,
)
}
pub async fn handle_incoming_bytes(&self, bytes: Bytes) {
let id = match id_from_transport_bytes(&bytes) {
Ok(res) => res,
Err(e) => {
// TODO: handle better
warn!("Failed to read id | {e:#}");
return;
}
};
let Some(channel) = self.channels.get(&id).await else {
// TODO: handle better
warn!("Failed to send response | No response channel found");
return;
};
if let Err(e) = channel.send(bytes).await {
// TODO: handle better
warn!("Failed to send response | Channel failure | {e:#}");
}
}
pub async fn send(
&self,
value: Bytes,
) -> Result<(), SendError<Bytes>> {
self.sender.send(value).await
}
pub fn set_connected(&self, connected: bool) {
self.connected.store(connected, atomic::Ordering::Relaxed);
}
pub fn connected(&self) -> bool {
self.connected.load(atomic::Ordering::Relaxed)
}
/// Polls connected 3 times (1s in between) before bailing.
pub async fn bail_if_not_connected(&self) -> anyhow::Result<()> {
for i in 0..3 {
if self.connected() {
return Ok(());
}
if i < 2 {
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
if let Some(e) = self.error().await {
Err(serror_into_anyhow_error(e))
} else {
Err(anyhow!("Server is not currently connected"))
}
}
pub async fn error(&self) -> Option<serror::Serror> {
self.error.read().await.clone()
}
pub async fn set_error(&self, e: anyhow::Error) {
let mut error = self.error.write().await;
*error = Some(e.into());
}
pub async fn clear_error(&self) {
let mut error = self.error.write().await;
*error = None;
}
pub fn cancel(&self) {
self.cancel.cancel();
}
}

17
bin/core/src/resource/server.rs
View File

@@ -15,8 +15,10 @@ use komodo_client::entities::{
use crate::{
config::core_config,
connection::PeripheryConnectionArgs,
helpers::query::get_system_info,
monitor::update_cache_for_server,
periphery::PeripheryClient,
state::{
action_states, db_client, periphery_connections,
server_status_cache,
@@ -148,6 +150,21 @@ impl super::KomodoResource for Server {
updated: &Self,
_update: &mut Update,
) -> anyhow::Result<()> {
if updated.config.enabled {
// Init periphery client to trigger reconnection
// if relevant parameters change.
let _ = PeripheryClient::new(
updated.id.clone(),
PeripheryConnectionArgs {
address: &updated.config.address,
private_key: &updated.config.private_key,
expected_public_key: &updated.config.public_key,
},
)
.await;
} else {
periphery_connections().remove(&updated.id).await;
}
update_cache_for_server(updated, true).await;
Ok(())
}

2
bin/core/src/state.rs
View File

@@ -22,6 +22,7 @@ use octorust::auth::{
use crate::{
auth::jwt::JwtClient,
config::core_config,
connection::PeripheryConnections,
helpers::{
action_state::ActionStates, all_resources::AllResourcesById,
},
@@ -29,7 +30,6 @@ use crate::{
CachedDeploymentStatus, CachedRepoStatus, CachedServerStatus,
CachedStackStatus, History,
},
periphery::PeripheryConnections,
};
static DB_CLIENT: OnceLock<database::Client> = OnceLock::new();