[nihav.git] / dispatch.rs

use std::sync::{Arc, Barrier};
use std::sync::atomic::*;
use std::thread;

use nihav_core::codecs::{DecoderError, DecoderResult};

use super::{FrameDecoder, PictureInfo, Shareable};

#[derive(Clone,Copy,Debug,PartialEq)]
pub enum FrameDecodingStatus {
    Ok,
    NotReady,
    Error,
    NotFound,
}

struct FrameState {
    pinfo:      PictureInfo,
    mb_pos:     AtomicUsize,
    error:      AtomicBool,
    complete:   AtomicBool,
    output:     AtomicBool,
    worker:     Option<thread::JoinHandle<DecoderResult<()>>>,
    result:     DecoderResult<()>,
    num_refs:   usize,
    ref_frames: Vec<u32>,
}

impl FrameState {
    fn get_id(&self) -> u32 { self.pinfo.full_id }
    fn get_user_id(&self) -> u32 { self.pinfo.user_id }
    fn is_working(&self) -> bool {
        self.worker.is_some() &&
        !self.complete.load(Ordering::Relaxed) &&
        !self.error.load(Ordering::Relaxed)
    }
    fn is_output_candidate(&self) -> bool {
        !self.output.load(Ordering::Relaxed) &&
            (self.complete.load(Ordering::Relaxed) || self.error.load(Ordering::Relaxed))
    }
}

pub struct ThreadDispatcher {
        fstate:         Vec<FrameState>,
    pub max_threads:    usize,
        cur_threads:    usize,
}

impl ThreadDispatcher {
    pub fn new() -> Self {
        Self {
            fstate:         Vec::new(),
            max_threads:    3,
            cur_threads:    0,
        }
    }
    pub fn can_decode_more(&self) -> bool {
        let out_cand = self.fstate.iter().filter(|state| state.is_output_candidate()).count();
        if out_cand > self.max_threads {
            return false;
        }
        if (self.cur_threads < self.max_threads) || (self.max_threads == 0) {
            true
        } else {
            let real_workers = self.fstate.iter().fold(0usize,
                    |acc, state| acc + (state.is_working() as usize));
            real_workers < self.max_threads
        }
    }
    fn cleanup(&mut self) {
        for state in self.fstate.iter_mut() {
            if state.worker.is_some() && !state.is_working() {
                let mut ret = None;
                std::mem::swap(&mut state.worker, &mut ret);
                if let Some(handle) = ret {
                    state.result = handle.join().unwrap();
                }
                self.cur_threads -= 1;
            }
        }
    }
    fn unref_frame(&mut self, id: u32) {
        let mut toremove = Vec::new();
        for state in self.fstate.iter() {
            if state.num_refs == 0 && state.output.load(Ordering::Relaxed) {
                toremove.push(state.get_id());
            }
        }
        if let Some(idx) = self.find_by_id(id) {
            let mut ref_frm = Vec::new();
            std::mem::swap(&mut ref_frm, &mut self.fstate[idx].ref_frames);
            for state in self.fstate.iter_mut() {
                if ref_frm.contains(&state.get_id()) {
                    assert!(state.num_refs >= 2);
                    state.num_refs -= 2;
                }
            }
            if self.fstate[idx].num_refs == 0 && self.fstate[idx].output.load(Ordering::Relaxed) {
                self.remove_frame(id);
            }
        }
        for &id in toremove.iter() {
            self.remove_frame(id);
        }
    }
    fn find_by_id(&self, id: u32) -> Option<usize> {
        self.fstate.iter().position(|x| x.get_id() == id)
    }
    fn set_completed(&self, id: u32) {
        if let Some(idx) = self.find_by_id(id) {
            self.fstate[idx].complete.store(true, Ordering::Relaxed);
        }
    }
    fn set_error(&self, id: u32) {
        if let Some(idx) = self.find_by_id(id) {
            self.fstate[idx].error.store(true, Ordering::Relaxed);
        }
    }
    pub fn update_pos(&self, id: u32, mb_pos: usize) {
        if let Some(idx) = self.find_by_id(id) {
            self.fstate[idx].mb_pos.store(mb_pos, Ordering::Relaxed);
        }
    }
    pub fn check_pos(&self, id: u32, mb_pos: usize) -> FrameDecodingStatus {
        if let Some(idx) = self.find_by_id(id) {
            let state = &self.fstate[idx];
            if !state.error.load(Ordering::Relaxed) {
                if state.complete.load(Ordering::Relaxed) || mb_pos < state.mb_pos.load(Ordering::Relaxed) {
                    FrameDecodingStatus::Ok
                } else {
                    FrameDecodingStatus::NotReady
                }
            } else {
                FrameDecodingStatus::Error
            }
        } else {
            FrameDecodingStatus::NotFound
        }
    }
    fn remove_frame(&mut self, id: u32) {
        if let Some(idx) = self.find_by_id(id) {
            self.fstate.remove(idx);
        }
    }
    /*fn print_state(&self) {
        print!("   state:");
        for state in self.fstate.iter() {
            print!(" s{}b{}r{}{}{}{}", state.get_id(),
                state.mb_pos.load(Ordering::Relaxed), state.num_refs,
                if state.error.load(Ordering::Relaxed) { "E" } else {""},
                if state.complete.load(Ordering::Relaxed) {"C"} else {""},
                if state.output.load(Ordering::Relaxed) {"O"} else {""});
        }
        println!();
    }*/
    pub fn has_output(&self) -> bool {
        for state in self.fstate.iter() {
            if state.is_output_candidate() {
                return true;
            }
        }
        false
    }
}

