rockraft 0.1.6

use std::fmt::Debug;
use std::io;
use std::io::Error;
use std::marker::PhantomData;
use std::ops::Bound;
use std::ops::RangeBounds;
use std::sync::Arc;

use crate::raft::types::{decode, encode};
use byteorder::BigEndian;
use byteorder::ReadBytesExt as _;
use byteorder::WriteBytesExt as _;
use openraft::OptionalSend;
use openraft::RaftLogReader;
use openraft::RaftTypeConfig;
use openraft::alias::EntryOf;
use openraft::alias::LogIdOf;
use openraft::alias::VoteOf;
use openraft::entry::RaftEntry;
use openraft::storage::IOFlushed;
use openraft::storage::RaftLogStorage;
use rocksdb::BoundColumnFamily;
use rocksdb::DB;
use rocksdb::Direction;
use rocksdb::IteratorMode;
use tokio::task::spawn_blocking;

use super::StoreMeta;
use super::keys::LOG_DATA_FAMILY;
use super::keys::LOG_META_FAMILY;
use super::meta::LastPurged;
use crate::raft::types::Entry;
use crate::raft::types::LogId;
use crate::raft::types::LogState;
use crate::raft::types::RaftCodec;
use crate::raft::types::TypeConfig;
use crate::raft::types::read_logs_err;

#[derive(Debug, Clone)]
pub struct RocksLogStore<C>
where
  C: RaftTypeConfig,
{
  pub(crate) db: Arc<DB>,
  _p: PhantomData<C>,
}

impl RocksLogStore<TypeConfig> {
  pub fn create(db: Arc<DB>) -> Result<Self, Error> {
    db.cf_handle(LOG_META_FAMILY)
      .ok_or_else(|| Error::other("column family `_log_meta` not found"))?;
    db.cf_handle(LOG_DATA_FAMILY)
      .ok_or_else(|| Error::other("column family `_log_data` not found"))?;

    Ok(Self {
      db,
      _p: Default::default(),
    })
  }

  fn cf_meta(&self) -> Arc<BoundColumnFamily<'_>> {
    self.db.cf_handle(LOG_META_FAMILY).unwrap()
  }

  fn cf_logs(&self) -> Arc<BoundColumnFamily<'_>> {
    self.db.cf_handle(LOG_DATA_FAMILY).unwrap()
  }

  /// Get a store metadata.
  ///
  /// It returns `None` if the store does not have such a metadata stored.
  fn get_meta<M: StoreMeta>(&self) -> Result<Option<M::Value>, io::Error> {
    let bytes = self
      .db
      .get_cf(&self.cf_meta(), M::KEY)
      .map_err(M::read_err)?;

    let Some(bytes) = bytes else {
      return Ok(None);
    };

    let entry = M::Value::decode_from(bytes.as_ref()).map_err(read_logs_err)?;

    Ok(Some(entry))
  }

  /// Save a store metadata.
  fn put_meta<M: StoreMeta>(&self, value: &M::Value) -> Result<(), io::Error> {
    let encode_valude = value.encode_to()?;

    self
      .db
      .put_cf(&self.cf_meta(), M::KEY, encode_valude)
      .map_err(|e| M::write_err(value, e))?;

    Ok(())
  }

  fn delete_meta<M: StoreMeta>(&self) -> Result<(), io::Error> {
    self
      .db
      .delete_cf(&self.cf_meta(), M::KEY)
      .map_err(|e| M::delete_err(e))?;

    Ok(())
  }

  fn get_last_purged_log_id(&self) -> Result<Option<LogId>, io::Error> {
    self.get_meta::<LastPurged>()
  }

  fn set_last_purged_log_id(&self, log_id: &LogId) -> Result<(), io::Error> {
    self.put_meta::<LastPurged>(log_id)
  }

  fn get_vote(&mut self) -> Result<Option<VoteOf<TypeConfig>>, io::Error> {
    self.get_meta::<super::meta::Vote>()
  }

  #[allow(dead_code)]
  fn set_vote(&mut self, vote: &VoteOf<TypeConfig>) -> Result<(), io::Error> {
    self.put_meta::<super::meta::Vote>(vote)
  }

  fn set_committed(&self, committed: &Option<LogIdOf<TypeConfig>>) -> Result<(), io::Error> {
    if let Some(committed) = committed {
      self.put_meta::<super::meta::Committed>(committed)?;
    } else {
      self.delete_meta::<super::meta::Committed>()?;
    }
    Ok(())
  }

  fn get_committed(&self) -> Result<Option<LogIdOf<TypeConfig>>, io::Error> {
    self.get_meta::<super::meta::Committed>()
  }
}

