// Copyright © 2021 Mozilla Foundation
//
// This program is made available under an ISC-style license. See the
// accompanying file LICENSE for details
use crossbeam_queue::ArrayQueue;
use mio::Token;
use std::cell::UnsafeCell;
use std::collections::VecDeque;
use std::io::{self, Error, ErrorKind, Result};
use std::marker::PhantomPinned;
use std::mem::ManuallyDrop;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Weak};
use crate::ipccore::EventLoopHandle;
// This provides a safe-ish method for a thread to allocate
// stack storage space for a result, then pass a (wrapped)
// pointer to that location to another thread via
// a CompletionWriter to eventually store a result into.
struct Completion<T> {
item: UnsafeCell<Option<T>>,
writer: AtomicBool,
_pin: PhantomPinned, // disable rustc's no-alias
}
impl<T> Completion<T> {
fn new() -> Self {
Completion {
item: UnsafeCell::new(None),
writer: AtomicBool::new(false),
_pin: PhantomPinned,
}
}
// Wait until the writer completes, then return the result.
// This is intended to be a single-use function, once the writer
// has completed any further attempts to wait will return None.
fn wait(&self) -> Option<T> {
// Wait for the writer to complete or be dropped.
while self.writer.load(Ordering::Acquire) {
std::thread::park();
}
unsafe { (*self.item.get()).take() }
}
// Create a writer for the other thread to store the
// expected result into.
fn writer(&self) -> CompletionWriter<T> {
assert!(!self.writer.load(Ordering::Relaxed));
self.writer.store(true, Ordering::Release);
CompletionWriter {
ptr: self as *const _ as *mut _,
waiter: std::thread::current(),
}
}
}
impl<T> Drop for Completion<T> {
fn drop(&mut self) {
// Wait for the outstanding writer to complete before
// dropping, since the CompletionWriter references
// memory owned by this object.
while self.writer.load(Ordering::Acquire) {
std::thread::park();
}
}
}
struct CompletionWriter<T> {
ptr: *mut Completion<T>, // Points to a Completion on another thread's stack
waiter: std::thread::Thread, // Identifies thread waiting for completion
}
impl<T> CompletionWriter<T> {
fn set(self, value: T) {
// Store the result into the Completion's memory.
// Since `set` consumes `self`, rely on `Drop` to
// mark the writer as done and wake the Completion's
// thread.
unsafe {
assert!((*self.ptr).writer.load(Ordering::Relaxed));
*(*self.ptr).item.get() = Some(value);
}
}
}
impl<T> Drop for CompletionWriter<T> {
fn drop(&mut self) {
// Mark writer as complete - if `set` was not called,
// the waiter will receive `None`.
unsafe {
(*self.ptr).writer.store(false, Ordering::Release);
}
// Wake the Completion's thread.
self.waiter.unpark();
}
}
// Safety: CompletionWriter holds a pointer to a Completion
// residing on another thread's stack. The Completion always
// waits for an outstanding writer if present, and CompletionWriter
// releases the waiter and wakes the Completion's thread on drop,
// so this pointer will always be live for the duration of a
// CompletionWriter.
unsafe impl<T> Send for CompletionWriter<T> {}
// RPC message handler. Implemented by ClientHandler (for Client)
// and ServerHandler (for Server).
pub(crate) trait Handler {
type In;
type Out;
// Consume a request
fn consume(&mut self, request: Self::In) -> Result<()>;
// Produce a response
fn produce(&mut self) -> Result<Option<Self::Out>>;
}
// Client RPC definition. This supplies the expected message
// request and response types.
pub trait Client {
type ServerMessage;
type ClientMessage;
}
// Server RPC definition. This supplies the expected message
// request and response types. `process` is passed inbound RPC
// requests by the ServerHandler to be responded to by the server.
pub trait Server {
type ServerMessage;
type ClientMessage;
fn process(&mut self, req: Self::ServerMessage) -> Self::ClientMessage;
}
// RPC Client Proxy implementation.
type ProxyRequest<Request, Response> = (Request, CompletionWriter<Response>);
// RPC Proxy that may be `clone`d for use by multiple owners/threads.
// A Proxy `call` arranges for the supplied request to be transmitted
// to the associated Server via RPC and blocks awaiting the response
// via the associated `Completion`.
// A ClientHandler normally lives until the last Proxy is dropped, but if the ClientHandler
// encounters an internal error, `requests` will fail to upgrade, allowing
// the proxy to report an error.
#[derive(Debug)]
pub struct Proxy<Request, Response> {
handle: Option<(EventLoopHandle, Token)>,
requests: ManuallyDrop<RequestQueueSender<ProxyRequest<Request, Response>>>,
}
impl<Request, Response> Proxy<Request, Response> {
fn new(requests: RequestQueueSender<ProxyRequest<Request, Response>>) -> Self {
Self {
handle: None,
requests: ManuallyDrop::new(requests),
}
}
pub fn call(&self, request: Request) -> Result<Response> {
let response = Completion::new();
self.requests.push((request, response.writer()))?;
self.wake_connection();
match response.wait() {
Some(resp) => Ok(resp),
None => Err(Error::new(ErrorKind::Other, "proxy recv error")),
}
}
pub(crate) fn connect_event_loop(&mut self, handle: EventLoopHandle, token: Token) {
self.handle = Some((handle, token));
}
fn wake_connection(&self) {
let (handle, token) = self
.handle
.as_ref()
.expect("proxy not connected to event loop");
handle.wake_connection(*token);
}
}
impl<Request, Response> Clone for Proxy<Request, Response> {
fn clone(&self) -> Self {
let mut clone = Self::new((*self.requests).clone());
let (handle, token) = self
.handle
.as_ref()
.expect("proxy not connected to event loop");
clone.connect_event_loop(handle.clone(), *token);
clone
}
}
impl<Request, Response> Drop for Proxy<Request, Response> {
fn drop(&mut self) {
trace!("Proxy drop, waking EventLoop");
// Must drop `requests` before waking the connection, otherwise
// the wake may be processed before the (last) weak reference is
// dropped.
let last_proxy = self.requests.live_proxies();
unsafe {
ManuallyDrop::drop(&mut self.requests);
}
if last_proxy == 1 && self.handle.is_some() {
self.wake_connection()
}
}
}
const RPC_CLIENT_INITIAL_PROXIES: usize = 32; // Initial proxy pre-allocation per client.
// Client-specific Handler implementation.
// The IPC EventLoop Driver calls this to execute client-specific
// RPC handling. Serialized messages sent via a Proxy are queued
// for transmission when `produce` is called.
// Deserialized messages are passed via `consume` to
// trigger response completion by sending the response via a channel
// connected to a ProxyResponse.
pub(crate) struct ClientHandler<C: Client> {
in_flight: VecDeque<CompletionWriter<C::ClientMessage>>,
requests: Arc<RequestQueue<ProxyRequest<C::ServerMessage, C::ClientMessage>>>,
}
impl<C: Client> ClientHandler<C> {
fn new(
requests: Arc<RequestQueue<ProxyRequest<C::ServerMessage, C::ClientMessage>>>,
) -> ClientHandler<C> {
ClientHandler::<C> {
in_flight: VecDeque::with_capacity(RPC_CLIENT_INITIAL_PROXIES),
requests,
}
}
}
impl<C: Client> Handler for ClientHandler<C> {
type In = C::ClientMessage;
type Out = C::ServerMessage;
fn consume(&mut self, response: Self::In) -> Result<()> {
trace!("ClientHandler::consume");
if let Some(response_writer) = self.in_flight.pop_front() {
response_writer.set(response);
} else {
return Err(Error::new(ErrorKind::Other, "request/response mismatch"));
}
Ok(())
}
fn produce(&mut self) -> Result<Option<Self::Out>> {
trace!("ClientHandler::produce");
// If the weak count is zero, no proxies are attached and
// no further proxies can be attached since every proxy
// after the initial one is cloned from an existing instance.
self.requests.check_live_proxies()?;
// Try to get a new message
match self.requests.pop() {
Some((request, response_writer)) => {
trace!(" --> received request");
self.in_flight.push_back(response_writer);
Ok(Some(request))
}
None => {
trace!(" --> no request");
Ok(None)
}
}
}
}
#[derive(Debug)]
pub(crate) struct RequestQueue<T> {
queue: ArrayQueue<T>,
}
impl<T> RequestQueue<T> {
pub(crate) fn new(size: usize) -> Self {
RequestQueue {
queue: ArrayQueue::new(size),
}
}
pub(crate) fn pop(&self) -> Option<T> {
self.queue.pop()
}
pub(crate) fn new_sender(self: &Arc<Self>) -> RequestQueueSender<T> {
RequestQueueSender {
inner: Arc::downgrade(self),
}
}
pub(crate) fn check_live_proxies(self: &Arc<Self>) -> Result<()> {
if Arc::weak_count(self) == 0 {
return Err(io::ErrorKind::ConnectionAborted.into());
}
Ok(())
}
}
pub(crate) struct RequestQueueSender<T> {
inner: Weak<RequestQueue<T>>,
}
impl<T> RequestQueueSender<T> {
pub(crate) fn push(&self, request: T) -> Result<()> {
if let Some(consumer) = self.inner.upgrade() {
if consumer.queue.push(request).is_err() {
debug!("Proxy[{:p}]: call failed - CH::requests full", self);
return Err(io::ErrorKind::ConnectionAborted.into());
}
return Ok(());
}
debug!("Proxy[{:p}]: call failed - CH::requests dropped", self);
Err(Error::new(ErrorKind::Other, "proxy send error"))
}
pub(crate) fn live_proxies(&self) -> usize {
Weak::weak_count(&self.inner)
}
}
impl<T> Clone for RequestQueueSender<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
impl<T> std::fmt::Debug for RequestQueueSender<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RequestQueueProducer")
.field("inner", &self.inner.as_ptr())
.finish()
}
}
#[allow(clippy::type_complexity)]
pub(crate) fn make_client<C: Client>(
) -> Result<(ClientHandler<C>, Proxy<C::ServerMessage, C::ClientMessage>)> {
let requests = Arc::new(RequestQueue::new(RPC_CLIENT_INITIAL_PROXIES));
let proxy_req = requests.new_sender();
let handler = ClientHandler::new(requests);
Ok((handler, Proxy::new(proxy_req)))
}
// Server-specific Handler implementation.
// The IPC EventLoop Driver calls this to execute server-specific
// RPC handling. Deserialized messages are passed via `consume` to the
// associated `server` for processing. Server responses are then queued
// for RPC to the associated client when `produce` is called.
pub(crate) struct ServerHandler<S: Server> {
server: S,
in_flight: VecDeque<S::ClientMessage>,
}
impl<S: Server> Handler for ServerHandler<S> {
type In = S::ServerMessage;
type Out = S::ClientMessage;
fn consume(&mut self, message: Self::In) -> Result<()> {
trace!("ServerHandler::consume");
let response = self.server.process(message);
self.in_flight.push_back(response);
Ok(())
}
fn produce(&mut self) -> Result<Option<Self::Out>> {
trace!("ServerHandler::produce");
// Return the ready response
match self.in_flight.pop_front() {
Some(res) => {
trace!(" --> received response");
Ok(Some(res))
}
None => {
trace!(" --> no response ready");
Ok(None)
}
}
}
}
const RPC_SERVER_INITIAL_CLIENTS: usize = 32; // Initial client allocation per server.
pub(crate) fn make_server<S: Server>(server: S) -> ServerHandler<S> {
ServerHandler::<S> {
server,
in_flight: VecDeque::with_capacity(RPC_SERVER_INITIAL_CLIENTS),
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn basic() {
let queue = Arc::new(RequestQueue::new(1));
let producer = queue.new_sender();
assert!(queue.pop().is_none());
producer.push(1).unwrap();
assert!(queue.pop().is_some());
assert!(queue.pop().is_none());
}
#[test]
fn queue_dropped() {
let queue = Arc::new(RequestQueue::new(1));
let producer = queue.new_sender();
drop(queue);
assert!(producer.push(1).is_err());
}
#[test]
fn queue_full() {
let queue = Arc::new(RequestQueue::new(1));
let producer = queue.new_sender();
producer.push(1).unwrap();
assert!(producer.push(2).is_err());
}
#[test]
fn queue_producer_clone() {
let queue = Arc::new(RequestQueue::new(1));
let producer = queue.new_sender();
let producer2 = producer.clone();
producer.push(1).unwrap();
assert!(producer2.push(2).is_err());
}
#[test]
fn queue_producer_drop() {
let queue = Arc::new(RequestQueue::new(1));
let producer = queue.new_sender();
let producer2 = producer.clone();
drop(producer);
assert!(producer2.push(2).is_ok());
}
#[test]
fn queue_producer_weak() {
let queue = Arc::new(RequestQueue::new(1));
let producer = queue.new_sender();
let producer2 = producer.clone();
drop(queue);
assert!(producer2.push(2).is_err());
}
#[test]
fn queue_producer_shutdown() {
let queue = Arc::new(RequestQueue::new(1));
let producer = queue.new_sender();
let producer2 = producer.clone();
producer.push(1).unwrap();
assert!(Arc::weak_count(&queue) == 2);
drop(producer);
assert!(Arc::weak_count(&queue) == 1);
drop(producer2);
assert!(Arc::weak_count(&queue) == 0);
}
}
