ceph/src/jaegertracing/thrift/lib/rs/src/transport/framed.rs

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements. See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership. The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License. You may obtain a copy of the License at
   8 //
   9 //   http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing,
  12 // software distributed under the License is distributed on an
  13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 // KIND, either express or implied. See the License for the
  15 // specific language governing permissions and limitations
  16 // under the License.
  17
  18 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
  19 use std::cmp;
  20 use std::io;
  21 use std::io::{Read, Write};
  22
  23 use super::{TReadTransport, TReadTransportFactory, TWriteTransport, TWriteTransportFactory};
  24
  25 /// Default capacity of the read buffer in bytes.
  26 const READ_CAPACITY: usize = 4096;
  27
  28 /// Default capacity of the write buffer in bytes.
  29 const WRITE_CAPACITY: usize = 4096;
  30
  31 /// Transport that reads framed messages.
  32 ///
  33 /// A `TFramedReadTransport` maintains a fixed-size internal read buffer.
  34 /// On a call to `TFramedReadTransport::read(...)` one full message - both
  35 /// fixed-length header and bytes - is read from the wrapped channel and
  36 /// buffered. Subsequent read calls are serviced from the internal buffer
  37 /// until it is exhausted, at which point the next full message is read
  38 /// from the wrapped channel.
  39 ///
  40 /// # Examples
  41 ///
  42 /// Create and use a `TFramedReadTransport`.
  43 ///
  44 /// ```no_run
  45 /// use std::io::Read;
  46 /// use thrift::transport::{TFramedReadTransport, TTcpChannel};
  47 ///
  48 /// let mut c = TTcpChannel::new();
  49 /// c.open("localhost:9090").unwrap();
  50 ///
  51 /// let mut t = TFramedReadTransport::new(c);
  52 ///
  53 /// t.read(&mut vec![0u8; 1]).unwrap();
  54 /// ```
  55 #[derive(Debug)]
  56 pub struct TFramedReadTransport<C>
  57 where
  58     C: Read,
  59 {
  60     buf: Vec<u8>,
  61     pos: usize,
  62     cap: usize,
  63     chan: C,
  64 }
  65
  66 impl<C> TFramedReadTransport<C>
  67 where
  68     C: Read,
  69 {
  70     /// Create a `TFramedReadTransport` with a default-sized
  71     /// internal read buffer that wraps the given `TIoChannel`.
  72     pub fn new(channel: C) -> TFramedReadTransport<C> {
  73         TFramedReadTransport::with_capacity(READ_CAPACITY, channel)
  74     }
  75
  76     /// Create a `TFramedTransport` with an internal read buffer
  77     /// of size `read_capacity` that wraps the given `TIoChannel`.
  78     pub fn with_capacity(read_capacity: usize, channel: C) -> TFramedReadTransport<C> {
  79         TFramedReadTransport {
  80             buf: vec![0; read_capacity], // FIXME: do I actually have to do this?
  81             pos: 0,
  82             cap: 0,
  83             chan: channel,
  84         }
  85     }
  86 }
  87
  88 impl<C> Read for TFramedReadTransport<C>
  89 where
  90     C: Read,
  91 {
  92     fn read(&mut self, b: &mut [u8]) -> io::Result<usize> {
  93         if self.cap - self.pos == 0 {
  94             let message_size = self.chan.read_i32::<BigEndian>()? as usize;
  95
  96             let buf_capacity = cmp::max(message_size, READ_CAPACITY);
  97             self.buf.resize(buf_capacity, 0);
  98
  99             self.chan.read_exact(&mut self.buf[..message_size])?;
 100             self.cap = message_size as usize;
 101             self.pos = 0;
 102         }
 103
 104         let nread = cmp::min(b.len(), self.cap - self.pos);
 105         b[..nread].clone_from_slice(&self.buf[self.pos..self.pos + nread]);
 106         self.pos += nread;
 107
 108         Ok(nread)
 109     }
 110 }
 111
 112 /// Factory for creating instances of `TFramedReadTransport`.
 113 #[derive(Default)]
 114 pub struct TFramedReadTransportFactory;
 115
 116 impl TFramedReadTransportFactory {
 117     pub fn new() -> TFramedReadTransportFactory {
 118         TFramedReadTransportFactory {}
 119     }
 120 }
 121
 122 impl TReadTransportFactory for TFramedReadTransportFactory {
 123     /// Create a `TFramedReadTransport`.
 124     fn create(&self, channel: Box<dyn Read + Send>) -> Box<dyn TReadTransport + Send> {
 125         Box::new(TFramedReadTransport::new(channel))
 126     }
 127 }
 128
 129 /// Transport that writes framed messages.
 130 ///
 131 /// A `TFramedWriteTransport` maintains a fixed-size internal write buffer. All
 132 /// writes are made to this buffer and are sent to the wrapped channel only
 133 /// when `TFramedWriteTransport::flush()` is called. On a flush a fixed-length
 134 /// header with a count of the buffered bytes is written, followed by the bytes
 135 /// themselves.
 136 ///
 137 /// # Examples
 138 ///
 139 /// Create and use a `TFramedWriteTransport`.
 140 ///
 141 /// ```no_run
 142 /// use std::io::Write;
 143 /// use thrift::transport::{TFramedWriteTransport, TTcpChannel};
 144 ///
 145 /// let mut c = TTcpChannel::new();
 146 /// c.open("localhost:9090").unwrap();
 147 ///
 148 /// let mut t = TFramedWriteTransport::new(c);
 149 ///
 150 /// t.write(&[0x00]).unwrap();
 151 /// t.flush().unwrap();
 152 /// ```
 153 #[derive(Debug)]
 154 pub struct TFramedWriteTransport<C>
 155 where
 156     C: Write,
 157 {
 158     buf: Vec<u8>,
 159     channel: C,
 160 }
 161
 162 impl<C> TFramedWriteTransport<C>
 163 where
 164     C: Write,
 165 {
 166     /// Create a `TFramedWriteTransport` with default-sized internal
 167     /// write buffer that wraps the given `TIoChannel`.
 168     pub fn new(channel: C) -> TFramedWriteTransport<C> {
 169         TFramedWriteTransport::with_capacity(WRITE_CAPACITY, channel)
 170     }
 171
 172     /// Create a `TFramedWriteTransport` with an internal write buffer
 173     /// of size `write_capacity` that wraps the given `TIoChannel`.
 174     pub fn with_capacity(write_capacity: usize, channel: C) -> TFramedWriteTransport<C> {
 175         TFramedWriteTransport {
 176             buf: Vec::with_capacity(write_capacity),
 177             channel,
 178         }
 179     }
 180 }
 181
 182 impl<C> Write for TFramedWriteTransport<C>
 183 where
 184     C: Write,
 185 {
 186     fn write(&mut self, b: &[u8]) -> io::Result<usize> {
 187         let current_capacity = self.buf.capacity();
 188         let available_space = current_capacity - self.buf.len();
 189         if b.len() > available_space {
 190             let additional_space = cmp::max(b.len() - available_space, current_capacity);
 191             self.buf.reserve(additional_space);
 192         }
 193
 194         self.buf.extend_from_slice(b);
 195         Ok(b.len())
 196     }
 197
 198     fn flush(&mut self) -> io::Result<()> {
 199         let message_size = self.buf.len();
 200
 201         if let 0 = message_size {
 202             return Ok(());
 203         } else {
 204             self.channel.write_i32::<BigEndian>(message_size as i32)?;
 205         }
 206
 207         // will spin if the underlying channel can't be written to
 208         let mut byte_index = 0;
 209         while byte_index < message_size {
 210             let nwrite = self.channel.write(&self.buf[byte_index..message_size])?;
 211             byte_index = cmp::min(byte_index + nwrite, message_size);
 212         }
 213
 214         let buf_capacity = cmp::min(self.buf.capacity(), WRITE_CAPACITY);
 215         self.buf.resize(buf_capacity, 0);
 216         self.buf.clear();
 217
 218         self.channel.flush()
 219     }
 220 }
 221
 222 /// Factory for creating instances of `TFramedWriteTransport`.
 223 #[derive(Default)]
 224 pub struct TFramedWriteTransportFactory;
 225
 226 impl TFramedWriteTransportFactory {
 227     pub fn new() -> TFramedWriteTransportFactory {
 228         TFramedWriteTransportFactory {}
 229     }
 230 }
 231
 232 impl TWriteTransportFactory for TFramedWriteTransportFactory {
 233     /// Create a `TFramedWriteTransport`.
 234     fn create(&self, channel: Box<dyn Write + Send>) -> Box<dyn TWriteTransport + Send> {
 235         Box::new(TFramedWriteTransport::new(channel))
 236     }
 237 }
 238
 239 #[cfg(test)]
 240 mod tests {
 241     use super::*;
 242     use transport::mem::TBufferChannel;
 243
 244     // FIXME: test a forced reserve
 245
 246     #[test]
 247     fn must_read_message_smaller_than_initial_buffer_size() {
 248         let c = TBufferChannel::with_capacity(10, 10);
 249         let mut t = TFramedReadTransport::with_capacity(8, c);
 250
 251         t.chan.set_readable_bytes(&[
 252             0x00, 0x00, 0x00, 0x04, /* message size */
 253             0x00, 0x01, 0x02, 0x03, /* message body */
 254         ]);
 255
 256         let mut buf = vec![0; 8];
 257
 258         // we've read exactly 4 bytes
 259         assert_eq!(t.read(&mut buf).unwrap(), 4);
 260         assert_eq!(&buf[..4], &[0x00, 0x01, 0x02, 0x03]);
 261     }
 262
 263     #[test]
 264     fn must_read_message_greater_than_initial_buffer_size() {
 265         let c = TBufferChannel::with_capacity(10, 10);
 266         let mut t = TFramedReadTransport::with_capacity(2, c);
 267
 268         t.chan.set_readable_bytes(&[
 269             0x00, 0x00, 0x00, 0x04, /* message size */
 270             0x00, 0x01, 0x02, 0x03, /* message body */
 271         ]);
 272
 273         let mut buf = vec![0; 8];
 274
 275         // we've read exactly 4 bytes
 276         assert_eq!(t.read(&mut buf).unwrap(), 4);
 277         assert_eq!(&buf[..4], &[0x00, 0x01, 0x02, 0x03]);
 278     }
 279
 280     #[test]
 281     fn must_read_multiple_messages_in_sequence_correctly() {
 282         let c = TBufferChannel::with_capacity(10, 10);
 283         let mut t = TFramedReadTransport::with_capacity(2, c);
 284
 285         //
 286         // 1st message
 287         //
 288
 289         t.chan.set_readable_bytes(&[
 290             0x00, 0x00, 0x00, 0x04, /* message size */
 291             0x00, 0x01, 0x02, 0x03, /* message body */
 292         ]);
 293
 294         let mut buf = vec![0; 8];
 295
 296         // we've read exactly 4 bytes
 297         assert_eq!(t.read(&mut buf).unwrap(), 4);
 298         assert_eq!(&buf, &[0x00, 0x01, 0x02, 0x03, 0x00, 0x00, 0x00, 0x00]);
 299
 300         //
 301         // 2nd message
 302         //
 303
 304         t.chan.set_readable_bytes(&[
 305             0x00, 0x00, 0x00, 0x01, /* message size */
 306             0x04, /* message body */
 307         ]);
 308
 309         let mut buf = vec![0; 8];
 310
 311         // we've read exactly 1 byte
 312         assert_eq!(t.read(&mut buf).unwrap(), 1);
 313         assert_eq!(&buf, &[0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
 314     }
 315
 316     #[test]
 317     fn must_write_message_smaller_than_buffer_size() {
 318         let mem = TBufferChannel::with_capacity(0, 0);
 319         let mut t = TFramedWriteTransport::with_capacity(20, mem);
 320
 321         let b = vec![0; 10];
 322
 323         // should have written 10 bytes
 324         assert_eq!(t.write(&b).unwrap(), 10);
 325     }
 326
 327     #[test]
 328     fn must_return_zero_if_caller_calls_write_with_empty_buffer() {
 329         let mem = TBufferChannel::with_capacity(0, 10);
 330         let mut t = TFramedWriteTransport::with_capacity(10, mem);
 331
 332         let expected: [u8; 0] = [];
 333
 334         assert_eq!(t.write(&[]).unwrap(), 0);
 335         assert_eq_transport_written_bytes!(t, expected);
 336     }
 337
 338     #[test]
 339     fn must_write_to_inner_transport_on_flush() {
 340         let mem = TBufferChannel::with_capacity(10, 10);
 341         let mut t = TFramedWriteTransport::new(mem);
 342
 343         let b: [u8; 5] = [0x00, 0x01, 0x02, 0x03, 0x04];
 344         assert_eq!(t.write(&b).unwrap(), 5);
 345         assert_eq_transport_num_written_bytes!(t, 0);
 346
 347         assert!(t.flush().is_ok());
 348
 349         let expected_bytes = [
 350             0x00, 0x00, 0x00, 0x05, /* message size */
 351             0x00, 0x01, 0x02, 0x03, 0x04, /* message body */
 352         ];
 353
 354         assert_eq_transport_written_bytes!(t, expected_bytes);
 355     }
 356
 357     #[test]
 358     fn must_write_message_greater_than_buffer_size_00() {
 359         let mem = TBufferChannel::with_capacity(0, 10);
 360
 361         // IMPORTANT: DO **NOT** CHANGE THE WRITE_CAPACITY OR THE NUMBER OF BYTES TO BE WRITTEN!
 362         // these lengths were chosen to be just long enough
 363         // that doubling the capacity is a **worse** choice than
 364         // simply resizing the buffer to b.len()
 365
 366         let mut t = TFramedWriteTransport::with_capacity(1, mem);
 367         let b = [0x00, 0x01, 0x02];
 368
 369         // should have written 3 bytes
 370         assert_eq!(t.write(&b).unwrap(), 3);
 371         assert_eq_transport_num_written_bytes!(t, 0);
 372
 373         assert!(t.flush().is_ok());
 374
 375         let expected_bytes = [
 376             0x00, 0x00, 0x00, 0x03, /* message size */
 377             0x00, 0x01, 0x02, /* message body */
 378         ];
 379
 380         assert_eq_transport_written_bytes!(t, expected_bytes);
 381     }
 382
 383     #[test]
 384     fn must_write_message_greater_than_buffer_size_01() {
 385         let mem = TBufferChannel::with_capacity(0, 10);
 386
 387         // IMPORTANT: DO **NOT** CHANGE THE WRITE_CAPACITY OR THE NUMBER OF BYTES TO BE WRITTEN!
 388         // these lengths were chosen to be just long enough
 389         // that doubling the capacity is a **better** choice than
 390         // simply resizing the buffer to b.len()
 391
 392         let mut t = TFramedWriteTransport::with_capacity(2, mem);
 393         let b = [0x00, 0x01, 0x02];
 394
 395         // should have written 3 bytes
 396         assert_eq!(t.write(&b).unwrap(), 3);
 397         assert_eq_transport_num_written_bytes!(t, 0);
 398
 399         assert!(t.flush().is_ok());
 400
 401         let expected_bytes = [
 402             0x00, 0x00, 0x00, 0x03, /* message size */
 403             0x00, 0x01, 0x02, /* message body */
 404         ];
 405
 406         assert_eq_transport_written_bytes!(t, expected_bytes);
 407     }
 408
 409     #[test]
 410     fn must_return_error_if_nothing_can_be_written_to_inner_transport_on_flush() {
 411         let mem = TBufferChannel::with_capacity(0, 0);
 412         let mut t = TFramedWriteTransport::with_capacity(1, mem);
 413
 414         let b = vec![0; 10];
 415
 416         // should have written 10 bytes
 417         assert_eq!(t.write(&b).unwrap(), 10);
 418
 419         // let's flush
 420         let r = t.flush();
 421
 422         // this time we'll error out because the flush can't write to the underlying channel
 423         assert!(r.is_err());
 424     }
 425
 426     #[test]
 427     fn must_write_successfully_after_flush() {
 428         // IMPORTANT: write capacity *MUST* be greater
 429         // than message sizes used in this test + 4-byte frame header
 430         let mem = TBufferChannel::with_capacity(0, 10);
 431         let mut t = TFramedWriteTransport::with_capacity(5, mem);
 432
 433         // write and flush
 434         let first_message: [u8; 5] = [0x00, 0x01, 0x02, 0x03, 0x04];
 435         assert_eq!(t.write(&first_message).unwrap(), 5);
 436         assert!(t.flush().is_ok());
 437
 438         let mut expected = Vec::new();
 439         expected.write_all(&[0x00, 0x00, 0x00, 0x05]).unwrap(); // message size
 440         expected.extend_from_slice(&first_message);
 441
 442         // check the flushed bytes
 443         assert_eq!(t.channel.write_bytes(), expected);
 444
 445         // reset our underlying transport
 446         t.channel.empty_write_buffer();
 447
 448         let second_message: [u8; 3] = [0x05, 0x06, 0x07];
 449         assert_eq!(t.write(&second_message).unwrap(), 3);
 450         assert!(t.flush().is_ok());
 451
 452         expected.clear();
 453         expected.write_all(&[0x00, 0x00, 0x00, 0x03]).unwrap(); // message size
 454         expected.extend_from_slice(&second_message);
 455
 456         // check the flushed bytes
 457         assert_eq!(t.channel.write_bytes(), expected);
 458     }
 459 }