github, server: stream outputs from S3 if possible, with range requests

jclulow · jclulow · commit 7abf37c07a4f · 2025-03-04T17:50:38.000-08:00
diff --git a/bin/src/main.rs b/bin/src/main.rs
@@ -990,6 +990,8 @@ async fn do_job_copy(mut l: Level<Stuff>) -> Result<()> {
 async fn do_job_sign(mut l: Level<Stuff>) -> Result<()> {
     l.usage_args(Some("JOB SRC"));
 
+    l.optflag("x", "", "display extended information about signed URL");
+
     let a = args!(l);
 
     if a.args().len() != 2 {
@@ -998,6 +1000,7 @@ async fn do_job_sign(mut l: Level<Stuff>) -> Result<()> {
 
     let job = a.args()[0].as_str();
     let src = a.args()[1].as_str();
+    let ext = a.opts().opt_present("x");
 
     let c = l.context().user();
     for o in c.job_outputs_get().job(job).send().await?.into_inner() {
@@ -1016,7 +1019,15 @@ async fn do_job_sign(mut l: Level<Stuff>) -> Result<()> {
             .await?
             .into_inner();
 
-        println!("{}", su.url);
+        if ext {
+            let JobOutputSignedUrlResult { available, size, url } = su;
+
+            println!("url:          {url}");
+            println!("size:         {size}");
+            println!("available:    {available}");
+        } else {
+            println!("{}", su.url);
+        }
         return Ok(());
     }
 
diff --git a/client/openapi.json b/client/openapi.json
@@ -3232,6 +3232,18 @@
       "JobOutputSignedUrlResult": {
         "type": "object",
         "properties": {
+          "available": {
+            "description": "Has this file been uploaded to the object store yet?",
+            "default": false,
+            "type": "boolean"
+          },
+          "size": {
+            "description": "Size of the file in bytes.",
+            "default": 0,
+            "type": "integer",
+            "format": "uint64",
+            "minimum": 0
+          },
           "url": {
             "type": "string"
           }
diff --git a/download/src/kinds/url.rs b/download/src/kinds/url.rs
@@ -35,35 +35,40 @@ pub async fn stream_from_url(
     range: Option<PotentialRange>,
     head_only: bool,
     content_type: String,
+    assume_file_size: Option<u64>,
 ) -> Result<Option<Response<Body>>> {
     let log = log.new(o!("download" => "url"));
     let file_size: u64;
     let (mut body, want_bytes, crange) = if let Some(range) = range {
-        /*
-         * If this is a range request, we first need to determine the total file
-         * size.
-         */
-        let head = client.head(&url).send().await?;
-
-        /*
-         * Beware the content_length() method on the response: on a HEAD request
-         * it is (unhelpfully!) zero, because the body of the request should
-         * ultimately be empty.  Parse the actual Content-Length value
-         * ourselves:
-         */
-        file_size = head
-            .headers()
-            .get(CONTENT_LENGTH)
-            .and_then(|hv| hv.to_str().ok())
-            .and_then(|cl| cl.parse::<u64>().ok())
-            .ok_or_else(|| anyhow!("no content-length in response"))?;
+        if let Some(assumed) = assume_file_size {
+            file_size = assumed;
+        } else {
+            /*
+             * If this is a range request, we first need to determine the total
+             * file size.
+             */
+            let head = client.head(&url).send().await?;
 
-        if head.status() == StatusCode::NOT_FOUND {
             /*
-             * If the backend returns 404, we want to be able to pass that on to
-             * the client.
+             * Beware the content_length() method on the response: on a HEAD
+             * request it is (unhelpfully!) zero, because the body of the
+             * request should ultimately be empty.  Parse the actual
+             * Content-Length value ourselves:
              */
-            return Ok(None);
+            file_size = head
+                .headers()
+                .get(CONTENT_LENGTH)
+                .and_then(|hv| hv.to_str().ok())
+                .and_then(|cl| cl.parse::<u64>().ok())
+                .ok_or_else(|| anyhow!("no content-length in response"))?;
+
+            if head.status() == StatusCode::NOT_FOUND {
+                /*
+                 * If the backend returns 404, we want to be able to pass that
+                 * on to the client.
+                 */
+                return Ok(None);
+            }
         }
 
         match range.single_range(file_size) {
diff --git a/github/server/src/http.rs b/github/server/src/http.rs
@@ -1,5 +1,5 @@
 /*
- * Copyright 2024 Oxide Computer Company
+ * Copyright 2025 Oxide Computer Company
  */
 
 use anyhow::{anyhow, bail, Result};
@@ -203,6 +203,7 @@ async fn artefact(
     let app = rc.context();
     let path = path.into_inner();
     let query = query.into_inner();
+    let pr = rc.range();
 
     let load = match path.load(&rc) {
         Ok(Some(load)) => load,
@@ -222,6 +223,7 @@ async fn artefact(
                 &path.output,
                 &path.name,
                 query.format.as_deref(),
+                pr,
             )
             .await
         }
@@ -910,6 +912,7 @@ async fn published_file_common(
         pr,
         head_only,
         guess_mime_type(&path.name),
+        None,
     )
     .await
     {
diff --git a/github/server/src/variety/basic.rs b/github/server/src/variety/basic.rs
@@ -1,11 +1,12 @@
 /*
- * Copyright 2024 Oxide Computer Company
+ * Copyright 2025 Oxide Computer Company
  */
 
 use crate::{App, FlushOut, FlushState};
 use anyhow::{bail, Result};
 use buildomat_client::types::{DependSubmit, JobOutput};
 use buildomat_common::*;
+use buildomat_download::PotentialRange;
 use buildomat_github_database::{types::*, Database};
 use buildomat_jobsh::variety::basic::{output_sse, output_table, BasicConfig};
 use chrono::SecondsFormat;
@@ -19,6 +20,7 @@ use serde::{Deserialize, Serialize};
 use slog::{debug, error, info, o, trace, warn, Logger};
 use std::collections::{HashMap, VecDeque};
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::io::{AsyncSeekExt, AsyncWriteExt};
 
 const KILOBYTE: f64 = 1024.0;
@@ -939,6 +941,7 @@ pub(crate) async fn artefact(
     output: &str,
     name: &str,
     format: Option<&str>,
+    pr: Option<PotentialRange>,
 ) -> Result<Option<hyper::Response<Body>>> {
     let p: BasicPrivate = cr.get_private()?;
 
@@ -951,10 +954,6 @@ pub(crate) async fn artefact(
     if let Some(id) = &p.buildomat_id {
         let bm = app.buildomat(&app.db.load_repository(cs.repo)?);
 
-        let backend =
-            bm.job_output_download().job(id).output(output).send().await?;
-        let cl = backend.content_length().unwrap();
-
         /*
          * To try and help out the browser in deciding whether to display or
          * immediately download a particular file, we'll try to guess the
@@ -974,6 +973,10 @@ pub(crate) async fn artefact(
                 bail!("cannot reformat a file that is not plain text");
             }
 
+            let backend =
+                bm.job_output_download().job(id).output(output).send().await?;
+            let cl = backend.content_length().unwrap();
+
             if cl > MAX_RENDERED_LOG {
                 bail!("file too large for reformat");
             }
@@ -1048,15 +1051,69 @@ pub(crate) async fn artefact(
             ));
         }
 
-        return Ok(Some(
-            hyper::Response::builder()
-                .status(hyper::StatusCode::OK)
-                .header(hyper::header::CONTENT_TYPE, ct)
-                .header(hyper::header::CONTENT_LENGTH, cl)
-                .body(Body::wrap(StreamBody::new(
-                    backend.into_inner_stream().map_ok(|b| Frame::data(b)),
-                )))?,
-        ));
+        /*
+         * To improve efficiency, we can try to fetch the file directly from the
+         * object store instead of forwarding the request through the API
+         * server.
+         */
+        let url = bm
+            .job_output_signed_url()
+            .job(id)
+            .output(output)
+            .body_map(|body| body.content_type(ct.clone()).expiry_seconds(120))
+            .send()
+            .await?;
+        if url.available {
+            let client = reqwest::ClientBuilder::new()
+                .timeout(Duration::from_secs(3600))
+                .tcp_keepalive(Duration::from_secs(60))
+                .connect_timeout(Duration::from_secs(15))
+                .build()?;
+
+            /*
+             * The file is available for streaming directly from the object
+             * store.
+             */
+            match buildomat_download::stream_from_url(
+                &app.log,
+                format!("pre-signed job output: job {id}, output {output}"),
+                &client,
+                url.url.clone(),
+                pr,
+                false,
+                ct,
+                Some(url.size),
+            )
+            .await
+            {
+                Ok(res) => return Ok(res),
+                Err(e) => {
+                    bail!(
+                        "pre-signed job output: \
+                        job {id} output {output}: {e:?}"
+                    );
+                }
+            }
+        } else {
+            /*
+             * Fall back to a regular request through the core API server.
+             * XXX There is currently no good way to pass the range request to
+             * the backend here.
+             */
+            let backend =
+                bm.job_output_download().job(id).output(output).send().await?;
+            let cl = backend.content_length().unwrap();
+
+            return Ok(Some(
+                hyper::Response::builder()
+                    .status(hyper::StatusCode::OK)
+                    .header(hyper::header::CONTENT_TYPE, ct)
+                    .header(hyper::header::CONTENT_LENGTH, cl)
+                    .body(Body::wrap(StreamBody::new(
+                        backend.into_inner_stream().map_ok(|b| Frame::data(b)),
+                    )))?,
+            ));
+        }
     }
 
     Ok(None)
diff --git a/server/src/api/user.rs b/server/src/api/user.rs
@@ -361,7 +361,7 @@ pub(crate) async fn job_output_download(
     let owner = c.require_user(log, &rqctx.request).await?;
     let t = c.load_job_for_user(log, &owner, p.job()?).await?;
 
-    let o = c.load_job_output(log, &t, p.output()?).await.or_500()?;
+    let (o, _) = c.load_job_output(log, &t, p.output()?).await.or_500()?;
 
     let info = format!("job {} output {} path {:?}", t.id, o.id, o.path);
     c.file_response(log, info, t.id, o.id, pr, false).await
@@ -385,7 +385,7 @@ pub(crate) async fn job_output_head(
     let owner = c.require_user(log, &rqctx.request).await?;
     let t = c.load_job_for_user(log, &owner, p.job()?).await?;
 
-    let o = c.load_job_output(log, &t, p.output()?).await.or_500()?;
+    let (o, _) = c.load_job_output(log, &t, p.output()?).await.or_500()?;
 
     let info = format!("job {} output {} path {:?}", t.id, o.id, o.path);
     c.file_response(log, info, t.id, o.id, pr, true).await
@@ -401,6 +401,16 @@ pub(crate) struct JobOutputSignedUrl {
 #[derive(Serialize, JsonSchema)]
 pub(crate) struct JobOutputSignedUrlResult {
     url: String,
+    /**
+     * Has this file been uploaded to the object store yet?
+     */
+    #[serde(default)]
+    available: bool,
+    /**
+     * Size of the file in bytes.
+     */
+    #[serde(default)]
+    size: u64,
 }
 
 #[endpoint {
@@ -429,7 +439,7 @@ pub(crate) async fn job_output_signed_url(
     let owner = c.require_user(log, &rqctx.request).await?;
     let t = c.load_job_for_user(log, &owner, p.job()?).await?;
 
-    let o = c.load_job_output(log, &t, p.output()?).await.or_500()?;
+    let (o, f) = c.load_job_output(log, &t, p.output()?).await.or_500()?;
     let psu = c
         .file_presigned_url(
             t.id,
@@ -447,7 +457,11 @@ pub(crate) async fn job_output_signed_url(
         t.id, o.id, o.path, psu.info; "params" => ?b,
     );
 
-    Ok(HttpResponseOk(JobOutputSignedUrlResult { url: psu.url }))
+    Ok(HttpResponseOk(JobOutputSignedUrlResult {
+        url: psu.url,
+        size: f.size.0,
+        available: f.time_archived.is_some(),
+    }))
 }
 
 #[derive(Deserialize, JsonSchema)]
@@ -507,7 +521,7 @@ pub(crate) async fn job_output_publish(
     let owner = c.require_user(log, &rqctx.request).await?;
     let t = c.load_job_for_user(log, &owner, p.job()?).await?;
 
-    let o = c.load_job_output(log, &t, p.output()?).await.or_500()?;
+    let (o, _) = c.load_job_output(log, &t, p.output()?).await.or_500()?;
 
     info!(
         log,
diff --git a/server/src/archive/jobs.rs b/server/src/archive/jobs.rs
@@ -622,19 +622,30 @@ impl LoadedArchivedJob {
             .collect::<Result<Vec<_>>>()
     }
 
-    pub fn job_output(&self, id: db::JobFileId) -> Result<db::JobOutput> {
+    pub fn job_output(
+        &self,
+        id: db::JobFileId,
+    ) -> Result<(db::JobOutput, db::JobFile)> {
         let job = self.id;
 
         self.job
             .outputs
             .iter()
             .find(|f| f.file.id().ok() == Some(id))
             .map(|f| {
-                Ok(db::JobOutput {
-                    job,
-                    path: f.path.clone(),
-                    id: f.file.id()?,
-                })
+                Ok((
+                    db::JobOutput {
+                        job,
+                        id: f.file.id().unwrap(),
+                        path: f.path.clone(),
+                    },
+                    db::JobFile {
+                        job,
+                        id: f.file.id().unwrap(),
+                        size: db::DataSize(f.file.size),
+                        time_archived: Some(f.file.time_archived()?),
+                    },
+                ))
             })
             .ok_or_else(|| anyhow!("file {id} for job {job} not in archive"))?
     }
diff --git a/server/src/db/mod.rs b/server/src/db/mod.rs
@@ -1440,6 +1440,10 @@ impl Database {
         self.sql.tx(|h| self.i_job_outputs(h, job))
     }
 
+    pub fn job_file(&self, job: JobId, file: JobFileId) -> DBResult<JobFile> {
+        self.sql.tx(|h| h.get_row(JobFile::find(job, file)))
+    }
+
     pub fn job_file_opt(
         &self,
         job: JobId,
diff --git a/server/src/main.rs b/server/src/main.rs
