introduce sanity check when creating bufferdatasink

Differential Revision: D70190912

Pull Request resolved: #8709

Author: Gasoonjia

devtools/etdump/buffer_data_sink.cpp

@@ -11,10 +11,22 @@
 
 using ::executorch::runtime::Error;
 using ::executorch::runtime::Result;
+using ::executorch::runtime::Span;
 
 namespace executorch {
 namespace etdump {
 
+Result<BufferDataSink> BufferDataSink::create(
+    Span<uint8_t> buffer,
+    size_t alignment) noexcept {
+  // Check if alignment is a power of two and greater than 0
+  if (alignment == 0 || (alignment & (alignment - 1)) != 0) {
+    return Error::InvalidArgument;
+  }
+
+  return BufferDataSink(buffer, alignment);
+}
+
 Result<size_t> BufferDataSink::write(const void* ptr, size_t length) {
   if (length == 0) {
     return offset_;

devtools/etdump/buffer_data_sink.h

@@ -24,17 +24,20 @@ namespace etdump {
 class BufferDataSink : public DataSinkBase {
  public:
   /**
-   * Constructs a BufferDataSink with a given buffer.
+   * Creates a BufferDataSink with a given span buffer.
    *
    * @param[in] buffer A Span object representing the buffer where data will be
    * stored.
    * @param[in] alignment The alignment requirement for the buffer. It must be
-   * a power of two. Default is 64.
+   * a power of two and greater than zero. Default is 64.
+   * @return A Result object containing either:
+   *         - A BufferDataSink object if succees, or
+   *         - An error code indicating the failure reason, if any issue
+   *           occurs during the creation process.
    */
-  explicit BufferDataSink(
+  static ::executorch::runtime::Result<BufferDataSink> create(
       ::executorch::runtime::Span<uint8_t> buffer,
-      size_t alignment = 64)
-      : debug_buffer_(buffer), offset_(0), alignment_(alignment) {}
+      size_t alignment = 64) noexcept;
 
   // Uncopiable and unassignable to avoid double assignment and free of the
   // internal buffer.
@@ -77,6 +80,19 @@ class BufferDataSink : public DataSinkBase {
   size_t get_used_bytes() const override;
 
  private:
+  /**
+   * Constructs a BufferDataSink with a given buffer.
+   *
+   * @param[in] buffer A Span object representing the buffer where data will be
+   * stored.
+   * @param[in] alignment The alignment requirement for the buffer. It must be
+   * a power of two. Default is 64.
+   */
+  explicit BufferDataSink(
+      ::executorch::runtime::Span<uint8_t> buffer,
+      size_t alignment)
+      : debug_buffer_(buffer), offset_(0), alignment_(alignment) {}
+
   // A Span object representing the buffer used for storing debug data.
   ::executorch::runtime::Span<uint8_t> debug_buffer_;
 

devtools/etdump/tests/buffer_data_sink_test.cpp

@@ -16,6 +16,7 @@
 using namespace ::testing;
 using ::executorch::aten::ScalarType;
 using ::executorch::aten::Tensor;
+using ::executorch::etdump::BufferDataSink;
 using ::executorch::runtime::Error;
 using ::executorch::runtime::Result;
 using ::executorch::runtime::Span;
@@ -29,7 +30,11 @@ class BufferDataSinkTest : public ::testing::Test {
     buffer_size_ = 128; // Small size for testing
     buffer_ptr_ = malloc(buffer_size_);
     buffer_ = Span<uint8_t>(static_cast<uint8_t*>(buffer_ptr_), buffer_size_);
-    data_sink_ = std::make_unique<executorch::etdump::BufferDataSink>(buffer_);
+    Result<BufferDataSink> buffer_data_sink_ret =
+        BufferDataSink::create(buffer_);
+    ASSERT_EQ(buffer_data_sink_ret.error(), Error::Ok);
+    buffer_data_sink_ =
+        std::make_unique<BufferDataSink>(std::move(buffer_data_sink_ret.get()));
   }
 
   void TearDown() override {
@@ -39,11 +44,11 @@ class BufferDataSinkTest : public ::testing::Test {
   size_t buffer_size_;
   void* buffer_ptr_;
   Span<uint8_t> buffer_;
-  std::unique_ptr<executorch::etdump::BufferDataSink> data_sink_;
+  std::unique_ptr<BufferDataSink> buffer_data_sink_;
 };
 
 TEST_F(BufferDataSinkTest, StorageSizeCheck) {
-  Result<size_t> ret = data_sink_->get_storage_size();
+  Result<size_t> ret = buffer_data_sink_->get_storage_size();
   ASSERT_EQ(ret.error(), Error::Ok);
 
   size_t storage_size = ret.get();
@@ -55,7 +60,7 @@ TEST_F(BufferDataSinkTest, WriteOneTensorAndCheckData) {
   Tensor tensor = tf.make({1, 4}, {1.0, 2.0, 3.0, 4.0});
 
   Result<size_t> ret =
-      data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
+      buffer_data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
   ASSERT_EQ(ret.error(), Error::Ok);
 
   size_t offset = ret.get();
@@ -75,9 +80,10 @@ TEST_F(BufferDataSinkTest, WriteMultiTensorsAndCheckData) {
   std::vector<Tensor> tensors = {
       tf.make({1, 4}, {1.0, 2.0, 3.0, 4.0}),
       tf.make({1, 4}, {5.0, 6.0, 7.0, 8.0})};
+
   for (const auto& tensor : tensors) {
     Result<size_t> ret =
-        data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
+        buffer_data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
     ASSERT_EQ(ret.error(), Error::Ok);
 
     size_t offset = ret.get();
@@ -94,8 +100,9 @@ TEST_F(BufferDataSinkTest, WriteMultiTensorsAndCheckData) {
 TEST_F(BufferDataSinkTest, PointerAlignmentCheck) {
   TensorFactory<ScalarType::Float> tf;
   Tensor tensor = tf.make({1, 4}, {1.0, 2.0, 3.0, 4.0});
+
   Result<size_t> ret =
-      data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
+      buffer_data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
   ASSERT_EQ(ret.error(), Error::Ok);
 
   size_t offset = ret.get();
@@ -112,12 +119,32 @@ TEST_F(BufferDataSinkTest, WriteUntilOverflow) {
   // Write tensors until we run out of space
   for (size_t i = 0; i < 2; i++) {
     Result<size_t> ret =
-        data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
+        buffer_data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
     ASSERT_EQ(ret.error(), Error::Ok);
   }
 
   // Attempting to write another tensor should raise an error
   Result<size_t> ret =
-      data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
+      buffer_data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
   ASSERT_EQ(ret.error(), Error::OutOfResources);
 }
+
+TEST_F(BufferDataSinkTest, illegalAlignment) {
+  // Create a buffer_data_sink_ with legal alignment that is a power of 2 and
+  // greater than 0
+  for (size_t i = 1; i <= 128; i <<= 1) {
+    Result<BufferDataSink> buffer_data_sink_ret =
+        BufferDataSink::create(buffer_, i);
+    ASSERT_EQ(buffer_data_sink_ret.error(), Error::Ok);
+  }
+
+  // Create a buffer_data_sink_ with illegal alignment that is not a power of 2
+  // or greater than 0
+  std::vector<size_t> illegal_alignments = {0, 3, 5, 7, 100, 127};
+
+  for (size_t i = 0; i < illegal_alignments.size(); i++) {
+    Result<BufferDataSink> buffer_data_sink_ret =
+        BufferDataSink::create(buffer_, illegal_alignments[i]);
+    ASSERT_EQ(buffer_data_sink_ret.error(), Error::InvalidArgument);
+  }
+}