ytsaurus_job/lib.rs
1//! Write [YTsaurus](https://ytsaurus.tech) MapReduce jobs in Rust.
2//!
3//! A YTsaurus job is an ordinary executable. The cluster runs it once per chunk
4//! of input, feeds it rows on fd 0, and collects output tables from fds 1, 4, 7…
5//! The wire format is [YSON](https://ytsaurus.tech/docs/en/user-guide/storage/yson),
6//! normally binary. This crate handles that protocol so a job can be written as
7//! a loop over rows.
8//!
9//! # A complete mapper
10//!
11//! ```no_run
12//! use serde::{Deserialize, Serialize};
13//! use ytsaurus_job::{Event, JobReader, JobWriter};
14//!
15//! #[derive(Deserialize)]
16//! struct Input<'a> {
17//! #[serde(borrow)]
18//! url: &'a str,
19//! size: i64,
20//! }
21//!
22//! #[derive(Serialize)]
23//! struct Output<'a> {
24//! host: &'a str,
25//! size: i64,
26//! }
27//!
28//! fn main() {
29//! ytsaurus_job::run(|| {
30//! let mut reader = JobReader::from_stdin();
31//! let mut writer = JobWriter::descriptors(1)?;
32//!
33//! while let Some(event) = reader.next_event()? {
34//! let Event::Row(row) = event else { continue };
35//! let input: Input = row.parse()?;
36//! let host = input.url.split('/').next().unwrap_or("");
37//! writer.write(0, &Output { host, size: input.size })?;
38//! }
39//!
40//! writer.finish()
41//! })
42//! }
43//! ```
44//!
45//! Build it for the cluster with `scripts/build-worker.sh`, then launch it with
46//! the `yt` CLI — see `docs/writing-a-job.md`.
47//!
48//! # Memory
49//!
50//! The input stream is usually much larger than the job's memory limit.
51//! [`JobReader`] never accumulates it: it holds one buffer (1 MiB by default)
52//! and hands out rows that borrow from it. A row is only copied if you ask for
53//! an owned type when decoding it.
54//!
55//! # Control records
56//!
57//! When the operation enables them, YTsaurus interleaves control records with
58//! the data: `<table_index=N>#`, `<row_index=N>#`, `<range_index=N>#` and
59//! `<key_switch=%true>#`. [`JobReader`] consumes the first three and reflects
60//! them on each [`Row`]; `key_switch` surfaces as [`Event::KeySwitch`], or is
61//! turned into per-key iterators by [`JobReader::groups`].
62
63#![warn(missing_docs)]
64
65/// Errors a job can fail with.
66pub mod error;
67/// Reading the input stream.
68pub mod reader;
69/// Reading Skiff input streams.
70pub mod skiff;
71/// Writing Skiff output streams.
72pub mod skiff_writer;
73/// Custom job statistics.
74pub mod statistics;
75/// Shared-format worker reader and writer.
76pub mod worker;
77/// Writing output tables.
78pub mod writer;
79
80pub use crate::error::{JobError, Result};
81pub use crate::reader::{Event, Group, GroupKey, Groups, JobReader, Row};
82pub use crate::skiff::{SkiffJobReader, SkiffRow};
83pub use crate::skiff_writer::SkiffJobWriter;
84pub use crate::statistics::JobStatistics;
85pub use crate::worker::{WorkerEvent, WorkerReader, WorkerRow, WorkerWriter};
86pub use crate::writer::{JobWriter, TableId, table_descriptor};
87pub use ytsaurus_format::DataFormat;
88
89pub use ytsaurus_yson as yson;
90
91/// Installs a panic hook that reports panics in a form a human can act on.
92///
93/// A job's stderr is shown in the operation UI, so this is where a failing job
94/// explains itself. The default hook already prints the message and location;
95/// this one labels it so it is obvious in the UI that the job — not the
96/// infrastructure — is at fault, and reminds the reader that backtraces need
97/// `RUST_BACKTRACE`, which cannot be set after the fact on a cluster.
98///
99/// [`run`] calls this for you.
100pub fn install_panic_hook() {
101 let default_hook = std::panic::take_hook();
102 std::panic::set_hook(Box::new(move |info| {
103 eprintln!("─────────────────────────────────────────────");
104 eprintln!("ytsaurus-job: the job panicked and will fail.");
105 if std::env::var_os("RUST_BACKTRACE").is_none() {
106 eprintln!("Set RUST_BACKTRACE=1 in the operation spec's environment for a backtrace.");
107 }
108 eprintln!("─────────────────────────────────────────────");
109 default_hook(info);
110 }));
111}
112
113/// Runs a job body, reporting failures the way YTsaurus expects.
114///
115/// Installs [`install_panic_hook`], runs `job`, and on error prints the whole
116/// error chain to stderr and exits with a non-zero status. YTsaurus decides
117/// whether a job succeeded from its exit code, and shows stderr in the
118/// operation UI, so this is the difference between a diagnosable failure and a
119/// job that just says "exit code 1".
120///
121/// Note that `job` is responsible for calling [`JobWriter::finish`]; buffered
122/// output that is never flushed is missing output.
123pub fn run<F, E>(job: F) -> !
124where
125 F: FnOnce() -> std::result::Result<(), E>,
126 E: std::fmt::Display,
127{
128 install_panic_hook();
129
130 match job() {
131 Ok(()) => std::process::exit(0),
132 Err(e) => {
133 eprintln!("ytsaurus-job: the job failed: {e}");
134 std::process::exit(1);
135 }
136 }
137}
138
139/// The variable YTsaurus sets in every job's environment.
140///
141/// Verified on a cluster, not assumed: a job printed its environment and this
142/// was in it. The Go SDK's `mapreduce.InsideJob` tests the same variable.
143const JOB_ID_ENV: &str = "YT_JOB_ID";
144
145/// Whether this process is running as a job on a cluster.
146///
147/// This is what lets one binary be both the launcher and the job: the cluster
148/// starts the same executable with `YT_JOB_ID` set, so the program can tell
149/// which role it is playing.
150///
151/// ```no_run
152/// fn main() {
153/// // Inside a job this never returns.
154/// ytsaurus_job::run_if_inside_job(my_mapper);
155///
156/// // Only reached on your machine: upload this binary and start the
157/// // operation that will run it.
158/// }
159/// # fn my_mapper() -> ytsaurus_job::Result<()> { Ok(()) }
160/// ```
161#[must_use]
162pub fn is_inside_job() -> bool {
163 inside_job(std::env::var_os(JOB_ID_ENV))
164}
165
166/// The job's ID, when running inside one.
167///
168/// Worth putting in a log line: it is how a message on stderr is tied back to a
169/// job in the operation's UI.
170#[must_use]
171pub fn job_id() -> Option<String> {
172 std::env::var(JOB_ID_ENV).ok().filter(|id| !id.is_empty())
173}
174
175/// Which job of its task this is, counting from zero.
176///
177/// A map job rarely cares — its share of the work arrives on fd 0. A **vanilla**
178/// job has no input at all, so this is how it knows which part of the work is
179/// its own:
180///
181/// ```no_run
182/// let shard = ytsaurus_job::job_cookie().unwrap_or(0);
183/// // ... process the shard'th slice of whatever this operation is doing
184/// ```
185///
186/// Stable across a restart: a job that fails and is retried comes back with the
187/// same cookie, which is what makes it usable as a shard number.
188#[must_use]
189pub fn job_cookie() -> Option<u64> {
190 std::env::var("YT_JOB_COOKIE").ok()?.parse().ok()
191}
192
193/// Runs `job` if this process is a job, and returns otherwise.
194///
195/// The whole of the one-binary pattern:
196///
197/// ```no_run
198/// use ytsaurus_job::{Event, JobReader, JobWriter};
199///
200/// fn main() {
201/// ytsaurus_job::run_if_inside_job(mapper);
202/// launch(); // only your machine gets here
203/// }
204///
205/// fn mapper() -> ytsaurus_job::Result<()> {
206/// let mut reader = JobReader::from_stdin();
207/// let mut writer = JobWriter::descriptors(1)?;
208/// while let Some(event) = reader.next_event()? {
209/// let Event::Row(row) = event else { continue };
210/// writer.write_raw(0, row.raw())?;
211/// }
212/// writer.finish()
213/// }
214/// # fn launch() {}
215/// ```
216///
217/// Inside a job this behaves exactly like [`run`] and never returns; the
218/// process exits with the job's status.
219pub fn run_if_inside_job<F, E>(job: F)
220where
221 F: FnOnce() -> std::result::Result<(), E>,
222 E: std::fmt::Display,
223{
224 if is_inside_job() {
225 run(job);
226 }
227}
228
229/// The decision itself, split out so it can be tested without touching the
230/// process environment — which is global, and in edition 2024 unsafe to write.
231fn inside_job(job_id: Option<std::ffi::OsString>) -> bool {
232 job_id.is_some_and(|id| !id.is_empty())
233}
234
235#[cfg(test)]
236mod tests {
237 use super::*;
238
239 #[test]
240 fn a_job_id_means_we_are_inside_a_job() {
241 assert!(inside_job(Some("55aff293-7ef14284-3fe0384-3e07".into())));
242 }
243
244 #[test]
245 fn no_job_id_means_we_are_not() {
246 assert!(!inside_job(None));
247 }
248
249 #[test]
250 fn an_empty_job_id_does_not_count() {
251 // `YT_JOB_ID=` in a shell is not a job. Treating it as one would run
252 // the job body on a developer's machine, reading their terminal as if
253 // it were an input stream.
254 assert!(!inside_job(Some(std::ffi::OsString::new())));
255 }
256}