impl RaftLogReader<TypeConfig> for RocksLogStore<TypeConfig> {
  async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + Debug + OptionalSend>(
    &mut self,
    range: RB,
  ) -> Result<Vec<Entry>, io::Error> {
    let start = match range.start_bound() {
      Bound::Included(x) => id_to_bin(*x),
      Bound::Excluded(x) => id_to_bin(*x + 1),
      Bound::Unbounded => id_to_bin(0),
    };

    let mut entries = Vec::new();

    for item in self.db.iterator_cf(
      &self.cf_logs(),
      IteratorMode::From(&start, Direction::Forward),
    ) {
      let (key, val) = item.map_err(read_logs_err)?;

      let id = bin_to_id(&key)?;

      if !range.contains(&id) {
        break;
      }

      let entry: Entry = decode(&val).map_err(read_logs_err)?;

      entries.push(entry);
    }

    Ok(entries)
  }

  async fn read_vote(&mut self) -> Result<Option<VoteOf<TypeConfig>>, io::Error> {
    self.get_vote()
  }
}

impl RaftLogStorage<TypeConfig> for RocksLogStore<TypeConfig> {
  type LogReader = Self;

  async fn get_log_state(&mut self) -> Result<LogState, io::Error> {
    let last = self
      .db
      .iterator_cf(&self.cf_logs(), rocksdb::IteratorMode::End)
      .next();

    let last_log_id = match last {
      None => None,
      Some(res) => {
        let (_log_index, val) = res.map_err(read_logs_err)?;
        let entry: Entry = decode(val.as_ref()).map_err(read_logs_err)?;
        Some(entry.log_id)
      }
    };

    let last_purged_log_id = self.get_last_purged_log_id()?;

    let last_log_id = match last_log_id {
      None => last_purged_log_id,
      Some(x) => Some(x),
    };

    Ok(LogState {
      last_purged_log_id,
      last_log_id,
    })
  }

  async fn save_committed(
    &mut self,
    committed: Option<LogIdOf<TypeConfig>>,
  ) -> Result<(), io::Error> {
    self.set_committed(&committed)
  }

  async fn read_committed(&mut self) -> Result<Option<LogId>, io::Error> {
    self.get_committed()
  }

  async fn get_log_reader(&mut self) -> Self::LogReader {
    self.clone()
  }

  async fn save_vote(&mut self, vote: &VoteOf<TypeConfig>) -> Result<(), io::Error> {
    self.put_meta::<super::meta::Vote>(vote)?;

    // Vote must be persisted to disk before returning.
    let db = self.db.clone();
    spawn_blocking(move || db.flush_wal(true))
      .await
      .map_err(read_logs_err)?
      .map_err(read_logs_err)?;

    Ok(())
  }

  async fn append<I>(
    &mut self,
    entries: I,
    callback: IOFlushed<TypeConfig>,
  ) -> Result<(), io::Error>
  where
    I: IntoIterator<Item = EntryOf<TypeConfig>> + Send,
  {
    for entry in entries {
      let id = id_to_bin(entry.index());
      self
        .db
        .put_cf(&self.cf_logs(), id, encode(&entry).map_err(read_logs_err)?)
        .map_err(read_logs_err)?;
    }

    // Make sure the logs are persisted to disk before invoking the callback.
    //
    // But the above `pub_cf()` must be called in this function, not in another task.
    // Because when the function returns, it requires the log entries can be read.
    let db = self.db.clone();
    let handle = spawn_blocking(move || {
      let res = db.flush_wal(true).map_err(Error::other);
      callback.io_completed(res);
    });
    drop(handle);

    // Return now, and the callback will be invoked later when IO is done.
    Ok(())
  }

  async fn truncate_after(&mut self, log_id: Option<LogIdOf<TypeConfig>>) -> Result<(), io::Error> {
    tracing::debug!("truncate_after: [{:?}, +oo)", log_id);

    // If log_id is None, truncate everything
    let from = match log_id {
      Some(ref id) => id_to_bin(id.index()),
      None => id_to_bin(0),
    };
    let to = id_to_bin(u64::MAX);
    self
      .db
      .delete_range_cf(&self.cf_logs(), &from, &to)
      .map_err(read_logs_err)?;

    // Truncating does not need to be persisted.
    Ok(())
  }

  async fn purge(&mut self, log_id: LogIdOf<TypeConfig>) -> Result<(), io::Error> {
    tracing::debug!("delete_log: [0, {:?}]", log_id);

    // Write the last-purged log id before purging the logs.
    // The logs at and before last-purged log id will be ignored by openraft.
    // Therefore, there is no need to do it in a transaction.
    self.set_last_purged_log_id(&log_id)?;

    let from = id_to_bin(0);
    let to = id_to_bin(log_id.index() + 1);
    self
      .db
      .delete_range_cf(&self.cf_logs(), &from, &to)
      .map_err(read_logs_err)?;

    // Purging does not need to be persistent.
    Ok(())
  }
}

fn id_to_bin(id: u64) -> Vec<u8> {
  let mut buf = Vec::with_capacity(8);
  buf.write_u64::<BigEndian>(id).unwrap();
  buf
}

