100 lines
3.7 KiB
Rust
100 lines
3.7 KiB
Rust
//! The connection side of the per-agent resource-cap KV bucket: opening it for
|
|
//! writing or reading, and the hive's direct read. The bucket's names, subjects
|
|
//! and stored type are `swarm_wire::agent_limits`.
|
|
|
|
/// Open the caps bucket for writing, creating it if nothing has yet.
|
|
///
|
|
/// **Controller-side only.** `history: 1`: a hive wants the current value.
|
|
#[cfg(feature = "kv")]
|
|
pub async fn open_or_create(
|
|
client: &async_nats::Client,
|
|
) -> Result<async_nats::jetstream::kv::Store, crate::Error> {
|
|
let js = async_nats::jetstream::new(client.clone());
|
|
crate::kv::open_or_create(
|
|
&js,
|
|
swarm_wire::agent_limits::BUCKET,
|
|
"agent-limits",
|
|
"Resource caps the swarm controller sets for each agent".to_owned(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Open the caps bucket for reading, or `None` when nothing has created it
|
|
/// yet, which reads the same as no caps.
|
|
///
|
|
/// **Controller-side**, for reads that must not create the bucket.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// [`crate::Error::OpenBucket`] for every failure other than the bucket not
|
|
/// existing.
|
|
#[cfg(feature = "kv")]
|
|
pub async fn open(
|
|
client: &async_nats::Client,
|
|
) -> Result<Option<async_nats::jetstream::kv::Store>, crate::Error> {
|
|
let js = async_nats::jetstream::new(client.clone());
|
|
match js.get_key_value(swarm_wire::agent_limits::BUCKET).await {
|
|
Ok(store) => Ok(Some(store)),
|
|
Err(e) if crate::wanted::is_bucket_not_found(&e) => Ok(None),
|
|
Err(source) => Err(crate::Error::OpenBucket {
|
|
bucket: swarm_wire::agent_limits::BUCKET.to_owned(),
|
|
source,
|
|
}),
|
|
}
|
|
}
|
|
|
|
/// What reading one agent's key found.
|
|
#[cfg(feature = "agent-limits")]
|
|
#[derive(Debug)]
|
|
pub enum Read {
|
|
/// The bucket answered; an absent or deleted key has every field unset.
|
|
Current(swarm_wire::agent_limits::AgentCaps),
|
|
/// Nothing answers for the bucket's stream: no cap has been written yet,
|
|
/// or `JetStream` is not serving it right now. Not an unset value, since
|
|
/// the second case says nothing about the agent's caps.
|
|
NoBucket,
|
|
}
|
|
|
|
/// Read `agent`'s current caps with one direct get.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// A failed request or an undecodable value. A request the queue refused for
|
|
/// permissions reaches the client as a timeout, not as a refusal.
|
|
#[cfg(feature = "agent-limits")]
|
|
pub async fn read(
|
|
client: &async_nats::Client,
|
|
agent: &str,
|
|
) -> Result<Read, Box<dyn std::error::Error + Send + Sync>> {
|
|
use async_nats::jetstream::stream::DirectGetErrorKind;
|
|
use swarm_wire::agent_limits::{AgentCaps, STREAM, key_subject};
|
|
|
|
let js = async_nats::jetstream::new(client.clone());
|
|
let stream = js.get_stream_no_info(STREAM).await?;
|
|
let message = match stream.direct_get_last_for_subject(key_subject(agent)).await {
|
|
Ok(message) => message,
|
|
Err(e) if matches!(e.kind(), DirectGetErrorKind::NotFound) => {
|
|
return Ok(Read::Current(AgentCaps::default()));
|
|
}
|
|
Err(e) if is_no_responders(&e) => return Ok(Read::NoBucket),
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
// A KV delete or purge is stored as an empty marker carrying this header.
|
|
let removed = message
|
|
.headers
|
|
.get("KV-Operation")
|
|
.is_some_and(|op| matches!(op.as_str(), "DEL" | "PURGE"));
|
|
if removed {
|
|
return Ok(Read::Current(AgentCaps::default()));
|
|
}
|
|
Ok(Read::Current(serde_json::from_slice(&message.payload)?))
|
|
}
|
|
|
|
#[cfg(feature = "agent-limits")]
|
|
fn is_no_responders(e: &async_nats::jetstream::stream::DirectGetError) -> bool {
|
|
matches!(
|
|
e.kind(),
|
|
async_nats::jetstream::stream::DirectGetErrorKind::ErrorResponse(status, _)
|
|
if status == async_nats::StatusCode::NO_RESPONDERS
|
|
)
|
|
}
|