Resource management in Storage module.

Memory management implemented and added.

Bug: b:159361496
Bug: b:157943192
Change-Id: Iedfb204680c89867188180902ce1054a75e9fd24
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2552719Reviewed-by: Zach Trudo <zatrudo@google.com>
Commit-Queue: Leonid Baraz <lbaraz@chromium.org>
Cr-Commit-Position: refs/heads/master@{#830255}

chrome/browser/BUILD.gn

@@ -1219,6 +1219,10 @@ static_library("browser") {
     "policy/messaging_layer/public/report_queue.h",
     "policy/messaging_layer/public/report_queue_configuration.cc",
     "policy/messaging_layer/public/report_queue_configuration.h",
+    "policy/messaging_layer/storage/resources/memory_resource_impl.cc",
+    "policy/messaging_layer/storage/resources/memory_resource_impl.h",
+    "policy/messaging_layer/storage/resources/resource_interface.cc",
+    "policy/messaging_layer/storage/resources/resource_interface.h",
     "policy/messaging_layer/storage/storage.cc",
     "policy/messaging_layer/storage/storage.h",
     "policy/messaging_layer/storage/storage_configuration.h",

chrome/browser/policy/messaging_layer/storage/resources/memory_resource_impl.cc

+// Copyright 2020 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 "chrome/browser/policy/messaging_layer/storage/resources/memory_resource_impl.h"
+
+#include <atomic>
+#include <cstdint>
+
+#include "base/check_op.h"
+#include "base/no_destructor.h"
+
+namespace reporting {
+
+// TODO(b/159361496): Set total memory allowance based on the platform
+// (or policy?).
+MemoryResourceImpl::MemoryResourceImpl()
+    : total_(16u * 1024LLu * 1024LLu),  // 16 MiB
+      used_(0) {}
+
+MemoryResourceImpl::~MemoryResourceImpl() = default;
+
+bool MemoryResourceImpl::Reserve(uint64_t size) {
+  uint64_t old_used = used_.fetch_add(size);
+  if (old_used + size > total_) {
+    used_.fetch_sub(size);
+    return false;
+  }
+  return true;
+}
+
+void MemoryResourceImpl::Discard(uint64_t size) {
+  DCHECK_LE(size, used_.load());
+  used_.fetch_sub(size);
+}
+
+uint64_t MemoryResourceImpl::GetTotal() {
+  return total_;
+}
+
+uint64_t MemoryResourceImpl::GetUsed() {
+  return used_.load();
+}
+
+void MemoryResourceImpl::Test_SetTotal(uint64_t test_total) {
+  total_ = test_total;
+}
+
+ResourceInterface* GetMemoryResource() {
+  static base::NoDestructor<MemoryResourceImpl> memory;
+  return memory.get();
+}
+
+}  // namespace reporting

chrome/browser/policy/messaging_layer/storage/resources/memory_resource_impl.h

+// Copyright 2020 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 CHROME_BROWSER_POLICY_MESSAGING_LAYER_STORAGE_RESOURCES_MEMORY_RESOURCE_IMPL_H_
+#define CHROME_BROWSER_POLICY_MESSAGING_LAYER_STORAGE_RESOURCES_MEMORY_RESOURCE_IMPL_H_
+
+#include <atomic>
+#include <cstdint>
+
+#include "chrome/browser/policy/messaging_layer/storage/resources/resource_interface.h"
+
+namespace reporting {
+
+// Interface to resources management by Storage module.
+// Must be implemented by the caller base on the platform limitations.
+// All APIs are non-blocking.
+class MemoryResourceImpl : public ResourceInterface {
+ public:
+  MemoryResourceImpl();
+  ~MemoryResourceImpl() override;
+
+  // Implementation of ResourceInterface methods.
+  bool Reserve(uint64_t size) override;
+  void Discard(uint64_t size) override;
+  uint64_t GetTotal() override;
+  uint64_t GetUsed() override;
+  void Test_SetTotal(uint64_t test_total) override;
+
+ private:
+  uint64_t total_;
+  std::atomic<uint64_t> used_;
+};
+
+}  // namespace reporting
+
+#endif  // CHROME_BROWSER_POLICY_MESSAGING_LAYER_STORAGE_RESOURCES_MEMORY_RESOURCE_IMPL_H_

chrome/browser/policy/messaging_layer/storage/resources/resource_interface.cc

+// Copyright 2020 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 <cstdint>
+
+#include "chrome/browser/policy/messaging_layer/storage/resources/resource_interface.h"
+
+namespace reporting {
+
+ScopedReservation::ScopedReservation(uint64_t size,
+                                     ResourceInterface* resource_interface)
+    : resource_interface_(resource_interface) {
+  if (!resource_interface->Reserve(size)) {
+    return;
+  }
+  size_ = size;
+}
+
+ScopedReservation::ScopedReservation(ScopedReservation&& other)
+    : resource_interface_(other.resource_interface_),
+      size_(std::move(other.size_)) {}
+
+bool ScopedReservation::reserved() const {
+  return size_.has_value();
+}
+
+ScopedReservation::~ScopedReservation() {
+  if (reserved()) {
+    resource_interface_->Reserve(size_.value());
+  }
+}
+
+}  // namespace reporting

chrome/browser/policy/messaging_layer/storage/resources/resource_interface.h

+// Copyright 2020 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 CHROME_BROWSER_POLICY_MESSAGING_LAYER_STORAGE_RESOURCES_RESOURCE_INTERFACE_H_
+#define CHROME_BROWSER_POLICY_MESSAGING_LAYER_STORAGE_RESOURCES_RESOURCE_INTERFACE_H_
+
+#include <cstdint>
+
+#include "base/optional.h"
+
+namespace reporting {
+
+// Interface to resources management by Storage module.
+// Must be implemented by the caller base on the platform limitations.
+// All APIs are non-blocking.
+class ResourceInterface {
+ public:
+  virtual ~ResourceInterface() = default;
+
+  // Needs to be called before attempting to allocate specified size.
+  // Returns true if requested amount can be allocated.
+  // After that the caller can actually allocate it or must call
+  // |Discard| if decided not to allocate.
+  virtual bool Reserve(uint64_t size) = 0;
+
+  // Reverts reservation.
+  // Must be called after the specified amount is released.
+  virtual void Discard(uint64_t size) = 0;
+
+  // Returns total amount.
+  virtual uint64_t GetTotal() = 0;
+
+  // Returns current used amount.
+  virtual uint64_t GetUsed() = 0;
+
+  // Test only: Sets non-default usage limit.
+  virtual void Test_SetTotal(uint64_t test_total) = 0;
+
+ protected:
+  ResourceInterface() = default;
+};
+
+// Moveable RAII class used for scoped Reserve-Discard.
+//
+// Usage:
+//  {
+//    ScopedReservation reservation(1024u, GetMemoryResource());
+//    if (!reservation.reserved()) {
+//      // Allocation failed.
+//      return;
+//    }
+//    // Allocation succeeded.
+//    ...
+//  }   // Automatically discarded.
+//
+// Can be handed over to another owner.
+class ScopedReservation {
+ public:
+  ScopedReservation(uint64_t size, ResourceInterface* resource_interface);
+  ScopedReservation(ScopedReservation&& other);
+  ScopedReservation(const ScopedReservation& other) = delete;
+  ScopedReservation& operator=(const ScopedReservation& other) = delete;
+  ~ScopedReservation();
+
+  bool reserved() const;
+
+ private:
+  ResourceInterface* const resource_interface_;
+  base::Optional<uint64_t> size_;
+};
+
+ResourceInterface* GetMemoryResource();
+ResourceInterface* GetDiskResource();
+
+}  // namespace reporting
+
+#endif  // CHROME_BROWSER_POLICY_MESSAGING_LAYER_STORAGE_RESOURCES_RESOURCE_INTERFACE_H_

chrome/browser/policy/messaging_layer/storage/storage_queue.cc

@@ -33,6 +33,7 @@
 #include "base/task/post_task.h"
 #include "base/task_runner.h"
 #include "chrome/browser/policy/messaging_layer/encryption/encryption_module.h"
+#include "chrome/browser/policy/messaging_layer/storage/resources/resource_interface.h"
 #include "chrome/browser/policy/messaging_layer/storage/storage_configuration.h"
 #include "chrome/browser/policy/messaging_layer/util/status.h"
 #include "chrome/browser/policy/messaging_layer/util/status_macros.h"
@@ -54,11 +55,6 @@ const base::FilePath::CharType METADATA_NAME[] = FILE_PATH_LITERAL("META");
 // make it harder to differ between kinds of records).
 constexpr size_t FRAME_SIZE = 16u;
 
-// Size of the buffer to read data to. Must be multiple of FRAME_SIZE
-constexpr size_t MAX_BUFFER_SIZE = 1024u * 1024u;  // 1 MiB
-static_assert(MAX_BUFFER_SIZE % FRAME_SIZE == 0u,
-              "Buffer size not multiple of frame size");
-
 // Helper functions for FRAME_SIZE alignment support.
 size_t RoundUpToFrameSize(size_t size) {
   return (size + FRAME_SIZE - 1) / FRAME_SIZE * FRAME_SIZE;
@@ -310,10 +306,14 @@ Status StorageQueue::ScanLastFile() {
     return Status(error::DATA_LOSS, base::StrCat({"Error opening file: '",
                                                   last_file->name(), "'"}));
   }
+  const size_t max_buffer_size =
+      RoundUpToFrameSize(options_.max_record_size()) +
+      RoundUpToFrameSize(sizeof(RecordHeader));
   uint32_t pos = 0;
   for (;;) {
     // Read the header
-    auto read_result = last_file->Read(pos, sizeof(RecordHeader));
+    auto read_result =
+        last_file->Read(pos, sizeof(RecordHeader), max_buffer_size);
     if (read_result.status().error_code() == error::OUT_OF_RANGE) {
       // End of file detected.
       break;
@@ -335,7 +335,7 @@ Status StorageQueue::ScanLastFile() {
         *reinterpret_cast<const RecordHeader*>(read_result.ValueOrDie().data());
     // Read the data (rounded to frame size).
     const size_t data_size = RoundUpToFrameSize(header.record_size);
-    read_result = last_file->Read(pos, data_size);
+    read_result = last_file->Read(pos, data_size, max_buffer_size);
     if (!read_result.ok()) {
       // Error detected.
       LOG(ERROR) << "Error reading file " << last_file->name()
@@ -550,7 +550,10 @@ Status StorageQueue::RestoreMetadata(
                                                     /*size=*/it->second.second);
   RETURN_IF_ERROR(meta_file->Open(/*read_only=*/true));
   // Read generation id.
-  auto read_result = meta_file->Read(/*pos=*/0, sizeof(generation_id_));
+  constexpr size_t max_buffer_size =
+      sizeof(generation_id_) + crypto::kSHA256Length;
+  auto read_result =
+      meta_file->Read(/*pos=*/0, sizeof(generation_id_), max_buffer_size);
   if (!read_result.ok() ||
       read_result.ValueOrDie().size() != sizeof(generation_id_)) {
     return Status(error::DATA_LOSS,
@@ -560,8 +563,8 @@ Status StorageQueue::RestoreMetadata(
   const uint64_t generation_id =
       *reinterpret_cast<const uint64_t*>(read_result.ValueOrDie().data());
   // Read last record digest.
-  read_result =
-      meta_file->Read(/*pos=*/sizeof(generation_id_), crypto::kSHA256Length);
+  read_result = meta_file->Read(/*pos=*/sizeof(generation_id_),
+                                crypto::kSHA256Length, max_buffer_size);
   if (!read_result.ok() ||
       read_result.ValueOrDie().size() != crypto::kSHA256Length) {
     return Status(error::DATA_LOSS,
@@ -861,8 +864,11 @@ class StorageQueue::ReadContext : public TaskRunnerContext<Status> {
 
     // Read from the current file at the current offset.
     RETURN_IF_ERROR(current_file_->second->Open(/*read_only=*/true));
-    auto read_result =
-        current_file_->second->Read(current_pos_, sizeof(RecordHeader));
+    const size_t max_buffer_size =
+        RoundUpToFrameSize(storage_queue->options_.max_record_size()) +
+        RoundUpToFrameSize(sizeof(RecordHeader));
+    auto read_result = current_file_->second->Read(
+        current_pos_, sizeof(RecordHeader), max_buffer_size);
     RETURN_IF_ERROR(read_result.status());
     auto header_data = read_result.ValueOrDie();
     if (header_data.empty()) {
@@ -891,7 +897,8 @@ class StorageQueue::ReadContext : public TaskRunnerContext<Status> {
     const size_t data_size = RoundUpToFrameSize(header.record_size);
     // From this point on, header in memory is no longer used and can be
     // overwritten when reading rest of the data.
-    read_result = current_file_->second->Read(current_pos_, data_size);
+    read_result =
+        current_file_->second->Read(current_pos_, data_size, max_buffer_size);
     RETURN_IF_ERROR(read_result.status());
     current_pos_ += read_result.ValueOrDie().size();
     if (read_result.ValueOrDie().size() != data_size) {
@@ -1022,13 +1029,23 @@ class StorageQueue::WriteContext : public TaskRunnerContext<Status> {
 
   void SerializeAndEncryptWrappedRecord(WrappedRecord wrapped_record) {
     // Serialize wrapped record into a string.
+    ScopedReservation scoped_reservation(wrapped_record.ByteSizeLong(),
+                                         GetMemoryResource());
+    if (!scoped_reservation.reserved()) {
+      Schedule(&ReadContext::Response, base::Unretained(this),
+               Status(error::RESOURCE_EXHAUSTED,
+                      "Not enough memory for the write buffer"));
+      return;
+    }
+
     std::string buffer;
     if (!wrapped_record.SerializeToString(&buffer)) {
       Schedule(&ReadContext::Response, base::Unretained(this),
                Status(error::DATA_LOSS, "Cannot serialize record"));
       return;
     }
-    wrapped_record.Clear();  // Release wrapped record memory.
+    // Release wrapped record memory, so scoped reservation may act.
+    wrapped_record.Clear();
 
     // Encrypt the result.
     storage_queue_->encryption_module_->EncryptRecord(
@@ -1046,14 +1063,23 @@ class StorageQueue::WriteContext : public TaskRunnerContext<Status> {
     }
 
     // Serialize encrypted record.
+    ScopedReservation scoped_reservation(
+        encrypted_record_result.ValueOrDie().ByteSizeLong(),
+        GetMemoryResource());
+    if (!scoped_reservation.reserved()) {
+      Schedule(&ReadContext::Response, base::Unretained(this),
+               Status(error::RESOURCE_EXHAUSTED,
+                      "Not enough memory for the write buffer"));
+      return;
+    }
     std::string buffer;
     if (!encrypted_record_result.ValueOrDie().SerializeToString(&buffer)) {
       Schedule(&ReadContext::Response, base::Unretained(this),
                Status(error::DATA_LOSS, "Cannot serialize EncryptedRecord"));
       return;
     }
-    encrypted_record_result.ValueOrDie()
-        .Clear();  // Release encrypted record memory.
+    // Release encrypted record memory, so scoped reservation may act.
+    encrypted_record_result.ValueOrDie().Clear();
 
     // Write into storage on sequntial task runner.
     Schedule(&WriteContext::WriteRecord, base::Unretained(this), buffer);
@@ -1293,6 +1319,7 @@ void StorageQueue::SingleFile::Close() {
   handle_.reset();
   is_readonly_ = base::nullopt;
   buffer_.reset();
+  GetMemoryResource()->Discard(buffer_size_);
 }
 
 Status StorageQueue::SingleFile::Delete() {
@@ -1305,25 +1332,31 @@ Status StorageQueue::SingleFile::Delete() {
   return Status::StatusOK();
 }
 
-StatusOr<base::StringPiece> StorageQueue::SingleFile::Read(uint32_t pos,
-                                                           uint32_t size) {
+StatusOr<base::StringPiece> StorageQueue::SingleFile::Read(
+    uint32_t pos,
+    uint32_t size,
+    size_t max_buffer_size) {
   if (!handle_) {
     return Status(error::UNAVAILABLE, base::StrCat({"File not open ", name()}));
   }
-  if (size > MAX_BUFFER_SIZE) {
+  if (size > max_buffer_size) {
     return Status(error::RESOURCE_EXHAUSTED, "Too much data to read");
   }
   if (size_ == 0) {
     // Empty file, return EOF right away.
     return Status(error::OUT_OF_RANGE, "End of file");
   }
-  const size_t buffer_size =
-      std::min(MAX_BUFFER_SIZE, RoundUpToFrameSize(size_));
+  buffer_size_ = std::min(max_buffer_size, RoundUpToFrameSize(size_));
   // If no buffer yet, allocate.
   // TODO(b/157943192): Add buffer management - consider adding an UMA for
   // tracking the average + peak memory the Storage module is consuming.
   if (!buffer_) {
-    buffer_ = std::make_unique<char[]>(buffer_size);
+    // Register with resource management.
+    if (!GetMemoryResource()->Reserve(buffer_size_)) {
+      return Status(error::RESOURCE_EXHAUSTED,
+                    "Not enough memory for the read buffer");
+    }
+    buffer_ = std::make_unique<char[]>(buffer_size_);
     data_start_ = data_end_ = 0;
     file_position_ = 0;
   }
@@ -1334,7 +1367,7 @@ StatusOr<base::StringPiece> StorageQueue::SingleFile::Read(uint32_t pos,
   }
   // If expected data size does not fit into the buffer, move what's left to the
   // start.
-  if (data_start_ + size > buffer_size) {
+  if (data_start_ + size > buffer_size_) {
     DCHECK_GT(data_start_, 0u);  // Cannot happen if 0.
     memmove(buffer_.get(), buffer_.get() + data_start_,
             data_end_ - data_start_);
@@ -1346,7 +1379,7 @@ StatusOr<base::StringPiece> StorageQueue::SingleFile::Read(uint32_t pos,
     // Read as much as possible.
     const int32_t result =
         handle_->Read(pos, reinterpret_cast<char*>(buffer_.get() + data_end_),
-                      buffer_size - data_end_);
+                      buffer_size_ - data_end_);
     if (result < 0) {
       return Status(
           error::DATA_LOSS,
@@ -1359,7 +1392,7 @@ StatusOr<base::StringPiece> StorageQueue::SingleFile::Read(uint32_t pos,
     }
     pos += result;
     data_end_ += result;
-    DCHECK_LE(data_end_, buffer_size);
+    DCHECK_LE(data_end_, buffer_size_);
     actual_size += result;
   }
   if (actual_size > size) {

chrome/browser/policy/messaging_layer/storage/storage_queue.h

@@ -156,7 +156,12 @@ class StorageQueue : public base::RefCountedThreadSafe<StorageQueue> {
     // Attempts to read |size| bytes from position |pos| and returns
     // reference to the data that were actually read (no more than |size|).
     // End of file is indicated by empty data.
-    StatusOr<base::StringPiece> Read(uint32_t pos, uint32_t size);
+    // |max_buffer_size| specifies the largest allowed buffer, which
+    // must accommodate the largest possible data block plus header and
+    // overhead.
+    StatusOr<base::StringPiece> Read(uint32_t pos,
+                                     uint32_t size,
+                                     size_t max_buffer_size);
 
     // Appends data to the file.
     StatusOr<uint32_t> Append(base::StringPiece data);
@@ -192,6 +197,7 @@ class StorageQueue : public base::RefCountedThreadSafe<StorageQueue> {
     size_t data_start_ = 0;
     size_t data_end_ = 0;
     uint64_t file_position_ = 0;
+    size_t buffer_size_ = 0;
     std::unique_ptr<char[]> buffer_;
   };
 

chrome/browser/policy/messaging_layer/storage/storage_queue_unittest.cc

@@ -981,7 +981,8 @@ INSTANTIATE_TEST_SUITE_P(VaryingFileSize,
 // 1) Options object with a bad path.
 // 2) Have bad prefix files in the directory.
 // 3) Attempt to create file with duplicated file extension.
-// 4) Other negative tests.
+// 4) Disk and memory limit exceeded.
+// 5) Other negative tests.
 
 }  // namespace
 }  // namespace reporting