JMAP Registry API implementation - part 2
This commit is contained in:
mdecimus
@@ -25,7 +25,9 @@ use nlp::classifier::reservoir::SampleReservoir;
|
||||
use nlp::classifier::train::{CcfhTrainer, FhTrainer};
|
||||
use nlp::tokenizers::types::TypesTokenizer;
|
||||
use nlp::tokenizers::{stream::WordStemTokenizer, types::TokenType};
|
||||
use registry::schema::prelude::ObjectType;
|
||||
use registry::schema::prelude::{ObjectType, Property};
|
||||
use registry::schema::structs::SpamTrainingSample;
|
||||
use registry::types::EnumImpl;
|
||||
use std::time::Instant;
|
||||
use std::{
|
||||
borrow::Cow,
|
||||
@@ -33,15 +35,17 @@ use std::{
|
||||
hash::{Hash, RandomState},
|
||||
sync::Arc,
|
||||
};
|
||||
use store::ahash::AHashSet;
|
||||
use store::rand::seq::SliceRandom;
|
||||
use store::write::{BlobLink, now};
|
||||
use store::write::{BlobLink, RegistryClass, now};
|
||||
use store::{
|
||||
Deserialize, IterateParams, Serialize, U32_LEN, U64_LEN, ValueKey,
|
||||
Deserialize, IterateParams, Serialize, ValueKey,
|
||||
write::{
|
||||
AlignedBytes, Archive, Archiver, BatchBuilder, BlobOp, ValueClass,
|
||||
key::DeserializeBigEndian,
|
||||
},
|
||||
};
|
||||
use store::{SerializeInfallible, U16_LEN};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use trc::{AddContext, SpamEvent};
|
||||
use types::blob_hash::BlobHash;
|
||||
@@ -63,13 +67,25 @@ pub trait SpamClassifier {
|
||||
) -> impl Future<Output = Tokens<'x>> + Send;
|
||||
}
|
||||
|
||||
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Clone, PartialEq, Eq, Debug)]
|
||||
#[derive(
|
||||
rkyv::Archive,
|
||||
rkyv::Deserialize,
|
||||
rkyv::Serialize,
|
||||
Clone,
|
||||
PartialEq,
|
||||
Eq,
|
||||
Debug,
|
||||
PartialOrd,
|
||||
Ord,
|
||||
Hash,
|
||||
)]
|
||||
pub struct TrainingSample {
|
||||
hash: BlobHash,
|
||||
account_id: u32,
|
||||
}
|
||||
|
||||
struct TrainingTask {
|
||||
id: u64,
|
||||
sample: TrainingSample,
|
||||
is_spam: bool,
|
||||
is_replay: bool,
|
||||
@@ -80,7 +96,7 @@ struct TrainingTask {
|
||||
pub struct SpamTrainer {
|
||||
pub trainer: SpamTrainerClass,
|
||||
pub reservoir: SampleReservoir<TrainingSample>,
|
||||
pub last_sample_expiry: u64,
|
||||
pub last_id: u64,
|
||||
}
|
||||
|
||||
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
|
||||
@@ -136,7 +152,7 @@ impl SpamClassifier for Server {
|
||||
)))),
|
||||
},
|
||||
reservoir: SampleReservoir::default(),
|
||||
last_sample_expiry: 0,
|
||||
last_id: 0,
|
||||
}
|
||||
};
|
||||
|
||||
@@ -169,77 +185,79 @@ impl SpamClassifier for Server {
|
||||
|
||||
// Fetch blob hashes for samples
|
||||
let mut samples = Vec::new();
|
||||
let mut duplicate_samples = Vec::new();
|
||||
let mut remove_entries = false;
|
||||
let from_key = ValueKey {
|
||||
account_id: 0,
|
||||
collection: 0,
|
||||
document_id: 0,
|
||||
class: ValueClass::Blob(BlobOp::SpamSample {
|
||||
hash: BlobHash::default(),
|
||||
until: trainer.last_sample_expiry + 1,
|
||||
}),
|
||||
};
|
||||
let to_key = ValueKey {
|
||||
account_id: u32::MAX,
|
||||
collection: u8::MAX,
|
||||
document_id: u32::MAX,
|
||||
class: ValueClass::Blob(BlobOp::SpamSample {
|
||||
hash: BlobHash::new_max(),
|
||||
until: u64::MAX,
|
||||
}),
|
||||
};
|
||||
let object_id = ObjectType::SpamTrainingSample.to_id();
|
||||
let from_key = ValueKey::from(ValueClass::Registry(RegistryClass::Id {
|
||||
object_id,
|
||||
item_id: trainer.last_id + 1,
|
||||
}));
|
||||
let to_key = ValueKey::from(ValueClass::Registry(RegistryClass::Id {
|
||||
object_id,
|
||||
item_id: u64::MAX,
|
||||
}));
|
||||
let mut seen_samples = AHashSet::new();
|
||||
let mut spam_count = 0;
|
||||
let mut ham_count = 0;
|
||||
self.store()
|
||||
.iterate(
|
||||
IterateParams::new(from_key, to_key).ascending(),
|
||||
IterateParams::new(from_key, to_key).descending(),
|
||||
|key, value| {
|
||||
let until = key.deserialize_be_u64(1)?;
|
||||
let account_id = key.deserialize_be_u32(U64_LEN + 1)?;
|
||||
let hash = BlobHash::try_from_hash_slice(
|
||||
key.get(U64_LEN + U32_LEN + 1..).ok_or_else(|| {
|
||||
trc::Error::corrupted_key(key, value.into(), trc::location!())
|
||||
})?,
|
||||
)
|
||||
.unwrap();
|
||||
let (Some(is_spam), Some(hold)) = (value.first(), value.get(1)) else {
|
||||
return Err(trc::Error::corrupted_key(
|
||||
key,
|
||||
value.into(),
|
||||
trc::location!(),
|
||||
));
|
||||
let id = key.deserialize_be_u64(U16_LEN + 1)?;
|
||||
let sample = SpamTrainingSample::deserialize(value)?;
|
||||
|
||||
let until = sample.expires_at.timestamp() as u64;
|
||||
let do_remove = sample.delete_after_use;
|
||||
let is_spam = sample.is_spam;
|
||||
let sample = TrainingSample {
|
||||
hash: sample.blob_id.hash,
|
||||
account_id: sample
|
||||
.account_id
|
||||
.map(|a| a.document_id())
|
||||
.unwrap_or(u32::MAX),
|
||||
};
|
||||
|
||||
let do_remove = *hold == 0;
|
||||
let is_spam = *is_spam == 1;
|
||||
let sample = TrainingSample { hash, account_id };
|
||||
if seen_samples.insert(sample.clone()) {
|
||||
// Add to reservoir
|
||||
if !do_remove {
|
||||
trainer.reservoir.update_reservoir(
|
||||
&sample,
|
||||
is_spam,
|
||||
config.reservoir_capacity,
|
||||
);
|
||||
} else {
|
||||
trainer.reservoir.update_counts(is_spam);
|
||||
}
|
||||
|
||||
// Add to reservoir
|
||||
if !do_remove {
|
||||
trainer.reservoir.update_reservoir(
|
||||
&sample,
|
||||
samples.push(TrainingTask {
|
||||
id,
|
||||
sample,
|
||||
is_spam,
|
||||
config.reservoir_capacity,
|
||||
);
|
||||
is_replay: false,
|
||||
remove: do_remove.then_some(until),
|
||||
});
|
||||
|
||||
remove_entries |= do_remove;
|
||||
|
||||
// Update trainer stats
|
||||
if is_spam {
|
||||
spam_count += 1;
|
||||
} else {
|
||||
ham_count += 1;
|
||||
}
|
||||
} else {
|
||||
trainer.reservoir.update_counts(is_spam);
|
||||
duplicate_samples.push(TrainingTask {
|
||||
id,
|
||||
sample,
|
||||
is_spam,
|
||||
is_replay: false,
|
||||
remove: Some(until),
|
||||
});
|
||||
remove_entries = true;
|
||||
}
|
||||
|
||||
samples.push(TrainingTask {
|
||||
sample,
|
||||
is_spam,
|
||||
is_replay: false,
|
||||
remove: do_remove.then_some(until),
|
||||
});
|
||||
|
||||
remove_entries |= do_remove;
|
||||
|
||||
// Update trainer stats
|
||||
trainer.last_sample_expiry = until;
|
||||
if is_spam {
|
||||
spam_count += 1;
|
||||
} else {
|
||||
ham_count += 1;
|
||||
if trainer.last_id == 0 {
|
||||
trainer.last_id = id;
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
@@ -284,6 +302,7 @@ impl SpamClassifier for Server {
|
||||
.reservoir
|
||||
.replay_samples((spam_count - ham_count) as usize, false)
|
||||
.map(|sample| TrainingTask {
|
||||
id: 0,
|
||||
sample: sample.clone(),
|
||||
is_spam: false,
|
||||
is_replay: true,
|
||||
@@ -297,6 +316,7 @@ impl SpamClassifier for Server {
|
||||
.reservoir
|
||||
.replay_samples((ham_count - spam_count) as usize, true)
|
||||
.map(|sample| TrainingTask {
|
||||
id: 0,
|
||||
sample: sample.clone(),
|
||||
is_spam: true,
|
||||
is_replay: true,
|
||||
@@ -536,18 +556,26 @@ impl SpamClassifier for Server {
|
||||
// Remove samples marked for deletion
|
||||
if remove_entries {
|
||||
let mut batch = BatchBuilder::new();
|
||||
for sample in samples {
|
||||
for sample in samples.into_iter().chain(duplicate_samples.into_iter()) {
|
||||
if let Some(until) = sample.remove {
|
||||
batch
|
||||
.with_account_id(sample.sample.account_id)
|
||||
.clear(BlobOp::Link {
|
||||
hash: sample.sample.hash.clone(),
|
||||
hash: sample.sample.hash,
|
||||
to: BlobLink::Temporary { until },
|
||||
})
|
||||
.clear(BlobOp::SpamSample {
|
||||
hash: sample.sample.hash,
|
||||
until,
|
||||
});
|
||||
.clear(ValueClass::Registry(RegistryClass::Id {
|
||||
object_id,
|
||||
item_id: sample.id,
|
||||
}));
|
||||
if sample.sample.account_id != u32::MAX {
|
||||
batch.clear(ValueClass::Registry(RegistryClass::Index {
|
||||
index_id: Property::AccountId.to_id(),
|
||||
object_id,
|
||||
item_id: sample.id,
|
||||
key: sample.sample.account_id.serialize(),
|
||||
}));
|
||||
}
|
||||
if batch.is_large_batch() {
|
||||
self.store()
|
||||
.write(batch.build_all())
|
||||
|
||||
Reference in New Issue
Block a user