use std::collections::BTreeMap;
use std::num::NonZeroU64;
use std::time::Duration;
use async_trait::async_trait;
use onetaskgraph_plugin_api::{
Capabilities, DependencyEdge, Direction, Health, ItemWrite, Label, NativeId, Page, PageRequest,
Project, ProjectQuery, SourceError, SourceName, Task, TaskQuery, TaskSource, WriteSupport,
};
use serde::Deserialize;
use serde_json::{Value, json};
use super::connection::{Connection, Peer};
use super::wire::{
DeleteParams, DependencyParams, EngineIdentity, IdParams, InitializeParams, InitializeResult,
LabelParams, PROTOCOL_VERSION, ProjectQueryParams, ProjectResult, ProjectWriteParams, Request,
TaskQueryParams, TaskResult, TaskWriteParams, WriteResult,
};
const HANDSHAKE_ID: &str = "0";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RequestDeadline(NonZeroU64);
impl RequestDeadline {
pub const DEFAULT: Self = Self(NonZeroU64::new(30_000).expect("non-zero default"));
#[must_use]
pub const fn from_millis(milliseconds: NonZeroU64) -> Self {
Self(milliseconds)
}
#[must_use]
pub const fn milliseconds(self) -> NonZeroU64 {
self.0
}
fn duration(self) -> Duration {
Duration::from_millis(self.0.get())
}
}
pub struct SubprocessSource {
kind: &'static str,
capabilities: Capabilities,
writes: WriteSupport,
connection: Connection,
}
impl std::fmt::Debug for SubprocessSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SubprocessSource")
.field("kind", &self.kind)
.finish_non_exhaustive()
}
}
impl SubprocessSource {
pub fn connect(
program: &str,
args: &[String],
name: &SourceName,
config: &Value,
secrets: BTreeMap<String, String>,
) -> Result<Self, SourceError> {
Self::connect_with_deadline(
program,
args,
name,
config,
secrets,
RequestDeadline::DEFAULT,
)
}
pub fn connect_with_deadline(
program: &str,
args: &[String],
name: &SourceName,
config: &Value,
secrets: BTreeMap<String, String>,
deadline: RequestDeadline,
) -> Result<Self, SourceError> {
Self::adopt(
Peer::spawn(program, args, deadline.duration())?,
name,
config,
secrets,
)
}
pub fn over(
to_plugin: impl std::io::Write + Send + 'static,
from_plugin: impl std::io::Read + Send + 'static,
name: &SourceName,
config: &Value,
secrets: BTreeMap<String, String>,
) -> Result<Self, SourceError> {
Self::over_with_request_deadline(
to_plugin,
from_plugin,
name,
config,
secrets,
RequestDeadline::DEFAULT,
)
}
pub fn over_with_request_deadline(
to_plugin: impl std::io::Write + Send + 'static,
from_plugin: impl std::io::Read + Send + 'static,
name: &SourceName,
config: &Value,
secrets: BTreeMap<String, String>,
deadline: RequestDeadline,
) -> Result<Self, SourceError> {
Self::adopt(
Peer::over(to_plugin, from_plugin, deadline.duration()),
name,
config,
secrets,
)
}
fn adopt(
mut peer: Peer,
name: &SourceName,
config: &Value,
secrets: BTreeMap<String, String>,
) -> Result<Self, SourceError> {
let result = Self::handshake(&mut peer, name, config, secrets);
let InitializeResult {
protocol_version,
kind,
capabilities,
writes,
} = match result {
Ok(result) => result,
Err(error) => return Err(with_diagnostics(error, &mut peer)),
};
let kind = kind.into_string();
if protocol_version != Some(PROTOCOL_VERSION) {
return Err(SourceError::Config {
message: match protocol_version {
Some(spoken) => format!(
"the {kind:?} plugin was asked for protocol version \
{PROTOCOL_VERSION} and answered in version {spoken}; the two are \
incompatible and this engine does not guess between them"
),
None => format!(
"the {kind:?} plugin did not say which protocol version it \
answered in; this engine speaks version {PROTOCOL_VERSION} and \
does not guess"
),
},
});
}
Ok(Self {
kind: String::leak(kind),
capabilities,
writes: writes.unwrap_or(WriteSupport::Unsupported),
connection: Connection::adopt(peer),
})
}
fn handshake(
peer: &mut Peer,
name: &SourceName,
config: &Value,
secrets: BTreeMap<String, String>,
) -> Result<InitializeResult, SourceError> {
let params = InitializeParams {
protocol_version: PROTOCOL_VERSION,
engine: EngineIdentity {
name: "onetaskgraph".to_owned(),
version: env!("CARGO_PKG_VERSION").to_owned(),
},
source_name: name.as_str().to_owned(),
config: config.clone(),
secrets,
};
let request = Request {
id: HANDSHAKE_ID.to_owned(),
method: "initialize".to_owned(),
params: serde_json::to_value(¶ms).expect("a handshake is plain data"),
};
let line = peer.exchange(
&serde_json::to_string(&request).expect("a handshake request is plain data"),
)?;
let response: super::wire::Response =
serde_json::from_str(&line).map_err(|error| SourceError::Malformed {
message: format!(
"the plugin's handshake answer is not a response envelope: {error}"
),
})?;
if response.id != HANDSHAKE_ID {
return Err(SourceError::Malformed {
message: format!(
"the plugin answered the handshake with an envelope addressed to {:?} \
rather than to {HANDSHAKE_ID:?}",
response.id
),
});
}
let outcome = response.outcome().ok_or_else(|| SourceError::Malformed {
message: "the plugin's handshake answer carried both a result and an error, or \
neither"
.to_owned(),
})?;
let result = outcome?;
serde_json::from_value(result).map_err(|error| SourceError::Malformed {
message: format!("the plugin's handshake answer is not an initialize result: {error}"),
})
}
async fn ask<T: for<'de> Deserialize<'de>>(
&self,
method: &str,
params: Value,
) -> Result<T, SourceError> {
let result = self.connection.call(method, params).await?;
serde_json::from_value(result).map_err(|error| SourceError::Malformed {
message: format!(
"the plugin's answer to {method} is not the shape it promises: {error}"
),
})
}
}
fn with_diagnostics(error: SourceError, peer: &mut Peer) -> SourceError {
let said = peer.said();
if said.is_empty() {
return error;
}
let message = format!("{error}; the plugin wrote: {said}");
match error {
SourceError::RateLimited { .. } => error,
SourceError::Config { .. } => SourceError::Config { message },
SourceError::Auth { .. } => SourceError::Auth { message },
SourceError::Refused { .. } => SourceError::Refused { message },
SourceError::Malformed { .. } => SourceError::Malformed { message },
SourceError::Unavailable { .. } => SourceError::Unavailable { message },
}
}
#[async_trait]
impl TaskSource for SubprocessSource {
fn kind(&self) -> &'static str {
self.kind
}
fn capabilities(&self) -> Capabilities {
self.capabilities.clone()
}
async fn health(&self) -> Result<Health, SourceError> {
self.ask("health", json!({})).await
}
async fn get_task(&self, id: &NativeId) -> Result<Option<Task>, SourceError> {
let result: TaskResult = self
.ask("get_task", params(&IdParams { id: id.clone() }))
.await?;
Ok(result.task)
}
async fn get_project(&self, id: &NativeId) -> Result<Option<Project>, SourceError> {
let result: ProjectResult = self
.ask("get_project", params(&IdParams { id: id.clone() }))
.await?;
Ok(result.project)
}
async fn query_tasks(
&self,
query: &TaskQuery,
page: &PageRequest,
) -> Result<Page<Task>, SourceError> {
self.ask(
"query_tasks",
params(&TaskQueryParams {
query: query.clone(),
page: page.clone(),
}),
)
.await
}
async fn query_projects(
&self,
query: &ProjectQuery,
page: &PageRequest,
) -> Result<Page<Project>, SourceError> {
self.ask(
"query_projects",
params(&ProjectQueryParams {
query: query.clone(),
page: page.clone(),
}),
)
.await
}
async fn labels(&self, page: &PageRequest) -> Result<Page<Label>, SourceError> {
self.ask("labels", params(&LabelParams { page: page.clone() }))
.await
}
async fn task_dependencies(
&self,
id: &NativeId,
direction: Direction,
page: &PageRequest,
) -> Result<Page<DependencyEdge>, SourceError> {
self.ask(
"task_dependencies",
params(&DependencyParams {
id: id.clone(),
direction,
page: page.clone(),
}),
)
.await
}
async fn project_dependencies(
&self,
id: &NativeId,
direction: Direction,
page: &PageRequest,
) -> Result<Page<DependencyEdge>, SourceError> {
self.ask(
"project_dependencies",
params(&DependencyParams {
id: id.clone(),
direction,
page: page.clone(),
}),
)
.await
}
fn writes(&self) -> WriteSupport {
self.writes
}
async fn write_task(&self, write: &ItemWrite<Task>) -> Result<NativeId, SourceError> {
let result: WriteResult = self
.ask(
"write_task",
params(&TaskWriteParams {
write: write.clone(),
}),
)
.await?;
Ok(result.id)
}
async fn write_project(&self, write: &ItemWrite<Project>) -> Result<NativeId, SourceError> {
let result: WriteResult = self
.ask(
"write_project",
params(&ProjectWriteParams {
write: write.clone(),
}),
)
.await?;
Ok(result.id)
}
async fn delete_task(&self, id: &NativeId) -> Result<(), SourceError> {
let _: IgnoredResult = self
.ask("delete_task", params(&DeleteParams { id: id.clone() }))
.await?;
Ok(())
}
async fn delete_project(&self, id: &NativeId) -> Result<(), SourceError> {
let _: IgnoredResult = self
.ask("delete_project", params(&DeleteParams { id: id.clone() }))
.await?;
Ok(())
}
}
#[derive(serde::Deserialize)]
struct IgnoredResult {}
fn params<T: serde::Serialize>(value: &T) -> Value {
serde_json::to_value(value).expect("method parameters are plain data")
}