Manage key length conversion error correctly

This commit is contained in:
Clément Renault 2024-12-02 11:02:49 +01:00
parent 5e218f3f4d
commit d5c07ef7b3
No known key found for this signature in database
GPG Key ID: F250A4C4E3AE5F5F
3 changed files with 48 additions and 17 deletions

View File

@ -3,6 +3,7 @@ use std::convert::Infallible;
use std::fmt::Write;
use std::{io, str};
use bstr::BString;
use heed::{Error as HeedError, MdbError};
use rayon::ThreadPoolBuildError;
use rhai::EvalAltResult;
@ -62,14 +63,9 @@ pub enum InternalError {
#[error(transparent)]
Store(#[from] MdbError),
#[error("Cannot delete {key:?} from database {database_name}: {error}")]
StoreDeletion { database_name: &'static str, key: Box<[u8]>, error: heed::Error },
StoreDeletion { database_name: &'static str, key: BString, error: heed::Error },
#[error("Cannot insert {key:?} and value with length {value_length} into database {database_name}: {error}")]
StorePut {
database_name: &'static str,
key: Box<[u8]>,
value_length: usize,
error: heed::Error,
},
StorePut { database_name: &'static str, key: BString, value_length: usize, error: heed::Error },
#[error(transparent)]
Utf8(#[from] str::Utf8Error),
#[error("An indexation process was explicitly aborted")]

View File

@ -10,7 +10,7 @@ use bbqueue::BBBuffer;
use bytemuck::{checked, CheckedBitPattern, NoUninit};
use flume::{RecvTimeoutError, SendError};
use heed::types::Bytes;
use heed::BytesDecode;
use heed::{BytesDecode, MdbError};
use memmap2::{Mmap, MmapMut};
use roaring::RoaringBitmap;
@ -23,7 +23,7 @@ use crate::index::db_name;
use crate::index::main_key::{GEO_FACETED_DOCUMENTS_IDS_KEY, GEO_RTREE_KEY};
use crate::update::new::KvReaderFieldId;
use crate::vector::Embedding;
use crate::{CboRoaringBitmapCodec, DocumentId, Index};
use crate::{CboRoaringBitmapCodec, DocumentId, Index, InternalError};
/// Creates a tuple of senders/receiver to be used by
/// the extractors and the writer loop.
@ -524,7 +524,14 @@ impl<'b> ExtractorBbqueueSender<'b> {
}
fn write_key_value(&self, database: Database, key: &[u8], value: &[u8]) -> crate::Result<()> {
let key_length = NonZeroU16::new(key.len().try_into().unwrap()).unwrap();
let key_length = key.len().try_into().ok().and_then(NonZeroU16::new).ok_or_else(|| {
InternalError::StorePut {
database_name: database.database_name(),
key: key.into(),
value_length: value.len(),
error: MdbError::BadValSize.into(),
}
})?;
self.write_key_value_with(database, key_length, value.len(), |key_buffer, value_buffer| {
key_buffer.copy_from_slice(key);
value_buffer.copy_from_slice(value);
@ -587,7 +594,13 @@ impl<'b> ExtractorBbqueueSender<'b> {
}
fn delete_entry(&self, database: Database, key: &[u8]) -> crate::Result<()> {
let key_length = NonZeroU16::new(key.len().try_into().unwrap()).unwrap();
let key_length = key.len().try_into().ok().and_then(NonZeroU16::new).ok_or_else(|| {
InternalError::StoreDeletion {
database_name: database.database_name(),
key: key.into(),
error: MdbError::BadValSize.into(),
}
})?;
self.delete_entry_with(database, key_length, |buffer| {
buffer.copy_from_slice(key);
Ok(())
@ -702,8 +715,15 @@ pub struct WordDocidsSender<'a, 'b, D> {
impl<D: DatabaseType> WordDocidsSender<'_, '_, D> {
pub fn write(&self, key: &[u8], bitmap: &RoaringBitmap) -> crate::Result<()> {
let key_length = NonZeroU16::new(key.len().try_into().unwrap()).unwrap();
let value_length = CboRoaringBitmapCodec::serialized_size(bitmap);
let key_length = key.len().try_into().ok().and_then(NonZeroU16::new).ok_or_else(|| {
InternalError::StorePut {
database_name: D::DATABASE.database_name(),
key: key.into(),
value_length,
error: MdbError::BadValSize.into(),
}
})?;
self.sender.write_key_value_with(
D::DATABASE,
key_length,
@ -731,7 +751,6 @@ impl FacetDocidsSender<'_, '_> {
let (facet_kind, key) = FacetKind::extract_from_key(key);
let database = Database::from(facet_kind);
let key_length = NonZeroU16::new(key.len().try_into().unwrap()).unwrap();
let value_length = CboRoaringBitmapCodec::serialized_size(bitmap);
let value_length = match facet_kind {
// We must take the facet group size into account
@ -739,6 +758,14 @@ impl FacetDocidsSender<'_, '_> {
FacetKind::Number | FacetKind::String => value_length + 1,
FacetKind::Null | FacetKind::Empty | FacetKind::Exists => value_length,
};
let key_length = key.len().try_into().ok().and_then(NonZeroU16::new).ok_or_else(|| {
InternalError::StorePut {
database_name: database.database_name(),
key: key.into(),
value_length,
error: MdbError::BadValSize.into(),
}
})?;
self.sender.write_key_value_with(
database,
@ -862,12 +889,20 @@ impl GeoSender<'_, '_> {
}
pub fn set_geo_faceted(&self, bitmap: &RoaringBitmap) -> crate::Result<()> {
let key = GEO_FACETED_DOCUMENTS_IDS_KEY.as_bytes();
let key_length = NonZeroU16::new(key.len().try_into().unwrap()).unwrap();
let database = Database::Main;
let value_length = bitmap.serialized_size();
let key = GEO_FACETED_DOCUMENTS_IDS_KEY.as_bytes();
let key_length = key.len().try_into().ok().and_then(NonZeroU16::new).ok_or_else(|| {
InternalError::StorePut {
database_name: database.database_name(),
key: key.into(),
value_length,
error: MdbError::BadValSize.into(),
}
})?;
self.0.write_key_value_with(
Database::Main,
database,
key_length,
value_length,
|key_buffer, value_buffer| {

View File

@ -432,7 +432,7 @@ where
if let Err(error) = database.put(wtxn, &key, &value) {
return Err(Error::InternalError(InternalError::StorePut {
database_name,
key,
key: bstr::BString::from(&key[..]),
value_length: value.len(),
error,
}));