pub fn queue_decoding(disp: &mut Shareable<ThreadDispatcher>, mut fdec: FrameDecoder, initial_ref_frames: &[u32], ref_frames: &[u32]) {
    let barrier = Arc::new(Barrier::new(2));
    let starter = Arc::clone(&barrier);

    let pinfo = fdec.cur_pic.clone();
    let pic_id = pinfo.full_id;
    let shared_disp = Arc::clone(disp);
    let worker = thread::Builder::new().name("frame ".to_string() + &pic_id.to_string()).spawn(move || {
            barrier.wait();

            let mut slices = Vec::new();
            std::mem::swap(&mut slices, &mut fdec.slices);
            let mut cur_mb = 0;
            for (hdr, hdr_size, refs, nal) in slices.iter() {
                if hdr.first_mb_in_slice != cur_mb {
                    if let Ok(rd) = shared_disp.read() {
                        rd.set_error(pic_id);
                    } else {
                        panic!("can't set error");
                    }
                    return Err(DecoderError::InvalidData);
                }
                match fdec.decode_slice(hdr, *hdr_size, refs, nal) {
                    Ok(pos) => cur_mb = pos,
                    Err(err) => {
                        if let Ok(rd) = shared_disp.read() {
                            rd.set_error(pic_id);
                        } else {
                            panic!("can't set error");
                        }
                        return Err(err);
                    },
                };
            }

            if cur_mb == fdec.num_mbs {
                if let Ok(rd) = shared_disp.read() {
                    rd.set_completed(pic_id);
                } else {
                    panic!("can't set status");
                }
            }

            DecoderResult::Ok(())
        }).unwrap();
    let new_state = FrameState {
        pinfo,
        mb_pos:     AtomicUsize::new(0),
        error:      AtomicBool::new(false),
        complete:   AtomicBool::new(false),
        output:     AtomicBool::new(false),
        worker:     Some(worker),
        result:     DecoderResult::Err(DecoderError::Bug),
        num_refs:   0,
        ref_frames: initial_ref_frames.to_vec(),
    };
    if let Ok(ref mut ds) = disp.write() {
        let new_id = new_state.get_id();
        if ds.find_by_id(new_id).is_some() {
            ds.remove_frame(new_id);
        }
        ds.cleanup();
        ds.fstate.push(new_state);
        for state in ds.fstate.iter_mut() {
            if ref_frames.contains(&state.get_id()) {
                state.num_refs += 1;
            }
            if initial_ref_frames.contains(&state.get_id()) {
                state.num_refs += 1;
            }
        }
        ds.cur_threads += 1;
        starter.wait();
    } else {
        panic!("cannot invoke thread dispatcher");
    }
}

