Add Extent map
When installing OTA, we stream operations as data becomes available. For
each operation, update_engine needs to look at the merge sequence and
figure out which blocks should be converted to COW_XOR. To do that, we
need a data structure for querying a disjoint set of extents.
Test: th
Bug: 177104308
Change-Id: I14ec9d72a32e859b65516b894ca94d0153aa7e17
diff --git a/payload_consumer/extent_map.h b/payload_consumer/extent_map.h
new file mode 100644
index 0000000..aeb7066
--- /dev/null
+++ b/payload_consumer/extent_map.h
@@ -0,0 +1,90 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#ifndef UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
+#define UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
+
+#include <functional>
+#include <map>
+#include <utility>
+
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+// Data structure for storing a disjoint set of extents.
+// Currently the only usecase is for VABCPartitionWriter to keep track of which
+// block belongs to which merge operation. Therefore this class only contains
+// the minimal set of functions needed.
+template <typename T, typename Comparator = ExtentLess>
+class ExtentMap {
+ public:
+ bool AddExtent(const Extent& extent, T&& value) {
+ if (Get(extent)) {
+ return false;
+ }
+ const auto& [it, inserted] = map_.insert({extent, std::forward<T>(value)});
+ if (inserted) {
+ set_.AddExtent(extent);
+ }
+ return inserted;
+ }
+
+ size_t size() const { return map_.size(); }
+
+ // Return a pointer to entry which is intersecting |extent|. If T is already
+ // a pointer type, return T on success. This function always return
+ // |nullptr| on failure. Therefore you cannot store nullptr as an entry.
+ std::optional<T> Get(const Extent& extent) const {
+ const auto it = map_.find(extent);
+ if (it == map_.end()) {
+ LOG_IF(WARNING, set_.OverlapsWithExtent(extent))
+ << "Looking up a partially intersecting extent isn't supported by "
+ "this data structure.";
+ return {};
+ }
+ return {it->second};
+ }
+
+ // Return a set of extents that are contained in this extent map.
+ // If |extent| is completely covered by this extent map, a vector of itself
+ // will be returned.
+ // If only a subset of |extent| is covered by this extent map, a vector of
+ // parts in this map will be returned.
+ // If |extent| has no intersection with this map, an empty vector will be
+ // returned.
+ // E.g. extent map contains [0,5] and [10,15], GetIntersectingExtents([3, 12])
+ // would return [3,5] and [10,12]
+ std::vector<Extent> GetIntersectingExtents(const Extent& extent) const {
+ return set_.GetIntersectingExtents(extent);
+ }
+
+ // Complement of |GetIntersectingExtents|, return vector of extents which are
+ // part of |extent| but not covered by this map.
+ std::vector<Extent> GetNonIntersectingExtents(const Extent& extent) const {
+ return FilterExtentRanges({extent}, set_);
+ }
+
+ private:
+ // Get a range of exents that potentially intersect with parameter |extent|
+ std::map<Extent, T, Comparator> map_;
+ ExtentRanges set_;
+};
+} // namespace chromeos_update_engine
+
+#endif // UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
diff --git a/payload_consumer/extent_map_unittest.cc b/payload_consumer/extent_map_unittest.cc
new file mode 100644
index 0000000..e7972cd
--- /dev/null
+++ b/payload_consumer/extent_map_unittest.cc
@@ -0,0 +1,122 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include <gtest/gtest.h>
+#include <optional>
+
+#include "update_engine/payload_consumer/extent_map.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/extent_utils.h"
+
+namespace chromeos_update_engine {
+
+class ExtentMapTest : public ::testing::Test {
+ public:
+ ExtentMap<int> map_;
+};
+
+TEST_F(ExtentMapTest, QueryExactExtent) {
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
+ auto ret = map_.Get(ExtentForRange(0, 5));
+ ASSERT_NE(ret, std::nullopt);
+ ASSERT_EQ(*ret, 7);
+}
+
+TEST_F(ExtentMapTest, QuerySubset) {
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
+ auto ret = map_.Get(ExtentForRange(1, 2));
+ ASSERT_EQ(ret, std::nullopt);
+}
+
+TEST_F(ExtentMapTest, QueryTouching) {
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
+ auto ret = map_.Get(ExtentForRange(3, 2));
+ ASSERT_EQ(ret, std::nullopt);
+ ret = map_.Get(ExtentForRange(4, 1));
+ ASSERT_EQ(ret, std::nullopt);
+ ret = map_.Get(ExtentForRange(5, 5));
+ ASSERT_EQ(ret, std::nullopt);
+ ret = map_.Get(ExtentForRange(5, 6));
+ ASSERT_EQ(ret, std::nullopt);
+}
+
+TEST_F(ExtentMapTest, GetIntersectingExtents) {
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 7));
+ auto ret = std::vector<Extent>{};
+ ret = map_.GetIntersectingExtents(ExtentForRange(2, 10));
+ ASSERT_EQ(ret.size(), 2U);
+ ASSERT_EQ(ret[0].start_block(), 2U);
+ ASSERT_EQ(ret[0].num_blocks(), 3U);
+
+ ASSERT_EQ(ret[1].start_block(), 10U);
+ ASSERT_EQ(ret[1].num_blocks(), 2U);
+
+ ret = map_.GetIntersectingExtents(ExtentForRange(2, 17));
+ ASSERT_EQ(ret.size(), 2U);
+ ASSERT_EQ(ret[0].start_block(), 2U);
+ ASSERT_EQ(ret[0].num_blocks(), 3U);
+
+ ASSERT_EQ(ret[1].start_block(), 10U);
+ ASSERT_EQ(ret[1].num_blocks(), 5U);
+
+ ret = map_.GetIntersectingExtents(ExtentForRange(2, 2));
+ ASSERT_EQ(ret, std::vector<Extent>{ExtentForRange(2, 2)});
+
+ ret = map_.GetIntersectingExtents(ExtentForRange(10, 5));
+ ASSERT_EQ(ret, std::vector<Extent>{ExtentForRange(10, 5)});
+
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(20, 5), 7));
+ ret = map_.GetIntersectingExtents(ExtentForRange(0, 30));
+ ASSERT_EQ(ret.size(), 3U);
+ ASSERT_EQ(ret[0].start_block(), 0U);
+ ASSERT_EQ(ret[0].num_blocks(), 5U);
+
+ ASSERT_EQ(ret[1].start_block(), 10U);
+ ASSERT_EQ(ret[1].num_blocks(), 5U);
+
+ ASSERT_EQ(ret[2].start_block(), 20U);
+ ASSERT_EQ(ret[2].num_blocks(), 5U);
+}
+
+TEST_F(ExtentMapTest, GetNonIntersectingExtents) {
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 7));
+ ASSERT_TRUE(map_.AddExtent(ExtentForRange(20, 5), 7));
+
+ auto ret = std::vector<Extent>{};
+ ret = map_.GetNonIntersectingExtents(ExtentForRange(2, 13));
+
+ ASSERT_EQ(ret.size(), 1U);
+ ASSERT_EQ(ret[0].start_block(), 5U);
+ ASSERT_EQ(ret[0].num_blocks(), 5U);
+
+ ret = map_.GetNonIntersectingExtents(ExtentForRange(7, 20));
+ ASSERT_EQ(ret.size(), 3U);
+ ASSERT_EQ(ret[0].start_block(), 7U);
+ ASSERT_EQ(ret[0].num_blocks(), 3U);
+
+ ASSERT_EQ(ret[1].start_block(), 15U);
+ ASSERT_EQ(ret[1].num_blocks(), 5U);
+
+ ASSERT_EQ(ret[2].start_block(), 25U);
+ ASSERT_EQ(ret[2].num_blocks(), 2U);
+}
+
+} // namespace chromeos_update_engine