<commit_before>use super::super::trees::*;
use super::btree::level::LEVEL_M as BLOCK_SIZE;
use super::btree::storage;
use super::tree::*;
use super::super::sm::client::SMClient;
use crate::client::AsyncClient;
use crate::ram::types::Id;
use crate::ram::types::RandValue;
use bifrost::conshash::ConsistentHashing;
use bifrost::utils::time::get_time;
use bifrost_plugins::hash_ident;
use futures::future::BoxFuture;
use futures::prelude::*;
use lightning::map::Map;
use lightning::map::{HashMap, ObjectMap};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::cell::RefCell;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::time::Duration;

pub type EntryKeyBlock = [EntryKey; BLOCK_SIZE];
pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(LSM_TREE_RPC_SERVICE) as u64;

#[derive(Clone, Serialize, Deserialize)]
pub struct Boundary {
    upper: EntryKey,
    lower: EntryKey,
}

#[derive(Clone, Serialize, Deserialize)]
pub enum OpResult<T> {
    Successful(T),
    Failed,
    NotFound,
    OutOfBound,
    Migrating,
}

#[derive(Clone, Copy, Serialize, Deserialize)]
pub struct ServCursor {
    cursor_id: u64,
}

pub struct DistLSMTree {
    id: Id,
    tree: LSMTree,
    prop: RwLock<DistProp>,
}

struct DistProp {
    boundary: Boundary,
    migration: Option<Migration>,
}

struct Migration {
    pivot: EntryKey,
}

pub struct CursorMemo {
    tree_cursor: LSMTreeCursor,
    expires: i64,
}

service! {
    rpc crate_tree(id: Id, boundary: Boundary);
    rpc load_tree(id: Id, boundary: Boundary);
    rpc insert(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc delete(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc seek(id: Id, entry: EntryKey, ordering: Ordering, cursor_lifetime: u16)
        -> OpResult<(ServCursor, Option<EntryKey>)>;
    rpc renew_cursor(cursor: ServCursor, time: u16) -> bool;
    rpc dispose_cursor(cursor: ServCursor) -> bool;
    rpc cursor_next(cursor: ServCursor) -> Option<EntryKeyBlock>;
}

pub struct LSMTreeService {
    client: Arc<AsyncClient>,
    cursor_counter: AtomicUsize,
    cursors: ObjectMap<Arc<RefCell<CursorMemo>>>,
    trees: Arc<HashMap<Id, Arc<DistLSMTree>>>,
}

impl Service for LSMTreeService {
    fn crate_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::create(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn load_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::recover(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn insert(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.insert(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn delete(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.delete(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn seek(
        &self,
        id: Id,
        entry: EntryKey,
        ordering: Ordering,
        cursor_lifetime: u16,
    ) -> BoxFuture<OpResult<(ServCursor, Option<EntryKey>)>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            let mut tree_cursor = tree.seek(&entry, ordering);
            let cursor_id = self.cursor_counter.fetch_add(1, Relaxed);
            let expires = get_time() + cursor_lifetime as i64;
            let first_entry = tree_cursor.next();
            let cursor_memo = CursorMemo {
                tree_cursor,
                expires,
            };
            self.cursors
                .insert(&(cursor_id), Arc::new(RefCell::new(cursor_memo)));
            let cursor = ServCursor {
                cursor_id: cursor_id as u64,
            };
            OpResult::Successful((cursor, first_entry))
        })
    }

    fn renew_cursor(&self, cursor: ServCursor, time: u16) -> BoxFuture<bool> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                cursor.borrow_mut().expires = get_time() + time as i64;
                true
            } else {
                false
            },
        )
        .boxed()
    }

    fn dispose_cursor(&self, cursor: ServCursor) -> BoxFuture<bool> {
        future::ready(self.cursors.remove(&(cursor.cursor_id as usize)).is_some()).boxed()
    }

    fn cursor_next(&self, cursor: ServCursor) -> BoxFuture<Option<EntryKeyBlock>> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                let mut cursor_ref = cursor.borrow_mut();
                Some(entry_block_from_cursor_memo(&mut *cursor_ref))
            } else {
                None
            },
        )
        .boxed()
    }
}

