架构杂谈:TiKV弹性设计及实现

/// A TiKV key-value store
pub trait KvEngine:
    Peekable
    + SyncMutable
    + Iterable
    + WriteBatchExt
    + DBOptionsExt
    + CFNamesExt
    + CFOptionsExt
    + ImportExt
    + SstExt
    + TablePropertiesExt
    + CompactExt
    + RangePropertiesExt
    + MvccPropertiesExt
    + TtlPropertiesExt
    + PerfContextExt
    + MiscExt
    + Send
    + Sync
    + Clone
    + Debug
    + Unpin
    + 'static
{
    /// A consistent read-only snapshot of the database
    type Snapshot: Snapshot;
    /// Create a snapshot
    fn snapshot(&self) -> Self::Snapshot;
    /// Syncs any writes to disk
    fn sync(&self) -> Result<()>;
    /// Flush metrics to prometheus
    /// `instance` is the label of the metric to flush.
    fn flush_metrics(&self, _instance: &str) {}
    /// Reset internal statistics
    fn reset_statistics(&self) {}
    /// Cast to a concrete engine type
    /// This only exists as a temporary hack during refactor.
    /// It cannot be used forever.
    fn bad_downcast<T: 'static>(&self) -> &T;
}

pub trait RaftEngine: Clone + Sync + Send + 'static {
    type LogBatch: RaftLogBatch;
    fn log_batch(&self, capacity: usize) -> Self::LogBatch;
    /// Synchronize the Raft engine.
    fn sync(&self) -> Result<()>;
    fn get_raft_state(&self, raft_group_id: u64)
        -> Result<Option<RaftLocalState>>;
    fn get_entry(&self, raft_group_id: u64, index: u64)
        -> Result<Option<Entry>>;
    /// Append some log entries and return written bytes.
    fn append(&self, raft_group_id: u64, entries: Vec<Entry>)
        -> Result<usize>;
    fn put_raft_state(&self, raft_group_id: u64
        , state: &RaftLocalState) -> Result<()>;
    // 省略部分其他方法
}


pub trait RaftLogBatch: Send {
    fn append(&mut self, raft_group_id: u64
        , entries: Vec<Entry>) -> Result<()>;
    /// Remove Raft logs in [`from`, `to`) 
    /// which will be overwritten later.
    fn cut_logs(&mut self, raft_group_id: u64
        , from: u64, to: u64);
    fn put_raft_state(&mut self, raft_group_id: u64
        , state: &RaftLocalState) -> Result<()>;
    fn is_empty(&self) -> bool;
}

#[derive(Clone, Debug)]
pub struct Engines<K, R> {
    pub kv: K,
    pub raft: R,
}


impl<K: KvEngine, R: RaftEngine> Engines<K, R> {
    pub fn new(kv_engine: K, raft_engine: R) -> Self {
        Engines {
            kv: kv_engine,
            raft: raft_engine,
        }
    }
    // 省略部分其他方法
}

//基于RocksDB来提供RocksEngine以实现KvEngine trait；
pub struct RocksEngine {
    db: Arc<DB>,
    shared_block_cache: bool,
}


impl KvEngine for RocksEngine {
    type Snapshot = RocksSnapshot;
    fn snapshot(&self) -> RocksSnapshot {
        RocksSnapshot::new(self.db.clone())
    }
    // 省略部分其他方法
}

use raft_engine::{EntryExt, Error as RaftEngineError,
    LogBatch, RaftLogEngine as RawRaftEngine};


pub struct RaftLogEngine(RawRaftEngine<Entry, EntryExtTyped>);


impl RaftEngine for RaftLogEngine {
    type LogBatch = RaftLogBatch;


    fn log_batch(&self, _capacity: usize) -> Self::LogBatch {
        RaftLogBatch::default()
    }
    // 省略部分其他方法
}

use crate::{RocksEngine, RocksWriteBatch};


use engine_traits::{Error, RaftEngine, RaftLogBatch, Result};


impl RaftEngine for RocksEngine {
    type LogBatch = RocksWriteBatch;


    fn log_batch(&self, capacity: usize) -> Self::LogBatch {
        RocksWriteBatch::with_capacity(self.as_inner().clone()
          , capacity)
    }


