use std::collections::HashMap;
use std::sync::Arc;
use opendal::{Operator, services::GooseFs};
use url::Url;
use crate::object_store::opendal_store::OpendalStore;
use crate::object_store::{
DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore,
ObjectStoreParams, ObjectStoreProvider, StorageOptions,
};
use lance_core::error::{Error, Result};
const DEFAULT_GOOSEFS_PORT: u16 = 9200;
const STORAGE_OPTION_KEYS: &[&str] = &[
"goosefs_master_addr",
"goosefs_root",
"goosefs_write_type",
"goosefs_block_size",
"goosefs_chunk_size",
"goosefs_auth_type",
"goosefs_auth_username",
];
#[derive(Default, Debug)]
pub struct GooseFsStoreProvider;
fn parse_decimal_exponent(s: &str) -> Option<i32> {
if s.is_empty() {
return None;
}
let (negative, digits) = if let Some(rest) = s.strip_prefix('+') {
(false, rest)
} else if let Some(rest) = s.strip_prefix('-') {
(true, rest)
} else {
(false, s)
};
if digits.is_empty() || !digits.chars().all(|c| c.is_ascii_digit()) {
return None;
}
match digits.parse::<i32>() {
Ok(n) => Some(if negative { -n } else { n }),
Err(_) => Some(if negative { i32::MIN } else { i32::MAX }),
}
}
impl GooseFsStoreProvider {
fn validate_storage_option_keys(storage_options: &StorageOptions) -> Result<()> {
let mut wrong_case = Vec::new();
for key in storage_options.0.keys() {
let lower = key.to_ascii_lowercase();
if let Some(&canonical) = STORAGE_OPTION_KEYS
.iter()
.find(|&&canonical| canonical == lower.as_str())
&& key != canonical
{
wrong_case.push((key.clone(), canonical));
}
}
if wrong_case.is_empty() {
return Ok(());
}
wrong_case.sort_by(|a, b| a.1.cmp(b.1));
let details = wrong_case
.iter()
.map(|(got, want)| format!("`{got}` (use `{want}`)"))
.collect::<Vec<_>>()
.join(", ");
Err(Error::invalid_input(format!(
"GooseFS storage_options keys must be lowercase; got {details}"
)))
}
fn resolve_master_addr(url: &Url, storage_options: &StorageOptions) -> Result<String> {
if let Some(addr) = storage_options
.0
.get("goosefs_master_addr")
.filter(|v| !v.is_empty())
{
return Ok(addr.clone());
}
if let Ok(addr) = std::env::var("GOOSEFS_MASTER_ADDR")
&& !addr.is_empty()
{
return Ok(addr);
}
let host = url.host_str().ok_or_else(|| {
Error::invalid_input(
"GooseFS URL must contain a master address (host), e.g. goosefs://host:port/path",
)
})?;
let port = url.port().unwrap_or(DEFAULT_GOOSEFS_PORT);
Ok(format!("{}:{}", host, port))
}
fn resolve_option(
storage_options: &StorageOptions,
option_key: &str,
env_key: &str,
) -> Option<String> {
storage_options
.0
.get(option_key)
.cloned()
.or_else(|| std::env::var(env_key).ok())
.filter(|v| !v.is_empty())
}
fn resolve_root(storage_options: &StorageOptions) -> String {
Self::resolve_option(storage_options, "goosefs_root", "GOOSEFS_ROOT")
.unwrap_or_else(|| "/".to_string())
}
fn parse_space_size(option_key: &str, value: &str) -> Result<u64> {
let trimmed = value.trim();
if trimmed.is_empty() {
return Err(Error::invalid_input(format!(
"{option_key} must be a size such as `4MB` or `4194304`, got `{value}`"
)));
}
let suffix_start = trimmed
.char_indices()
.rev()
.find(|(_, c)| c.is_ascii_digit())
.map(|(i, c)| i + c.len_utf8())
.unwrap_or(0);
let (number, suffix) = trimmed.split_at(suffix_start);
if number.is_empty() {
return Err(Error::invalid_input(format!(
"invalid {option_key} `{value}`: missing numeric coefficient"
)));
}
let multiplier = match suffix.to_ascii_lowercase().as_str() {
"" | "b" => 1u64,
"k" | "kb" => 1024,
"m" | "mb" => 1024 * 1024,
"g" | "gb" => 1024 * 1024 * 1024,
"t" | "tb" => 1024u64.pow(4),
"p" | "pb" => 1024u64.pow(5),
other => {
return Err(Error::invalid_input(format!(
"invalid {option_key} `{value}`: unrecognized size suffix `{other}` \
(supported: b, k/kb, m/mb, g/gb, t/tb, p/pb; units are binary, 1KB=1024)"
)));
}
};
let overflow = || {
Error::invalid_input(format!(
"invalid {option_key} `{value}`: size overflows u64"
))
};
if let Ok(n) = number.parse::<u64>() {
return n.checked_mul(multiplier).ok_or_else(overflow);
}
if number.chars().all(|c| c.is_ascii_digit()) {
return Err(overflow());
}
Self::parse_decimal_times_unit(option_key, value, number, multiplier)
}
fn parse_decimal_times_unit(
option_key: &str,
value: &str,
number: &str,
multiplier: u64,
) -> Result<u64> {
let invalid_number = || {
Error::invalid_input(format!(
"invalid {option_key} `{value}`: `{number}` is not a valid number"
))
};
let overflow = || {
Error::invalid_input(format!(
"invalid {option_key} `{value}`: size overflows u64"
))
};
let negative = || {
Error::invalid_input(format!(
"invalid {option_key} `{value}`: size must be a non-negative number"
))
};
if number.starts_with('-') {
return Err(negative());
}
let s = number.strip_prefix('+').unwrap_or(number);
if s.is_empty() {
return Err(invalid_number());
}
let (mantissa, exp) = if let Some(e_idx) = s.find(['e', 'E']) {
let mantissa = &s[..e_idx];
let exp_str = &s[e_idx + 1..];
if mantissa.is_empty() {
return Err(invalid_number());
}
let exp = parse_decimal_exponent(exp_str).ok_or_else(invalid_number)?;
(mantissa, exp)
} else {
(s, 0i32)
};
let mut parts = mantissa.split('.');
let int_str = parts.next().unwrap_or("");
let frac_raw = parts.next();
if parts.next().is_some() {
return Err(invalid_number());
}
if !int_str.chars().all(|c| c.is_ascii_digit()) {
return Err(invalid_number());
}
let frac_raw = match frac_raw {
Some(frac) if frac.chars().all(|c| c.is_ascii_digit()) => frac,
Some(_) => return Err(invalid_number()),
None => "",
};
if int_str.is_empty() && frac_raw.is_empty() {
return Err(invalid_number());
}
let frac_str = frac_raw.trim_end_matches('0');
let int_part: u128 = if int_str.is_empty() {
0
} else {
int_str.parse().map_err(|_| overflow())?
};
let frac_digits = u32::try_from(frac_str.len()).map_err(|_| overflow())?;
let frac_part: u128 = if frac_str.is_empty() {
0
} else {
frac_str.parse().map_err(|_| overflow())?
};
let significand = if frac_digits == 0 {
int_part
} else {
let pow10 = 10u128.checked_pow(frac_digits).ok_or_else(overflow)?;
int_part
.checked_mul(pow10)
.and_then(|v| v.checked_add(frac_part))
.ok_or_else(overflow)?
};
if significand == 0 {
return Ok(0);
}
let mut num = significand
.checked_mul(u128::from(multiplier))
.ok_or_else(overflow)?;
let scale = i64::from(frac_digits) - i64::from(exp);
if scale > 0 {
match u32::try_from(scale)
.ok()
.and_then(|s| 10u128.checked_pow(s))
{
Some(den) => num /= den,
None => num = 0,
}
} else if scale < 0 {
let raise = u32::try_from(scale.unsigned_abs()).map_err(|_| overflow())?;
let factor = 10u128.checked_pow(raise).ok_or_else(overflow)?;
num = num.checked_mul(factor).ok_or_else(overflow)?;
}
u64::try_from(num).map_err(|_| overflow())
}
fn resolve_space_size(
storage_options: &StorageOptions,
option_key: &str,
env_key: &str,
) -> Result<Option<u64>> {
let Some(raw) = Self::resolve_option(storage_options, option_key, env_key) else {
return Ok(None);
};
Self::parse_space_size(option_key, &raw).map(Some)
}
}
#[async_trait::async_trait]
impl ObjectStoreProvider for GooseFsStoreProvider {
async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result<ObjectStore> {
let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE);
let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default());
Self::validate_storage_option_keys(&storage_options)?;
let master_addr = Self::resolve_master_addr(&base_path, &storage_options)?;
let root = Self::resolve_root(&storage_options);
let mut config_map: HashMap<String, String> = HashMap::new();
config_map.insert("master_addr".to_string(), master_addr);
config_map.insert("root".to_string(), root);
if let Some(wt) =
Self::resolve_option(&storage_options, "goosefs_write_type", "GOOSEFS_WRITE_TYPE")
{
config_map.insert("write_type".to_string(), wt);
}
if let Some(bs) =
Self::resolve_space_size(&storage_options, "goosefs_block_size", "GOOSEFS_BLOCK_SIZE")?
{
config_map.insert("block_size".to_string(), bs.to_string());
}
if let Some(cs) =
Self::resolve_space_size(&storage_options, "goosefs_chunk_size", "GOOSEFS_CHUNK_SIZE")?
{
config_map.insert("chunk_size".to_string(), cs.to_string());
}
if let Some(at) =
Self::resolve_option(&storage_options, "goosefs_auth_type", "GOOSEFS_AUTH_TYPE")
{
config_map.insert("auth_type".to_string(), at);
}
if let Some(au) = Self::resolve_option(
&storage_options,
"goosefs_auth_username",
"GOOSEFS_AUTH_USERNAME",
) {
config_map.insert("auth_username".to_string(), au);
}
let operator = Operator::from_iter::<GooseFs>(config_map).map_err(|e| {
Error::invalid_input(format!("Failed to create GooseFS operator: {:?}", e))
})?;
let opendal_store = Arc::new(OpendalStore::new(operator));
Ok(ObjectStore {
scheme: "goosefs".to_string(),
inner: opendal_store,
local_dir_operations: None,
block_size,
max_iop_size: *DEFAULT_MAX_IOP_SIZE,
use_constant_size_upload_parts: params.use_constant_size_upload_parts,
list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(false),
io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
download_retry_count: storage_options.download_retry_count(),
io_tracker: Default::default(),
store_prefix: self
.calculate_object_store_prefix(&base_path, params.storage_options())?,
paginated_lister: None,
})
}
fn calculate_object_store_prefix(
&self,
url: &Url,
storage_options: Option<&HashMap<String, String>>,
) -> Result<String> {
let opts = StorageOptions(storage_options.cloned().unwrap_or_default());
let root = Self::resolve_root(&opts);
if root == "/" {
Ok(format!("{}${}", url.scheme(), url.authority()))
} else {
Ok(format!("{}${}#{}", url.scheme(), url.authority(), root))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[test]
fn test_goosefs_extract_path_basic() {
let provider = GooseFsStoreProvider;
let url = Url::parse("goosefs://10.0.0.1:9200/data/embeddings.lance").unwrap();
let path = provider.extract_path(&url).unwrap();
assert_eq!(path.to_string(), "data/embeddings.lance");
}
#[test]
fn test_goosefs_extract_path_root() {
let provider = GooseFsStoreProvider;
let url = Url::parse("goosefs://10.0.0.1:9200/").unwrap();
let path = provider.extract_path(&url).unwrap();
assert_eq!(path.to_string(), "");
}
#[test]
fn test_goosefs_extract_path_deep() {
let provider = GooseFsStoreProvider;
let url = Url::parse("goosefs://master:9200/a/b/c/d.lance").unwrap();
let path = provider.extract_path(&url).unwrap();
assert_eq!(path.to_string(), "a/b/c/d.lance");
}
#[test]
fn test_goosefs_extract_path_percent_decoded() {
let provider = GooseFsStoreProvider;
let url = Url::parse("goosefs://master:9200/dir/with%20space/f.lance").unwrap();
let path = provider.extract_path(&url).unwrap();
assert_eq!(path.to_string(), "dir/with space/f.lance");
}
#[test]
fn test_calculate_object_store_prefix_default_root() {
let provider = GooseFsStoreProvider;
let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap();
let prefix = provider.calculate_object_store_prefix(&url, None).unwrap();
assert_eq!(prefix, "goosefs$10.0.0.1:9200");
}
#[test]
fn test_calculate_object_store_prefix_with_hostname() {
let provider = GooseFsStoreProvider;
let url = Url::parse("goosefs://myhost:9200/data").unwrap();
let prefix = provider.calculate_object_store_prefix(&url, None).unwrap();
assert_eq!(prefix, "goosefs$myhost:9200");
}
#[test]
fn test_prefix_shared_across_datasets_same_master() {
let provider = GooseFsStoreProvider;
let url_a = Url::parse("goosefs://10.0.0.1:9200/repro/a.lance").unwrap();
let url_b = Url::parse("goosefs://10.0.0.1:9200/repro/b.lance").unwrap();
let pa = provider
.calculate_object_store_prefix(&url_a, None)
.unwrap();
let pb = provider
.calculate_object_store_prefix(&url_b, None)
.unwrap();
assert_eq!(pa, pb, "same master must share one cache prefix");
assert_ne!(
provider.extract_path(&url_a).unwrap(),
provider.extract_path(&url_b).unwrap(),
"distinct URLs must yield distinct object keys",
);
}
#[test]
fn test_prefix_isolated_across_masters() {
let provider = GooseFsStoreProvider;
let u1 = Url::parse("goosefs://host-a:9200/x.lance").unwrap();
let u2 = Url::parse("goosefs://host-b:9200/x.lance").unwrap();
assert_ne!(
provider.calculate_object_store_prefix(&u1, None).unwrap(),
provider.calculate_object_store_prefix(&u2, None).unwrap(),
);
}
#[test]
fn test_prefix_includes_custom_root() {
let provider = GooseFsStoreProvider;
let url = Url::parse("goosefs://host:9200/x.lance").unwrap();
let default_prefix = provider.calculate_object_store_prefix(&url, None).unwrap();
let custom_opts: HashMap<String, String> =
HashMap::from([("goosefs_root".to_string(), "/tenant-a".to_string())]);
let custom_prefix = provider
.calculate_object_store_prefix(&url, Some(&custom_opts))
.unwrap();
assert_eq!(default_prefix, "goosefs$host:9200");
assert_eq!(custom_prefix, "goosefs$host:9200#/tenant-a");
assert_ne!(default_prefix, custom_prefix);
}
#[test]
fn test_resolve_master_addr_from_url() {
let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap();
let storage_options = StorageOptions(HashMap::new());
let addr = GooseFsStoreProvider::resolve_master_addr(&url, &storage_options).unwrap();
assert_eq!(addr, "10.0.0.1:9200");
}
#[test]
fn test_resolve_master_addr_default_port() {
let url = Url::parse("goosefs://10.0.0.1/data").unwrap();
let storage_options = StorageOptions(HashMap::new());
let addr = GooseFsStoreProvider::resolve_master_addr(&url, &storage_options).unwrap();
assert_eq!(addr, "10.0.0.1:9200");
}
#[test]
fn test_resolve_master_addr_from_storage_options() {
let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap();
let storage_options = StorageOptions(HashMap::from([(
"goosefs_master_addr".to_string(),
"10.0.0.2:9200,10.0.0.3:9200".to_string(),
)]));
let addr = GooseFsStoreProvider::resolve_master_addr(&url, &storage_options).unwrap();
assert_eq!(addr, "10.0.0.2:9200,10.0.0.3:9200");
}
#[test]
fn test_resolve_root_defaults_to_slash() {
let opts = StorageOptions(HashMap::new());
assert_eq!(GooseFsStoreProvider::resolve_root(&opts), "/");
}
#[test]
fn test_resolve_root_from_storage_options() {
let opts = StorageOptions(HashMap::from([(
"goosefs_root".to_string(),
"/tenant-a".to_string(),
)]));
assert_eq!(GooseFsStoreProvider::resolve_root(&opts), "/tenant-a");
}
#[test]
fn test_validate_storage_option_keys_accepts_lowercase() {
let opts = StorageOptions(HashMap::from([
(
"goosefs_master_addr".to_string(),
"10.0.0.1:9200,10.0.0.2:9200".to_string(),
),
("goosefs_auth_type".to_string(), "simple".to_string()),
("goosefs_auth_username".to_string(), "lance".to_string()),
("allow_http".to_string(), "true".to_string()),
]));
GooseFsStoreProvider::validate_storage_option_keys(&opts).unwrap();
}
#[test]
fn test_validate_storage_option_keys_rejects_uppercase() {
let opts = StorageOptions(HashMap::from([(
"GOOSEFS_MASTER_ADDR".to_string(),
"10.0.0.1:9200,10.0.0.2:9200".to_string(),
)]));
let err = GooseFsStoreProvider::validate_storage_option_keys(&opts).unwrap_err();
let msg = err.to_string();
assert!(
msg.contains("must be lowercase"),
"expected lowercase-key error, got: {msg}"
);
assert!(
msg.contains("`GOOSEFS_MASTER_ADDR` (use `goosefs_master_addr`)"),
"expected canonical key hint, got: {msg}"
);
}
#[test]
fn test_validate_storage_option_keys_rejects_mixed_case() {
let opts = StorageOptions(HashMap::from([
(
"Goosefs_Master_Addr".to_string(),
"10.0.0.1:9200".to_string(),
),
("GOOSEFS_AUTH_TYPE".to_string(), "simple".to_string()),
]));
let err = GooseFsStoreProvider::validate_storage_option_keys(&opts).unwrap_err();
let msg = err.to_string();
assert!(msg.contains("must be lowercase"), "got: {msg}");
assert!(msg.contains("goosefs_master_addr"), "got: {msg}");
assert!(msg.contains("goosefs_auth_type"), "got: {msg}");
}
#[rstest]
#[case::bare_bytes("4096", 4096)]
#[case::suffix_b("4096b", 4096)]
#[case::suffix_k("4k", 4096)]
#[case::suffix_kb("4KB", 4096)]
#[case::suffix_m("4m", 4 * 1024 * 1024)]
#[case::suffix_mb("4MB", 4 * 1024 * 1024)]
#[case::suffix_gb("1GB", 1024 * 1024 * 1024)]
#[case::fractional_mb("1.5MB", 1_572_864)]
#[case::fractional_kb("2.5KB", 2560)]
#[case::scientific_kb("1e3KB", 1000 * 1024)]
#[case::leading_dot(".5KB", 512)]
#[case::truncate_toward_zero("1.9", 1)]
#[case::trimmed(" 4MB ", 4 * 1024 * 1024)]
fn test_parse_space_size_accepts_goosefs_suffixes(#[case] input: &str, #[case] expected: u64) {
assert_eq!(
GooseFsStoreProvider::parse_space_size("goosefs_chunk_size", input).unwrap(),
expected
);
}
#[rstest]
#[case::empty("")]
#[case::whitespace(" ")]
#[case::missing_number("MB")]
#[case::unknown_suffix("4MiB")]
#[case::si_mismatch_not_mib("4XB")]
#[case::negative("-4MB")]
#[case::two_dots("1.2.3")]
fn test_parse_space_size_rejects_invalid(#[case] input: &str) {
let err = GooseFsStoreProvider::parse_space_size("goosefs_chunk_size", input).unwrap_err();
let msg = err.to_string();
assert!(
msg.contains("goosefs_chunk_size"),
"error should name the option, got: {msg}"
);
assert!(
msg.contains(input.trim()) || input.trim().is_empty(),
"error should include the input, got: {msg}"
);
}
#[rstest]
#[case::u64_max_plus_one("18446744073709551616")]
#[case::u64_max_plus_one_bytes("18446744073709551616b")]
#[case::integer_mul_overflow("16384PB")]
#[case::fractional_mul_overflow("16384.0PB")]
#[case::scientific_overflow("1e20")]
#[case::scientific_u64_max_plus_one("1.8446744073709551616e19")]
fn test_parse_space_size_rejects_overflow(#[case] input: &str) {
let err = GooseFsStoreProvider::parse_space_size("goosefs_chunk_size", input).unwrap_err();
assert!(
matches!(err, Error::InvalidInput { .. }),
"expected InvalidInput, got: {err:?}"
);
let msg = err.to_string();
assert!(
msg.contains("overflows u64"),
"overflowing size should be rejected, got: {msg}"
);
assert!(
msg.contains(input),
"error should include the input, got: {msg}"
);
}
#[test]
fn test_parse_space_size_accepts_u64_max_integer() {
assert_eq!(
GooseFsStoreProvider::parse_space_size("goosefs_chunk_size", "18446744073709551615")
.unwrap(),
u64::MAX
);
assert_eq!(
GooseFsStoreProvider::parse_space_size(
"goosefs_chunk_size",
"1.8446744073709551615e19"
)
.unwrap(),
u64::MAX
);
}
#[test]
fn test_parse_space_size_preserves_values_beyond_f64_precision() {
assert_eq!(
GooseFsStoreProvider::parse_space_size("goosefs_chunk_size", "9007199254740993.0")
.unwrap(),
9007199254740993
);
assert_eq!(
GooseFsStoreProvider::parse_space_size("goosefs_chunk_size", "9007199254740993.0b")
.unwrap(),
9007199254740993
);
}
#[test]
fn test_parse_space_size_four_mb_is_binary_not_si() {
let bytes = GooseFsStoreProvider::parse_space_size("goosefs_chunk_size", "4MB").unwrap();
assert_eq!(bytes, 4_194_304, "GooseFS MB is 1024^2, not 10^6");
assert_ne!(bytes, 4_000_000);
}
#[test]
fn test_resolve_space_size_from_storage_options() {
let opts = StorageOptions(HashMap::from([(
"goosefs_chunk_size".to_string(),
"4MB".to_string(),
)]));
let bytes = GooseFsStoreProvider::resolve_space_size(
&opts,
"goosefs_chunk_size",
"GOOSEFS_CHUNK_SIZE",
)
.unwrap();
assert_eq!(bytes, Some(4 * 1024 * 1024));
}
}