impl LSMTreeService {
    pub fn new(client: &Arc<AsyncClient>, sm_client: &Arc<SMClient>) -> Self {
        let trees_map = Arc::new(HashMap::with_capacity(32));
        super::btree::storage::start_external_nodes_write_back(client);
        Self::start_tree_balancer(&trees_map, client);
        Self {
            client: client.clone(),
            cursor_counter: AtomicUsize::new(0),
            cursors: ObjectMap::with_capacity(64),
            trees: trees_map,
        }
    }

    pub fn start_tree_balancer(
        trees_map: &Arc<HashMap<Id, Arc<DistLSMTree>>>,
        client: &Arc<AsyncClient>,
        sm_client: &Arc<SMClient>,
    ) {
        let trees_map = trees_map.clone();
        let client = client.clone();
        let sm_client = sm_client.clone();
        tokio::spawn(async move {
            loop {
                for (_, dist_tree) in trees_map.entries() {
                    let tree = &dist_tree.tree;
                    tree.merge_levels();
                    if tree.oversized() {
                        // Tree oversized, need to migrate
                        let mid_key = tree.mid_key().unwrap();
                        let migration_target_id = Id::rand();
                        let migration_tree = LSMTree::create(&client, &migration_target_id).await;
                        {
                            let mut dist_tree_prop = dist_tree.prop.write();
                            dist_tree_prop.migration = Some(Migration {
                                pivot: mid_key.clone(),
                            });<commit_msg>record on state machine for split<commit_after>use super::super::trees::*;
use super::btree::level::LEVEL_M as BLOCK_SIZE;
use super::btree::storage;
use super::tree::*;
use super::super::sm::client::SMClient;
use crate::client::AsyncClient;
use crate::ram::types::Id;
use crate::ram::types::RandValue;
use bifrost::conshash::ConsistentHashing;
use bifrost::utils::time::get_time;
use bifrost_plugins::hash_ident;
use futures::future::BoxFuture;
use futures::prelude::*;
use lightning::map::Map;
use lightning::map::{HashMap, ObjectMap};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::cell::RefCell;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::time::Duration;

pub type EntryKeyBlock = [EntryKey; BLOCK_SIZE];
pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(LSM_TREE_RPC_SERVICE) as u64;

#[derive(Clone, Serialize, Deserialize)]
pub struct Boundary {
    upper: EntryKey,
    lower: EntryKey,
}

#[derive(Clone, Serialize, Deserialize)]
pub enum OpResult<T> {
    Successful(T),
    Failed,
    NotFound,
    OutOfBound,
    Migrating,
}

#[derive(Clone, Copy, Serialize, Deserialize)]
pub struct ServCursor {
    cursor_id: u64,
}

pub struct DistLSMTree {
    id: Id,
    tree: LSMTree,
    prop: RwLock<DistProp>,
}

struct DistProp {
    boundary: Boundary,
    migration: Option<Migration>,
}

struct Migration {
    pivot: EntryKey,
}

pub struct CursorMemo {
    tree_cursor: LSMTreeCursor,
    expires: i64,
}

service! {
    rpc crate_tree(id: Id, boundary: Boundary);
    rpc load_tree(id: Id, boundary: Boundary);
    rpc insert(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc delete(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc seek(id: Id, entry: EntryKey, ordering: Ordering, cursor_lifetime: u16)
        -> OpResult<(ServCursor, Option<EntryKey>)>;
    rpc renew_cursor(cursor: ServCursor, time: u16) -> bool;
    rpc dispose_cursor(cursor: ServCursor) -> bool;
    rpc cursor_next(cursor: ServCursor) -> Option<EntryKeyBlock>;
}

pub struct LSMTreeService {
    client: Arc<AsyncClient>,
    cursor_counter: AtomicUsize,
    cursors: ObjectMap<Arc<RefCell<CursorMemo>>>,
    trees: Arc<HashMap<Id, Arc<DistLSMTree>>>,
}

impl Service for LSMTreeService {
    fn crate_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::create(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn load_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::recover(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn insert(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.insert(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn delete(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.delete(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn seek(
        &self,
        id: Id,
        entry: EntryKey,
        ordering: Ordering,
        cursor_lifetime: u16,
    ) -> BoxFuture<OpResult<(ServCursor, Option<EntryKey>)>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            let mut tree_cursor = tree.seek(&entry, ordering);
            let cursor_id = self.cursor_counter.fetch_add(1, Relaxed);
            let expires = get_time() + cursor_lifetime as i64;
            let first_entry = tree_cursor.next();
            let cursor_memo = CursorMemo {
                tree_cursor,
                expires,
            };
            self.cursors
                .insert(&(cursor_id), Arc::new(RefCell::new(cursor_memo)));
            let cursor = ServCursor {
                cursor_id: cursor_id as u64,
            };
            OpResult::Successful((cursor, first_entry))
        })
    }

    fn renew_cursor(&self, cursor: ServCursor, time: u16) -> BoxFuture<bool> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                cursor.borrow_mut().expires = get_time() + time as i64;
                true
            } else {
                false
            },
        )
        .boxed()
    }

    fn dispose_cursor(&self, cursor: ServCursor) -> BoxFuture<bool> {
        future::ready(self.cursors.remove(&(cursor.cursor_id as usize)).is_some()).boxed()
    }

    fn cursor_next(&self, cursor: ServCursor) -> BoxFuture<Option<EntryKeyBlock>> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                let mut cursor_ref = cursor.borrow_mut();
                Some(entry_block_from_cursor_memo(&mut *cursor_ref))
            } else {
                None
            },
        )
        .boxed()
    }
}

impl LSMTreeService {
    pub fn new(client: &Arc<AsyncClient>, sm_client: &Arc<SMClient>) -> Self {
        let trees_map = Arc::new(HashMap::with_capacity(32));
        super::btree::storage::start_external_nodes_write_back(client);
        Self::start_tree_balancer(&trees_map, client, sm_client);
        Self {
            client: client.clone(),
            cursor_counter: AtomicUsize::new(0),
            cursors: ObjectMap::with_capacity(64),
            trees: trees_map,
        }
    }

    pub fn start_tree_balancer(
        trees_map: &Arc<HashMap<Id, Arc<DistLSMTree>>>,
        client: &Arc<AsyncClient>,
        sm_client: &Arc<SMClient>,
    ) {
        let trees_map = trees_map.clone();
        let client = client.clone();
        let sm_client = sm_client.clone();
        tokio::spawn(async move {
            loop {
                for (_, dist_tree) in trees_map.entries() {
                    let tree = &dist_tree.tree;
                    tree.merge_levels();
                    if tree.oversized() {
                        // Tree oversized, need to migrate
                        let mid_key = tree.mid_key().unwrap();
                        let migration_target_id = Id::rand();
                        let migration_tree = LSMTree::create(&client, &migration_target_id).await;
                        {
                            let mut dist_tree_prop = dist_tree.prop.write();
                            dist_tree_prop.migration = Some(Migration {
                                pivot: mid_key.clone(),
                            });
{'fraction_non_alphanumeric': 0.12866461185578096, 'fraction_numerical': 0.004014049172102358, 'mean_word_length': 2.2066191680073546, 'pattern_counts': {'":': 0, '<': 73, '<?xml version=': 0, '>': 103, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5877248', 'n_tokens_mistral': 4374, 'n_tokens_neox': 4239, 'n_words': 882}