Improve ExtentRanges::AddExtent() runtime from O(n) to O(log n)

Old algorithm iterates over all extents to detect if there's any
intersection, use binary search instead.

Test: Add ~200K extents from an actual OTA, runtime improves from 13s to
0.4s
Bug: 315215541

Change-Id: Ia99eda3a22e726a9c02bfe2a152dc0a5c16efcfc

6 files changed, 79 insertions, 25 deletions

diff --git a/payload_consumer/extent_map.h b/payload_consumer/extent_map.h
index a83bf0f7..27c67ec4 100644
--- a/payload_consumer/extent_map.h
+++ b/payload_consumer/extent_map.h
@@ -17,7 +17,6 @@
 #ifndef UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
 #define UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
 
-#include <functional>
 #include <map>
 #include <utility>
 #include <vector>
diff --git a/payload_generator/extent_ranges.cc b/payload_generator/extent_ranges.cc
index eecc8b32..32580b91 100644
--- a/payload_generator/extent_ranges.cc
+++ b/payload_generator/extent_ranges.cc
@@ -25,7 +25,6 @@
 
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/payload_constants.h"
-#include "update_engine/payload_generator/extent_utils.h"
 
 using std::vector;
 
@@ -84,9 +83,8 @@ void ExtentRanges::AddExtent(Extent extent) {
   ExtentSet::iterator begin_del = extent_set_.end();
   ExtentSet::iterator end_del = extent_set_.end();
   uint64_t del_blocks = 0;
-  for (ExtentSet::iterator it = extent_set_.begin(), e = extent_set_.end();
-       it != e;
-       ++it) {
+  const auto range = GetCandidateRange(extent);
+  for (ExtentSet::iterator it = range.begin(), e = range.end(); it != e; ++it) {
     const bool should_merge = merge_touching_extents_
                                   ? ExtentsOverlapOrTouch(*it, extent)
                                   : ExtentsOverlap(*it, extent);
@@ -291,12 +289,20 @@ vector<Extent> ExtentRanges::GetExtentsForBlockCount(uint64_t count) const {
 Range<ExtentRanges::ExtentSet::const_iterator> ExtentRanges::GetCandidateRange(
     const Extent& extent) const {
   auto lower_it = extent_set_.lower_bound(extent);
-  if (lower_it != extent_set_.begin()) {
+  while (lower_it != extent_set_.begin() &&
+         (lower_it == extent_set_.end() ||
+          lower_it->start_block() + lower_it->num_blocks() >
+              extent.start_block())) {
     --lower_it;
   }
 
-  const auto upper_it = extent_set_.upper_bound(
+  auto upper_it = extent_set_.upper_bound(
       ExtentForRange(extent.start_block() + extent.num_blocks(), 1));
+  while (upper_it != extent_set_.end() &&
+         upper_it->start_block() <=
+             extent.start_block() + extent.num_blocks()) {
+    ++upper_it;
+  }
   return {lower_it, upper_it};
 }
 
diff --git a/payload_generator/extent_ranges.h b/payload_generator/extent_ranges.h
index 08cf5fe1..bd468a11 100644
--- a/payload_generator/extent_ranges.h
+++ b/payload_generator/extent_ranges.h
@@ -17,13 +17,13 @@
 #ifndef UPDATE_ENGINE_PAYLOAD_GENERATOR_EXTENT_RANGES_H_
 #define UPDATE_ENGINE_PAYLOAD_GENERATOR_EXTENT_RANGES_H_
 
-#include <map>
 #include <set>
 #include <vector>
 
 #include <base/macros.h>
 
 #include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/extent_utils.h"
 #include "update_engine/update_metadata.pb.h"
 
 // An ExtentRanges object represents an unordered collection of extents (and
@@ -35,15 +35,6 @@
 
 namespace chromeos_update_engine {
 
-struct ExtentLess {
-  bool operator()(const Extent& x, const Extent& y) const {
-    if (x.start_block() == y.start_block()) {
-      return x.num_blocks() < y.num_blocks();
-    }
-    return x.start_block() < y.start_block();
-  }
-};
-
 Extent ExtentForRange(uint64_t start_block, uint64_t num_blocks);
 Extent ExtentForBytes(uint64_t block_size,
                       uint64_t start_bytes,
diff --git a/payload_generator/extent_ranges_unittest.cc b/payload_generator/extent_ranges_unittest.cc
index f7a4cd22..5f36aa3d 100644
--- a/payload_generator/extent_ranges_unittest.cc
+++ b/payload_generator/extent_ranges_unittest.cc
@@ -21,11 +21,11 @@
 #include <base/stl_util.h>
 #include <gtest/gtest.h>
 
-#include "update_engine/common/test_utils.h"
-#include "update_engine/payload_consumer/payload_constants.h"
 #include "update_engine/payload_generator/extent_utils.h"
+#include "update_engine/payload_consumer/payload_constants.h"
 
 using std::vector;
+using chromeos_update_engine::operator==;
 
 namespace chromeos_update_engine {
 
@@ -390,4 +390,35 @@ TEST(ExtentRangesTest, OverlapsWithExtent) {
   ASSERT_TRUE(ranges.OverlapsWithExtent(ExtentForRange(19, 1)));
 }
 
+TEST(ExtentRangesTest, AddExtentMergeStressTest) {
+  ExtentRanges ranges(true);
+  for (size_t i = 0; i < 1000000; i++) {
+    ranges.AddExtent(ExtentForRange(i, 1));
+  }
+  ASSERT_EQ(ranges.extent_set().size(), 1UL) << ranges.extent_set();
+}
+
+TEST(ExtentRangesTest, AddExtentNoMergeStressTest) {
+  ExtentRanges ranges(true);
+  for (size_t i = 0; i < 200000; i++) {
+    ranges.AddExtent(ExtentForRange(i * 2, 1));
+  }
+  ASSERT_EQ(ranges.extent_set().size(), 200000UL) << ranges.extent_set();
+}
+
+TEST(ExtentRangesTest, AddExtentTouching) {
+  ExtentRanges ranges(true);
+  ranges.AddExtent(ExtentForRange(5, 5));
+  ranges.AddExtent(ExtentForRange(25, 7));
+  ASSERT_EQ(ranges.extent_set().size(), 2UL) << ranges.extent_set();
+  ranges.AddExtent(ExtentForRange(0, 5));
+  ASSERT_EQ(ranges.extent_set().size(), 2UL) << ranges.extent_set();
+  ranges.AddExtent(ExtentForRange(10, 15));
+  ASSERT_EQ(ranges.extent_set().size(), 1UL) << ranges.extent_set();
+  ranges.AddExtent(ExtentForRange(32, 8));
+  ASSERT_EQ(ranges.extent_set().size(), 1UL) << ranges.extent_set();
+  ranges.AddExtent(ExtentForRange(45, 5));
+  ASSERT_EQ(ranges.extent_set().size(), 2UL) << ranges.extent_set();
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_utils.cc b/payload_generator/extent_utils.cc
index 612fd672..851db8a2 100644
--- a/payload_generator/extent_utils.cc
+++ b/payload_generator/extent_utils.cc
@@ -19,16 +19,13 @@
 #include <inttypes.h>
 
 #include <string>
-#include <utility>
 #include <vector>
 
 #include <base/logging.h>
 #include <base/macros.h>
 #include <base/strings/stringprintf.h>
 
-#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/payload_constants.h"
-#include "update_engine/payload_generator/annotated_operation.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 
 using std::string;
@@ -171,8 +168,7 @@ bool operator!=(const Extent& a, const Extent& b) noexcept {
 }
 
 std::ostream& operator<<(std::ostream& out, const Extent& extent) {
-  out << "[" << extent.start_block() << " - "
-      << extent.start_block() + extent.num_blocks() - 1 << "]";
+  out << "(" << extent.start_block() << " - " << extent.num_blocks() << ")";
   return out;
 }
 
@@ -202,4 +198,19 @@ std::ostream& operator<<(
   return PrintExtents(out, extents);
 }
 
+std::ostream& operator<<(std::ostream& out, const std::set<Extent>& extents) {
+  return PrintExtents(out, extents);
+}
+
+std::ostream& operator<<(std::ostream& out,
+                         const std::set<Extent, ExtentLess>& extents) {
+  return PrintExtents(out, extents);
+}
+
+std::ostream& operator<<(
+    std::ostream& out,
+    Range<std::set<Extent, ExtentLess>::const_iterator> range) {
+  return PrintExtents(out, range);
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_utils.h b/payload_generator/extent_utils.h
index f8d36e7f..52b6d1e0 100644
--- a/payload_generator/extent_utils.h
+++ b/payload_generator/extent_utils.h
@@ -23,13 +23,21 @@
 
 #include <base/logging.h>
 
-#include "google/protobuf/repeated_field.h"
+#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/payload_constants.h"
 #include "update_engine/update_metadata.pb.h"
 
 // Utility functions for manipulating Extents and lists of blocks.
 
 namespace chromeos_update_engine {
+struct ExtentLess {
+  constexpr bool operator()(const Extent& x, const Extent& y) const {
+    if (x.start_block() == y.start_block()) {
+      return x.num_blocks() < y.num_blocks();
+    }
+    return x.start_block() < y.start_block();
+  }
+};
 
 // |block| must either be the next block in the last extent or a block
 // in the next extent. This function will not handle inserting block
@@ -130,6 +138,14 @@ struct BlockIterator {
 
 std::ostream& operator<<(std::ostream& out, const Extent& extent);
 std::ostream& operator<<(std::ostream& out, const std::vector<Extent>& extent);
+std::ostream& operator<<(std::ostream& out, const std::set<Extent>& extents);
+
+std::ostream& operator<<(std::ostream& out,
+                         const std::set<Extent, ExtentLess>& extents);
+std::ostream& operator<<(
+    std::ostream& out,
+    Range<std::set<Extent, ExtentLess>::const_iterator> range);
+
 std::ostream& operator<<(
     std::ostream& out,
     const google::protobuf::RepeatedPtrField<Extent>& extent);