fn bin_to_id(buf: &[u8]) -> Result<u64, Error> {
  (&buf[0..8]).read_u64::<BigEndian>()
}

#[cfg(test)]
mod tests {
  use crate::utils::test::{append_entries, create_entry, create_log_id, create_test_log_store};
  use openraft::Vote;

  use super::*;

  #[tokio::test]
  async fn test_raft_log_operations() -> Result<(), io::Error> {
    let mut log_store = create_test_log_store();

    let entries: Vec<_> = (1..=10).map(|i| create_entry(1, 1, i)).collect();
    append_entries(&mut log_store, entries).await?;

    let all = log_store.try_get_log_entries(1..=10).await?;
    assert_eq!(all.len(), 10);
    assert_eq!(all[all.len() - 1].log_id.index, 10);

    let range = log_store.try_get_log_entries(3..=7).await?;
    assert_eq!(range.len(), 5);
    assert_eq!(range[0].log_id.index, 3);

    let more: Vec<_> = (11..=15).map(|i| create_entry(1, 1, i)).collect();
    append_entries(&mut log_store, more).await?;
    assert_eq!(log_store.try_get_log_entries(1..=15).await?.len(), 15);

    log_store
      .truncate_after(Some(create_log_id(1, 1, 11)))
      .await?;
    assert_eq!(log_store.try_get_log_entries(1..=15).await?.len(), 10);
    assert_eq!(log_store.try_get_log_entries(11..=15).await?.len(), 0);

    log_store.purge(create_log_id(1, 1, 5)).await?;
    assert_eq!(log_store.get_last_purged_log_id()?.unwrap().index, 5);
    let after_purge = log_store.try_get_log_entries(1..=10).await?;
    assert_eq!(after_purge.len(), 5);
    assert_eq!(after_purge[0].log_id.index, 6);

    let new: Vec<_> = (11..=13).map(|i| create_entry(2, 1, i)).collect();
    append_entries(&mut log_store, new).await?;
    let final_logs = log_store.try_get_log_entries(6..=13).await?;
    assert_eq!(final_logs.len(), 8);
    assert_eq!(final_logs[0].log_id.leader_id.term, 1);
    assert_eq!(final_logs[5].log_id.leader_id.term, 2);

    Ok(())
  }

  #[test]
  fn test_set_and_get_last_purged() -> Result<(), io::Error> {
    let log_store = create_test_log_store();

    assert!(log_store.get_last_purged_log_id()?.is_none());

    let log_id = create_log_id(1, 1, 100);
    log_store.set_last_purged_log_id(&log_id)?;

    let retrieved = log_store.get_last_purged_log_id()?.unwrap();
    assert_eq!(retrieved.leader_id.term, log_id.leader_id.term);
    assert_eq!(retrieved.leader_id.node_id, log_id.leader_id.node_id);
    assert_eq!(retrieved.index, log_id.index);

    let new_log_id = create_log_id(2, 2, 200);
    log_store.set_last_purged_log_id(&new_log_id).unwrap();

    let updated = log_store.get_last_purged_log_id()?.unwrap();
    assert_eq!(updated.index, new_log_id.index);

    Ok(())
  }

  #[test]
  fn test_set_and_get_committed() -> Result<(), io::Error> {
    let log_store = create_test_log_store();

    assert!(log_store.get_committed()?.is_none());

    let log_id = create_log_id(1, 1, 100);
    log_store.set_committed(&Some(log_id))?;

    let retrieved = log_store.get_committed()?.unwrap();
    assert_eq!(retrieved.leader_id.term, log_id.leader_id.term);
    assert_eq!(retrieved.leader_id.node_id, log_id.leader_id.node_id);
    assert_eq!(retrieved.index, log_id.index);

    let new_log_id = create_log_id(2, 2, 200);
    log_store.set_committed(&Some(new_log_id))?;

    let updated = log_store.get_committed()?.unwrap();
    assert_eq!(updated.index, new_log_id.index);

    log_store.set_committed(&None)?;
    assert!(log_store.get_committed()?.is_none());

    Ok(())
  }

  #[test]
  fn test_set_and_get_vote() -> Result<(), io::Error> {
    let mut log_store = create_test_log_store();

    assert!(log_store.get_vote()?.is_none());

    let vote = Vote::new(1, 1);
    log_store.set_vote(&vote)?;

    let retrieved = log_store.get_vote()?.unwrap();
    assert_eq!(retrieved.leader_id().term, vote.leader_id().term);
    assert_eq!(retrieved.leader_id().node_id, vote.leader_id().node_id);

    let new_vote = Vote::new(2, 2);
    log_store.set_vote(&new_vote)?;

    let updated = log_store.get_vote()?.unwrap();
    assert_eq!(updated.leader_id().term, 2);
    assert_eq!(updated.leader_id().node_id, 2);

    Ok(())
  }
}