1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244

use axum::{
    body::Body,
    extract::{Extension, Path},
    http::{header::HeaderName, HeaderValue, Response, StatusCode},
    response::{sse::Event as AxumSseEvent, Sse},
    routing::get,
    BoxError, Json, Router,
};
use dis_spawner::{
    event_stream::{event_stream, past_events},
    SessionLivedBackend, SessionLivedBackendBuilder, SPAWNER_GROUP,
};
use futures::{Stream, TryStreamExt};
use k8s_openapi::api::core::v1::Event as KubeEventResource;
use kube::{
    api::PostParams, runtime::watcher::Error as KubeWatcherError, Api, Client, ResourceExt,
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::{collections::HashMap, sync::Arc};
use tokio_stream::StreamExt;

pub async fn get_client() -> Result<Client, StatusCode> {
    Client::try_default().await.map_err(|error| {
        tracing::error!(%error, "Error getting client");
        StatusCode::INTERNAL_SERVER_ERROR
    })
}

pub fn backend_routes() -> Router {
    Router::new()
        .route("/:backend_id/status/stream", get(status_handler))
        .route("/:backend_id/status", get(last_status_handler))
        .route("/:backend_id/ready", get(ready_handler))
}

async fn ready_handler(
    Path((backend_id,)): Path<(String,)>,
    Extension(settings): Extension<Arc<ApiSettings>>,
) -> Result<Response<Body>, StatusCode> {
    let client = get_client().await?;
    let name = settings.backend_to_slab_name(&backend_id);

    let api = Api::<SessionLivedBackend>::namespaced(client, &settings.namespace);
    let slab = api.get(&name).await;

    match slab {
        Ok(slab) => {
            if slab.is_ready() {
                let url = settings
                    .backend_to_url(&backend_id)
                    .ok_or(StatusCode::INTERNAL_SERVER_ERROR)?;

                return Response::builder()
                    .status(StatusCode::FOUND)
                    .header(
                        HeaderName::from_static("location"),
                        HeaderValue::from_str(&url)
                            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
                    )
                    .body(Body::empty())
                    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR);
            } else {
                return Err(StatusCode::CONFLICT);
            }
        }
        Err(error) => {
            tracing::warn!(?error, "Error when looking up SessionLivedBackend.");
            return Err(StatusCode::NOT_FOUND);
        }
    }
}

fn event_to_json(event: &KubeEventResource) -> Value {
    json!({
        "state": event.action,
        "time": event.event_time,
    })
}

fn convert_stream<T>(stream: T) -> impl Stream<Item = Result<AxumSseEvent, BoxError>>
where
    T: Stream<Item = Result<KubeEventResource, KubeWatcherError>>,
{
    stream.map(|event| {
        let event: KubeEventResource = event.map_err(Box::new)?;

        Ok(AxumSseEvent::default()
            .json_data(event_to_json(&event))
            .map_err(Box::new)?)
    })
}

async fn last_status_handler(
    Path((backend_id,)): Path<(String,)>,
    Extension(settings): Extension<Arc<ApiSettings>>,
) -> Result<Json<Value>, StatusCode> {
    let client = settings.get_client().await?;

    let resource_name = settings.backend_to_slab_name(&backend_id);
    let mut events = past_events(client, &resource_name, &settings.namespace)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    events.sort_by_key(|d| d.event_time.clone());
    let last_event = events.last().ok_or(StatusCode::NO_CONTENT)?;

    Ok(Json(event_to_json(last_event)))
}

async fn status_handler(
    Path((backend_id,)): Path<(String,)>,
    Extension(settings): Extension<Arc<ApiSettings>>,
) -> Result<Sse<impl Stream<Item = Result<AxumSseEvent, BoxError>>>, StatusCode> {
    let client = settings.get_client().await?;

    let name = format!("{}{}", settings.service_prefix, backend_id);
    let events = event_stream(client, &name, &settings.namespace)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    let sse_events: _ = convert_stream(events).into_stream();

    Ok(Sse::new(sse_events))
}

pub struct ApiSettings {
    pub namespace: String,
    pub url_template: Option<String>,
    pub api_server_base: Option<String>,
    pub service_prefix: String,
}

impl ApiSettings {
    async fn get_client(&self) -> Result<Client, StatusCode> {
        Client::try_default().await.map_err(|error| {
            tracing::error!(%error, "Error getting client");
            StatusCode::INTERNAL_SERVER_ERROR
        })
    }

    pub fn backend_to_slab_name(&self, backend_id: &str) -> String {
        format!("{}{}", self.service_prefix, backend_id)
    }

    pub fn backend_to_url(&self, backend_id: &str) -> Option<String> {
        self.url_template
            .as_ref()
            .map(|d| d.replace("{}", &backend_id))
    }

    pub fn backend_api_path(&self, backend_id: &str, path: &str) -> Option<String> {
        let api_server_base = self.api_server_base.as_ref()?;

        Some(format!("{}/backend/{}/{}", api_server_base, backend_id, path))
    }

    pub fn get_init_result(&self, backend_id: &str) -> SpawnResult {
        let ready_url = self.backend_api_path(&backend_id, "ready");
        let status_url = self.backend_api_path(&backend_id, "status");

        SpawnResult {
            url: self.backend_to_url(backend_id),
            name: backend_id.to_string(),
            ready_url,
            status_url,
        }
    }

    pub fn slab_name_to_backend(&self, slab_name: &str) -> Option<String> {
        slab_name
            .strip_prefix(&self.service_prefix)
            .map(|t| t.to_string())
    }
}

#[derive(Serialize)]
pub struct SpawnResult {
    pub url: Option<String>,
    pub name: String,
    pub ready_url: Option<String>,
    pub status_url: Option<String>,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InitPayload {
    /// The container image used to create the container.
    image: String,

    /// HTTP port to expose on the container.
    port: Option<u16>,

    /// Environment variables to expose on the container.
    #[serde(default)]
    env: HashMap<String, String>,

    /// Duration of time (in seconds) before the pod is shut down.
    grace_period_seconds: Option<u32>,
}

pub async fn init_handler(
    Json(payload): Json<InitPayload>,
    Extension(settings): Extension<Arc<ApiSettings>>,
) -> Result<Json<SpawnResult>, StatusCode> {
    let slab = SessionLivedBackendBuilder::new(&payload.image)
        .with_env(payload.env)
        .with_port(payload.port)
        .with_grace_period(payload.grace_period_seconds)
        .build_prefixed(&settings.service_prefix);

    let client = Client::try_default().await.map_err(|error| {
        tracing::error!(%error, "Error getting client");
        StatusCode::INTERNAL_SERVER_ERROR
    })?;
    let api = Api::<SessionLivedBackend>::namespaced(client, &settings.namespace);

    let result = api
        .create(
            &PostParams {
                field_manager: Some(SPAWNER_GROUP.to_string()),
                ..PostParams::default()
            },
            &slab,
        )
        .await
        .map_err(|error| {
            tracing::error!(%error, "Error creating SessionLivedBackend.");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    let prefixed_name = result.name();
    let name = settings
        .slab_name_to_backend(&prefixed_name)
        .ok_or_else(|| {
            tracing::warn!("Couldn't strip prefix from name.");
            StatusCode::EXPECTATION_FAILED
        })?;

    let url = settings.backend_to_url(&name);

    tracing::info!(?url, %name, "Created SessionLivedBackend.");

    Ok(Json(settings.get_init_result(&name)))
}