Feature/new registry (#2612)

* wip

* overhaul boot process

* wip: new registry

* wip

* wip

* wip

* wip

* wip

* wip

* os registry complete

* ui fixes

* fixes

* fixes

* more fixes

* fix merkle archive

core/startos/src/logs.rs

@@ -4,13 +4,17 @@ use std::time::{Duration, UNIX_EPOCH};
 
 use axum::extract::ws::{self, WebSocket};
 use chrono::{DateTime, Utc};
-use clap::Parser;
+use clap::{Args, FromArgMatches, Parser};
 use color_eyre::eyre::eyre;
 use futures::stream::BoxStream;
-use futures::{FutureExt, Stream, StreamExt, TryStreamExt};
+use futures::{Future, FutureExt, Stream, StreamExt, TryStreamExt};
+use itertools::Itertools;
 use models::PackageId;
 use rpc_toolkit::yajrc::RpcError;
-use rpc_toolkit::{command, from_fn_async, CallRemote, Empty, HandlerExt, ParentHandler};
+use rpc_toolkit::{
+    from_fn_async, CallRemote, Context, Empty, HandlerArgs, HandlerExt, HandlerFor, ParentHandler,
+};
+use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use tokio::io::{AsyncBufReadExt, BufReader};
 use tokio::process::{Child, Command};
@@ -19,10 +23,10 @@ use tokio_tungstenite::tungstenite::Message;
 use tracing::instrument;
 
 use crate::context::{CliContext, RpcContext};
-use crate::core::rpc_continuations::{RequestGuid, RpcContinuation};
 use crate::error::ResultExt;
 use crate::lxc::ContainerId;
 use crate::prelude::*;
+use crate::rpc_continuations::{RequestGuid, RpcContinuation, RpcContinuations};
 use crate::util::serde::Reversible;
 use crate::util::Invoke;
 
@@ -211,7 +215,6 @@ fn deserialize_log_message<'de, D: serde::de::Deserializer<'de>>(
 pub enum LogSource {
     Kernel,
     Unit(&'static str),
-    System,
     Container(ContainerId),
 }
 
@@ -220,168 +223,195 @@ pub const SYSTEM_UNIT: &str = "startd";
 #[derive(Deserialize, Serialize, Parser)]
 #[serde(rename_all = "camelCase")]
 #[command(rename_all = "kebab-case")]
-pub struct LogsParam {
+pub struct PackageIdParams {
     id: PackageId,
+}
+
+#[derive(Deserialize, Serialize, Parser)]
+#[serde(rename_all = "camelCase")]
+#[command(rename_all = "kebab-case")]
+pub struct LogsParams<Extra: FromArgMatches + Args = Empty> {
+    #[command(flatten)]
+    #[serde(flatten)]
+    extra: Extra,
     #[arg(short = 'l', long = "limit")]
     limit: Option<usize>,
-    #[arg(short = 'c', long = "cursor")]
+    #[arg(short = 'c', long = "cursor", conflicts_with = "follow")]
     cursor: Option<String>,
-    #[arg(short = 'B', long = "before")]
+    #[arg(short = 'B', long = "before", conflicts_with = "follow")]
     #[serde(default)]
     before: bool,
+}
+
+#[derive(Deserialize, Serialize, Parser)]
+#[serde(rename_all = "camelCase")]
+#[command(rename_all = "kebab-case")]
+pub struct CliLogsParams<Extra: FromArgMatches + Args = Empty> {
+    #[command(flatten)]
+    #[serde(flatten)]
+    rpc_params: LogsParams<Extra>,
     #[arg(short = 'f', long = "follow")]
     #[serde(default)]
     follow: bool,
 }
 
-pub fn logs() -> ParentHandler<LogsParam> {
-    ParentHandler::<LogsParam>::new()
+#[allow(private_bounds)]
+pub fn logs<
+    C: Context + AsRef<RpcContinuations>,
+    Extra: FromArgMatches + Serialize + DeserializeOwned + Args + Send + Sync + 'static,
+>(
+    source: impl for<'a> LogSourceFn<'a, C, Extra>,
+) -> ParentHandler<C, LogsParams<Extra>> {
+    ParentHandler::new()
         .root_handler(
-            from_fn_async(cli_logs)
-                .no_display()
-                .with_inherited(|params, _| params),
-        )
-        .root_handler(
-            from_fn_async(logs_nofollow)
+            logs_nofollow::<C, Extra>(source.clone())
                 .with_inherited(|params, _| params)
                 .no_cli(),
         )
         .subcommand(
             "follow",
-            from_fn_async(logs_follow)
+            logs_follow::<C, Extra>(source)
                 .with_inherited(|params, _| params)
                 .no_cli(),
         )
 }
-pub async fn cli_logs(
-    ctx: CliContext,
-    _: Empty,
-    LogsParam {
-        id,
-        limit,
-        cursor,
-        before,
-        follow,
-    }: LogsParam,
-) -> Result<(), RpcError> {
-    if follow {
-        if cursor.is_some() {
-            return Err(RpcError::from(Error::new(
-                eyre!("The argument '--cursor <cursor>' cannot be used with '--follow'"),
-                crate::ErrorKind::InvalidRequest,
-            )));
-        }
-        if before {
-            return Err(RpcError::from(Error::new(
-                eyre!("The argument '--before' cannot be used with '--follow'"),
-                crate::ErrorKind::InvalidRequest,
-            )));
-        }
-        cli_logs_generic_follow(ctx, "package.logs.follow", Some(id), limit).await
-    } else {
-        cli_logs_generic_nofollow(ctx, "package.logs", Some(id), limit, cursor, before).await
-    }
-}
-pub async fn logs_nofollow(
-    ctx: RpcContext,
-    _: Empty,
-    LogsParam {
-        id,
-        limit,
-        cursor,
-        before,
+
+pub async fn cli_logs<RemoteContext, Extra>(
+    HandlerArgs {
+        context: ctx,
+        parent_method,
+        method,
+        params: CliLogsParams { rpc_params, follow },
         ..
-    }: LogsParam,
-) -> Result<LogResponse, Error> {
-    let container_id = ctx
-        .services
-        .get(&id)
-        .await
-        .as_ref()
-        .map(|x| x.container_id())
-        .ok_or_else(|| {
-            Error::new(
-                eyre!("No service found with id: {}", id),
-                ErrorKind::NotFound,
-            )
-        })??;
-    fetch_logs(LogSource::Container(container_id), limit, cursor, before).await
-}
-pub async fn logs_follow(
-    ctx: RpcContext,
-    _: Empty,
-    LogsParam { id, limit, .. }: LogsParam,
-) -> Result<LogFollowResponse, Error> {
-    let container_id = ctx
-        .services
-        .get(&id)
-        .await
-        .as_ref()
-        .map(|x| x.container_id())
-        .ok_or_else(|| {
-            Error::new(
-                eyre!("No service found with id: {}", id),
-                ErrorKind::NotFound,
-            )
-        })??;
-    follow_logs(ctx, LogSource::Container(container_id), limit).await
-}
+    }: HandlerArgs<CliContext, CliLogsParams<Extra>>,
+) -> Result<(), RpcError>
+where
+    CliContext: CallRemote<RemoteContext>,
+    Extra: FromArgMatches + Args + Serialize + Send + Sync,
+{
+    let method = parent_method
+        .into_iter()
+        .chain(method)
+        .chain(follow.then_some("follow"))
+        .join(".");
 
-pub async fn cli_logs_generic_nofollow(
-    ctx: CliContext,
-    method: &str,
-    id: Option<PackageId>,
-    limit: Option<usize>,
-    cursor: Option<String>,
-    before: bool,
-) -> Result<(), RpcError> {
-    let res = from_value::<LogResponse>(
-        ctx.call_remote(
-            method,
-            imbl_value::json!({
-                "id": id,
-                "limit": limit,
-                "cursor": cursor,
-                "before": before,
-            }),
-        )
-        .await?,
-    )?;
+    if follow {
+        let res = from_value::<LogFollowResponse>(
+            ctx.call_remote::<RemoteContext>(&method, to_value(&rpc_params)?)
+                .await?,
+        )?;
 
-    for entry in res.entries.iter() {
-        println!("{}", entry);
-    }
+        let mut stream = ctx.ws_continuation(res.guid).await?;
+        while let Some(log) = stream.try_next().await? {
+            if let Message::Text(log) = log {
+                println!("{}", serde_json::from_str::<LogEntry>(&log)?);
+            }
+        }
+    } else {
+        let res = from_value::<LogResponse>(
+            ctx.call_remote::<RemoteContext>(&method, to_value(&rpc_params)?)
+                .await?,
+        )?;
 
-    Ok(())
-}
-
-pub async fn cli_logs_generic_follow(
-    ctx: CliContext,
-    method: &str,
-    id: Option<PackageId>,
-    limit: Option<usize>,
-) -> Result<(), RpcError> {
-    let res = from_value::<LogFollowResponse>(
-        ctx.call_remote(
-            method,
-            imbl_value::json!({
-                "id": id,
-                "limit": limit,
-            }),
-        )
-        .await?,
-    )?;
-
-    let mut stream = ctx.ws_continuation(res.guid).await?;
-    while let Some(log) = stream.try_next().await? {
-        if let Message::Text(log) = log {
-            println!("{}", serde_json::from_str::<LogEntry>(&log)?);
+        for entry in res.entries.iter() {
+            println!("{}", entry);
         }
     }
 
     Ok(())
 }
 
+trait LogSourceFn<'a, Context, Extra>: Clone + Send + Sync + 'static {
+    type Fut: Future<Output = Result<LogSource, Error>> + Send + 'a;
+    fn call(&self, ctx: &'a Context, extra: Extra) -> Self::Fut;
+}
+
+impl<'a, C: Context, Extra, F, Fut> LogSourceFn<'a, C, Extra> for F
+where
+    F: Fn(&'a C, Extra) -> Fut + Clone + Send + Sync + 'static,
+    Fut: Future<Output = Result<LogSource, Error>> + Send + 'a,
+{
+    type Fut = Fut;
+    fn call(&self, ctx: &'a C, extra: Extra) -> Self::Fut {
+        self(ctx, extra)
+    }
+}
+
+fn logs_nofollow<C, Extra>(
+    f: impl for<'a> LogSourceFn<'a, C, Extra>,
+) -> impl HandlerFor<C, Params = Empty, InheritedParams = LogsParams<Extra>, Ok = LogResponse, Err = Error>
+where
+    C: Context,
+    Extra: FromArgMatches + Args + Send + Sync + 'static,
+{
+    from_fn_async(
+        move |HandlerArgs {
+                  context,
+                  inherited_params:
+                      LogsParams {
+                          extra,
+                          limit,
+                          cursor,
+                          before,
+                      },
+                  ..
+              }: HandlerArgs<C, Empty, LogsParams<Extra>>| {
+            let f = f.clone();
+            async move { fetch_logs(f.call(&context, extra).await?, limit, cursor, before).await }
+        },
+    )
+}
+
+fn logs_follow<
+    C: Context + AsRef<RpcContinuations>,
+    Extra: FromArgMatches + Args + Send + Sync + 'static,
+>(
+    f: impl for<'a> LogSourceFn<'a, C, Extra>,
+) -> impl HandlerFor<
+    C,
+    Params = Empty,
+    InheritedParams = LogsParams<Extra>,
+    Ok = LogFollowResponse,
+    Err = Error,
+> {
+    from_fn_async(
+        move |HandlerArgs {
+                  context,
+                  inherited_params: LogsParams { extra, limit, .. },
+                  ..
+              }: HandlerArgs<C, Empty, LogsParams<Extra>>| {
+            let f = f.clone();
+            async move {
+                let src = f.call(&context, extra).await?;
+                follow_logs(context, src, limit).await
+            }
+        },
+    )
+}
+
+async fn get_package_id(
+    ctx: &RpcContext,
+    PackageIdParams { id }: PackageIdParams,
+) -> Result<LogSource, Error> {
+    let container_id = ctx
+        .services
+        .get(&id)
+        .await
+        .as_ref()
+        .map(|x| x.container_id())
+        .ok_or_else(|| {
+            Error::new(
+                eyre!("No service found with id: {}", id),
+                ErrorKind::NotFound,
+            )
+        })??;
+    Ok(LogSource::Container(container_id))
+}
+
+pub fn package_logs() -> ParentHandler<RpcContext, LogsParams<PackageIdParams>> {
+    logs::<RpcContext, PackageIdParams>(get_package_id)
+}
+
 pub async fn journalctl(
     id: LogSource,
     limit: usize,
@@ -474,10 +504,6 @@ fn gen_journalctl_command(id: &LogSource) -> Command {
             cmd.arg("-u");
             cmd.arg(id);
         }
-        LogSource::System => {
-            cmd.arg("-u");
-            cmd.arg(SYSTEM_UNIT);
-        }
         LogSource::Container(_container_id) => {
             cmd.arg("-u").arg("container-runtime.service");
         }
@@ -533,8 +559,8 @@ pub async fn fetch_logs(
 }
 
 #[instrument(skip_all)]
-pub async fn follow_logs(
-    ctx: RpcContext,
+pub async fn follow_logs<Context: AsRef<RpcContinuations>>(
+    ctx: Context,
     id: LogSource,
     limit: Option<usize>,
 ) -> Result<LogFollowResponse, Error> {
@@ -556,23 +582,24 @@ pub async fn follow_logs(
     }
 
     let guid = RequestGuid::new();
-    ctx.add_continuation(
-        guid.clone(),
-        RpcContinuation::ws(
-            Box::new(move |socket| {
-                ws_handler(first_entry, stream, socket)
-                    .map(|x| match x {
-                        Ok(_) => (),
-                        Err(e) => {
-                            tracing::error!("Error in log stream: {}", e);
-                        }
-                    })
-                    .boxed()
-            }),
-            Duration::from_secs(30),
-        ),
-    )
-    .await;
+    ctx.as_ref()
+        .add(
+            guid.clone(),
+            RpcContinuation::ws(
+                Box::new(move |socket| {
+                    ws_handler(first_entry, stream, socket)
+                        .map(|x| match x {
+                            Ok(_) => (),
+                            Err(e) => {
+                                tracing::error!("Error in log stream: {}", e);
+                            }
+                        })
+                        .boxed()
+                }),
+                Duration::from_secs(30),
+            ),
+        )
+        .await;
     Ok(LogFollowResponse { start_cursor, guid })
 }