thrift/transport/TShortReadTransport.h

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one
0003  * or more contributor license agreements. See the NOTICE file
0004  * distributed with this work for additional information
0005  * regarding copyright ownership. The ASF licenses this file
0006  * to you under the Apache License, Version 2.0 (the
0007  * "License"); you may not use this file except in compliance
0008  * with the License. You may obtain a copy of the License at
0009  *
0010  *   http://www.apache.org/licenses/LICENSE-2.0
0011  *
0012  * Unless required by applicable law or agreed to in writing,
0013  * software distributed under the License is distributed on an
0014  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0015  * KIND, either express or implied. See the License for the
0016  * specific language governing permissions and limitations
0017  * under the License.
0018  */
0019
0020 #ifndef _THRIFT_TRANSPORT_TSHORTREADTRANSPORT_H_
0021 #define _THRIFT_TRANSPORT_TSHORTREADTRANSPORT_H_ 1
0022
0023 #include <cstdlib>
0024
0025 #include <thrift/transport/TTransport.h>
0026 #include <thrift/transport/TVirtualTransport.h>
0027
0028 namespace apache {
0029 namespace thrift {
0030 namespace transport {
0031 namespace test {
0032
0033 /**
0034  * This class is only meant for testing.  It wraps another transport.
0035  * Calls to read are passed through with some probability.  Otherwise,
0036  * the read amount is randomly reduced before being passed through.
0037  *
0038  */
0039 class TShortReadTransport : public TVirtualTransport<TShortReadTransport> {
0040 public:
0041   TShortReadTransport(std::shared_ptr<TTransport> transport, double full_prob,
0042                      std::shared_ptr<TConfiguration> config = nullptr)
0043     : TVirtualTransport(config), transport_(transport), fullProb_(full_prob) {
0044     }
0045
0046   bool isOpen() const override { return transport_->isOpen(); }
0047
0048   bool peek() override { return transport_->peek(); }
0049
0050   void open() override { transport_->open(); }
0051
0052   void close() override { transport_->close(); }
0053
0054   uint32_t read(uint8_t* buf, uint32_t len) {
0055     checkReadBytesAvailable(len);
0056     if (len == 0) {
0057       return 0;
0058     }
0059
0060     if (rand() / (double)RAND_MAX >= fullProb_) {
0061       len = 1 + rand() % len;
0062     }
0063     return transport_->read(buf, len);
0064   }
0065
0066   void write(const uint8_t* buf, uint32_t len) { transport_->write(buf, len); }
0067
0068   void flush() override {
0069     resetConsumedMessageSize();
0070     transport_->flush();
0071   }
0072
0073   const uint8_t* borrow(uint8_t* buf, uint32_t* len) { return transport_->borrow(buf, len); }
0074
0075   void consume(uint32_t len) {
0076     countConsumedMessageBytes(len);
0077     return transport_->consume(len);
0078   }
0079
0080   std::shared_ptr<TTransport> getUnderlyingTransport() { return transport_; }
0081
0082 protected:
0083   std::shared_ptr<TTransport> transport_;
0084   double fullProb_;
0085 };
0086 }
0087 }
0088 }
0089 } // apache::thrift::transport::test
0090
0091 #endif // #ifndef _THRIFT_TRANSPORT_TSHORTREADTRANSPORT_H_