pub fn wait_for_one(dispatch: &mut Shareable<ThreadDispatcher>) -> Result<PictureInfo, (DecoderError, u32)> {
    /*if let Ok(ref ds) = dispatch.read() {
        ds.print_state();
    }*/
    let start = std::time::Instant::now();
    'main_loop: loop {
        if std::time::Instant::now().duration_since(start) > std::time::Duration::from_millis(20000) { panic!(" too long!"); }
        if let Ok(ref ds) = dispatch.read() {
            let mut nw = 0;
            for state in ds.fstate.iter() {
                if state.is_working() {
                    nw += 1;
                }
                if state.is_output_candidate() {
                    break 'main_loop;
                }
            }
            if nw == 0 {
                return Err((DecoderError::NoFrame, 0));
            }
        } else {
            panic!("can't peek into status");
        }
        thread::yield_now();
    }
    if let Ok(ref mut ds) = dispatch.write() {
        ds.cleanup();
        let mut found = None;
        for state in ds.fstate.iter() {
            if state.is_output_candidate() {
                state.output.store(true, Ordering::Relaxed);
                if let DecoderResult::Err(err) = state.result {
                    let id = state.get_id();
                    let user_id = state.get_user_id();
                    ds.unref_frame(id);
                    return Err((err, user_id));
                } else {
                    found = Some(state.pinfo.clone());
                    break;
                }
            }
        }
        if let Some(ret) = found {
            ds.unref_frame(ret.full_id);
            Ok(ret)
        } else {
            unreachable!();
        }
    } else {
        panic!("can't grab status");
    }
}

