deviceplug / btleplug

Rust Cross-Platform Host-Side Bluetooth LE Access Library

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

AsyncWrite for Peripheral

robbym opened this issue · comments

I'm trying to implement AsyncWrite for writing to a characteristic. I created my own struct that holds a Peripheral, Characteristic, and a Future (the result of periph.write(&characteristic, data, WriteType::WithResponse). The struct looks like this:

#[pin_project]
pub struct CharWriteTask<'a, P>
where
    P: Peripheral,
{
    periph: P,
    characteristic: Characteristic,
    #[pin]
    data: Box<[u8]>,
    #[pin]
    future: Pin<Box<dyn Future<Output = Result<(), btleplug::Error>> + Send + 'a>>,
}

I am running into a lot of lifetime issues and even GPT-4 isn't much help :S.
I'm fairly new to dealing with async/await self-referencing structs, pinning, etc. Any help would be appreciated.

Ok, I have got something that's working, but I can't speak much to its correctness/soundness:

struct CharWriteTask<'a> {
    _periph: Pin<Box<Peripheral>>,
    _characteristic: Pin<Box<Characteristic>>,
    _buffer: Pin<Box<[u8]>>,
    future: Pin<Box<dyn futures::Future<Output = Result<(), btleplug::Error>> + Send + 'a>>,
}

impl<'a> CharWriteTask<'a> {
    fn new(
        periph: Peripheral,
        characteristic: Characteristic,
        buffer: &[u8],
        write_type: WriteType,
    ) -> CharWriteTask<'a> {
        let periph = Box::pin(periph);
        let characteristic = Box::pin(characteristic);
        let buffer = Pin::new(buffer.to_vec().into_boxed_slice());
        let periph_ptr = periph.as_ref().get_ref() as *const _;
        let characteristic_ptr = characteristic.as_ref().get_ref() as *const _;
        let buffer_ptr = buffer.as_ref().get_ref() as *const _;

        let future = unsafe {
            Peripheral::write(&*periph_ptr, &*characteristic_ptr, &*buffer_ptr, write_type)
        };

        CharWriteTask {
            _periph: periph,
            _characteristic: characteristic,
            _buffer: buffer,
            future,
        }
    }
}

impl<'a> Future for CharWriteTask<'a> {
    type Output = Result<(), btleplug::Error>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.future.poll_unpin(cx)
    }
}

#[pin_project]
struct PeripheralStream<'a> {
    periph: Peripheral,
    char_rx: Characteristic,
    char_tx: Characteristic,
    rx_stream: Pin<Box<dyn Stream<Item = ValueNotification> + Send>>,
    rx_buffer: VecDeque<u8>,
    #[pin]
    tx_write_task: Option<CharWriteTask<'a>>,
}

impl<'a> PeripheralStream<'a> {
    fn new(
        periph: Peripheral,
        service_uuid: Uuid,
        rx_char_uuid: Uuid,
        tx_char_uuid: Uuid,
    ) -> Pin<Box<dyn futures::Future<Output = Result<Self, anyhow::Error>> + 'a>> {
        Box::pin(async move {
            periph.connect().await?;
            periph.discover_services().await?;

            let mut char_rx = None;
            let mut char_tx = None;

            for service in periph.services() {
                if service.uuid == service_uuid {
                    for characteristic in service.characteristics {
                        if characteristic.uuid == rx_char_uuid
                            && characteristic.properties.contains(CharPropFlags::NOTIFY)
                        {
                            periph.subscribe(&characteristic).await?;
                            char_rx = Some(characteristic);
                        } else if characteristic.uuid == tx_char_uuid {
                            char_tx = Some(characteristic);
                        }
                    }
                }
            }

            let rx_stream = periph.notifications().await?;

            Ok(PeripheralStream {
                periph,
                char_rx: char_rx.unwrap(),
                char_tx: char_tx.unwrap(),
                rx_stream,
                rx_buffer: VecDeque::new(),
                tx_write_task: None,
            })
        })
    }
}

impl<'a> AsyncWrite for PeripheralStream<'a> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<io::Result<usize>> {
        let mut this = self.as_mut().project();

        if this.tx_write_task.is_none() {
            let write_task = CharWriteTask::new(
                this.periph.clone(),
                this.char_tx.clone(),
                buf,
                WriteType::WithoutResponse,
            );
            this.tx_write_task.set(Some(write_task));
        }

        match this.tx_write_task.as_mut().as_pin_mut().unwrap().poll(cx) {
            Poll::Ready(Ok(_)) => {
                this.tx_write_task.set(None);
                Poll::Ready(Ok(buf.len()))
            }
            Poll::Ready(Err(e)) => {
                this.tx_write_task.set(None);
                Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e)))
            }
            Poll::Pending => Poll::Pending,
        }
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<io::Result<()>> {
        std::task::Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        std::task::Poll::Ready(Ok(()))
    }
}

impl<'a> AsyncRead for PeripheralStream<'a> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let stream = self.get_mut();
        match stream.rx_stream.poll_next_unpin(cx) {
            std::task::Poll::Ready(Some(data)) => {
                stream.rx_buffer.extend(data.value);
            }
            _ => {}
        }
        let amount = stream.rx_buffer.len();
        if stream.rx_buffer.len() > 0 {
            let data = stream.rx_buffer.drain(0..amount).collect::<Vec<_>>();
            buf.put_slice(&data);
            return std::task::Poll::Ready(Ok(()));
        } else {
            return std::task::Poll::Pending;
        }
    }
}

This doesn't handle adjusting the write size depending on the characteristic buffer size or negotiated MTU size, so keep that in mind.