Thrift: Add writePtr(), wroteBytes() to TMemoryBuffer

Summary: This adds the discussed interface to TMemoryBuffer, as follows:

         - writePtr(size) returns a ptr you can write (size) bytes to
         - wroteBytes() lets it know you wrote that many bytes

         To do this, I refactored an:
         - ensureCanWrite(size) private func

Reviewed By: dreiss

Test Plan: works in my test environment

Revert: OK

TracCamp Project: Thrift

DiffCamp Revision: 8739


git-svn-id: https://svn.apache.org/repos/asf/incubator/thrift/trunk@665555 13f79535-47bb-0310-9956-ffa450edef68

lib/cpp/src/transport/TTransportUtils.cpp

@@ -297,30 +297,44 @@ uint32_t TMemoryBuffer::readAppendToString(std::string& str, uint32_t len) {
   return give;
 }
 
-void TMemoryBuffer::write(const uint8_t* buf, uint32_t len) {
+void TMemoryBuffer::ensureCanWrite(uint32_t len) {
   // Check available space
   uint32_t avail = bufferSize_ - wPos_;
-
-  // Grow the buffer
-  if (len > avail) {
-    if (!owner_) {
-      throw TTransportException("Insufficient space in external MemoryBuffer");
-    }
-    while (len > avail) {
-      bufferSize_ *= 2;
-      avail = bufferSize_ - wPos_;
-    }
-    buffer_ = (uint8_t*)std::realloc(buffer_, bufferSize_);
-    if (buffer_ == NULL) {
-      throw TTransportException("Out of memory.");
-    }
+  if (len <= avail) {
+    return;
   }
 
+  if (!owner_) {
+    throw TTransportException("Insufficient space in external MemoryBuffer");
+  }
+
+  // Grow the buffer as necessary
+  while (len > avail) {
+    bufferSize_ *= 2;
+    avail = bufferSize_ - wPos_;
+  }
+  buffer_ = (uint8_t*)std::realloc(buffer_, bufferSize_);
+  if (buffer_ == NULL) {
+    throw TTransportException("Out of memory.");
+  }
+}
+
+void TMemoryBuffer::write(const uint8_t* buf, uint32_t len) {
+  ensureCanWrite(len);
+
   // Copy into the buffer and increment wPos_
   memcpy(buffer_ + wPos_, buf, len);
   wPos_ += len;
 }
 
+void TMemoryBuffer::wroteBytes(uint32_t len) {
+  uint32_t avail = bufferSize_ - wPos_;
+  if (len > avail) {
+    throw TTransportException("Client wrote more bytes than size of buffer.");
+  }
+  wPos_ += len;
+}
+
 const uint8_t* TMemoryBuffer::borrow(uint8_t* buf, uint32_t* len) {
   if (wPos_-rPos_ >= *len) {
     *len = wPos_-rPos_;

lib/cpp/src/transport/TTransportUtils.h

@@ -315,6 +315,9 @@ class TMemoryBuffer : public TTransport {
     rPos_ = 0;
   }
 
+  // make sure there's at least 'len' bytes available for writing
+  void ensureCanWrite(uint32_t len);
+
  public:
   static const uint32_t defaultSize = 1024;
 
@@ -480,7 +483,7 @@ class TMemoryBuffer : public TTransport {
 
   void write(const uint8_t* buf, uint32_t len);
 
-  uint32_t available() {
+  uint32_t available() const {
     return wPos_ - rPos_;
   }
 
@@ -496,6 +499,20 @@ class TMemoryBuffer : public TTransport {
     swap(owner_,      that.owner_);
   }
 
+  // Returns a pointer to where the client can write data to append to
+  // the TMemoryBuffer, and ensures the buffer is big enough to accomodate a
+  // write of the provided length.  The returned pointer is very convenient for
+  // passing to read(), recv(), or similar. You must call wroteBytes() as soon
+  // as data is written or the buffer will not be aware that data has changed.
+  uint8_t* getWritePtr(uint32_t len) {
+    ensureCanWrite(len);
+    return buffer_ + wPos_;
+  }
+
+  // Informs the buffer that the client has written 'len' bytes into storage
+  // that had been provided by getWritePtr().
+  void wroteBytes(uint32_t len);
+
  private:
   // Data buffer
   uint8_t* buffer_;