#![deny(missing_docs)]
pub mod cache;
pub mod error;
mod utils;
use crate::error::RepodataError;
use std::{
collections::{BTreeMap, HashMap, HashSet},
io::{BufRead, BufReader, Cursor, Read, Seek},
path::{Path, PathBuf},
str::FromStr,
sync::Arc,
time::SystemTime,
};
use anyhow::{Context, Result};
use bytes::buf::Buf;
use fs_err::{self as fs};
use futures::{StreamExt, stream::FuturesUnordered};
use indexmap::IndexMap;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
#[cfg(feature = "s3")]
use opendal::layers::RetryLayer;
#[cfg(feature = "s3")]
use opendal::services::S3Config;
use opendal::{Configurator, Operator, services::FsConfig};
use rattler_conda_types::{
ChannelInfo, ChannelRelations, PackageRecord, PatchInstructions, Platform, RepoData, Shard,
ShardedRepodata, ShardedSubdirInfo, UrlOrPath, V3Packages, WhlPackageRecord,
package::{
CondaArchiveType, DistArchiveIdentifier, DistArchiveType, IndexJson, PackageFile,
RunExportsJson, WheelArchiveType,
},
};
pub use rattler_conda_types::{RepodataRevision, RepodataRevisionInfo};
pub use rattler_config::config::index::{
IndexChannelConfig, IndexConfig, PackageRevisionAssignment,
};
use rattler_digest::Sha256Hash;
use rattler_package_streaming::{
read,
seek::{self, stream_conda_content},
};
#[cfg(feature = "s3")]
use rattler_s3::ResolvedS3Credentials;
use retry_policies::{Jitter, RetryDecision, RetryPolicy, policies::ExponentialBackoff};
use serde::Serialize;
use sha2::{Digest, Sha256};
use tokio::sync::Semaphore;
use tracing::Instrument;
#[cfg(feature = "s3")]
use url::Url;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ChannelMetadata {
pub base_url: Option<String>,
pub channel_relations: Option<ChannelRelations>,
}
impl ChannelMetadata {
pub fn from_index_config(config: &IndexChannelConfig) -> Self {
Self {
base_url: config.base_url.clone(),
channel_relations: config
.channel_relations
.clone()
.filter(|relations| !relations.is_empty()),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PreconditionChecks {
#[default]
Enabled,
Disabled,
}
impl PreconditionChecks {
pub fn is_enabled(self) -> bool {
matches!(self, PreconditionChecks::Enabled)
}
}
#[derive(Debug, Clone)]
pub(crate) struct IndexedPackageRecord {
record: PackageRecord,
repodata_revision: RepodataRevision,
wheel_url: Option<UrlOrPath>,
}
#[derive(Debug, Clone, Default)]
pub struct SubdirIndexStats {
pub packages_added: usize,
pub packages_removed: usize,
pub retries: usize,
}
#[derive(Debug, Clone, Default)]
pub struct IndexStats {
pub subdirs: HashMap<Platform, SubdirIndexStats>,
}
const REPODATA_FROM_PACKAGES: &str = "repodata_from_packages.json";
const REPODATA: &str = "repodata.json";
const REPODATA_SHARDS: &str = "repodata_shards.msgpack.zst";
const ZSTD_REPODATA_COMPRESSION_LEVEL: i32 = 19;
const CACHE_CONTROL_IMMUTABLE: &str = "public, max-age=31536000, immutable";
const CACHE_CONTROL_REPODATA: &str = "public, max-age=300";
pub fn write_retry_policy() -> impl RetryPolicy {
ExponentialBackoff::builder()
.retry_bounds(
std::time::Duration::from_secs(10), std::time::Duration::from_secs(90), )
.jitter(Jitter::Bounded)
.build_with_total_retry_duration(std::time::Duration::from_secs(600)) }
pub fn package_record_from_index_json<T: Read>(
package_as_bytes: impl AsRef<[u8]>,
index_json_reader: &mut T,
) -> std::io::Result<PackageRecord> {
indexed_package_record_from_index_json(package_as_bytes, index_json_reader)
.map(|indexed| indexed.record)
}
fn indexed_package_record_from_index_json<T: Read>(
package_as_bytes: impl AsRef<[u8]>,
index_json_reader: &mut T,
) -> std::io::Result<IndexedPackageRecord> {
let index = IndexJson::from_reader(index_json_reader)?;
index
.validate()
.map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?;
let repodata_revision = index.required_repodata_revision();
let sha256_result =
rattler_digest::compute_bytes_digest::<rattler_digest::Sha256>(&package_as_bytes);
let md5_result = rattler_digest::compute_bytes_digest::<rattler_digest::Md5>(&package_as_bytes);
let size = package_as_bytes.as_ref().len();
let package_record = PackageRecord {
name: index.name,
version: index.version,
build: index.build,
build_number: index.build_number,
subdir: index.subdir.unwrap_or_else(|| "unknown".to_string()),
md5: Some(md5_result),
sha256: Some(sha256_result),
size: Some(size as u64),
arch: index.arch,
platform: index.platform,
depends: index.depends,
extra_depends: index.extra_depends,
constrains: index.constrains,
track_features: index.track_features,
features: index.features,
flags: index.flags,
noarch: index.noarch,
license: index.license,
license_family: index.license_family,
timestamp: index.timestamp,
python_site_packages_path: index.python_site_packages_path,
legacy_bz2_md5: None,
legacy_bz2_size: None,
purls: index.purls,
run_exports: None,
};
Ok(IndexedPackageRecord {
record: package_record,
repodata_revision,
wheel_url: None,
})
}
fn repodata_patch_from_conda_package_stream<'a>(
package: impl Read + Seek + 'a,
) -> anyhow::Result<rattler_conda_types::RepoDataPatch> {
let mut subdirs = HashMap::default();
let mut content_reader = stream_conda_content(package)?;
let entries = content_reader.entries()?;
for entry in entries {
let mut entry = entry?;
if !entry.header().entry_type().is_file() {
return Err(anyhow::anyhow!(
"Expected repodata patch package to be a file"
));
}
let mut buf = Vec::new();
entry.read_to_end(&mut buf)?;
let path = entry.path()?;
let components = path.components().collect::<Vec<_>>();
let subdir =
if components.len() == 2 && components[1].as_os_str() == "patch_instructions.json" {
let subdir_str = components[0]
.as_os_str()
.to_str()
.context("Could not convert OsStr to str")?;
let _ = Platform::from_str(subdir_str)?;
subdir_str.to_string()
} else {
return Err(anyhow::anyhow!(
"Expected files of form <subdir>/patch_instructions.json, but found {}",
path.display()
));
};
let instructions: PatchInstructions = serde_json::from_slice(&buf)?;
subdirs.insert(subdir, instructions);
}
Ok(rattler_conda_types::RepoDataPatch { subdirs })
}
pub fn package_record_from_tar_bz2(file: &Path) -> std::io::Result<PackageRecord> {
let reader = fs::File::open(file)?;
package_record_from_tar_bz2_reader(BufReader::new(reader))
}
pub fn package_record_from_tar_bz2_reader(reader: impl BufRead) -> std::io::Result<PackageRecord> {
let bytes = reader.bytes().collect::<Result<Vec<u8>, _>>()?;
let reader = Cursor::new(&bytes);
let mut archive = read::stream_tar_bz2(reader);
for entry in archive.entries()?.flatten() {
let mut entry = entry;
let path = entry.path()?;
if path.as_os_str().eq("info/index.json") {
return package_record_from_index_json(&bytes, &mut entry);
}
}
Err(std::io::Error::other("No index.json found"))
}
pub fn package_record_from_conda(file: &Path) -> std::io::Result<PackageRecord> {
let reader = fs::File::open(file)?;
package_record_from_conda_reader(BufReader::new(reader))
}
fn read_indexed_json_from_archive(
bytes: &Vec<u8>,
archive: &mut tar::Archive<impl Read>,
) -> std::io::Result<IndexedPackageRecord> {
let mut index_json = None;
let mut run_exports_json = None;
for entry in archive.entries()?.flatten() {
let mut entry = entry;
let path = entry.path()?;
if path.as_os_str().eq("info/index.json") {
index_json = Some(indexed_package_record_from_index_json(bytes, &mut entry)?);
} else if path.as_os_str().eq("info/run_exports.json") {
run_exports_json = Some(RunExportsJson::from_reader(&mut entry)?);
}
}
if let Some(mut index_json) = index_json {
index_json.record.run_exports = run_exports_json;
return Ok(index_json);
}
Err(std::io::Error::other("No index.json found"))
}
fn read_index_json_from_archive(
bytes: &Vec<u8>,
archive: &mut tar::Archive<impl Read>,
) -> std::io::Result<PackageRecord> {
read_indexed_json_from_archive(bytes, archive).map(|indexed| indexed.record)
}
pub fn package_record_from_conda_reader(reader: impl BufRead) -> std::io::Result<PackageRecord> {
let bytes = reader.bytes().collect::<Result<Vec<u8>, _>>()?;
let reader = Cursor::new(&bytes);
let mut archive = seek::stream_conda_info(reader).expect("Could not open conda file");
read_index_json_from_archive(&bytes, &mut archive)
}
fn indexed_package_record_from_tar_bz2_reader(
reader: impl BufRead,
) -> std::io::Result<IndexedPackageRecord> {
let bytes = reader.bytes().collect::<Result<Vec<u8>, _>>()?;
let reader = Cursor::new(&bytes);
let mut archive = read::stream_tar_bz2(reader);
for entry in archive.entries()?.flatten() {
let mut entry = entry;
let path = entry.path()?;
if path.as_os_str().eq("info/index.json") {
return indexed_package_record_from_index_json(&bytes, &mut entry);
}
}
Err(std::io::Error::other("No index.json found"))
}
fn indexed_package_record_from_conda_reader(
reader: impl BufRead,
) -> std::io::Result<IndexedPackageRecord> {
let bytes = reader.bytes().collect::<Result<Vec<u8>, _>>()?;
let reader = Cursor::new(&bytes);
let mut archive = seek::stream_conda_info(reader).expect("Could not open conda file");
read_indexed_json_from_archive(&bytes, &mut archive)
}
fn parse_package_buffer(
buffer: opendal::Buffer,
filename: &str,
) -> std::io::Result<IndexedPackageRecord> {
let reader = buffer.reader();
let archive_type = DistArchiveType::try_from(filename).unwrap();
match archive_type {
DistArchiveType::Conda(CondaArchiveType::TarBz2) => {
indexed_package_record_from_tar_bz2_reader(reader)
}
DistArchiveType::Conda(CondaArchiveType::Conda) => {
indexed_package_record_from_conda_reader(reader)
}
DistArchiveType::Wheel(WheelArchiveType::Whl) => Err(std::io::Error::other(
"Package type \".whl\" not yet supported.",
)),
}
}
async fn read_and_parse_package(
op: &Operator,
cache: &cache::PackageRecordCache,
subdir: Platform,
filename: &str,
) -> std::io::Result<IndexedPackageRecord> {
let file_path = format!("{subdir}/{filename}");
match cache.get_or_stat(op, &file_path).await {
Ok(cache::CacheResult::Hit(record)) => {
Ok(*record)
}
Ok(cache::CacheResult::Miss {
etag,
last_modified,
}) => {
let (buffer, final_metadata) = cache::read_package_with_retry(
op,
&file_path,
RepodataFileMetadata {
etag,
last_modified,
file_existed: true, precondition_checks: PreconditionChecks::Enabled, },
)
.await
.map_err(|e| std::io::Error::other(e.to_string()))?;
let record = parse_package_buffer(buffer, filename)?;
cache
.insert(
&file_path,
record.clone(),
final_metadata.etag,
final_metadata.last_modified,
)
.await;
Ok(record)
}
Err(e) => {
tracing::warn!("Cache stat failed for {file_path}: {e}, proceeding without cache");
let buffer = op
.read(&file_path)
.await
.map_err(|e| std::io::Error::other(e.to_string()))?;
parse_package_buffer(buffer, filename)
}
}
}
#[derive(Debug, Clone)]
pub struct RepodataFileMetadata {
pub etag: Option<String>,
pub last_modified: Option<opendal::raw::Timestamp>,
pub file_existed: bool,
pub precondition_checks: PreconditionChecks,
}
impl RepodataFileMetadata {
pub async fn new(
op: &Operator,
path: &str,
precondition_checks: PreconditionChecks,
) -> opendal::Result<Self> {
if !precondition_checks.is_enabled() {
return Ok(Self {
etag: None,
last_modified: None,
file_existed: false,
precondition_checks,
});
}
match op.stat(path).await {
Ok(metadata) => Ok(Self {
etag: metadata.etag().map(str::to_owned),
last_modified: metadata.last_modified(),
file_existed: true,
precondition_checks,
}),
Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(Self {
etag: None,
last_modified: None,
file_existed: false,
precondition_checks,
}),
Err(e) => Err(e),
}
}
}
#[derive(Debug, Clone)]
pub struct RepodataMetadataCollection {
pub repodata: RepodataFileMetadata,
pub repodata_from_packages: Option<RepodataFileMetadata>,
pub repodata_zst: Option<RepodataFileMetadata>,
pub repodata_shards: Option<RepodataFileMetadata>,
}
impl RepodataMetadataCollection {
pub async fn new(
op: &Operator,
subdir: Platform,
has_patch: bool,
write_zst: bool,
write_shards: bool,
precondition_checks: PreconditionChecks,
) -> opendal::Result<Self> {
let repodata =
RepodataFileMetadata::new(op, &format!("{subdir}/{REPODATA}"), precondition_checks)
.await?;
let repodata_from_packages = if has_patch {
Some(
RepodataFileMetadata::new(
op,
&format!("{subdir}/{REPODATA_FROM_PACKAGES}"),
precondition_checks,
)
.await?,
)
} else {
None
};
let repodata_zst = if write_zst {
Some(
RepodataFileMetadata::new(
op,
&format!("{subdir}/{REPODATA}.zst"),
precondition_checks,
)
.await?,
)
} else {
None
};
let repodata_shards = if write_shards {
Some(
RepodataFileMetadata::new(
op,
&format!("{subdir}/{REPODATA_SHARDS}"),
precondition_checks,
)
.await?,
)
} else {
None
};
Ok(Self {
repodata,
repodata_from_packages,
repodata_zst,
repodata_shards,
})
}
}
#[allow(clippy::too_many_arguments)]
async fn index_subdir(
subdir: Platform,
op: Operator,
force: bool,
write_zst: bool,
write_shards: bool,
repodata_revisions: Vec<RepodataRevisionInfo>,
package_revision_assignment: PackageRevisionAssignment,
channel_metadata: ChannelMetadata,
repodata_patch: Option<PatchInstructions>,
progress: Option<MultiProgress>,
semaphore: Arc<Semaphore>,
cache: cache::PackageRecordCache,
precondition_checks: PreconditionChecks,
) -> Result<SubdirIndexStats, RepodataError> {
let retry_policy = write_retry_policy();
let mut current_try = 0;
loop {
let request_start_time = SystemTime::now();
match index_subdir_inner(
subdir,
op.clone(),
force,
write_zst,
write_shards,
repodata_revisions.clone(),
package_revision_assignment,
channel_metadata.clone(),
repodata_patch.clone(),
progress.clone(),
semaphore.clone(),
cache.clone(),
precondition_checks,
)
.await
{
Ok(mut stats) => {
stats.retries = current_try;
return Ok(stats);
}
Err(e) => {
let is_retryable_condition_error = match &e {
RepodataError::Opendal(opendal_err) => {
matches!(
opendal_err.kind(),
opendal::ErrorKind::ConditionNotMatch | opendal::ErrorKind::Unexpected
) && {
let error_str = format!("{opendal_err:?}");
error_str.contains("ConditionalRequestConflict")
|| error_str.contains("status: 409")
|| opendal_err.kind() == opendal::ErrorKind::ConditionNotMatch
}
}
_ => false,
};
if is_retryable_condition_error {
match retry_policy.should_retry(request_start_time, current_try as u32) {
RetryDecision::Retry { execute_after } => {
let duration = execute_after
.duration_since(SystemTime::now())
.unwrap_or_default();
tracing::warn!(
"Detected concurrent modification of repodata for {} (attempt {}/max). \
Error: {:?}. Retrying in {:?}.",
subdir,
current_try + 1,
e,
duration
);
tokio::time::sleep(duration).await;
current_try += 1;
continue;
}
RetryDecision::DoNotRetry => {
tracing::error!(
"Max retries exceeded for {subdir}. Final error: {e:?}"
);
return Err(e);
}
}
}
return Err(e);
}
}
}
}
#[allow(clippy::too_many_arguments)]
async fn index_subdir_inner(
subdir: Platform,
op: Operator,
force: bool,
write_zst: bool,
write_shards: bool,
repodata_revisions: Vec<RepodataRevisionInfo>,
package_revision_assignment: PackageRevisionAssignment,
channel_metadata: ChannelMetadata,
repodata_patch: Option<PatchInstructions>,
progress: Option<MultiProgress>,
semaphore: Arc<Semaphore>,
cache: cache::PackageRecordCache,
precondition_checks: PreconditionChecks,
) -> Result<SubdirIndexStats, RepodataError> {
let metadata = RepodataMetadataCollection::new(
&op,
subdir,
repodata_patch.is_some(),
write_zst,
write_shards,
precondition_checks,
)
.await?;
let mut registered_packages: ahash::HashMap<DistArchiveIdentifier, IndexedPackageRecord> =
if force {
HashMap::default()
} else {
let (repodata_path, read_metadata) = if repodata_patch.is_some() {
(
format!("{subdir}/{REPODATA_FROM_PACKAGES}"),
metadata.repodata_from_packages.as_ref().unwrap(),
)
} else {
(format!("{subdir}/{REPODATA}"), &metadata.repodata)
};
match crate::utils::read_with_metadata_check(&op, &repodata_path, read_metadata).await {
Ok(bytes) => match serde_json::from_slice::<RepoData>(&bytes.to_vec()) {
Ok(repodata) => package_records_from_repodata(repodata),
Err(err) => {
tracing::warn!(
"Failed to parse {repodata_path}: {err}. Not reusing content from this file"
);
HashMap::default()
}
},
Err(err) if err.kind() == opendal::ErrorKind::NotFound => {
tracing::info!("Could not find {repodata_path}. Creating new one.");
HashMap::default()
}
Err(err) => return Err(err.into()),
}
};
let uploaded_packages: HashSet<DistArchiveIdentifier> = op
.list_with(&format!("{}/", subdir.as_str()))
.await?
.iter()
.filter_map(|entry| {
if entry.metadata().mode().is_file() {
let filename = entry.name().to_string();
DistArchiveIdentifier::try_from_filename(&filename)
} else {
None
}
})
.collect();
tracing::debug!(
"Found {} already uploaded packages in subdir {}.",
uploaded_packages.len(),
subdir
);
let packages_to_delete = registered_packages
.keys()
.cloned()
.collect::<HashSet<_>>()
.difference(&uploaded_packages)
.cloned()
.collect::<Vec<_>>();
tracing::debug!(
"Deleting {} packages from subdir {}.",
packages_to_delete.len(),
subdir
);
for filename in &packages_to_delete {
registered_packages.remove(filename);
}
let packages_to_add = uploaded_packages
.difference(®istered_packages.keys().cloned().collect::<HashSet<_>>())
.cloned()
.collect::<Vec<_>>();
tracing::info!(
"Adding {} packages to subdir {}.",
packages_to_add.len(),
subdir
);
let pb = if let Some(progress) = progress {
progress.add(ProgressBar::new(packages_to_add.len() as u64))
} else {
ProgressBar::hidden()
};
let sty = ProgressStyle::with_template(
"[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}",
)
.unwrap()
.progress_chars("##-");
pb.set_style(sty);
let mut tasks = FuturesUnordered::new();
for filename in packages_to_add.iter() {
let task = {
let op = op.clone();
let filename = filename.clone();
let pb = pb.clone();
let semaphore = semaphore.clone();
let cache = cache.clone();
async move {
let _permit = semaphore
.acquire()
.await
.expect("Semaphore was unexpectedly closed");
pb.set_message(format!(
"Indexing {} {}",
subdir.as_str(),
console::style(&filename).dim()
));
let record =
read_and_parse_package(&op, &cache, subdir, &filename.to_file_name()).await?;
pb.inc(1);
Ok::<(DistArchiveIdentifier, IndexedPackageRecord), std::io::Error>((
filename, record,
))
}
};
tasks.push(tokio::spawn(task));
}
let mut results = Vec::new();
while let Some(join_result) = tasks.next().await {
match join_result {
Ok(Ok(result)) => results.push(result),
Ok(Err(e)) => {
tasks.clear();
tracing::error!("Failed to process package: {}", e);
pb.abandon_with_message(format!(
"{} {}",
console::style("Failed to index").red(),
console::style(subdir.as_str()).dim()
));
return Err(RepodataError::Other(anyhow::anyhow!(e)));
}
Err(join_err) => {
tasks.clear();
tracing::error!("Task panicked: {}", join_err);
pb.abandon_with_message(format!(
"{} {}",
console::style("Failed to index").red(),
console::style(subdir.as_str()).dim()
));
return Err(join_err.into());
}
}
}
pb.finish_with_message(format!(
"{} {}",
console::style("Finished").green(),
subdir.as_str()
));
tracing::info!(
"Successfully added {} packages to subdir {}.",
results.len(),
subdir
);
for (filename, record) in results {
registered_packages.insert(filename, record);
}
let mut packages: IndexMap<DistArchiveIdentifier, PackageRecord, ahash::RandomState> =
IndexMap::default();
let mut conda_packages: IndexMap<DistArchiveIdentifier, PackageRecord, ahash::RandomState> =
IndexMap::default();
let mut v3 = V3Packages::default();
let latest_revision = latest_repodata_revision(&repodata_revisions);
for (filename, package) in registered_packages {
let revision =
package_revision_assignment.assign(package.repodata_revision, latest_revision);
insert_package_record_by_revision(
&mut packages,
&mut conda_packages,
&mut v3,
filename,
package,
revision,
)?;
}
let repodata_before_patches = RepoData {
info: Some(ChannelInfo {
subdir: Some(subdir.to_string()),
base_url: channel_metadata.base_url,
repodata_revisions: repodata_revisions_for_packages(&repodata_revisions, &v3),
channel_relations: channel_metadata.channel_relations,
}),
packages,
conda_packages,
v3,
removed: HashSet::default(),
version: Some(2),
};
write_repodata(
repodata_before_patches,
repodata_patch,
subdir,
op,
&metadata,
)
.await?;
Ok(SubdirIndexStats {
packages_added: packages_to_add.len(),
packages_removed: packages_to_delete.len(),
retries: 0, })
}
fn serialize_msgpack_zst<T>(val: &T) -> Result<Vec<u8>, RepodataError>
where
T: Serialize + ?Sized,
{
let msgpack = rmp_serde::to_vec_named(val)?;
let encoded = zstd::stream::encode_all(&msgpack[..], 0)?;
Ok(encoded)
}
fn latest_repodata_revision(revisions: &[RepodataRevisionInfo]) -> RepodataRevision {
revisions
.iter()
.map(|revision| revision.revision)
.max()
.unwrap_or(RepodataRevision::Legacy)
}
fn package_records_from_repodata(
repodata: RepoData,
) -> ahash::HashMap<DistArchiveIdentifier, IndexedPackageRecord> {
let mut packages = ahash::HashMap::default();
packages.extend(
repodata
.packages
.into_iter()
.chain(repodata.conda_packages)
.map(|(identifier, record)| {
(
identifier,
IndexedPackageRecord {
record,
repodata_revision: RepodataRevision::Legacy,
wheel_url: None,
},
)
}),
);
packages.extend(
repodata
.v3
.into_records_with_url()
.map(|(identifier, record, wheel_url)| {
(
identifier,
IndexedPackageRecord {
record,
repodata_revision: RepodataRevision::V3,
wheel_url,
},
)
}),
);
packages
}
#[allow(clippy::too_many_arguments)]
fn insert_package_record_by_revision(
packages: &mut IndexMap<DistArchiveIdentifier, PackageRecord, ahash::RandomState>,
conda_packages: &mut IndexMap<DistArchiveIdentifier, PackageRecord, ahash::RandomState>,
v3: &mut V3Packages,
filename: DistArchiveIdentifier,
package: IndexedPackageRecord,
revision: RepodataRevision,
) -> Result<(), RepodataError> {
let IndexedPackageRecord {
record, wheel_url, ..
} = package;
match revision {
RepodataRevision::Legacy => match filename.archive_type {
DistArchiveType::Conda(CondaArchiveType::TarBz2) => {
packages.insert(filename, record);
}
DistArchiveType::Conda(CondaArchiveType::Conda) => {
conda_packages.insert(filename, record);
}
_ => {
return Err(RepodataError::Other(anyhow::anyhow!(
"archive type '{:?}' is not supported in legacy repodata maps",
filename.archive_type
)));
}
},
RepodataRevision::V3 => match filename.archive_type {
DistArchiveType::Conda(CondaArchiveType::TarBz2) => {
v3.tar_bz2.insert(filename.identifier, record);
}
DistArchiveType::Conda(CondaArchiveType::Conda) => {
v3.conda.insert(filename.identifier, record);
}
DistArchiveType::Wheel(WheelArchiveType::Whl) => {
let url = wheel_url.ok_or_else(|| {
RepodataError::Other(anyhow::anyhow!(
"indexing new wheel packages into v3 repodata is not supported yet"
))
})?;
v3.whl.insert(
filename.identifier,
WhlPackageRecord {
package_record: record,
url,
},
);
}
},
RepodataRevision::Unknown(unsupported) => {
return Err(RepodataError::Other(anyhow::anyhow!(
"repodata revision v{unsupported} is not supported by this indexer"
)));
}
}
Ok(())
}
#[derive(Default)]
struct RevisionStats {
n_packages: u64,
oldest: Option<rattler_conda_types::utils::TimestampMs>,
newest: Option<rattler_conda_types::utils::TimestampMs>,
}
impl RevisionStats {
fn add(&mut self, record: &PackageRecord) {
self.n_packages += 1;
if let Some(timestamp) = record.timestamp {
self.oldest = Some(
self.oldest
.map_or(timestamp, |oldest| oldest.min(timestamp)),
);
self.newest = Some(
self.newest
.map_or(timestamp, |newest| newest.max(timestamp)),
);
}
}
}
fn repodata_revisions_for_packages(
configured: &[RepodataRevisionInfo],
v3: &V3Packages,
) -> Vec<RepodataRevisionInfo> {
let mut revisions = configured
.iter()
.filter(|revision| revision.revision != RepodataRevision::Legacy)
.map(|revision| (revision.revision, revision.clone()))
.collect::<BTreeMap<_, _>>();
let mut stats = BTreeMap::<RepodataRevision, RevisionStats>::new();
for (_, record) in v3.records() {
stats.entry(RepodataRevision::V3).or_default().add(record);
}
for (revision, revision_stats) in stats {
let info = revisions
.entry(revision)
.or_insert_with(|| RepodataRevisionInfo {
revision,
n_packages: None,
oldest: None,
newest: None,
});
if info.n_packages.is_none() {
info.n_packages = Some(revision_stats.n_packages);
}
if info.oldest.is_none() {
info.oldest = revision_stats.oldest;
}
if info.newest.is_none() {
info.newest = revision_stats.newest;
}
}
for revision in revisions.values_mut() {
if revision.n_packages.is_none() {
revision.n_packages = Some(0);
}
}
revisions.into_values().collect()
}
pub async fn write_repodata(
repodata: RepoData,
repodata_patch: Option<PatchInstructions>,
subdir: Platform,
op: Operator,
metadata: &RepodataMetadataCollection,
) -> Result<(), RepodataError> {
if let Some(repodata_from_packages_metadata) = &metadata.repodata_from_packages {
let unpatched_repodata_path = format!("{subdir}/{REPODATA_FROM_PACKAGES}");
tracing::info!("Writing unpatched repodata to {unpatched_repodata_path}");
let unpatched_repodata_bytes = serde_json::to_vec(&repodata)?;
crate::utils::write_with_metadata_check(
&op,
&unpatched_repodata_path,
unpatched_repodata_bytes,
repodata_from_packages_metadata,
Some(CACHE_CONTROL_REPODATA),
)
.await?;
}
let repodata = if let Some(instructions) = repodata_patch {
tracing::info!("Patching repodata");
let mut patched_repodata = repodata.clone();
patched_repodata.apply_patches(&instructions);
patched_repodata
} else {
repodata
};
let repodata_bytes = serde_json::to_vec(&repodata)?;
if let Some(repodata_zst_metadata) = &metadata.repodata_zst {
tracing::info!("Compressing repodata bytes");
let repodata_zst_bytes =
zstd::stream::encode_all(&repodata_bytes[..], ZSTD_REPODATA_COMPRESSION_LEVEL)?;
let repodata_zst_path = format!("{subdir}/{REPODATA}.zst");
tracing::info!("Writing zst repodata to {repodata_zst_path}");
crate::utils::write_with_metadata_check(
&op,
&repodata_zst_path,
repodata_zst_bytes,
repodata_zst_metadata,
Some(CACHE_CONTROL_REPODATA),
)
.await?;
}
let repodata_path = format!("{subdir}/{REPODATA}");
tracing::info!("Writing repodata to {repodata_path}");
crate::utils::write_with_metadata_check(
&op,
&repodata_path,
repodata_bytes,
&metadata.repodata,
Some(CACHE_CONTROL_REPODATA),
)
.await?;
if metadata.repodata_shards.is_some() {
tracing::info!("Creating sharded repodata");
let mut shards_by_package_names: HashMap<String, Shard> = HashMap::new();
let sharded_base_url = repodata
.info
.as_ref()
.and_then(|info| info.base_url.clone())
.unwrap_or_default();
let sharded_repodata_revisions = repodata
.info
.as_ref()
.map(|info| info.repodata_revisions.clone())
.unwrap_or_default();
let sharded_channel_relations = repodata
.info
.as_ref()
.and_then(|info| info.channel_relations.clone());
for (k, package_record) in repodata.conda_packages {
let package_name = package_record.name.as_normalized();
let shard = shards_by_package_names
.entry(package_name.into())
.or_default();
shard.conda_packages.insert(k, package_record);
}
for (k, package_record) in repodata.packages {
let package_name = package_record.name.as_normalized();
let shard = shards_by_package_names
.entry(package_name.into())
.or_default();
shard.packages.insert(k, package_record);
}
for (k, package_record) in repodata.v3.conda {
let package_name = package_record.name.as_normalized();
let shard = shards_by_package_names
.entry(package_name.into())
.or_default();
shard.v3.conda.insert(k, package_record);
}
for (k, package_record) in repodata.v3.tar_bz2 {
let package_name = package_record.name.as_normalized();
let shard = shards_by_package_names
.entry(package_name.into())
.or_default();
shard.v3.tar_bz2.insert(k, package_record);
}
for (k, package_record) in repodata.v3.whl {
let package_name = package_record.package_record.name.as_normalized();
let shard = shards_by_package_names
.entry(package_name.into())
.or_default();
shard.v3.whl.insert(k, package_record);
}
for package in repodata.removed {
let package_name = package.identifier.name.clone();
let shard = shards_by_package_names.entry(package_name).or_default();
shard.removed.insert(package);
}
let shards = shards_by_package_names
.iter()
.map(|(k, shard)| {
serialize_msgpack_zst(shard).map(|encoded| {
let mut hasher = Sha256::new();
hasher.update(&encoded);
let digest: Sha256Hash = hasher.finalize();
(k, (digest, encoded))
})
})
.collect::<Result<HashMap<_, _>, RepodataError>>()?;
let sharded_repodata = ShardedRepodata {
info: ShardedSubdirInfo {
subdir: subdir.to_string(),
base_url: sharded_base_url,
shards_base_url: "./shards/".into(),
created_at: Some(jiff::Timestamp::now()),
repodata_revisions: sharded_repodata_revisions,
channel_relations: sharded_channel_relations,
},
shards: shards
.iter()
.map(|(&k, (digest, _))| (k.clone(), *digest))
.collect(),
};
let mut tasks = FuturesUnordered::new();
for (_, (digest, encoded_shard)) in shards {
let op = op.clone();
let future = async move || {
let shard_path = format!("{subdir}/shards/{}.msgpack.zst", hex::encode(digest));
tracing::trace!("Writing repodata shard to {shard_path}");
match op
.write_with(&shard_path, encoded_shard)
.if_not_exists(true)
.cache_control(CACHE_CONTROL_IMMUTABLE)
.await
{
Err(e) if e.kind() == opendal::ErrorKind::ConditionNotMatch => {
tracing::trace!("{shard_path} already exists");
Ok(())
}
Ok(_metadata) => Ok(()),
Err(e) => Err(e),
}
};
tasks.push(tokio::spawn(future()));
}
while let Some(join_result) = tasks.next().await {
match join_result {
Ok(Ok(_)) => {}
Ok(Err(e)) => Err(e)?,
Err(join_err) => Err(join_err)?,
}
}
if let Some(repodata_shards_metadata) = &metadata.repodata_shards {
let repodata_shards_path = format!("{subdir}/{REPODATA_SHARDS}");
tracing::trace!("Writing repodata shards to {repodata_shards_path}");
let sharded_repodata_encoded = serialize_msgpack_zst(&sharded_repodata)?;
crate::utils::write_with_metadata_check(
&op,
&repodata_shards_path,
sharded_repodata_encoded,
repodata_shards_metadata,
Some(CACHE_CONTROL_REPODATA),
)
.await?;
}
}
Ok(())
}
pub struct IndexFsConfig {
pub channel: PathBuf,
pub target_platform: Option<Platform>,
pub repodata_patch: Option<String>,
pub write_zst: bool,
pub write_shards: bool,
pub repodata_revisions: Vec<RepodataRevisionInfo>,
pub package_revision_assignment: PackageRevisionAssignment,
pub force: bool,
pub max_parallel: usize,
pub multi_progress: Option<MultiProgress>,
}
pub async fn index_fs(config: IndexFsConfig) -> anyhow::Result<()> {
index_fs_with_channel_metadata(config, ChannelMetadata::default()).await
}
pub async fn index_fs_with_channel_metadata(
IndexFsConfig {
channel,
target_platform,
repodata_patch,
write_zst,
write_shards,
repodata_revisions,
package_revision_assignment,
force,
max_parallel,
multi_progress,
}: IndexFsConfig,
channel_metadata: ChannelMetadata,
) -> anyhow::Result<()> {
let mut config = FsConfig::default();
config.root = Some(channel.canonicalize()?.to_string_lossy().to_string());
let builder = config.into_builder();
let op = Operator::new(builder)?.finish();
index_with_channel_metadata(
target_platform,
op,
repodata_patch,
write_zst,
write_shards,
repodata_revisions,
package_revision_assignment,
force,
max_parallel,
multi_progress,
PreconditionChecks::Disabled,
channel_metadata,
)
.await
.map(|_| ())
}
#[cfg(feature = "s3")]
pub struct IndexS3Config {
pub channel: Url,
pub credentials: ResolvedS3Credentials,
pub target_platform: Option<Platform>,
pub repodata_patch: Option<String>,
pub write_zst: bool,
pub write_shards: bool,
pub repodata_revisions: Vec<RepodataRevisionInfo>,
pub package_revision_assignment: PackageRevisionAssignment,
pub force: bool,
pub max_parallel: usize,
pub multi_progress: Option<MultiProgress>,
pub precondition_checks: PreconditionChecks,
}
#[cfg(feature = "s3")]
fn s3_config(
credentials: &ResolvedS3Credentials,
channel: &Url,
) -> Result<S3Config, anyhow::Error> {
let mut s3_config = S3Config::default();
s3_config.root = Some(channel.path().to_string());
s3_config.bucket = channel
.host_str()
.ok_or(anyhow::anyhow!("No bucket in S3 URL"))?
.to_string();
s3_config.region = Some(credentials.region.clone());
s3_config.endpoint = Some(credentials.endpoint_url.to_string());
s3_config.secret_access_key = Some(credentials.secret_access_key.clone());
s3_config.access_key_id = Some(credentials.access_key_id.clone());
s3_config.session_token = credentials.session_token.clone();
s3_config.enable_virtual_host_style =
credentials.addressing_style == rattler_s3::S3AddressingStyle::VirtualHost;
Ok(s3_config)
}
#[cfg(feature = "s3")]
pub async fn index_s3(config: IndexS3Config) -> anyhow::Result<()> {
index_s3_with_channel_metadata(config, ChannelMetadata::default()).await
}
#[cfg(feature = "s3")]
pub async fn index_s3_with_channel_metadata(
IndexS3Config {
channel,
credentials,
target_platform,
repodata_patch,
write_zst,
write_shards,
repodata_revisions,
package_revision_assignment,
force,
max_parallel,
multi_progress,
precondition_checks,
}: IndexS3Config,
channel_metadata: ChannelMetadata,
) -> anyhow::Result<()> {
let s3_config = s3_config(&credentials, &channel)?;
let builder = s3_config.into_builder();
let op = Operator::new(builder)?.layer(RetryLayer::new()).finish();
index_with_channel_metadata(
target_platform,
op,
repodata_patch,
write_zst,
write_shards,
repodata_revisions,
package_revision_assignment,
force,
max_parallel,
multi_progress,
precondition_checks,
channel_metadata,
)
.await
.map(|_| ())
}
#[allow(clippy::too_many_arguments)]
pub async fn index(
target_platform: Option<Platform>,
op: Operator,
repodata_patch: Option<String>,
write_zst: bool,
write_shards: bool,
repodata_revisions: Vec<RepodataRevisionInfo>,
package_revision_assignment: PackageRevisionAssignment,
force: bool,
max_parallel: usize,
multi_progress: Option<MultiProgress>,
precondition_checks: PreconditionChecks,
) -> anyhow::Result<IndexStats> {
index_with_channel_metadata(
target_platform,
op,
repodata_patch,
write_zst,
write_shards,
repodata_revisions,
package_revision_assignment,
force,
max_parallel,
multi_progress,
precondition_checks,
ChannelMetadata::default(),
)
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn index_with_channel_metadata(
target_platform: Option<Platform>,
op: Operator,
repodata_patch: Option<String>,
write_zst: bool,
write_shards: bool,
repodata_revisions: Vec<RepodataRevisionInfo>,
package_revision_assignment: PackageRevisionAssignment,
force: bool,
max_parallel: usize,
multi_progress: Option<MultiProgress>,
precondition_checks: PreconditionChecks,
channel_metadata: ChannelMetadata,
) -> anyhow::Result<IndexStats> {
let entries = op.list_with("").await?;
let mut subdirs = if let Some(target_platform) = target_platform {
if !op.exists(&format!("{}/", target_platform.as_str())).await? {
tracing::debug!("Did not find {target_platform} subdir, creating.");
op.create_dir(&format!("{}/", target_platform.as_str()))
.await?;
}
HashSet::from([target_platform])
} else {
entries
.iter()
.filter_map(|entry| {
if entry.metadata().mode().is_dir() && entry.name() != "/" {
Some(entry.name().trim_end_matches('/').to_string())
} else {
None
}
})
.filter_map(|s| Platform::from_str(&s).ok())
.collect::<HashSet<_>>()
};
if !op
.exists(&format!("{}/", Platform::NoArch.as_str()))
.await?
{
tracing::debug!("Did not find noarch subdir, creating.");
op.create_dir(&format!("{}/", Platform::NoArch.as_str()))
.await?;
subdirs.insert(Platform::NoArch);
}
let repodata_patch = if let Some(path) = repodata_patch {
match DistArchiveType::try_from(path.clone()) {
Some(DistArchiveType::Conda(CondaArchiveType::Conda)) => {}
Some(
DistArchiveType::Conda(CondaArchiveType::TarBz2)
| DistArchiveType::Wheel(WheelArchiveType::Whl),
)
| None => {
return Err(anyhow::anyhow!(
"Only .conda packages are supported for repodata patches. Got: {path}",
));
}
}
let repodata_patch_path = format!("noarch/{path}");
let repodata_patch_bytes = op.read(&repodata_patch_path).await?.to_bytes();
let reader = Cursor::new(repodata_patch_bytes);
let repodata_patch = repodata_patch_from_conda_package_stream(reader)?;
Some(repodata_patch)
} else {
None
};
let semaphore = Semaphore::new(max_parallel);
let semaphore = Arc::new(semaphore);
let mut tasks: Vec<(Platform, _)> = Vec::new();
for subdir in subdirs.iter() {
let cache = cache::PackageRecordCache::new();
let task = index_subdir(
*subdir,
op.clone(),
force,
write_zst,
write_shards,
repodata_revisions.clone(),
package_revision_assignment,
channel_metadata.clone(),
repodata_patch
.as_ref()
.and_then(|p| p.subdirs.get(&subdir.to_string()).cloned()),
multi_progress.clone(),
semaphore.clone(),
cache,
precondition_checks,
)
.instrument(tracing::info_span!("index_subdir", subdir = %subdir));
tasks.push((*subdir, task));
}
let mut stats = IndexStats {
subdirs: HashMap::new(),
};
for (subdir, task) in tasks {
match task.await {
Ok(subdir_stats) => {
stats.subdirs.insert(subdir, subdir_stats);
}
Err(e) => {
tracing::error!("Failed to process subdir: {e}");
return Err(e.into());
}
}
}
Ok(stats)
}
pub async fn ensure_channel_initialized(op: &Operator) -> anyhow::Result<()> {
ensure_channel_initialized_with_channel_metadata(op, ChannelMetadata::default()).await
}
pub async fn ensure_channel_initialized_with_channel_metadata(
op: &Operator,
channel_metadata: ChannelMetadata,
) -> anyhow::Result<()> {
let noarch_repodata_path = format!("{}/{REPODATA}", Platform::NoArch.as_str());
if op.exists(&noarch_repodata_path).await? {
tracing::debug!("Channel already initialized");
return Ok(());
}
tracing::info!("Initializing channel with empty noarch/repodata.json");
let noarch_path = format!("{}/", Platform::NoArch.as_str());
if !op.exists(&noarch_path).await? {
op.create_dir(&noarch_path).await?;
}
let empty_repodata = RepoData {
info: Some(ChannelInfo {
subdir: Some(Platform::NoArch.to_string()),
base_url: channel_metadata.base_url,
repodata_revisions: Vec::new(),
channel_relations: channel_metadata.channel_relations,
}),
packages: IndexMap::default(),
conda_packages: IndexMap::default(),
v3: V3Packages::default(),
removed: HashSet::default(),
version: Some(2),
};
let repodata_bytes = serde_json::to_vec(&empty_repodata)?;
match op
.write_with(&noarch_repodata_path, repodata_bytes)
.if_not_exists(true)
.cache_control(CACHE_CONTROL_REPODATA)
.await
{
Ok(_) => {
tracing::info!("Successfully initialized channel");
Ok(())
}
Err(e) if e.kind() == opendal::ErrorKind::ConditionNotMatch => {
tracing::debug!("Channel already initialized by another process");
Ok(())
}
Err(e) => Err(e.into()),
}
}
pub async fn ensure_channel_initialized_fs(channel: &Path) -> anyhow::Result<()> {
ensure_channel_initialized_fs_with_channel_metadata(channel, ChannelMetadata::default()).await
}
pub async fn ensure_channel_initialized_fs_with_channel_metadata(
channel: &Path,
channel_metadata: ChannelMetadata,
) -> anyhow::Result<()> {
let mut config = FsConfig::default();
config.root = Some(channel.canonicalize()?.to_string_lossy().to_string());
let op = Operator::new(config.into_builder())?.finish();
ensure_channel_initialized_with_channel_metadata(&op, channel_metadata).await
}
#[cfg(feature = "s3")]
pub async fn ensure_channel_initialized_s3(
channel: &Url,
credentials: &ResolvedS3Credentials,
) -> anyhow::Result<()> {
ensure_channel_initialized_s3_with_channel_metadata(
channel,
credentials,
ChannelMetadata::default(),
)
.await
}
#[cfg(feature = "s3")]
pub async fn ensure_channel_initialized_s3_with_channel_metadata(
channel: &Url,
credentials: &ResolvedS3Credentials,
channel_metadata: ChannelMetadata,
) -> anyhow::Result<()> {
let s3_config = s3_config(credentials, channel)?;
let op = Operator::new(s3_config.into_builder())?
.layer(RetryLayer::new())
.finish();
ensure_channel_initialized_with_channel_metadata(&op, channel_metadata).await
}
#[cfg(test)]
mod tests {
use std::str::FromStr;
use indexmap::IndexMap;
use rattler_conda_types::Version;
use rattler_conda_types::{
PackageName, UrlOrPath, WhlPackageRecord, package::ArchiveIdentifier,
};
use super::*;
#[test]
fn package_records_from_repodata_preserves_v3_wheels() {
let identifier = ArchiveIdentifier::from_str("demo-1.0-py_0").unwrap();
let package_record = PackageRecord::new(
PackageName::new_unchecked("demo"),
Version::from_str("1.0").unwrap(),
"py_0".to_string(),
);
let wheel_url = UrlOrPath::Path("demo-1.0-py_0.whl".to_string());
let mut repodata = RepoData {
info: None,
packages: IndexMap::default(),
conda_packages: IndexMap::default(),
v3: V3Packages::default(),
removed: HashSet::default(),
version: None,
};
repodata.v3.whl.insert(
identifier.clone(),
WhlPackageRecord {
package_record,
url: wheel_url.clone(),
},
);
let records = package_records_from_repodata(repodata);
let dist_identifier = DistArchiveIdentifier::new(identifier.clone(), WheelArchiveType::Whl);
let indexed_record = records
.get(&dist_identifier)
.expect("v3 wheel should be preserved");
assert_eq!(indexed_record.repodata_revision, RepodataRevision::V3);
assert_eq!(indexed_record.wheel_url, Some(wheel_url));
let (_, indexed_record) = records
.into_iter()
.next()
.expect("v3 wheel should be present");
let mut packages = IndexMap::default();
let mut conda_packages = IndexMap::default();
let mut v3 = V3Packages::default();
insert_package_record_by_revision(
&mut packages,
&mut conda_packages,
&mut v3,
dist_identifier,
indexed_record,
RepodataRevision::V3,
)
.unwrap();
assert!(packages.is_empty());
assert!(conda_packages.is_empty());
assert!(v3.whl.contains_key(&identifier));
}
}