Add handle refcount, delete on full removal

profiles/src/store/mod.rs

@@ -559,7 +559,7 @@ fn count(
     tree: &sled::transaction::TransactionalTree,
     key: &str,
     f: impl Fn(u64) -> u64,
-) -> Result<(), sled::transaction::ConflictableTransactionError> {
+) -> Result<u64, sled::transaction::ConflictableTransactionError> {
     let count = match tree.get(key.as_bytes())? {
         Some(ivec) => {
             let s = String::from_utf8_lossy(&ivec);
@@ -569,11 +569,12 @@ fn count(
         None => 0,
     };
 
-    let count = (f)(count).to_string();
+    let count = (f)(count);
+    let count_string = count.to_string();
 
-    tree.insert(key.as_bytes(), count.as_bytes())?;
+    tree.insert(key.as_bytes(), count_string.as_bytes())?;
 
-    Ok(())
+    Ok(count)
 }
 
 impl fmt::Display for OwnerSource {

profiles/src/store/profile.rs

@@ -237,7 +237,7 @@ impl Store {
         domain: &str,
         published: DateTime<Utc>,
     ) -> Result<Profile, StoreError> {
-        let handle = handle.to_lowercase();
+        let handle_lower = handle.to_lowercase();
 
         let mut id;
         let mut stored_profile;
@@ -285,7 +285,7 @@ impl Store {
 
         // ensure handle uniqueness
         match self.handle_tree.compare_and_swap(
-            handle_id_key(&handle, domain),
+            handle_id_key(&handle_lower, domain),
             None as Option<&[u8]>,
             Some(id.as_bytes()),
         ) {
@@ -323,8 +323,8 @@ impl Store {
             Ok(())
         });
 
-        if let Err(e) = self.handle_index.insert(&handle) {
+        if let Err(e) = self.handle_index.insert(&handle_lower) {
-            log::error!("Failed to add {} to search index: {}", handle, e);
+            log::error!("Failed to add {} to search index: {}", handle_lower, e);
         }
 
         if let Err(e) = res {
@@ -333,7 +333,7 @@ impl Store {
                 let handle_tree = &trees[1];
 
                 profile_tree.remove(id_profile_key(id).as_bytes())?;
-                handle_tree.remove(handle_id_key(&handle, domain).as_bytes())?;
+                handle_tree.remove(handle_id_key(&handle_lower, domain).as_bytes())?;
 
                 Ok(())
             })?;
@@ -558,6 +558,14 @@ impl Store {
         stored_profile.updated = Some(now);
         let stored_profile_vec = serde_json::to_vec(&stored_profile)?;
 
+        if let Err(e) = self.handle_index.remove(&stored_profile.handle) {
+            log::error!(
+                "Failed to remove {} from handle index: {}",
+                stored_profile.handle,
+                e
+            );
+        }
+
         if self
             .profile_tree
             .compare_and_swap(

profiles/src/store/term_search.rs

@@ -1,10 +1,11 @@
-use super::StoreError;
+use super::{count, StoreError};
-use sled::{Db, Tree};
+use sled::{Db, Transactional, Tree};
 use std::collections::HashSet;
 
 #[derive(Clone, Debug)]
 pub struct TermSearch {
     trees: Vec<Tree>,
+    counts: Tree,
 }
 
 impl TermSearch {
@@ -15,11 +16,14 @@ impl TermSearch {
             trees.push(db.open_tree(&format!("/profiles/search/{}/{}", name, i))?);
         }
 
-        Ok(TermSearch { trees })
+        let counts = db.open_tree(&format!("/profiles/search/{}/counts", name))?;
+
+        Ok(TermSearch { trees, counts })
     }
 
     pub(crate) fn insert(&self, term: &str) -> Result<(), StoreError> {
-        let term_vec = term.to_lowercase().as_bytes().to_vec();
+        let term = term.to_lowercase();
+        let term_vec = term.as_bytes().to_vec();
 
         if term.is_empty() {
             return Err(StoreError::Empty);
@@ -31,17 +35,51 @@ impl TermSearch {
 
         let tree = &self.trees[term.len() - 1];
 
-        tree.transaction(|tree| {
+        [tree, &self.counts].transaction(|trees| {
-            for i in 0..term_vec.len() {
+            let tree = &trees[0];
-                let (start, end) = term_vec.split_at(i);
+            let counts = &trees[1];
 
-                let mut key = Vec::with_capacity(term_vec.len() + 1);
+            if count(counts, &term, |c| c.saturating_add(1))? == 1 {
-                key.extend(end);
+                for i in 0..term_vec.len() {
-                key.push(127);
+                    let key = key(&term_vec, i);
-                key.extend(start);
 
                     tree.insert(key.as_slice(), term_vec.as_slice())?;
                 }
+            }
+
+            Ok(())
+        })?;
+
+        Ok(())
+    }
+
+    pub(crate) fn remove(&self, term: &str) -> Result<(), StoreError> {
+        let term = term.to_lowercase();
+        let term_vec = term.as_bytes().to_vec();
+
+        if term.is_empty() {
+            return Err(StoreError::Empty);
+        }
+
+        if term.len() > self.trees.len() {
+            return Err(StoreError::TooLong);
+        }
+
+        let tree = &self.trees[term.len() - 1];
+
+        [tree, &self.counts].transaction(|trees| {
+            let tree = &trees[0];
+            let counts = &trees[1];
+
+            if count(counts, &term, |c| c.saturating_add(1))? == 0 {
+                for i in 0..term_vec.len() {
+                    let key = key(&term_vec, i);
+
+                    tree.remove(key.as_slice())?;
+                }
+
+                counts.remove(term.as_bytes())?;
+            }
 
             Ok(())
         })?;
@@ -67,6 +105,17 @@ impl TermSearch {
     }
 }
 
+fn key(term_vec: &[u8], i: usize) -> Vec<u8> {
+    let (start, end) = term_vec.split_at(i);
+
+    let mut key = Vec::with_capacity(term_vec.len() + 1);
+    key.extend(end);
+    key.push(127);
+    key.extend(start);
+
+    key
+}
+
 struct Deduplicate<T> {
     iter: T,
     seen: HashSet<sled::IVec>,