    fn sync(&self) -> Result<()> {
        self.sync_wal()
    }
    // 省略部分其他方法
}

/// Scheduler provides interface to schedule task
/// to underlying workers.
pub struct Scheduler<T: Display + Send> {
    counter: Arc<AtomicUsize>,
    sender: UnboundedSender<Msg<T>>,
    pending_capacity: usize,
    metrics_pending_task_count: IntGauge,
}


enum Msg<T: Display + Send> {
    Task(T),
    Timeout,
}


pub trait Runnable: Send {
    type Task: Display + Send + 'static;
    /// Runs a task.
    fn run(&mut self, _: Self::Task) {
        unimplemented!()
    }
    fn on_tick(&mut self) {}
    fn shutdown(&mut self) {}
}


/// A worker that can schedule time consuming tasks.
#[derive(Clone)]
pub struct Worker {
    pool: Arc<Mutex<Option<ThreadPool<
      yatp::task::future::TaskCell>>>>,
    remote: Remote<yatp::task::future::TaskCell>,
    pending_capacity: usize,
    counter: Arc<AtomicUsize>,
    stop: Arc<AtomicBool>,
    thread_count: usize,
}


pub struct Builder<S: Into<String>> {
    name: S,
    thread_count: usize,
    pending_capacity: usize,
}


impl<S: Into<String>> Builder<S> {
    pub fn thread_count(mut self, thread_count: usize)->Self{
        self.thread_count = thread_count;
        self
    }
    pub fn create(self) -> Worker {
        let pool = YatpPoolBuilder::new(
           DefaultTicker::default())
           .name_prefix(self.name)
           .thread_count(self.thread_count, self.thread_count)
           .build_single_level_pool();
        let remote = pool.remote().clone();
        let pool = Arc::new(Mutex::new(Some(pool)));
        Worker {
            remote,
            stop: Arc::new(AtomicBool::new(false)),
            pool,
            counter: Arc::new(AtomicUsize::new(0)),
            pending_capacity: self.pending_capacity,
            thread_count: self.thread_count,
        }
    }
}


impl Worker {
    pub fn new<S: Into<String>>(name: S) -> Worker {
        Builder::new(name).create()
    }
    pub fn start<R: Runnable + 'static, S: Into<String>>(
        &self,
        name: S,
        runner: R,
    ) -> Scheduler<R::Task> {
        let (tx, rx) = unbounded();
        self.start_impl(runner, rx
            , metrics_pending_task_count.clone());
        Scheduler::new(
            tx,
            self.counter.clone(),
            self.pending_capacity,
            metrics_pending_task_count,
        )
    }
    fn start_impl<R: Runnable + 'static>(
        &self,
        runner: R,
        mut receiver: UnboundedReceiver<Msg<R::Task>>,
        metrics_pending_task_count: IntGauge,
    ) {
        let counter = self.counter.clone();
        self.remote.spawn(async move {
            let mut handle = RunnableWrapper {inner: runner};
            while let Some(msg) = receiver.next().await {
                match msg {
                    Msg::Task(task) => {
                      handle.inner.run(task);
                      counter.fetch_sub(1, Ordering::SeqCst);
                      metrics_pending_task_count.dec();
                    }
                    Msg::Timeout => (),
                }
            }
        });
    }
    // 省略部分其他方法
}


impl<T: Display + Send> Scheduler<T> {
    /// Schedules a task to run.
    ///
    /// If the worker is stopped or number pending tasks
    /// exceeds capacity, an error will return.
    pub fn schedule(&self, task: T) -> Result<()
        , ScheduleError<T>> {
        debug!("scheduling task {}", task);
        if self.counter.load(Ordering::Acquire) >= 
            self.pending_capacity {
            return Err(ScheduleError::Full(task));
        }
        self.counter.fetch_add(1, Ordering::SeqCst);
        self.metrics_pending_task_count.inc();
        if let Err(e) = self.sender.unbounded_send(
            Msg::Task(task)) {
            if let Msg::Task(t) = e.into_inner() {
                self.counter.fetch_sub(1, Ordering::SeqCst);
                self.metrics_pending_task_count.dec();
                return Err(ScheduleError::Stopped(t));
            }
        }
        Ok(())
    }
    // 省略部分其他方法
}