forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
7.0 KiB
1//! HTTP client for the forge runner API (contract section 4).
2//!
3//! Runner calls (register, claim) authenticate with the runner secret. Job calls (logs,
4//! heartbeat, artifacts, finish) authenticate with the per-job token.
5
6use std::path::Path;
7use std::time::Duration;
8
9use forge_ci::proto::{
10 ClaimRequest, FinishRequest, HeartbeatResponse, JobAssignment, RegisterRequest,
11 RegisterResponse,
12};
13use reqwest::StatusCode;
14use reqwest::header::CONTENT_LENGTH;
15use serde::de::DeserializeOwned;
16
17/// The claim long-poll is 25s on the forge side, so 35s leaves room for the response.
18pub const CLAIM_TIMEOUT: Duration = Duration::from_secs(35);
19const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
20const LOG_TIMEOUT: Duration = Duration::from_secs(60);
21const ARTIFACT_TIMEOUT: Duration = Duration::from_secs(60 * 60);
22const HEADER_LOG_OFFSET: &str = "X-Forge-Log-Offset";
23const HEADER_LOG_SIZE: &str = "X-Forge-Log-Size";
24
25/// Errors from talking to the forge.
26#[derive(Debug, thiserror::Error)]
27pub enum ClientError {
28 #[error("http: {0}")]
29 Http(#[from] reqwest::Error),
30 #[error("forge answered {status}: {body}")]
31 Status { status: StatusCode, body: String },
32 #[error("local io: {0}")]
33 Io(#[from] std::io::Error),
34 #[error("protocol: {0}")]
35 Protocol(String),
36}
37
38/// The forge's answer to a log chunk.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum LogAck {
41 /// The chunk was stored (or was already stored, which is an idempotent retry).
42 Accepted,
43 /// The forge holds a different number of bytes. The runner must continue from this size.
44 Resync(u64),
45}
46
47/// A client for one forge base URL.
48#[derive(Clone, Debug)]
49pub struct ForgeClient {
50 http: reqwest::Client,
51 base: String,
52}
53
54impl ForgeClient {
55 /// Builds a client for `base_url`, for example `https://forge.example`.
56 pub fn new(base_url: &str) -> Result<Self, ClientError> {
57 let http = reqwest::Client::builder()
58 .user_agent(concat!("forge-runner/", env!("CARGO_PKG_VERSION")))
59 .connect_timeout(Duration::from_secs(10))
60 .build()?;
61 Ok(Self {
62 http,
63 base: base_url.trim_end_matches('/').to_owned(),
64 })
65 }
66
67 fn url(&self, path: &str) -> String {
68 format!("{}{}", self.base, path)
69 }
70
71 /// `POST /api/v1/ci/runners/register`.
72 pub async fn register(&self, req: &RegisterRequest) -> Result<RegisterResponse, ClientError> {
73 let resp = self
74 .http
75 .post(self.url("/api/v1/ci/runners/register"))
76 .json(req)
77 .timeout(REQUEST_TIMEOUT)
78 .send()
79 .await?;
80 json_or_error(resp).await
81 }
82
83 /// `POST /api/v1/ci/runners/claim`. Returns `None` on 204 (nothing to run).
84 pub async fn claim(
85 &self,
86 runner_secret: &str,
87 labels: &[String],
88 ) -> Result<Option<JobAssignment>, ClientError> {
89 let req = ClaimRequest {
90 labels: labels.to_vec(),
91 };
92 let resp = self
93 .http
94 .post(self.url("/api/v1/ci/runners/claim"))
95 .bearer_auth(runner_secret)
96 .json(&req)
97 .timeout(CLAIM_TIMEOUT)
98 .send()
99 .await?;
100 match resp.status() {
101 StatusCode::NO_CONTENT => Ok(None),
102 status if status.is_success() => Ok(Some(resp.json().await?)),
103 _ => Err(status_error(resp).await),
104 }
105 }
106
107 /// `POST /api/v1/ci/jobs/{id}/logs` with the byte offset of `body` in the job log.
108 pub async fn post_logs(
109 &self,
110 job_id: &str,
111 job_token: &str,
112 offset: u64,
113 body: Vec<u8>,
114 ) -> Result<LogAck, ClientError> {
115 let resp = self
116 .http
117 .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/logs")))
118 .bearer_auth(job_token)
119 .header(HEADER_LOG_OFFSET, offset.to_string())
120 .header(reqwest::header::CONTENT_TYPE, "text/plain; charset=utf-8")
121 .body(body)
122 .timeout(LOG_TIMEOUT)
123 .send()
124 .await?;
125 match resp.status() {
126 StatusCode::CONFLICT => {
127 let size = resp
128 .headers()
129 .get(HEADER_LOG_SIZE)
130 .and_then(|v| v.to_str().ok())
131 .and_then(|v| v.parse::<u64>().ok())
132 .ok_or_else(|| {
133 ClientError::Protocol("409 without a valid X-Forge-Log-Size".into())
134 })?;
135 Ok(LogAck::Resync(size))
136 }
137 status if status.is_success() => Ok(LogAck::Accepted),
138 _ => Err(status_error(resp).await),
139 }
140 }
141
142 /// `POST /api/v1/ci/jobs/{id}/heartbeat`.
143 pub async fn heartbeat(
144 &self,
145 job_id: &str,
146 job_token: &str,
147 ) -> Result<HeartbeatResponse, ClientError> {
148 let resp = self
149 .http
150 .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/heartbeat")))
151 .bearer_auth(job_token)
152 .timeout(REQUEST_TIMEOUT)
153 .send()
154 .await?;
155 json_or_error(resp).await
156 }
157
158 /// `PUT /api/v1/ci/jobs/{id}/artifacts/{name}`, streaming the file from disk.
159 pub async fn upload_artifact(
160 &self,
161 job_id: &str,
162 job_token: &str,
163 name: &str,
164 path: &Path,
165 ) -> Result<(), ClientError> {
166 let file = tokio::fs::File::open(path).await?;
167 let len = file.metadata().await?.len();
168 let stream = tokio_util::io::ReaderStream::new(file);
169 let resp = self
170 .http
171 .put(self.url(&format!("/api/v1/ci/jobs/{job_id}/artifacts/{name}")))
172 .bearer_auth(job_token)
173 .header(CONTENT_LENGTH, len)
174 .body(reqwest::Body::wrap_stream(stream))
175 .timeout(ARTIFACT_TIMEOUT)
176 .send()
177 .await?;
178 if resp.status().is_success() {
179 Ok(())
180 } else {
181 Err(status_error(resp).await)
182 }
183 }
184
185 /// `POST /api/v1/ci/jobs/{id}/finish`.
186 pub async fn finish(
187 &self,
188 job_id: &str,
189 job_token: &str,
190 req: &FinishRequest,
191 ) -> Result<(), ClientError> {
192 let resp = self
193 .http
194 .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/finish")))
195 .bearer_auth(job_token)
196 .json(req)
197 .timeout(REQUEST_TIMEOUT)
198 .send()
199 .await?;
200 if resp.status().is_success() {
201 Ok(())
202 } else {
203 Err(status_error(resp).await)
204 }
205 }
206}
207
208async fn json_or_error<T: DeserializeOwned>(resp: reqwest::Response) -> Result<T, ClientError> {
209 if resp.status().is_success() {
210 Ok(resp.json().await?)
211 } else {
212 Err(status_error(resp).await)
213 }
214}
215
216async fn status_error(resp: reqwest::Response) -> ClientError {
217 let status = resp.status();
218 let body: String = resp
219 .text()
220 .await
221 .unwrap_or_default()
222 .chars()
223 .take(512)
224 .collect();
225 ClientError::Status { status, body }
226}