IOBuf: Add goodSize() Allocation-Sizing Method

Summary:
For long-lived IOBufs, it can be useful to know whether their memory footprint
can be usefully shrunk. This diff exposes a method that lets users see the
smallest capacity an IOBuf could have while wrapping a particularly sized
object.

Without this, they have to do hacky math.

Reviewed By: yfeldblum

Differential Revision: D18440359

fbshipit-source-id: 1324700b58f0b2762cb3347b907ca5e04ebecb99

folly/io/IOBuf.cpp

@@ -312,6 +312,24 @@ unique_ptr<IOBuf> IOBuf::createChain(
   return out;
 }
 
+size_t IOBuf::goodSize(size_t minCapacity, CombinedOption combined) {
+  if (combined == CombinedOption::DEFAULT) {
+    combined = minCapacity <= kDefaultCombinedBufSize
+        ? CombinedOption::COMBINED
+        : CombinedOption::SEPARATE;
+  }
+  size_t overhead;
+  if (combined == CombinedOption::COMBINED) {
+    overhead = offsetof(HeapFullStorage, align);
+  } else {
+    // Pad minCapacity to a multiple of 8
+    minCapacity = (minCapacity + 7) & ~7;
+    overhead = sizeof(SharedInfo);
+  }
+  size_t goodSize = folly::goodMallocSize(minCapacity + overhead);
+  return goodSize - overhead;
+}
+
 IOBuf::IOBuf(
     TakeOwnershipOp,
     void* buf,

folly/io/IOBuf.h

@@ -228,6 +228,8 @@ class IOBuf {
   enum TakeOwnershipOp { TAKE_OWNERSHIP };
   enum CopyBufferOp { COPY_BUFFER };
 
+  enum class CombinedOption { DEFAULT, COMBINED, SEPARATE };
+
   typedef ByteRange value_type;
   typedef Iterator iterator;
   typedef Iterator const_iterator;
@@ -278,6 +280,20 @@ class IOBuf {
       size_t totalCapacity,
       std::size_t maxBufCapacity);
 
+  /**
+   * Uses folly::goodMallocSize() to figure out what the largest capacity would
+   * be that would trigger the same underlying allocation size as would be
+   * triggered by the given capacity.
+   *
+   * Note that IOBufs do this up-sizing for you: they will round up to the full
+   * allocation size and make that capacity available to you without your using
+   * this function. This just lets you introspect into that process, so you can
+   * for example figure out whether a given IOBuf can be usefully compacted.
+   */
+  static size_t goodSize(
+      size_t minCapacity,
+      CombinedOption combined = CombinedOption::DEFAULT);
+
   /**
    * Create a new IOBuf pointing to an existing data buffer.
    *

folly/io/test/IOBufTest.cpp

@@ -74,24 +74,24 @@ TEST(IOBuf, Simple) {
 }
 
 void testAllocSize(uint32_t requestedCapacity) {
+  auto expectedSize = IOBuf::goodSize(requestedCapacity);
   unique_ptr<IOBuf> iobuf(IOBuf::create(requestedCapacity));
   EXPECT_GE(iobuf->capacity(), requestedCapacity);
+  EXPECT_EQ(iobuf->capacity(), expectedSize);
 }
 
 TEST(IOBuf, AllocSizes) {
-  // Try with a small allocation size that should fit in the internal buffer
-  testAllocSize(28);
+  // cover small evil values exhaustively, including the
+  // kDefaultCombinedBufSize transition.
+  for (uint32_t i = 0; i < 1234; i++) {
+    testAllocSize(i);
+  }
 
-  // Try with a large allocation size that will require an external buffer.
+  // Try with large allocation sizes that will require an external buffer.
   testAllocSize(9000);
-
-  // 220 bytes is currently the cutoff
-  // (It would be nice to use the IOBuf::kMaxInternalDataSize constant,
-  // but it's private and it doesn't seem worth making it public just for this
-  // test code.)
-  testAllocSize(220);
-  testAllocSize(219);
-  testAllocSize(221);
+  testAllocSize(1048575);
+  testAllocSize(1048576);
+  testAllocSize(1048577);
 }
 
 void deleteArrayBuffer(void* buf, void* arg) {