mz_service/
local.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Process-local transport for the [client](crate::client) module.
11
12use std::fmt;
13use std::task::{Context, Poll};
14use std::thread::Thread;
15
16use anyhow::anyhow;
17use itertools::Itertools;
18use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
19
20use crate::client::{GenericClient, Partitionable, Partitioned};
21
22/// A client to a thread in the same process.
23///
24/// The thread is unparked on every call to [`send`](LocalClient::send) and on
25/// `Drop`.
26#[derive(Debug)]
27pub struct LocalClient<C, R> {
28    rx: UnboundedReceiver<R>,
29    tx: UnboundedSender<C>,
30    thread: Thread,
31}
32
33impl<C, R> GenericClient<C, R> for LocalClient<C, R>
34where
35    C: fmt::Debug + Send,
36    R: fmt::Debug + Send,
37{
38    async fn send(&mut self, cmd: C) -> Result<(), anyhow::Error> {
39        self.tx.send(cmd).map_err(|_| anyhow!("receiver dropped"))?;
40        self.thread.unpark();
41
42        Ok(())
43    }
44
45    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<R>, anyhow::Error>> {
46        self.rx.poll_recv(cx).map(Ok)
47    }
48}
49
50impl<C, R> LocalClient<C, R> {
51    /// Create a new instance of [`LocalClient`] from its parts.
52    pub fn new(rx: UnboundedReceiver<R>, tx: UnboundedSender<C>, thread: Thread) -> Self {
53        Self { rx, tx, thread }
54    }
55
56    /// Create a new partitioned local client from parts for each client.
57    pub fn new_partitioned(
58        rxs: Vec<UnboundedReceiver<R>>,
59        txs: Vec<UnboundedSender<C>>,
60        threads: Vec<Thread>,
61    ) -> Partitioned<Self, C, R>
62    where
63        (C, R): Partitionable<C, R>,
64    {
65        let clients = rxs
66            .into_iter()
67            .zip_eq(txs)
68            .zip_eq(threads)
69            .map(|((rx, tx), thread)| LocalClient::new(rx, tx, thread))
70            .collect();
71        Partitioned::new(clients)
72    }
73}
74
75impl<C, R> Drop for LocalClient<C, R> {
76    fn drop(&mut self) {
77        self.thread.unpark();
78    }
79}
mz_service/local.rs

mz_service/
local.rs
