proxmox-backup/src/client/pipe_to_stream.rs

// Implement simple flow control for h2 client
//
// See also: hyper/src/proto/h2/mod.rs

use bytes::Bytes;
use failure::*;
use futures::{try_ready, Async, Future, Poll};
use h2::SendStream;

pub struct PipeToSendStream {
    body_tx: SendStream<Bytes>,
    data: Option<Bytes>,
}

impl PipeToSendStream {
    pub fn new(data: Bytes, tx: SendStream<Bytes>) -> PipeToSendStream {
        PipeToSendStream {
            body_tx: tx,
            data: Some(data),
        }
    }
}

impl Future for PipeToSendStream {
    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if self.data != None {
                // just reserve 1 byte to make sure there's some
                // capacity available. h2 will handle the capacity
                // management for the actual body chunk.
                self.body_tx.reserve_capacity(1);

                if self.body_tx.capacity() == 0 {
                    loop {
                        match try_ready!(self.body_tx.poll_capacity().map_err(Error::from)) {
                            Some(0) => {}
                            Some(_) => break,
                            None => return Err(format_err!("protocol canceled")),
                        }
                    }
                } else {
                    if let Async::Ready(reason) = self.body_tx.poll_reset().map_err(Error::from)? {
                        return Err(format_err!("stream received RST_STREAM: {:?}", reason));
                    }
                }

                self.body_tx
                    .send_data(self.data.take().unwrap(), true)
                    .map_err(Error::from)?;

                return Ok(Async::Ready(()));
            } else {
                if let Async::Ready(reason) = self.body_tx.poll_reset().map_err(Error::from)? {
                    return Err(format_err!("stream received RST_STREAM: {:?}", reason));
                }
                return Ok(Async::Ready(()));
            }
        }
    }
}