Fix IOBufQueue::preallocate for callers who really wanted a hard max

Summary: D633912 broke unicorn because it relied on preallocate having a hard max.  Add an optional hard max parameter and rename maxHint to newAllocationSize.  Pass the hard max in for unicorn

Test Plan: folly and unicorn unit tests

Reviewed By: tudorb@fb.com

FB internal diff: D634894

diff --git a/folly/experimental/io/IOBufQueue.cpp b/folly/experimental/io/IOBufQueue.cpp
index 422121e7920bb8476a9cf13bf1d3a8607dc37800..6491d6d3ac62cc6012d8207942785ba095559d9a 100644 (file)
--- a/folly/experimental/io/IOBufQueue.cpp
+++ b/folly/experimental/io/IOBufQueue.cpp
@@ -161,23 +161,24 @@ IOBufQueue::wrapBuffer(const void* buf, size_t len, uint32_t blockSize) {
 }
 
 pair<void*,uint32_t>
-IOBufQueue::preallocate(uint32_t min, uint32_t maxHint) {
+IOBufQueue::preallocate(uint32_t min, uint32_t newAllocationSize,
+                        uint32_t max) {
   if (head_ != NULL) {
     // If there's enough space left over at the end of the queue, use that.
     IOBuf* last = head_->prev();
     if (!last->isSharedOne()) {
       uint32_t avail = last->tailroom();
       if (avail >= min) {
-        return make_pair(last->writableTail(), avail);
+        return make_pair(last->writableTail(), std::min(max, avail));
       }
     }
   }
   // Allocate a new buffer of the requested max size.
-  unique_ptr<IOBuf> newBuf(IOBuf::create(std::max(min, maxHint)));
+  unique_ptr<IOBuf> newBuf(IOBuf::create(std::max(min, newAllocationSize)));
   appendToChain(head_, std::move(newBuf));
   IOBuf* last = head_->prev();
   return make_pair(last->writableTail(),
-                   last->tailroom() /* may exceed maxHint */);
+                   std::min(max, last->tailroom()));
 }
 
 void

diff --git a/folly/experimental/io/IOBufQueue.h b/folly/experimental/io/IOBufQueue.h
index b9c1bac912e93349c7057622a5495e25d6eced91..cca5f98539dd3f87fc86da3c694e477c291e88d4 100644 (file)
--- a/folly/experimental/io/IOBufQueue.h
+++ b/folly/experimental/io/IOBufQueue.h
@@ -115,9 +115,10 @@ class IOBufQueue {
    * Obtain a writable block of contiguous bytes at the end of this
    * queue, allocating more space if necessary.  The amount of space
    * reserved will be at least min.  If min contiguous space is not
-   * available at the end of the queue, and IOBuf with size maxHint is
-   * appended to the chain and returned.  The actual available space
-   * may be larger than maxHint.
+   * available at the end of the queue, and IOBuf with size newAllocationSize
+   * is appended to the chain and returned.  The actual available space
+   * may be larger than newAllocationSize, but will be truncated to max,
+   * if specified.
    *
    * If the caller subsequently writes anything into the returned space,
    * it must call the postallocate() method.
@@ -130,7 +131,9 @@ class IOBufQueue {
    *       callback, tell the application how much of the buffer they've
    *       filled with data.
    */
-  std::pair<void*,uint32_t> preallocate(uint32_t min, uint32_t maxHint);
+  std::pair<void*,uint32_t> preallocate(
+    uint32_t min, uint32_t newAllocationSize,
+    uint32_t max = std::numeric_limits<uint32_t>::max());
 
   /**
    * Tell the queue that the caller has written data into the first n

diff --git a/folly/experimental/io/test/IOBufQueueTest.cpp b/folly/experimental/io/test/IOBufQueueTest.cpp
index e769ab5ce1d97c80dcf94dd161a49aa25f1697ad..47555e05be6d0651d794b2bad8c0114a1dbc8313 100644 (file)
--- a/folly/experimental/io/test/IOBufQueueTest.cpp
+++ b/folly/experimental/io/test/IOBufQueueTest.cpp
@@ -149,10 +149,11 @@ TEST(IOBufQueue, Split) {
 TEST(IOBufQueue, Preallocate) {
   IOBufQueue queue(clOptions);
   queue.append(string("Hello"));
-  pair<void*,uint32_t> writable = queue.preallocate(2, 64);
+  pair<void*,uint32_t> writable = queue.preallocate(2, 64, 64);
   checkConsistency(queue);
   EXPECT_NE((void*)NULL, writable.first);
   EXPECT_LE(2, writable.second);
+  EXPECT_GE(64, writable.second);
   memcpy(writable.first, SCL(", "));
   queue.postallocate(2);
   checkConsistency(queue);
@@ -160,15 +161,18 @@ TEST(IOBufQueue, Preallocate) {
   queue.append(SCL("World"));
   checkConsistency(queue);
   EXPECT_EQ(12, queue.front()->computeChainDataLength());
-  writable = queue.preallocate(1024, 4096);
+  // There are not 2048 bytes available, this will alloc a new buf
+  writable = queue.preallocate(2048, 4096);
   checkConsistency(queue);
-  EXPECT_LE(1024, writable.second);
+  EXPECT_LE(2048, writable.second);
+  // IOBuf allocates more than newAllocationSize, and we didn't cap it
+  EXPECT_GE(writable.second, 4096);
   queue.postallocate(writable.second);
   // queue has no empty space, make sure we allocate at least min, even if
-  // maxHint < min
-  writable = queue.preallocate(1024, 1);
+  // newAllocationSize < min
+  writable = queue.preallocate(1024, 1, 1024);
   checkConsistency(queue);
-  EXPECT_LE(1024, writable.second);
+  EXPECT_EQ(1024, writable.second);
 }
 
 TEST(IOBufQueue, Wrap) {