Add isManaged / makeManaged

Summary: Add a way to tell whether an IOBuf (either indvidually or the whole chain)
refers to memory that's outside of IOBuf's control (via WRAP_BUFFER). That is,
can you assume that clone() will extend the lifetime of the memory? Is it safe
to use clone() instead of copying the data (and avoid a copy)?

Reviewed By: @simpkins

Differential Revision: D2191883

folly/io/IOBuf.cpp

@@ -552,6 +552,19 @@ void IOBuf::unshareChained() {
   coalesceSlow();
 }
 
+void IOBuf::makeManagedChained() {
+  assert(isChained());
+
+  IOBuf* current = this;
+  while (true) {
+    current->makeManagedOne();
+    current = current->next_;
+    if (current == this) {
+      break;
+    }
+  }
+}
+
 void IOBuf::coalesceSlow() {
   // coalesceSlow() should only be called if we are part of a chain of multiple
   // IOBufs.  The caller should have already verified this.

folly/io/IOBuf.h

@@ -871,6 +871,33 @@ class IOBuf {
     }
   }
 
+  /**
+   * Return true if all IOBufs in this chain are managed by the usual
+   * refcounting mechanism (and so the lifetime of the underlying memory
+   * can be extended by clone()).
+   */
+  bool isManaged() const {
+    const IOBuf* current = this;
+    while (true) {
+      if (!current->isManagedOne()) {
+        return false;
+      }
+      current = current->next_;
+      if (current == this) {
+        return true;
+      }
+    }
+  }
+
+  /**
+   * Return true if this IOBuf is managed by the usual refcounting mechanism
+   * (and so the lifetime of the underlying memory can be extended by
+   * cloneOne()).
+   */
+  bool isManagedOne() const {
+    return sharedInfo();
+  }
+
   /**
    * Return true if other IOBufs are also pointing to the buffer used by this
    * IOBuf, and false otherwise.
@@ -946,6 +973,39 @@ class IOBuf {
     }
   }
 
+  /**
+   * Ensure that the memory that IOBufs in this chain refer to will continue to
+   * be allocated for as long as the IOBufs of the chain (or any clone()s
+   * created from this point onwards) is alive.
+   *
+   * This only has an effect for user-owned buffers (created with the
+   * WRAP_BUFFER constructor or wrapBuffer factory function), in which case
+   * those buffers are unshared.
+   */
+  void makeManaged() {
+    if (isChained()) {
+      makeManagedChained();
+    } else {
+      makeManagedOne();
+    }
+  }
+
+  /**
+   * Ensure that the memory that this IOBuf refers to will continue to be
+   * allocated for as long as this IOBuf (or any clone()s created from this
+   * point onwards) is alive.
+   *
+   * This only has an effect for user-owned buffers (created with the
+   * WRAP_BUFFER constructor or wrapBuffer factory function), in which case
+   * those buffers are unshared.
+   */
+  void makeManagedOne() {
+    if (!isManagedOne()) {
+      // We can call the internal function directly; unmanaged implies shared.
+      unshareOneSlow();
+    }
+  }
+
   /**
    * Coalesce this IOBuf chain into a single buffer.
    *
@@ -1158,6 +1218,7 @@ class IOBuf {
 
   void unshareOneSlow();
   void unshareChained();
+  void makeManagedChained();
   void coalesceSlow();
   void coalesceSlow(size_t maxLength);
   // newLength must be the entire length of the buffers between this and

folly/io/test/IOBufTest.cpp

@@ -1152,6 +1152,91 @@ TEST(IOBuf, CopyConstructorAndAssignmentOperator) {
   EXPECT_FALSE(buf->isShared());
 }
 
+namespace {
+// Use with string literals only
+std::unique_ptr<IOBuf> wrap(const char* str) {
+  return IOBuf::wrapBuffer(str, strlen(str));
+}
+
+std::unique_ptr<IOBuf> copy(const char* str) {
+  // At least 1KiB of tailroom, to ensure an external buffer
+  return IOBuf::copyBuffer(str, strlen(str), 0, 1024);
+}
+
+std::string toString(const folly::IOBuf& buf) {
+  std::string result;
+  result.reserve(buf.computeChainDataLength());
+  for (auto& b : buf) {
+    result.append(reinterpret_cast<const char*>(b.data()), b.size());
+  }
+  return result;
+}
+
+char* writableStr(folly::IOBuf& buf) {
+  return reinterpret_cast<char*>(buf.writableData());
+}
+
+}  // namespace
+
+TEST(IOBuf, Managed) {
+  auto hello = "hello";
+  auto buf1UP = wrap(hello);
+  auto buf1 = buf1UP.get();
+  EXPECT_FALSE(buf1->isManagedOne());
+  auto buf2UP = copy("world");
+  auto buf2 = buf2UP.get();
+  EXPECT_TRUE(buf2->isManagedOne());
+  auto buf3UP = wrap(hello);
+  auto buf3 = buf3UP.get();
+  auto buf4UP = buf2->clone();
+  auto buf4 = buf4UP.get();
+
+  // buf1 and buf3 share the same memory (but are unmanaged)
+  EXPECT_FALSE(buf1->isManagedOne());
+  EXPECT_FALSE(buf3->isManagedOne());
+  EXPECT_TRUE(buf1->isSharedOne());
+  EXPECT_TRUE(buf3->isSharedOne());
+  EXPECT_EQ(buf1->data(), buf3->data());
+
+  // buf2 and buf4 share the same memory (but are managed)
+  EXPECT_TRUE(buf2->isManagedOne());
+  EXPECT_TRUE(buf4->isManagedOne());
+  EXPECT_TRUE(buf2->isSharedOne());
+  EXPECT_TRUE(buf4->isSharedOne());
+  EXPECT_EQ(buf2->data(), buf4->data());
+
+  buf1->prependChain(std::move(buf2UP));
+  buf1->prependChain(std::move(buf3UP));
+  buf1->prependChain(std::move(buf4UP));
+
+  EXPECT_EQ("helloworldhelloworld", toString(*buf1));
+  EXPECT_FALSE(buf1->isManaged());
+
+  buf1->makeManaged();
+  EXPECT_TRUE(buf1->isManaged());
+
+  // buf1 and buf3 are now unshared (because they were unmanaged)
+  EXPECT_TRUE(buf1->isManagedOne());
+  EXPECT_TRUE(buf3->isManagedOne());
+  EXPECT_FALSE(buf1->isSharedOne());
+  EXPECT_FALSE(buf3->isSharedOne());
+  EXPECT_NE(buf1->data(), buf3->data());
+
+  // buf2 and buf4 are still shared
+  EXPECT_TRUE(buf2->isManagedOne());
+  EXPECT_TRUE(buf4->isManagedOne());
+  EXPECT_TRUE(buf2->isSharedOne());
+  EXPECT_TRUE(buf4->isSharedOne());
+  EXPECT_EQ(buf2->data(), buf4->data());
+
+  // And verify that the truth is what we expect: modify a byte in buf1 and
+  // buf2, see that the change from buf1 is *not* reflected in buf3, but the
+  // change from buf2 is reflected in buf4.
+  writableStr(*buf1)[0] = 'j';
+  writableStr(*buf2)[0] = 'x';
+  EXPECT_EQ("jelloxorldhelloxorld", toString(*buf1));
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   gflags::ParseCommandLineFlags(&argc, &argv, true);