mirror of
https://github.com/meilisearch/MeiliSearch
synced 2024-06-09 14:27:47 +02:00
177 lines
6.3 KiB
Rust
177 lines
6.3 KiB
Rust
use std::net::ToSocketAddrs;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::sync::{atomic, Arc, Mutex, RwLock};
|
|
|
|
use bus::{Bus, BusReader};
|
|
use crossbeam::channel::{unbounded, Receiver, Sender};
|
|
use ductile::{ChannelReceiver, ChannelSender, ChannelServer};
|
|
use log::info;
|
|
use meilisearch_types::tasks::Task;
|
|
|
|
use crate::batch::Batch;
|
|
use crate::{Consistency, FollowerMsg, LeaderMsg};
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct Leader {
|
|
task_ready_to_commit: Receiver<u32>,
|
|
broadcast_to_follower: Sender<LeaderMsg>,
|
|
|
|
cluster_size: Arc<AtomicUsize>,
|
|
|
|
batch_id: Arc<RwLock<u32>>,
|
|
}
|
|
|
|
impl Leader {
|
|
pub fn new(listen_on: impl ToSocketAddrs + Send + 'static) -> Leader {
|
|
let cluster_size = Arc::new(AtomicUsize::new(1));
|
|
let (process_batch_sender, process_batch_receiver) = unbounded();
|
|
let (task_finished_sender, task_finished_receiver) = unbounded();
|
|
|
|
let cs = cluster_size.clone();
|
|
std::thread::spawn(move || {
|
|
Self::listener(listen_on, cs, process_batch_receiver, task_finished_sender)
|
|
});
|
|
|
|
Leader {
|
|
task_ready_to_commit: task_finished_receiver,
|
|
broadcast_to_follower: process_batch_sender,
|
|
cluster_size,
|
|
batch_id: Arc::default(),
|
|
}
|
|
}
|
|
|
|
/// Takes all the necessary channels to chat with the scheduler and give them
|
|
/// to each new followers
|
|
fn listener(
|
|
listen_on: impl ToSocketAddrs,
|
|
cluster_size: Arc<AtomicUsize>,
|
|
broadcast_to_follower: Receiver<LeaderMsg>,
|
|
task_finished: Sender<u32>,
|
|
) {
|
|
let listener: ChannelServer<LeaderMsg, FollowerMsg> =
|
|
ChannelServer::bind(listen_on).unwrap();
|
|
|
|
// We're going to broadcast all the batches to all our follower
|
|
let bus: Bus<LeaderMsg> = Bus::new(10);
|
|
let bus = Arc::new(Mutex::new(bus));
|
|
let b = bus.clone();
|
|
|
|
std::thread::spawn(move || loop {
|
|
let msg = broadcast_to_follower.recv().expect("Main thread is dead");
|
|
b.lock().unwrap().broadcast(msg);
|
|
});
|
|
|
|
for (sender, receiver, _addr) in listener {
|
|
let task_finished = task_finished.clone();
|
|
let cs = cluster_size.clone();
|
|
|
|
let process_batch = bus.lock().unwrap().add_rx();
|
|
|
|
std::thread::spawn(move || {
|
|
Self::follower(sender, receiver, cs, process_batch, task_finished)
|
|
});
|
|
}
|
|
}
|
|
|
|
/// Allow a follower to chat with the scheduler
|
|
fn follower(
|
|
sender: ChannelSender<LeaderMsg>,
|
|
receiver: ChannelReceiver<FollowerMsg>,
|
|
cluster_size: Arc<AtomicUsize>,
|
|
mut broadcast_to_follower: BusReader<LeaderMsg>,
|
|
task_finished: Sender<u32>,
|
|
) {
|
|
let size = cluster_size.fetch_add(1, atomic::Ordering::Relaxed) + 1;
|
|
info!("A new follower joined the cluster. {} members.", size);
|
|
|
|
// send messages to the follower
|
|
std::thread::spawn(move || loop {
|
|
let msg = broadcast_to_follower.recv().expect("Main thread died");
|
|
if sender.send(msg).is_err() {
|
|
// the follower died, the logging and cluster size update should be done
|
|
// in the other thread
|
|
break;
|
|
}
|
|
});
|
|
|
|
// receive messages from the follower
|
|
loop {
|
|
match receiver.recv() {
|
|
Err(_) => break,
|
|
Ok(msg) => match msg {
|
|
FollowerMsg::ReadyToCommit(id) => {
|
|
task_finished.send(id).expect("Can't reach the main thread")
|
|
}
|
|
FollowerMsg::RegisterNewTask(_) => todo!(),
|
|
},
|
|
}
|
|
}
|
|
|
|
// if we exited from the previous loop it means the follower is down and should
|
|
// be removed from the cluster
|
|
let size = cluster_size.fetch_sub(1, atomic::Ordering::Relaxed) - 1;
|
|
info!("A follower left the cluster. {} members.", size);
|
|
}
|
|
|
|
pub fn starts_batch(&self, batch: Batch) {
|
|
let mut batch_id = self.batch_id.write().unwrap();
|
|
|
|
assert!(
|
|
*batch_id % 2 == 0,
|
|
"Tried to start processing a batch before commiting the previous one"
|
|
);
|
|
info!("Send the batch to process to the followers");
|
|
*batch_id += 1;
|
|
|
|
self.broadcast_to_follower
|
|
.send(LeaderMsg::StartBatch { id: *batch_id, batch })
|
|
.expect("Can't reach the cluster");
|
|
}
|
|
|
|
pub fn commit(&self, consistency_level: Consistency) {
|
|
info!("Wait until enough followers are ready to commit a batch");
|
|
|
|
let mut batch_id = self.batch_id.write().unwrap();
|
|
|
|
// if zero nodes needs to be sync we can commit right away and early exit
|
|
if consistency_level != Consistency::One {
|
|
// else, we wait till enough nodes are ready to commit
|
|
for ready_to_commit in self
|
|
.task_ready_to_commit
|
|
.iter()
|
|
// we need to filter out the messages from the old batches
|
|
.filter(|id| *id == *batch_id)
|
|
.enumerate()
|
|
// we do a +2 because enumerate starts at 1 and we must includes ourselves in the count
|
|
.map(|(id, _)| id + 2)
|
|
{
|
|
// TODO: if the last node dies we're stuck on the iterator
|
|
|
|
// we need to reload the cluster size everytime in case a node dies
|
|
let cluster_size = self.cluster_size.load(atomic::Ordering::Relaxed);
|
|
|
|
info!("{ready_to_commit} nodes are ready to commit for a cluster size of {cluster_size}");
|
|
match consistency_level {
|
|
Consistency::Two if ready_to_commit >= 1 => break,
|
|
Consistency::Quorum if ready_to_commit >= (cluster_size / 2) => break,
|
|
Consistency::All if ready_to_commit == cluster_size => break,
|
|
_ => (),
|
|
}
|
|
}
|
|
}
|
|
|
|
info!("Tells all the follower to commit");
|
|
|
|
self.broadcast_to_follower.send(LeaderMsg::Commit(*batch_id)).unwrap();
|
|
|
|
*batch_id += 1;
|
|
}
|
|
|
|
pub fn register_new_task(&self, task: Task, update_file: Option<Vec<u8>>) {
|
|
info!("Tells all the follower to register a new task");
|
|
self.broadcast_to_follower
|
|
.send(LeaderMsg::RegisterNewTask { task, update_file })
|
|
.expect("Main thread is dead");
|
|
}
|
|
}
|