[Chromecast] Create buffer pool for IOBuffers

IOBufferPool is a class that allocates fixed-size IOBuffers. There is
an optional limit on the maximum number of buffers that will be allocated.
Buffers from the pool can be used as usual; when they are deleted, the memory
will be returned to the pool to be reused. Buffers can be safely used and
freed even after the IOBufferPool instance is destroyed.

IOBufferPool is optionally threadsafe (defaults to non-threadsafe). When
thread safety is enabled, the pool and allocated buffers can be used from
any sequence. If thread safety is not enabled, the pool and buffers may only
be used from the sequence that created IOBufferPool.

Change-Id: I901c70bbb0f7dca1395886246209e0aeb6805151
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1805866
Commit-Queue: Kenneth MacKay <kmackay@chromium.org>
Reviewed-by: Luke Halliwell <halliwell@chromium.org>
Reviewed-by: Yuchen Liu <yucliu@chromium.org>
Cr-Commit-Position: refs/heads/master@{#697012}

chromecast/net/BUILD.gn

@@ -49,6 +49,18 @@ cast_source_set("net") {
   }
 }
 
+cast_source_set("io_buffer_pool") {
+  sources = [
+    "io_buffer_pool.cc",
+    "io_buffer_pool.h",
+  ]
+
+  deps = [
+    "//base",
+    "//net",
+  ]
+}
+
 cast_source_set("small_message_socket") {
   sources = [
     "small_message_socket.cc",
@@ -86,9 +98,11 @@ cast_source_set("test_support") {
 test("cast_net_unittests") {
   sources = [
     "fake_stream_socket_unittest.cc",
+    "io_buffer_pool_unittest.cc",
   ]
 
   deps = [
+    ":io_buffer_pool",
     ":test_support",
     "//base",
     "//base/test:run_all_unittests",

chromecast/net/io_buffer_pool.cc

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chromecast/net/io_buffer_pool.h"
+
+#include <new>
+
+#include "base/synchronization/lock.h"
+#include "net/base/io_buffer.h"
+
+namespace chromecast {
+
+// The IOBufferPool allocates IOBuffers and the associated data as a single
+// contiguous buffer. The buffer is laid out like this:
+// |------------Wrapper----------|---data buffer---|
+// |--IOBuffer--|--Internal ptr--|---data buffer---|
+//
+// The contiguous buffer is allocated as a character array, and then a Wrapper
+// instance is placement-newed into it. We return a pointer to the IOBuffer
+// within the Wrapper.
+//
+// When the IOBuffer is deleted (in operator delete), we get a pointer to the
+// beginning of storage for the IOBuffer, which is the same memory location
+// as the Wrapper instance (since the Wrapper has no vtable or base class, this
+// should be true for any compiler). We can therefore cast the "deleted"
+// pointer to a Wrapper* and then reclaim the buffer.
+//
+// All the actual data and logic for the buffer pool is held in the Internal
+// class, which is refcounted with 1 ref for the IOBufferPool owner, and 1 ref
+// for each buffer currently in use (ie, not in the free list). The Internal
+// instance is only deleted when its internal refcount drops to 0; this allows
+// buffers allocated from the pool to be safely used and deleted even after the
+// pool has been destroyed.
+//
+// Note that locking in the Internal methods is optional since it is only needed
+// when threadsafe operation is requested.
+
+class IOBufferPool::Internal {
+ public:
+  Internal(size_t buffer_size, size_t max_buffers, bool threadsafe);
+
+  size_t num_allocated() const {
+    base::AutoLockMaybe lock(lock_ptr_);
+    return num_allocated_;
+  }
+
+  size_t num_free() const {
+    base::AutoLockMaybe lock(lock_ptr_);
+    return num_free_;
+  }
+
+  void Preallocate(size_t num_buffers);
+
+  void OwnerDestroyed();
+
+  scoped_refptr<net::IOBuffer> GetBuffer();
+
+ private:
+  class Buffer;
+  class Wrapper;
+  union Storage;
+
+  ~Internal();
+
+  void Reclaim(Wrapper* wrapper);
+
+  const size_t buffer_size_;
+  const size_t max_buffers_;
+
+  mutable base::Lock lock_;
+  base::Lock* const lock_ptr_;
+
+  Storage* free_buffers_;
+  size_t num_allocated_;
+  size_t num_free_;
+
+  int refs_;
+
+  DISALLOW_COPY_AND_ASSIGN(Internal);
+};
+
+class IOBufferPool::Internal::Buffer : public net::IOBuffer {
+ public:
+  explicit Buffer(char* data) : net::IOBuffer(data) {}
+
+ private:
+  friend class Wrapper;
+
+  ~Buffer() override { data_ = nullptr; }
+  static void operator delete(void* ptr);
+
+  DISALLOW_COPY_AND_ASSIGN(Buffer);
+};
+
+class IOBufferPool::Internal::Wrapper {
+ public:
+  Wrapper(char* data, IOBufferPool::Internal* pool)
+      : buffer_(data), pool_(pool) {}
+
+  ~Wrapper() = delete;
+  static void operator delete(void*) = delete;
+
+  Buffer* buffer() { return &buffer_; }
+
+  void Reclaim() { pool_->Reclaim(this); }
+
+ private:
+  Buffer buffer_;
+  IOBufferPool::Internal* const pool_;
+
+  DISALLOW_COPY_AND_ASSIGN(Wrapper);
+};
+
+union IOBufferPool::Internal::Storage {
+  Storage* next;  // Pointer to next free buffer.
+  Wrapper wrapper;
+};
+
+void IOBufferPool::Internal::Buffer::operator delete(void* ptr) {
+  Wrapper* wrapper = reinterpret_cast<Wrapper*>(ptr);
+  wrapper->Reclaim();
+}
+
+IOBufferPool::Internal::Internal(size_t buffer_size,
+                                 size_t max_buffers,
+                                 bool threadsafe)
+    : buffer_size_(buffer_size),
+      max_buffers_(max_buffers),
+      lock_ptr_(threadsafe ? &lock_ : nullptr),
+      free_buffers_(nullptr),
+      num_allocated_(0),
+      num_free_(0),
+      refs_(1) {  // 1 ref for the owner.
+}
+
+IOBufferPool::Internal::~Internal() {
+  while (free_buffers_) {
+    char* data = reinterpret_cast<char*>(free_buffers_);
+    free_buffers_ = free_buffers_->next;
+    delete[] data;
+  }
+}
+
+void IOBufferPool::Internal::Preallocate(size_t num_buffers) {
+  // We assume that this is uncontended in normal usage, so just lock for the
+  // entire method.
+  base::AutoLockMaybe lock(lock_ptr_);
+  if (num_buffers > max_buffers_) {
+    num_buffers = max_buffers_;
+  }
+  if (num_allocated_ >= num_buffers) {
+    return;
+  }
+  size_t num_extra_buffers = num_buffers - num_allocated_;
+  num_free_ += num_extra_buffers;
+  num_allocated_ += num_extra_buffers;
+  while (num_extra_buffers > 0) {
+    char* ptr = new char[sizeof(Storage) + buffer_size_];
+    Storage* storage = reinterpret_cast<Storage*>(ptr);
+    storage->next = free_buffers_;
+    free_buffers_ = storage;
+
+    --num_extra_buffers;
+  }
+  // No need to add refs here, since the newly allocated buffers are not in use.
+}
+
+void IOBufferPool::Internal::OwnerDestroyed() {
+  bool deletable;
+  {
+    base::AutoLockMaybe lock(lock_ptr_);
+    --refs_;  // Remove the owner's ref.
+    deletable = (refs_ == 0);
+  }
+
+  if (deletable) {
+    delete this;
+  }
+}
+
+scoped_refptr<net::IOBuffer> IOBufferPool::Internal::GetBuffer() {
+  char* ptr = nullptr;
+
+  {
+    base::AutoLockMaybe lock(lock_ptr_);
+    if (free_buffers_) {
+      ptr = reinterpret_cast<char*>(free_buffers_);
+      free_buffers_ = free_buffers_->next;
+      --num_free_;
+    } else {
+      if (num_allocated_ == max_buffers_)
+        return nullptr;
+      ++num_allocated_;
+    }
+    ++refs_;  // Add a ref for the now in-use buffer.
+  }
+
+  if (!ptr) {
+    ptr = new char[sizeof(Storage) + buffer_size_];
+  }
+
+  char* data = ptr + sizeof(Storage);
+  Wrapper* wrapper = new (ptr) Wrapper(data, this);
+  return scoped_refptr<net::IOBuffer>(wrapper->buffer());
+}
+
+void IOBufferPool::Internal::Reclaim(Wrapper* wrapper) {
+  Storage* storage = reinterpret_cast<Storage*>(wrapper);
+  bool deletable;
+  {
+    base::AutoLockMaybe lock(lock_ptr_);
+    storage->next = free_buffers_;
+    free_buffers_ = storage;
+    ++num_free_;
+    --refs_;  // Remove a ref since this buffer is no longer in use.
+    deletable = (refs_ == 0);
+  }
+
+  if (deletable) {
+    delete this;
+  }
+}
+
+IOBufferPool::IOBufferPool(size_t buffer_size,
+                           size_t max_buffers,
+                           bool threadsafe)
+    : buffer_size_(buffer_size),
+      max_buffers_(max_buffers),
+      internal_(new Internal(buffer_size, max_buffers, threadsafe)) {}
+
+IOBufferPool::IOBufferPool(size_t buffer_size)
+    : IOBufferPool(buffer_size, static_cast<size_t>(-1)) {}
+
+IOBufferPool::~IOBufferPool() {
+  internal_->OwnerDestroyed();
+}
+
+size_t IOBufferPool::NumAllocatedForTesting() const {
+  return internal_->num_allocated();
+}
+
+size_t IOBufferPool::NumFreeForTesting() const {
+  return internal_->num_free();
+}
+
+void IOBufferPool::Preallocate(size_t num_buffers) {
+  internal_->Preallocate(num_buffers);
+}
+
+scoped_refptr<net::IOBuffer> IOBufferPool::GetBuffer() {
+  return internal_->GetBuffer();
+}
+
+}  // namespace chromecast

chromecast/net/io_buffer_pool.h

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROMECAST_NET_IO_BUFFER_POOL_H_
+#define CHROMECAST_NET_IO_BUFFER_POOL_H_
+
+#include <stddef.h>
+
+#include "base/macros.h"
+#include "base/memory/ref_counted.h"
+
+namespace net {
+class IOBuffer;
+}  // namespace net
+
+namespace chromecast {
+
+// Buffer pool to allocate ::net::IOBuffers of a constant size. When a buffer
+// from this pool is destroyed, it is returned to the pool to be reused (rather
+// than constantly reallocating memory). If the buffer is destroyed after the
+// pool has been destroyed, the buffer's memory will be freed correctly.
+class IOBufferPool : public base::RefCountedThreadSafe<IOBufferPool> {
+ public:
+  // |buffer_size| is the usable size of each buffer that will be returned
+  // by GetBuffer(). |max_buffers| is the maximum number of buffers that this
+  // pool will allocate. If |threadsafe| is true then the pool and any buffers
+  // allocated from it may be used safely on any thread; otherwise all access
+  // to pool and buffers must be made on the creating sequence.
+  IOBufferPool(size_t buffer_size, size_t max_buffers, bool threadsafe = false);
+  // If |max_buffers| is not specified, the maximum value of size_t is used.
+  explicit IOBufferPool(size_t buffer_size);
+
+  size_t buffer_size() const { return buffer_size_; }
+  size_t max_buffers() const { return max_buffers_; }
+
+  // Ensures that at least |num_buffers| are allocated. If |num_buffers| is
+  // greater than |max_buffers|, makes sure that |max_buffers| buffers have been
+  // allocated.
+  void Preallocate(size_t num_buffers);
+
+  // Returns an IOBuffer from the pool, or |nullptr| if there are no buffers in
+  // the pool and |max buffers| has been reached. The pool does not keep any
+  // references to the buffer. It is safe to use the returned buffer after the
+  // pool has been destroyed; if the buffer is destroyed after the pool is
+  // destroyed, the buffer's memory will be freed correctly.
+  scoped_refptr<net::IOBuffer> GetBuffer();
+
+  // Returns the number of buffers that have ever been allocated by the pool.
+  size_t NumAllocatedForTesting() const;
+  // Returns the number of buffers that are currently free in the pool.
+  size_t NumFreeForTesting() const;
+
+ private:
+  class Internal;
+
+  friend class base::RefCountedThreadSafe<IOBufferPool>;
+  ~IOBufferPool();
+
+  const size_t buffer_size_;
+  const size_t max_buffers_;
+  Internal* internal_;  // Manages its own lifetime.
+
+  DISALLOW_COPY_AND_ASSIGN(IOBufferPool);
+};
+
+}  // namespace chromecast
+
+#endif  // CHROMECAST_NET_IO_BUFFER_POOL_H_

chromecast/net/io_buffer_pool_unittest.cc

+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chromecast/net/io_buffer_pool.h"
+
+#include "base/memory/ref_counted.h"
+#include "net/base/io_buffer.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace chromecast {
+
+namespace {
+const size_t kDefaultBufferSize = 256;
+}  // namespace
+
+TEST(IOBufferPoolTest, ZeroMaxBuffers) {
+  auto pool = base::MakeRefCounted<IOBufferPool>(kDefaultBufferSize, 0);
+  EXPECT_EQ(nullptr, pool->GetBuffer());
+  EXPECT_EQ(kDefaultBufferSize, pool->buffer_size());
+  EXPECT_EQ(0u, pool->max_buffers());
+  EXPECT_EQ(0u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(0u, pool->NumFreeForTesting());
+}
+
+TEST(IOBufferPoolTest, OneMaxBuffer) {
+  auto pool = base::MakeRefCounted<IOBufferPool>(kDefaultBufferSize, 1);
+  scoped_refptr<net::IOBuffer> buffer = pool->GetBuffer();
+  EXPECT_NE(nullptr, buffer.get());
+  EXPECT_EQ(nullptr, pool->GetBuffer());
+  EXPECT_EQ(1u, pool->max_buffers());
+  EXPECT_EQ(1u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(0u, pool->NumFreeForTesting());
+}
+
+TEST(IOBufferPoolTest, SeveralMaxBuffers) {
+  auto pool = base::MakeRefCounted<IOBufferPool>(kDefaultBufferSize, 4);
+  scoped_refptr<net::IOBuffer> buffer1 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer2 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer3 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer4 = pool->GetBuffer();
+  EXPECT_NE(nullptr, buffer1.get());
+  EXPECT_NE(nullptr, buffer2.get());
+  EXPECT_NE(nullptr, buffer3.get());
+  EXPECT_NE(nullptr, buffer4.get());
+  EXPECT_EQ(nullptr, pool->GetBuffer());
+  EXPECT_EQ(4u, pool->max_buffers());
+  EXPECT_EQ(4u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(0u, pool->NumFreeForTesting());
+}
+
+TEST(IOBufferPoolTest, Reclaim) {
+  auto pool = base::MakeRefCounted<IOBufferPool>(kDefaultBufferSize, 3);
+  scoped_refptr<net::IOBuffer> buffer1 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer2 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer3 = pool->GetBuffer();
+  EXPECT_NE(nullptr, buffer1.get());
+  EXPECT_NE(nullptr, buffer2.get());
+  EXPECT_NE(nullptr, buffer3.get());
+  EXPECT_EQ(3u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(0u, pool->NumFreeForTesting());
+  buffer1 = nullptr;
+  EXPECT_EQ(3u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(1u, pool->NumFreeForTesting());
+  buffer2 = nullptr;
+  EXPECT_EQ(3u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(2u, pool->NumFreeForTesting());
+  buffer3 = nullptr;
+  EXPECT_EQ(3u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(3u, pool->NumFreeForTesting());
+  buffer1 = pool->GetBuffer();
+  EXPECT_EQ(3u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(2u, pool->NumFreeForTesting());
+}
+
+TEST(IOBufferPoolTest, DestroyBufferAfterPool) {
+  auto pool = base::MakeRefCounted<IOBufferPool>(kDefaultBufferSize, 3);
+  scoped_refptr<net::IOBuffer> buffer1 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer2 = pool->GetBuffer();
+  EXPECT_NE(nullptr, buffer1.get());
+  EXPECT_NE(nullptr, buffer2.get());
+  EXPECT_EQ(2u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(0u, pool->NumFreeForTesting());
+  buffer1 = nullptr;
+  EXPECT_EQ(2u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(1u, pool->NumFreeForTesting());
+
+  pool = nullptr;
+  buffer2 = nullptr;  // Expect no crash and no memory errors.
+}
+
+TEST(IOBufferPoolTest, Preallocate) {
+  auto pool = base::MakeRefCounted<IOBufferPool>(kDefaultBufferSize, 3);
+  pool->Preallocate(8);
+
+  EXPECT_EQ(3u, pool->max_buffers());
+  EXPECT_EQ(3u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(3u, pool->NumFreeForTesting());
+
+  scoped_refptr<net::IOBuffer> buffer1 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer2 = pool->GetBuffer();
+  scoped_refptr<net::IOBuffer> buffer3 = pool->GetBuffer();
+
+  EXPECT_NE(nullptr, buffer1.get());
+  EXPECT_NE(nullptr, buffer2.get());
+  EXPECT_NE(nullptr, buffer3.get());
+  EXPECT_EQ(nullptr, pool->GetBuffer());
+
+  EXPECT_EQ(3u, pool->NumAllocatedForTesting());
+  EXPECT_EQ(0u, pool->NumFreeForTesting());
+}
+
+}  // namespace chromecast