Skip to content

Replace all usages of BTreeMap with RBTree. #403

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 30, 2021
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 42 additions & 28 deletions drivers/android/process.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
// SPDX-License-Identifier: GPL-2.0

use alloc::{collections::btree_map::BTreeMap, sync::Arc, vec::Vec};
use alloc::{sync::Arc, vec::Vec};
use core::{
mem::{swap, MaybeUninit},
mem::{take, MaybeUninit},
ops::Range,
pin::Pin,
};
@@ -14,6 +14,7 @@ use kernel::{
linked_list::List,
pages::Pages,
prelude::*,
rbtree::RBTree,
sync::{Guard, Mutex, Ref},
user_ptr::{UserSlicePtr, UserSlicePtrReader},
Error,
@@ -62,11 +63,11 @@ impl Mapping {
pub(crate) struct ProcessInner {
is_manager: bool,
is_dead: bool,
threads: BTreeMap<i32, Arc<Thread>>,
threads: RBTree<i32, Arc<Thread>>,
ready_threads: List<Arc<Thread>>,
work: List<DeliverToReadListAdapter>,
mapping: Option<Mapping>,
nodes: BTreeMap<usize, Arc<Node>>,
nodes: RBTree<usize, Arc<Node>>,

delivered_deaths: List<Arc<NodeDeath>>,

@@ -85,11 +86,11 @@ impl ProcessInner {
Self {
is_manager: false,
is_dead: false,
threads: BTreeMap::new(),
threads: RBTree::new(),
ready_threads: List::new(),
work: List::new(),
mapping: None,
nodes: BTreeMap::new(),
nodes: RBTree::new(),
requested_thread_count: 0,
max_threads: 0,
started_thread_count: 0,
@@ -260,25 +261,25 @@ impl NodeRefInfo {
}

struct ProcessNodeRefs {
by_handle: BTreeMap<u32, NodeRefInfo>,
by_global_id: BTreeMap<u64, u32>,
by_handle: RBTree<u32, NodeRefInfo>,
by_global_id: RBTree<u64, u32>,
}

impl ProcessNodeRefs {
fn new() -> Self {
Self {
by_handle: BTreeMap::new(),
by_global_id: BTreeMap::new(),
by_handle: RBTree::new(),
by_global_id: RBTree::new(),
}
}
}

pub(crate) struct Process {
ctx: Ref<Context>,

// TODO: For now this a mutex because we have allocations in BTreeMap and RangeAllocator while
// holding the lock. We may want to split up the process state at some point to use a spin lock
// for the other fields; we can also get rid of allocations in BTreeMap once we replace it.
// TODO: For now this a mutex because we have allocations in RangeAllocator while holding the
// lock. We may want to split up the process state at some point to use a spin lock for the
// other fields.
// TODO: Make this private again.
pub(crate) inner: Mutex<ProcessInner>,

@@ -348,6 +349,7 @@ impl Process {

// Allocate a new `Thread` without holding any locks.
let ta = Thread::new(id, self.clone())?;
let node = RBTree::try_allocate_node(id, ta.clone())?;

let mut inner = self.inner.lock();

@@ -356,9 +358,7 @@ impl Process {
return Ok(thread.clone());
}

// TODO: We need a better solution here. Since this allocates, we cannot do it while
// holding a spinlock because it could block. It could panic too.
inner.threads.insert(id, ta.clone());
inner.threads.insert(node);
Ok(ta)
}

@@ -407,14 +407,14 @@ impl Process {

// Allocate the node before reacquiring the lock.
let node = Arc::try_new(Node::new(ptr, cookie, flags, self.clone()))?;
let rbnode = RBTree::try_allocate_node(ptr, node.clone())?;

let mut inner = self.inner.lock();
if let Some(node) = inner.get_existing_node_ref(ptr, cookie, strong, thread)? {
return Ok(node);
}

// TODO: Avoid allocation while holding lock.
inner.nodes.insert(ptr, node.clone());
inner.nodes.insert(rbnode);
Ok(inner.new_node_ref(node, strong, thread))
}

@@ -423,9 +423,25 @@ impl Process {
node_ref: NodeRef,
is_mananger: bool,
) -> Result<u32> {
{
let mut refs = self.node_refs.lock();

// Do a lookup before inserting.
if let Some(handle_ref) = refs.by_global_id.get(&node_ref.node.global_id) {
let handle = *handle_ref;
let info = refs.by_handle.get_mut(&handle).unwrap();
info.node_ref.absorb(node_ref);
return Ok(handle);
}
}

// Reserve memory for tree nodes.
let reserve1 = RBTree::try_reserve_node()?;
let reserve2 = RBTree::try_reserve_node()?;

let mut refs = self.node_refs.lock();

// Do a lookup before inserting.
// Do a lookup again as node may have been inserted before the lock was reacquired.
if let Some(handle_ref) = refs.by_global_id.get(&node_ref.node.global_id) {
let handle = *handle_ref;
let info = refs.by_handle.get_mut(&handle).unwrap();
@@ -449,9 +465,10 @@ impl Process {
if inner.is_dead {
return Err(Error::ESRCH);
}
// TODO: Two allocations below.
refs.by_global_id.insert(node_ref.node.global_id, target);
refs.by_handle.insert(target, NodeRefInfo::new(node_ref));
refs.by_global_id
.insert(reserve1.into_node(node_ref.node.global_id, target));
refs.by_handle
.insert(reserve2.into_node(target, NodeRefInfo::new(node_ref)));
Ok(target)
}

@@ -853,8 +870,7 @@ impl FileOperations for Process {
// Drop all references. We do this dance with `swap` to avoid destroying the references
// while holding the lock.
let mut refs = obj.node_refs.lock();
let mut node_refs = BTreeMap::new();
swap(&mut refs.by_handle, &mut node_refs);
let mut node_refs = take(&mut refs.by_handle);
drop(refs);

// Remove all death notifications from the nodes (that belong to a different process).
@@ -870,10 +886,8 @@ impl FileOperations for Process {

// Do similar dance for the state lock.
let mut inner = obj.inner.lock();
let mut threads = BTreeMap::new();
let mut nodes = BTreeMap::new();
swap(&mut inner.threads, &mut threads);
swap(&mut inner.nodes, &mut nodes);
let threads = take(&mut inner.threads);
let nodes = take(&mut inner.nodes);
drop(inner);

// Release all threads.
7 changes: 6 additions & 1 deletion rust/kernel/rbtree.rs
Original file line number Diff line number Diff line change
@@ -197,7 +197,6 @@ struct Node<K, V> {
/// Ok(())
/// }
/// ```
#[derive(Default)]
pub struct RBTree<K, V> {
root: bindings::rb_root,
_p: PhantomData<Node<K, V>>,
@@ -407,6 +406,12 @@ impl<K, V> RBTree<K, V> {
}
}

impl<K, V> Default for RBTree<K, V> {
fn default() -> Self {
Self::new()
}
}

impl<K, V> Drop for RBTree<K, V> {
fn drop(&mut self) {
// SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`.