expose share count of IOBuf

Summary: Exposing this counter allows us to fairly divide the memory usage of shared IOBuf buffers among all the owners in memory-usage tracking code.

Reviewed By: simpkins

Differential Revision: D14361976

fbshipit-source-id: 73c63c473b117198786143b421c6fcb8aed684d5

folly/io/IOBuf.cpp

@@ -1134,6 +1134,21 @@ IOBuf::FillIovResult IOBuf::fillIov(struct iovec* iov, size_t len) const {
   return {0, 0};
 }
 
+uint32_t IOBuf::approximateShareCountOne() const {
+  if (UNLIKELY(!sharedInfo())) {
+    return 1U;
+  }
+  if (LIKELY(!(flags() & kFlagMaybeShared))) {
+    return 1U;
+  }
+  auto shareCount = sharedInfo()->refcount.load(std::memory_order_acquire);
+  if (shareCount < 2) {
+    // we're the last one left
+    clearFlags(kFlagMaybeShared);
+  }
+  return shareCount;
+}
+
 size_t IOBufHash::operator()(const IOBuf& buf) const noexcept {
   folly::hash::SpookyHashV2 hasher;
   hasher.Init(0, 0);

folly/io/IOBuf.h

@@ -960,6 +960,21 @@ class IOBuf {
     return sharedInfo();
   }
 
+  /**
+   * For most of the  use-cases where it seems like a good idea to call this
+   * function, what you really want is `isSharedOne()`.
+   *
+   * If this IOBuf is managed by the usual refcounting mechanism (ie
+   * `isManagedOne()` returns `true`): this returns the reference count as it
+   * was when recently observed by this thread.
+   *
+   * If this IOBuf is *not* managed by the usual refcounting mechanism then the
+   * result of this function is not defined.
+   *
+   * This only checks the current IOBuf, and not other IOBufs in the chain.
+   */
+  uint32_t approximateShareCountOne() const;
+
   /**
    * Return true if other IOBufs are also pointing to the buffer used by this
    * IOBuf, and false otherwise.

folly/io/test/IOBufTest.cpp

@@ -531,6 +531,8 @@ TEST(IOBuf, Chaining) {
   EXPECT_TRUE(iob1->isShared());
 
   EXPECT_TRUE(iob1->isSharedOne());
+  // Also verify the share count is consistent:
+  EXPECT_LT(1, iob1->approximateShareCountOne());
   EXPECT_TRUE(iob2ptr->isSharedOne());
   EXPECT_TRUE(iob3ptr->isSharedOne());
   EXPECT_TRUE(iob4ptr->isSharedOne());
@@ -540,6 +542,8 @@ TEST(IOBuf, Chaining) {
   chainClone->unshare();
   EXPECT_FALSE(chainClone->isShared());
   EXPECT_FALSE(iob1->isShared());
+  // Also verify the share count is consistent:
+  EXPECT_EQ(1, iob1->approximateShareCountOne());
 
   // Make sure the unshared result still has the same data
   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());