|
| 1 | +use camino::Utf8PathBuf; |
| 2 | +use dropshot::{test_util::LogContext, ConfigLogging, ConfigLoggingLevel}; |
| 3 | +use omicron_test_utils::dev; |
| 4 | +use slog::Logger; |
| 5 | +use std::io::Write; |
| 6 | + |
| 7 | +// Creates a string identifier for the current DB schema and version. |
| 8 | +// |
| 9 | +// The goal here is to allow to create different "seed" directories |
| 10 | +// for each revision of the DB. |
| 11 | +fn digest_unique_to_schema() -> String { |
| 12 | + let schema = include_str!("../../schema/crdb/dbinit.sql"); |
| 13 | + let crdb_version = include_str!("../../tools/cockroachdb_version"); |
| 14 | + let mut ctx = ring::digest::Context::new(&ring::digest::SHA256); |
| 15 | + ctx.update(&schema.as_bytes()); |
| 16 | + ctx.update(&crdb_version.as_bytes()); |
| 17 | + let digest = ctx.finish(); |
| 18 | + hex::encode(digest.as_ref()) |
| 19 | +} |
| 20 | + |
| 21 | +enum SeedDirectoryStatus { |
| 22 | + Created, |
| 23 | + Existing, |
| 24 | +} |
| 25 | + |
| 26 | +async fn ensure_seed_directory_exists( |
| 27 | + log: &Logger, |
| 28 | +) -> (Utf8PathBuf, SeedDirectoryStatus) { |
| 29 | + let base_seed_dir = Utf8PathBuf::from_path_buf(std::env::temp_dir()) |
| 30 | + .expect("Not a UTF-8 path") |
| 31 | + .join("crdb-base"); |
| 32 | + std::fs::create_dir_all(&base_seed_dir).unwrap(); |
| 33 | + let desired_seed_dir = base_seed_dir.join(digest_unique_to_schema()); |
| 34 | + |
| 35 | + if desired_seed_dir.exists() { |
| 36 | + return (desired_seed_dir, SeedDirectoryStatus::Existing); |
| 37 | + } |
| 38 | + |
| 39 | + // The directory didn't exist when we started, so try to create it. |
| 40 | + // |
| 41 | + // Nextest will execute it just once, but it is possible for a user to start |
| 42 | + // up multiple nextest processes to be running at the same time. So we |
| 43 | + // should consider it possible for another caller to create this seed |
| 44 | + // directory before we finish setting it up ourselves. |
| 45 | + let tmp_seed_dir = |
| 46 | + camino_tempfile::Utf8TempDir::new_in(base_seed_dir).unwrap(); |
| 47 | + dev::test_setup_database_seed(log, tmp_seed_dir.path()).await; |
| 48 | + |
| 49 | + // If we can successfully perform the rename, there was either no |
| 50 | + // contention or we won a creation race. |
| 51 | + // |
| 52 | + // If we couldn't perform the rename, the directory might already exist. |
| 53 | + // Check that this is the error we encountered -- otherwise, we're |
| 54 | + // struggling. |
| 55 | + if let Err(err) = std::fs::rename(tmp_seed_dir.path(), &desired_seed_dir) { |
| 56 | + if !desired_seed_dir.exists() { |
| 57 | + panic!("Cannot rename seed directory for CockroachDB: {err}"); |
| 58 | + } |
| 59 | + } |
| 60 | + |
| 61 | + (desired_seed_dir, SeedDirectoryStatus::Created) |
| 62 | +} |
| 63 | + |
| 64 | +#[tokio::main] |
| 65 | +async fn main() { |
| 66 | + // TODO: dropshot is v heavyweight for this, we should be able to pull in a |
| 67 | + // smaller binary |
| 68 | + let logctx = LogContext::new( |
| 69 | + "crdb_seeding", |
| 70 | + &ConfigLogging::StderrTerminal { level: ConfigLoggingLevel::Info }, |
| 71 | + ); |
| 72 | + let (dir, status) = ensure_seed_directory_exists(&logctx.log).await; |
| 73 | + match status { |
| 74 | + SeedDirectoryStatus::Created => { |
| 75 | + slog::info!(logctx.log, "Created seed directory: `{dir}`"); |
| 76 | + } |
| 77 | + SeedDirectoryStatus::Existing => { |
| 78 | + slog::info!(logctx.log, "Using existing seed directory: `{dir}`"); |
| 79 | + } |
| 80 | + } |
| 81 | + if let Ok(env_path) = std::env::var("NEXTEST_ENV") { |
| 82 | + let mut file = std::fs::File::create(&env_path) |
| 83 | + .expect("failed to open NEXTEST_ENV file"); |
| 84 | + writeln!(file, "CRDB_SEED_DIR={dir}") |
| 85 | + .expect("failed to write to NEXTEST_ENV file"); |
| 86 | + } else { |
| 87 | + slog::warn!( |
| 88 | + logctx.log, |
| 89 | + "NEXTEST_ENV not set (is this script running under nextest?)" |
| 90 | + ); |
| 91 | + } |
| 92 | +} |
0 commit comments