fix: Fix merge storageStats_ in IoStatistics (#12263)

Summary: Pull Request resolved: #12263

Differential Revision: D69172665

Pulled By: kewang1024

velox/common/io/CMakeLists.txt

@@ -15,3 +15,7 @@
 velox_add_library(velox_common_io IoStatistics.cpp)
 
 velox_link_libraries(velox_common_io Folly::folly glog::glog)
+
+if(${VELOX_BUILD_TESTING})
+  add_subdirectory(tests)
+endif()

velox/common/io/IoStatistics.cpp

@@ -149,7 +149,11 @@ void IoStatistics::merge(const IoStatistics& other) {
     const auto& otherStorageStats = other.storageStats();
     std::lock_guard<std::mutex> storageStatsLock(storageStatsMutex_);
     for (auto& item : otherStorageStats) {
-      storageStats_[item.first].merge(item.second);
+      if (storageStats_.count(item.first) == 0) {
+        storageStats_.emplace(item);
+      } else {
+        storageStats_[item.first].merge(item.second);
+      }
     }
   }
 }

velox/common/io/tests/CMakeLists.txt

@@ -0,0 +1,23 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# 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.
+
+add_executable(velox_ioStatistics_test IoStatisticsTest.cpp)
+
+add_test(NAME velox_ioStatistics_test COMMAND velox_ioStatistics_test)
+
+target_link_libraries(
+  velox_ioStatistics_test
+  velox_common_base
+  velox_common_io
+  GTest::gtest GTest::gtest_main)

velox/common/io/tests/IoStatisticsTest.cpp

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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 "velox/common/io/IoStatistics.h"
+#include <gtest/gtest.h>
+#include "velox/common/base/RuntimeMetrics.h"
+
+namespace facebook::velox::io {
+
+TEST(IoStatisticsTest, merge) {
+  IoStatistics ioStats;
+  ioStats.addStorageStats(
+      "remoteRead", RuntimeCounter(100, RuntimeCounter::Unit::kBytes));
+
+  IoStatistics ioStats2;
+  ioStats2.addStorageStats(
+      "localRead", RuntimeCounter(100, RuntimeCounter::Unit::kBytes));
+  ioStats2.addStorageStats(
+      "remoteRead", RuntimeCounter(200, RuntimeCounter::Unit::kBytes));
+
+  ioStats.merge(ioStats2);
+
+  RuntimeMetric localReadResult;
+  localReadResult.unit = RuntimeCounter::Unit::kBytes;
+  localReadResult.sum = 100;
+  localReadResult.count = 1;
+  localReadResult.min = 100;
+  localReadResult.max = 100;
+
+  RuntimeMetric remoteReadResult;
+  remoteReadResult.unit = RuntimeCounter::Unit::kBytes;
+  remoteReadResult.sum = 300;
+  remoteReadResult.count = 2;
+  remoteReadResult.min = 100;
+  remoteReadResult.max = 200;
+
+  EXPECT_EQ(
+      ioStats.storageStats().at("localRead").toString(),
+      localReadResult.toString());
+  EXPECT_EQ(
+      ioStats.storageStats().at("remoteRead").toString(),
+      remoteReadResult.toString());
+}
+} // namespace facebook::velox::io