[core] check source usage before remove

2 files changed, 55 insertions, 0 deletions

diff --git a/src/mbgl/style/style.cpp b/src/mbgl/style/style.cpp
index a03a18fdaf..0a212b977f 100644
--- a/src/mbgl/style/style.cpp
+++ b/src/mbgl/style/style.cpp
@@ -158,7 +158,30 @@ void Style::addSource(std::unique_ptr<Source> source) {
     sources.emplace_back(std::move(source));
 }
 
+struct SourceIdUsageEvaluator {
+    const std::string& sourceId;
+
+    bool operator()(BackgroundLayer&) { return false; }
+    bool operator()(CustomLayer&) { return false; }
+
+    template <class LayerType>
+    bool operator()(LayerType& layer) {
+        return layer.getSourceID() == sourceId;
+    }
+};
+
 std::unique_ptr<Source> Style::removeSource(const std::string& id) {
+    // Check if source is in use
+    SourceIdUsageEvaluator sourceIdEvaluator {id};
+    auto layerIt = std::find_if(layers.begin(), layers.end(), [&](const auto& layer) {
+        return layer->accept(sourceIdEvaluator);
+    });
+
+    if (layerIt != layers.end()) {
+        Log::Warning(Event::General, "Source '%s' is in use, cannot remove", id.c_str());
+        return nullptr;
+    }
+
     auto it = std::find_if(sources.begin(), sources.end(), [&](const auto& source) {
         return source->getID() == id;
     });
diff --git a/test/style/style.test.cpp b/test/style/style.test.cpp
index b529abad4a..841c7b291b 100644
--- a/test/style/style.test.cpp
+++ b/test/style/style.test.cpp
@@ -1,10 +1,12 @@
 #include <mbgl/test/util.hpp>
 #include <mbgl/test/stub_file_source.hpp>
+#include <mbgl/test/fixture_log_observer.hpp>
 
 #include <mbgl/style/style.hpp>
 #include <mbgl/style/source_impl.hpp>
 #include <mbgl/style/sources/vector_source.hpp>
 #include <mbgl/style/layer.hpp>
+#include <mbgl/style/layers/line_layer.hpp>
 #include <mbgl/util/io.hpp>
 #include <mbgl/util/run_loop.hpp>
 #include <mbgl/util/default_thread_pool.hpp>
@@ -67,3 +69,33 @@ TEST(Style, DuplicateSource) {
         // Expected
     }
 }
+
+TEST(Style, RemoveSourceInUse) {
+    util::RunLoop loop;
+
+    auto log = new FixtureLogObserver();
+    Log::setObserver(std::unique_ptr<Log::Observer>(log));
+
+    ThreadPool threadPool{ 1 };
+    StubFileSource fileSource;
+    Style style { threadPool, fileSource, 1.0 };
+
+    style.setJSON(util::read_file("test/fixtures/resources/style-unused-sources.json"));
+
+    style.addSource(std::make_unique<VectorSource>("sourceId", "mapbox://mapbox.mapbox-terrain-v2"));
+    style.addLayer(std::make_unique<LineLayer>("layerId", "sourceId"));
+
+    // Should not remove the source
+    auto removed = style.removeSource("sourceId");
+    ASSERT_EQ(nullptr, removed);
+    ASSERT_NE(nullptr, style.getSource("sourceId"));
+
+    const FixtureLogObserver::LogMessage logMessage {
+            EventSeverity::Warning,
+            Event::General,
+            int64_t(-1),
+            "Source 'sourceId' is in use, cannot remove",
+    };
+
+    EXPECT_EQ(log->count(logMessage), 1u);
+}