blob: c30ccd9ed620e6e95b866f5e6df043b9823d659c [file] [log] [blame]
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use std::cmp;
use std::io;
use std::io::{Read, Write};
use super::{TReadTransport, TReadTransportFactory, TWriteTransport, TWriteTransportFactory};
/// Default capacity of the read buffer in bytes.
const READ_CAPACITY: usize = 4096;
/// Default capacity of the write buffer in bytes.
const WRITE_CAPACITY: usize = 4096;
/// Transport that reads framed messages.
///
/// A `TFramedReadTransport` maintains a fixed-size internal read buffer.
/// On a call to `TFramedReadTransport::read(...)` one full message - both
/// fixed-length header and bytes - is read from the wrapped channel and
/// buffered. Subsequent read calls are serviced from the internal buffer
/// until it is exhausted, at which point the next full message is read
/// from the wrapped channel.
///
/// # Examples
///
/// Create and use a `TFramedReadTransport`.
///
/// ```no_run
/// use std::io::Read;
/// use thrift::transport::{TFramedReadTransport, TTcpChannel};
///
/// let mut c = TTcpChannel::new();
/// c.open("localhost:9090").unwrap();
///
/// let mut t = TFramedReadTransport::new(c);
///
/// t.read(&mut vec![0u8; 1]).unwrap();
/// ```
#[derive(Debug)]
pub struct TFramedReadTransport<C>
where
C: Read,
{
buf: Vec<u8>,
pos: usize,
cap: usize,
chan: C,
}
impl<C> TFramedReadTransport<C>
where
C: Read,
{
/// Create a `TFramedReadTransport` with a default-sized
/// internal read buffer that wraps the given `TIoChannel`.
pub fn new(channel: C) -> TFramedReadTransport<C> {
TFramedReadTransport::with_capacity(READ_CAPACITY, channel)
}
/// Create a `TFramedTransport` with an internal read buffer
/// of size `read_capacity` that wraps the given `TIoChannel`.
pub fn with_capacity(read_capacity: usize, channel: C) -> TFramedReadTransport<C> {
TFramedReadTransport {
buf: vec![0; read_capacity], // FIXME: do I actually have to do this?
pos: 0,
cap: 0,
chan: channel,
}
}
}
impl<C> Read for TFramedReadTransport<C>
where
C: Read,
{
fn read(&mut self, b: &mut [u8]) -> io::Result<usize> {
if self.cap - self.pos == 0 {
let message_size = self.chan.read_i32::<BigEndian>()? as usize;
let buf_capacity = cmp::max(message_size, READ_CAPACITY);
self.buf.resize(buf_capacity, 0);
self.chan.read_exact(&mut self.buf[..message_size])?;
self.cap = message_size as usize;
self.pos = 0;
}
let nread = cmp::min(b.len(), self.cap - self.pos);
b[..nread].clone_from_slice(&self.buf[self.pos..self.pos + nread]);
self.pos += nread;
Ok(nread)
}
}
/// Factory for creating instances of `TFramedReadTransport`.
#[derive(Default)]
pub struct TFramedReadTransportFactory;
impl TFramedReadTransportFactory {
pub fn new() -> TFramedReadTransportFactory {
TFramedReadTransportFactory {}
}
}
impl TReadTransportFactory for TFramedReadTransportFactory {
/// Create a `TFramedReadTransport`.
fn create(&self, channel: Box<dyn Read + Send>) -> Box<dyn TReadTransport + Send> {
Box::new(TFramedReadTransport::new(channel))
}
}
/// Transport that writes framed messages.
///
/// A `TFramedWriteTransport` maintains a fixed-size internal write buffer. All
/// writes are made to this buffer and are sent to the wrapped channel only
/// when `TFramedWriteTransport::flush()` is called. On a flush a fixed-length
/// header with a count of the buffered bytes is written, followed by the bytes
/// themselves.
///
/// # Examples
///
/// Create and use a `TFramedWriteTransport`.
///
/// ```no_run
/// use std::io::Write;
/// use thrift::transport::{TFramedWriteTransport, TTcpChannel};
///
/// let mut c = TTcpChannel::new();
/// c.open("localhost:9090").unwrap();
///
/// let mut t = TFramedWriteTransport::new(c);
///
/// t.write(&[0x00]).unwrap();
/// t.flush().unwrap();
/// ```
#[derive(Debug)]
pub struct TFramedWriteTransport<C>
where
C: Write,
{
buf: Vec<u8>,
channel: C,
}
impl<C> TFramedWriteTransport<C>
where
C: Write,
{
/// Create a `TFramedWriteTransport` with default-sized internal
/// write buffer that wraps the given `TIoChannel`.
pub fn new(channel: C) -> TFramedWriteTransport<C> {
TFramedWriteTransport::with_capacity(WRITE_CAPACITY, channel)
}
/// Create a `TFramedWriteTransport` with an internal write buffer
/// of size `write_capacity` that wraps the given `TIoChannel`.
pub fn with_capacity(write_capacity: usize, channel: C) -> TFramedWriteTransport<C> {
TFramedWriteTransport {
buf: Vec::with_capacity(write_capacity),
channel,
}
}
}
impl<C> Write for TFramedWriteTransport<C>
where
C: Write,
{
fn write(&mut self, b: &[u8]) -> io::Result<usize> {
let current_capacity = self.buf.capacity();
let available_space = current_capacity - self.buf.len();
if b.len() > available_space {
let additional_space = cmp::max(b.len() - available_space, current_capacity);
self.buf.reserve(additional_space);
}
self.buf.extend_from_slice(b);
Ok(b.len())
}
fn flush(&mut self) -> io::Result<()> {
let message_size = self.buf.len();
if let 0 = message_size {
return Ok(());
} else {
self.channel.write_i32::<BigEndian>(message_size as i32)?;
}
// will spin if the underlying channel can't be written to
let mut byte_index = 0;
while byte_index < message_size {
let nwrite = self.channel.write(&self.buf[byte_index..message_size])?;
byte_index = cmp::min(byte_index + nwrite, message_size);
}
let buf_capacity = cmp::min(self.buf.capacity(), WRITE_CAPACITY);
self.buf.resize(buf_capacity, 0);
self.buf.clear();
self.channel.flush()
}
}
/// Factory for creating instances of `TFramedWriteTransport`.
#[derive(Default)]
pub struct TFramedWriteTransportFactory;
impl TFramedWriteTransportFactory {
pub fn new() -> TFramedWriteTransportFactory {
TFramedWriteTransportFactory {}
}
}
impl TWriteTransportFactory for TFramedWriteTransportFactory {
/// Create a `TFramedWriteTransport`.
fn create(&self, channel: Box<dyn Write + Send>) -> Box<dyn TWriteTransport + Send> {
Box::new(TFramedWriteTransport::new(channel))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::transport::mem::TBufferChannel;
// FIXME: test a forced reserve
#[test]
fn must_read_message_smaller_than_initial_buffer_size() {
let c = TBufferChannel::with_capacity(10, 10);
let mut t = TFramedReadTransport::with_capacity(8, c);
t.chan.set_readable_bytes(&[
0x00, 0x00, 0x00, 0x04, /* message size */
0x00, 0x01, 0x02, 0x03, /* message body */
]);
let mut buf = vec![0; 8];
// we've read exactly 4 bytes
assert_eq!(t.read(&mut buf).unwrap(), 4);
assert_eq!(&buf[..4], &[0x00, 0x01, 0x02, 0x03]);
}
#[test]
fn must_read_message_greater_than_initial_buffer_size() {
let c = TBufferChannel::with_capacity(10, 10);
let mut t = TFramedReadTransport::with_capacity(2, c);
t.chan.set_readable_bytes(&[
0x00, 0x00, 0x00, 0x04, /* message size */
0x00, 0x01, 0x02, 0x03, /* message body */
]);
let mut buf = vec![0; 8];
// we've read exactly 4 bytes
assert_eq!(t.read(&mut buf).unwrap(), 4);
assert_eq!(&buf[..4], &[0x00, 0x01, 0x02, 0x03]);
}
#[test]
fn must_read_multiple_messages_in_sequence_correctly() {
let c = TBufferChannel::with_capacity(10, 10);
let mut t = TFramedReadTransport::with_capacity(2, c);
//
// 1st message
//
t.chan.set_readable_bytes(&[
0x00, 0x00, 0x00, 0x04, /* message size */
0x00, 0x01, 0x02, 0x03, /* message body */
]);
let mut buf = vec![0; 8];
// we've read exactly 4 bytes
assert_eq!(t.read(&mut buf).unwrap(), 4);
assert_eq!(&buf, &[0x00, 0x01, 0x02, 0x03, 0x00, 0x00, 0x00, 0x00]);
//
// 2nd message
//
t.chan.set_readable_bytes(&[
0x00, 0x00, 0x00, 0x01, /* message size */
0x04, /* message body */
]);
let mut buf = vec![0; 8];
// we've read exactly 1 byte
assert_eq!(t.read(&mut buf).unwrap(), 1);
assert_eq!(&buf, &[0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
}
#[test]
fn must_write_message_smaller_than_buffer_size() {
let mem = TBufferChannel::with_capacity(0, 0);
let mut t = TFramedWriteTransport::with_capacity(20, mem);
let b = vec![0; 10];
// should have written 10 bytes
assert_eq!(t.write(&b).unwrap(), 10);
}
#[test]
fn must_return_zero_if_caller_calls_write_with_empty_buffer() {
let mem = TBufferChannel::with_capacity(0, 10);
let mut t = TFramedWriteTransport::with_capacity(10, mem);
let expected: [u8; 0] = [];
assert_eq!(t.write(&[]).unwrap(), 0);
assert_eq_transport_written_bytes!(t, expected);
}
#[test]
fn must_write_to_inner_transport_on_flush() {
let mem = TBufferChannel::with_capacity(10, 10);
let mut t = TFramedWriteTransport::new(mem);
let b: [u8; 5] = [0x00, 0x01, 0x02, 0x03, 0x04];
assert_eq!(t.write(&b).unwrap(), 5);
assert_eq_transport_num_written_bytes!(t, 0);
assert!(t.flush().is_ok());
let expected_bytes = [
0x00, 0x00, 0x00, 0x05, /* message size */
0x00, 0x01, 0x02, 0x03, 0x04, /* message body */
];
assert_eq_transport_written_bytes!(t, expected_bytes);
}
#[test]
fn must_write_message_greater_than_buffer_size_00() {
let mem = TBufferChannel::with_capacity(0, 10);
// IMPORTANT: DO **NOT** CHANGE THE WRITE_CAPACITY OR THE NUMBER OF BYTES TO BE WRITTEN!
// these lengths were chosen to be just long enough
// that doubling the capacity is a **worse** choice than
// simply resizing the buffer to b.len()
let mut t = TFramedWriteTransport::with_capacity(1, mem);
let b = [0x00, 0x01, 0x02];
// should have written 3 bytes
assert_eq!(t.write(&b).unwrap(), 3);
assert_eq_transport_num_written_bytes!(t, 0);
assert!(t.flush().is_ok());
let expected_bytes = [
0x00, 0x00, 0x00, 0x03, /* message size */
0x00, 0x01, 0x02, /* message body */
];
assert_eq_transport_written_bytes!(t, expected_bytes);
}
#[test]
fn must_write_message_greater_than_buffer_size_01() {
let mem = TBufferChannel::with_capacity(0, 10);
// IMPORTANT: DO **NOT** CHANGE THE WRITE_CAPACITY OR THE NUMBER OF BYTES TO BE WRITTEN!
// these lengths were chosen to be just long enough
// that doubling the capacity is a **better** choice than
// simply resizing the buffer to b.len()
let mut t = TFramedWriteTransport::with_capacity(2, mem);
let b = [0x00, 0x01, 0x02];
// should have written 3 bytes
assert_eq!(t.write(&b).unwrap(), 3);
assert_eq_transport_num_written_bytes!(t, 0);
assert!(t.flush().is_ok());
let expected_bytes = [
0x00, 0x00, 0x00, 0x03, /* message size */
0x00, 0x01, 0x02, /* message body */
];
assert_eq_transport_written_bytes!(t, expected_bytes);
}
#[test]
fn must_return_error_if_nothing_can_be_written_to_inner_transport_on_flush() {
let mem = TBufferChannel::with_capacity(0, 0);
let mut t = TFramedWriteTransport::with_capacity(1, mem);
let b = vec![0; 10];
// should have written 10 bytes
assert_eq!(t.write(&b).unwrap(), 10);
// let's flush
let r = t.flush();
// this time we'll error out because the flush can't write to the underlying channel
assert!(r.is_err());
}
#[test]
fn must_write_successfully_after_flush() {
// IMPORTANT: write capacity *MUST* be greater
// than message sizes used in this test + 4-byte frame header
let mem = TBufferChannel::with_capacity(0, 10);
let mut t = TFramedWriteTransport::with_capacity(5, mem);
// write and flush
let first_message: [u8; 5] = [0x00, 0x01, 0x02, 0x03, 0x04];
assert_eq!(t.write(&first_message).unwrap(), 5);
assert!(t.flush().is_ok());
let mut expected = Vec::new();
expected.write_all(&[0x00, 0x00, 0x00, 0x05]).unwrap(); // message size
expected.extend_from_slice(&first_message);
// check the flushed bytes
assert_eq!(t.channel.write_bytes(), expected);
// reset our underlying transport
t.channel.empty_write_buffer();
let second_message: [u8; 3] = [0x05, 0x06, 0x07];
assert_eq!(t.write(&second_message).unwrap(), 3);
assert!(t.flush().is_ok());
expected.clear();
expected.write_all(&[0x00, 0x00, 0x00, 0x03]).unwrap(); // message size
expected.extend_from_slice(&second_message);
// check the flushed bytes
assert_eq!(t.channel.write_bytes(), expected);
}
}