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

use super::{wait_for_output, Channel, Example};
use indicatif::MultiProgress;
use rayon::prelude::IntoParallelIterator;
use rayon::prelude::ParallelIterator;
use std::fs::create_dir_all;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;

/// Collect examples in the repository and run them to produce `.rrd` files.
#[derive(argh::FromArgs)]
#[argh(subcommand, name = "rrd")]
pub struct Rrd {
    #[argh(positional, description = "directory to output `rrd` files into")]
    output_dir: PathBuf,

    #[argh(option, description = "include only examples in this channel")]
    channel: Channel,

    #[argh(option, description = "run only these examples")]
    examples: Vec<String>,
}

impl Rrd {
    pub fn run(self) -> anyhow::Result<()> {
        create_dir_all(&self.output_dir)?;

        let workspace_root = re_build_tools::cargo_metadata()?.workspace_root;
        let mut examples = if self.examples.is_empty() {
            self.channel.examples(workspace_root)?
        } else {
            Channel::Nightly
                .examples(workspace_root)?
                .into_iter()
                .filter(|example| self.examples.contains(&example.name))
                .collect()
        };
        examples.sort_by(|a, b| a.name.cmp(&b.name));

        let progress = MultiProgress::new();
        let results: Vec<anyhow::Result<PathBuf>> = examples
            .into_par_iter()
            .map(|example| example.build(&progress, &self.output_dir))
            .collect();

        let mut failed = false;
        for result in results {
            match result {
                Ok(rrd_path) => {
                    if let Ok(metadata) = std::fs::metadata(&rrd_path) {
                        println!(
                            "Output: {} ({})",
                            rrd_path.display(),
                            re_format::format_bytes(metadata.len() as _)
                        );
                    } else {
                        eprintln!("Missing rrd at {}", rrd_path.display());
                        failed = true;
                    }
                }
                Err(err) => {
                    eprintln!("{err}");
                    failed = true;
                }
            }
        }
        if failed {
            anyhow::bail!("Failed to run some examples");
        }

        Ok(())
    }
}

impl Example {
    fn build(self, progress: &MultiProgress, output_dir: &Path) -> anyhow::Result<PathBuf> {
        let tempdir = tempfile::tempdir()?;

        let initial_rrd_path = tempdir.path().join(&self.name).with_extension("rrd");

        {
            let mut cmd = Command::new("python3");
            cmd.arg("-m").arg(&self.name);
            cmd.arg("--save").arg(&initial_rrd_path);
            cmd.args(self.script_args);

            // Configure flushing so that:
            // * the resulting file size is deterministic
            // * the file is chunked into small batches for better streaming
            cmd.env("RERUN_FLUSH_TICK_SECS", 1_000_000_000.to_string());
            cmd.env("RERUN_FLUSH_NUM_BYTES", (128 * 1024).to_string());

            wait_for_output(cmd, &self.name, progress)?;
        }

        // Now run compaction on the result:
        let final_rrd_path = output_dir.join(&self.name).with_extension("rrd");

        let mut cmd = Command::new("python3");
        cmd.arg("-m").arg("rerun");
        cmd.arg("rrd");
        cmd.arg("compact");
        // Small chunks for better streaming:
        cmd.arg("--max-bytes").arg((128 * 1024).to_string());
        cmd.arg(&initial_rrd_path);
        cmd.arg("-o").arg(&final_rrd_path);

        wait_for_output(cmd, &format!("{} compaction", self.name), progress)?;

        Ok(final_rrd_path)
    }
}