pub fn clear_threads(dispatch: &mut Shareable<ThreadDispatcher>) {
    /*if let Ok(ref ds) = dispatch.read() {
        ds.print_state();
    }*/
    let mut to_wait = Vec::new();
    if let Ok(ref mut ds) = dispatch.write() {
        while let Some(state) = ds.fstate.pop() {
            if let Some(handle) = state.worker {
                to_wait.push(handle);
            }
        }
        ds.cur_threads = 0;
    } else {
        panic!("can't grab status");
    }
    while let Some(handle) = to_wait.pop() {
        let _ = handle.join();
    }
}
Commit	Line	Data
	1	use std::sync::{Arc, Barrier};
	2	use std::sync::atomic::*;
	3	use std::thread;
	4
	5	use nihav_core::codecs::{DecoderError, DecoderResult};
	6
	7	use super::{FrameDecoder, PictureInfo, Shareable};
	8
	9	#[derive(Clone,Copy,Debug,PartialEq)]
	10	pub enum FrameDecodingStatus {
	11	Ok,
	12	NotReady,
	13	Error,
	14	NotFound,
	15	}
	16
	17	struct FrameState {
	18	pinfo: PictureInfo,
	19	mb_pos: AtomicUsize,
	20	error: AtomicBool,
	21	complete: AtomicBool,
	22	output: AtomicBool,
	23	worker: Option<thread::JoinHandle<DecoderResult<()>>>,
	24	result: DecoderResult<()>,
	25	num_refs: usize,
	26	ref_frames: Vec<u32>,
	27	}
	28
	29	impl FrameState {
	30	fn get_id(&self) -> u32 { self.pinfo.full_id }
	31	fn get_user_id(&self) -> u32 { self.pinfo.user_id }
	32	fn is_working(&self) -> bool {
	33	self.worker.is_some() &&
	34	!self.complete.load(Ordering::Relaxed) &&
	35	!self.error.load(Ordering::Relaxed)
	36	}
	37	fn is_output_candidate(&self) -> bool {
	38	!self.output.load(Ordering::Relaxed) &&
	39	(self.complete.load(Ordering::Relaxed) \|\| self.error.load(Ordering::Relaxed))
	40	}
	41	}
	42
	43	pub struct ThreadDispatcher {
	44	fstate: Vec<FrameState>,
	45	pub max_threads: usize,
	46	cur_threads: usize,
	47	}
	48
	49	impl ThreadDispatcher {
	50	pub fn new() -> Self {
	51	Self {
	52	fstate: Vec::new(),
	53	max_threads: 3,
	54	cur_threads: 0,
	55	}
	56	}
	57	pub fn can_decode_more(&self) -> bool {
	58	let out_cand = self.fstate.iter().filter(\|state\| state.is_output_candidate()).count();
	59	if out_cand > self.max_threads {
	60	return false;
	61	}
	62	if (self.cur_threads < self.max_threads) \|\| (self.max_threads == 0) {
	63	true
	64	} else {
	65	let real_workers = self.fstate.iter().fold(0usize,
	66	\|acc, state\| acc + (state.is_working() as usize));
	67	real_workers < self.max_threads
	68	}
	69	}
	70	fn cleanup(&mut self) {
	71	for state in self.fstate.iter_mut() {
	72	if state.worker.is_some() && !state.is_working() {
	73	let mut ret = None;
	74	std::mem::swap(&mut state.worker, &mut ret);
	75	if let Some(handle) = ret {
	76	state.result = handle.join().unwrap();
	77	}
	78	self.cur_threads -= 1;
	79	}
	80	}
	81	}
	82	fn unref_frame(&mut self, id: u32) {
	83	let mut toremove = Vec::new();
	84	for state in self.fstate.iter() {
	85	if state.num_refs == 0 && state.output.load(Ordering::Relaxed) {
	86	toremove.push(state.get_id());
	87	}
	88	}
	89	if let Some(idx) = self.find_by_id(id) {
	90	let mut ref_frm = Vec::new();
	91	std::mem::swap(&mut ref_frm, &mut self.fstate[idx].ref_frames);
	92	for state in self.fstate.iter_mut() {
	93	if ref_frm.contains(&state.get_id()) {
	94	assert!(state.num_refs >= 2);
	95	state.num_refs -= 2;
	96	}
	97	}
	98	if self.fstate[idx].num_refs == 0 && self.fstate[idx].output.load(Ordering::Relaxed) {
	99	self.remove_frame(id);
	100	}
	101	}
	102	for &id in toremove.iter() {
	103	self.remove_frame(id);
	104	}
	105	}
	106	fn find_by_id(&self, id: u32) -> Option<usize> {
	107	self.fstate.iter().position(\|x\| x.get_id() == id)
	108	}
	109	fn set_completed(&self, id: u32) {
	110	if let Some(idx) = self.find_by_id(id) {
	111	self.fstate[idx].complete.store(true, Ordering::Relaxed);
	112	}
	113	}
	114	fn set_error(&self, id: u32) {
	115	if let Some(idx) = self.find_by_id(id) {
	116	self.fstate[idx].error.store(true, Ordering::Relaxed);
	117	}
	118	}
	119	pub fn update_pos(&self, id: u32, mb_pos: usize) {
	120	if let Some(idx) = self.find_by_id(id) {
	121	self.fstate[idx].mb_pos.store(mb_pos, Ordering::Relaxed);
	122	}
	123	}
	124	pub fn check_pos(&self, id: u32, mb_pos: usize) -> FrameDecodingStatus {
	125	if let Some(idx) = self.find_by_id(id) {
	126	let state = &self.fstate[idx];
	127	if !state.error.load(Ordering::Relaxed) {
	128	if state.complete.load(Ordering::Relaxed) \|\| mb_pos < state.mb_pos.load(Ordering::Relaxed) {
	129	FrameDecodingStatus::Ok
	130	} else {
	131	FrameDecodingStatus::NotReady
	132	}
	133	} else {
	134	FrameDecodingStatus::Error
	135	}
	136	} else {
	137	FrameDecodingStatus::NotFound
	138	}
	139	}
	140	fn remove_frame(&mut self, id: u32) {
	141	if let Some(idx) = self.find_by_id(id) {
	142	self.fstate.remove(idx);
	143	}
	144	}
	145	/*fn print_state(&self) {
	146	print!(" state:");
	147	for state in self.fstate.iter() {
	148	print!(" s{}b{}r{}{}{}{}", state.get_id(),
	149	state.mb_pos.load(Ordering::Relaxed), state.num_refs,
	150	if state.error.load(Ordering::Relaxed) { "E" } else {""},
	151	if state.complete.load(Ordering::Relaxed) {"C"} else {""},
	152	if state.output.load(Ordering::Relaxed) {"O"} else {""});
	153	}
	154	println!();
	155	}*/
	156	pub fn has_output(&self) -> bool {
	157	for state in self.fstate.iter() {
	158	if state.is_output_candidate() {
	159	return true;
	160	}
	161	}
	162	false
	163	}
	164	}
	165
	166	pub fn queue_decoding(disp: &mut Shareable<ThreadDispatcher>, mut fdec: FrameDecoder, initial_ref_frames: &[u32], ref_frames: &[u32]) {
	167	let barrier = Arc::new(Barrier::new(2));
	168	let starter = Arc::clone(&barrier);
	169
	170	let pinfo = fdec.cur_pic.clone();
	171	let pic_id = pinfo.full_id;
	172	let shared_disp = Arc::clone(disp);
	173	let worker = thread::Builder::new().name("frame ".to_string() + &pic_id.to_string()).spawn(move \|\| {
	174	barrier.wait();
	175
	176	let mut slices = Vec::new();
	177	std::mem::swap(&mut slices, &mut fdec.slices);
	178	let mut cur_mb = 0;
	179	for (hdr, hdr_size, refs, nal) in slices.iter() {
	180	if hdr.first_mb_in_slice != cur_mb {
	181	if let Ok(rd) = shared_disp.read() {
	182	rd.set_error(pic_id);
	183	} else {
	184	panic!("can't set error");
	185	}
	186	return Err(DecoderError::InvalidData);
	187	}
	188	match fdec.decode_slice(hdr, *hdr_size, refs, nal) {
	189	Ok(pos) => cur_mb = pos,
	190	Err(err) => {
	191	if let Ok(rd) = shared_disp.read() {
	192	rd.set_error(pic_id);
	193	} else {
	194	panic!("can't set error");
	195	}
	196	return Err(err);
	197	},
	198	};
	199	}
	200
	201	if cur_mb == fdec.num_mbs {
	202	if let Ok(rd) = shared_disp.read() {
	203	rd.set_completed(pic_id);
	204	} else {
	205	panic!("can't set status");
	206	}
	207	}
	208
	209	DecoderResult::Ok(())
	210	}).unwrap();
	211	let new_state = FrameState {
	212	pinfo,
	213	mb_pos: AtomicUsize::new(0),
	214	error: AtomicBool::new(false),
	215	complete: AtomicBool::new(false),
	216	output: AtomicBool::new(false),
	217	worker: Some(worker),
	218	result: DecoderResult::Err(DecoderError::Bug),
	219	num_refs: 0,
	220	ref_frames: initial_ref_frames.to_vec(),
	221	};
	222	if let Ok(ref mut ds) = disp.write() {
	223	let new_id = new_state.get_id();
	224	if ds.find_by_id(new_id).is_some() {
	225	ds.remove_frame(new_id);
	226	}
	227	ds.cleanup();
	228	ds.fstate.push(new_state);
	229	for state in ds.fstate.iter_mut() {
	230	if ref_frames.contains(&state.get_id()) {
	231	state.num_refs += 1;
	232	}
	233	if initial_ref_frames.contains(&state.get_id()) {
	234	state.num_refs += 1;
	235	}
	236	}
	237	ds.cur_threads += 1;
	238	starter.wait();
	239	} else {
	240	panic!("cannot invoke thread dispatcher");
	241	}
	242	}
	243
	244	pub fn wait_for_one(dispatch: &mut Shareable<ThreadDispatcher>) -> Result<PictureInfo, (DecoderError, u32)> {
	245	/*if let Ok(ref ds) = dispatch.read() {
	246	ds.print_state();
	247	}*/
	248	let start = std::time::Instant::now();
	249	'main_loop: loop {
	250	if std::time::Instant::now().duration_since(start) > std::time::Duration::from_millis(20000) { panic!(" too long!"); }
	251	if let Ok(ref ds) = dispatch.read() {
	252	let mut nw = 0;
	253	for state in ds.fstate.iter() {
	254	if state.is_working() {
	255	nw += 1;
	256	}
	257	if state.is_output_candidate() {
	258	break 'main_loop;
	259	}
	260	}
	261	if nw == 0 {
	262	return Err((DecoderError::NoFrame, 0));
	263	}
	264	} else {
	265	panic!("can't peek into status");
	266	}
	267	thread::yield_now();
	268	}
	269	if let Ok(ref mut ds) = dispatch.write() {
	270	ds.cleanup();
	271	let mut found = None;
	272	for state in ds.fstate.iter() {
	273	if state.is_output_candidate() {
	274	state.output.store(true, Ordering::Relaxed);
	275	if let DecoderResult::Err(err) = state.result {
	276	let id = state.get_id();
	277	let user_id = state.get_user_id();
	278	ds.unref_frame(id);
	279	return Err((err, user_id));
	280	} else {
	281	found = Some(state.pinfo.clone());
	282	break;
	283	}
	284	}
	285	}
	286	if let Some(ret) = found {
	287	ds.unref_frame(ret.full_id);
	288	Ok(ret)
	289	} else {
	290	unreachable!();
	291	}
	292	} else {
	293	panic!("can't grab status");
	294	}
	295	}
	296
	297	pub fn clear_threads(dispatch: &mut Shareable<ThreadDispatcher>) {
	298	/*if let Ok(ref ds) = dispatch.read() {
	299	ds.print_state();
	300	}*/
	301	let mut to_wait = Vec::new();
	302	if let Ok(ref mut ds) = dispatch.write() {
	303	while let Some(state) = ds.fstate.pop() {
	304	if let Some(handle) = state.worker {
	305	to_wait.push(handle);
	306	}
	307	}
	308	ds.cur_threads = 0;
	309	} else {
	310	panic!("can't grab status");
	311	}
	312	while let Some(handle) = to_wait.pop() {
	313	let _ = handle.join();
	314	}
	315	}