[Celestica] Leh800bcls: Resolve multi-switch empty fb303 counters in IS_OSS Split Agent mode

Author: gang-tao

fboss/agent/hw/test/BUCK

@@ -759,6 +759,8 @@ cpp_library(
     ],
     exported_deps = [
         ":hw_test_thrift_handler_h",
+        "//fb303:service_data",
+        "//fb303:thread_cached_service_data",
         "//folly:synchronized",
         "//folly/logging:log_handler",
         "//folly/logging:logging",

fboss/agent/hw/test/HwTestLogCaptureThriftHandler.cpp

@@ -2,6 +2,8 @@
 
 #include "fboss/agent/hw/test/HwTestThriftHandler.h"
 
+#include <fb303/ServiceData.h>
+#include <fb303/ThreadCachedServiceData.h>
 #include <folly/Synchronized.h>
 #include <folly/logging/LogHandler.h>
 #include <folly/logging/LogHandlerConfig.h>
@@ -87,6 +89,18 @@ void HwTestThriftHandler::getMatchingLogMessages(
   out = static_cast<CaptureLogHandler*>(handler.get())->matching(*substring);
 }
 
+void HwTestThriftHandler::getFb303RegexCounters(
+    std::map<std::string, int64_t>& counters,
+    std::unique_ptr<std::string> regex) {
+  facebook::fb303::ThreadCachedServiceData::get()->publishStats();
+  counters = facebook::fb303::fbData->getRegexCounters(*regex);
+}
+
+int64_t HwTestThriftHandler::getFb303Counter(std::unique_ptr<std::string> key) {
+  facebook::fb303::ThreadCachedServiceData::get()->publishStats();
+  return facebook::fb303::fbData->getCounterIfExists(*key).value_or(0);
+}
+
 } // namespace utility
 } // namespace fboss
 } // namespace facebook

fboss/agent/hw/test/HwTestThriftHandler.h

@@ -258,6 +258,12 @@ class HwTestThriftHandler : public AgentHwTestCtrlSvIf {
       std::vector<std::string>& out,
       std::unique_ptr<std::string> substring) override;
 
+  void getFb303RegexCounters(
+      std::map<std::string, int64_t>& counters,
+      std::unique_ptr<std::string> regex) override;
+
+  int64_t getFb303Counter(std::unique_ptr<std::string> key) override;
+
  private:
   HwSwitch* hwSwitch_;
   // Captures log messages in this process for tests to read over RPC.

fboss/agent/if/agent_hw_test_ctrl.thrift

@@ -240,4 +240,8 @@ service AgentHwTestCtrl {
   // installLogCapture() must be called before the log is emitted.
   void installLogCapture();
   list<string> getMatchingLogMessages(1: string substring);
+
+  // fb303 cross-process utils for multi-switch testing
+  map<string, i64> getFb303RegexCounters(1: string regex);
+  i64 getFb303Counter(1: string key);
 }

fboss/agent/test/AgentEnsemble.cpp

@@ -881,8 +881,8 @@ std::map<std::string, int64_t> AgentEnsemble::getFb303CountersByRegex(
       client->getChannelShared()};
   monitoringClient.sync_getRegexCounters(counters, regex);
 #else
-  // TODO: This needs to be updated to support multi-switch.
-  counters = facebook::fb303::fbData->getRegexCounters(regex);
+  auto switchID = scopeResolver().scope(portId).switchId();
+  counters = queryHwAgentFb303RegexCounters(switchID, regex);
 #endif
   return counters;
 }
@@ -908,8 +908,7 @@ int64_t AgentEnsemble::getFb303Counter(
       client->getChannelShared()};
   counter = monitoringClient.sync_getCounter(key);
 #else
-  // TODO: This needs to be updated to support multi-switch.
-  counter = facebook::fb303::fbData->getCounter(key);
+  counter = queryHwAgentFb303Counter(switchID, key);
 #endif
   return counter;
 }
@@ -959,12 +958,51 @@ std::map<std::string, int64_t> AgentEnsemble::getFb303RegexCounters(
       client->getChannelShared()};
   monitoringClient.sync_getRegexCounters(counters, regex);
 #else
-  // TODO: This needs to be updated to support multi-switch.
-  counters = facebook::fb303::fbData->getRegexCounters(regex);
+  counters = queryHwAgentFb303RegexCounters(switchID, regex);
 #endif
   return counters;
 }
 
+std::map<std::string, int64_t> AgentEnsemble::queryHwAgentFb303RegexCounters(
+    const SwitchID& switchID,
+    const std::string& regex) {
+  std::map<std::string, int64_t> counters;
+  if (!getSw()->isRunModeMultiSwitch()) {
+    counters = facebook::fb303::fbData->getRegexCounters(regex);
+  } else {
+    try {
+      auto hwTestClient = getHwAgentTestClient(switchID);
+      hwTestClient->sync_getFb303RegexCounters(counters, regex);
+    } catch (const std::exception& ex) {
+      XLOG(ERR)
+          << "Failed to fetch remote fb303 counters via HwTestCtrl for switch "
+          << switchID << ": " << ex.what();
+      throw;
+    }
+  }
+  return counters;
+}
+
+int64_t AgentEnsemble::queryHwAgentFb303Counter(
+    const SwitchID& switchID,
+    const std::string& key) {
+  int64_t counter{0};
+  if (!getSw()->isRunModeMultiSwitch()) {
+    counter = facebook::fb303::fbData->getCounter(key);
+  } else {
+    try {
+      auto hwTestClient = getHwAgentTestClient(switchID);
+      counter = hwTestClient->sync_getFb303Counter(key);
+    } catch (const std::exception& ex) {
+      XLOG(ERR)
+          << "Failed to fetch remote fb303 counter via HwTestCtrl for switch "
+          << switchID << ": " << ex.what();
+      throw;
+    }
+  }
+  return counter;
+}
+
 std::string AgentEnsemble::getHwDebugDump() {
   std::string out{};
   ThriftHandler(getSw()).getHwDebugDump(out);

fboss/agent/test/AgentEnsemble.h

@@ -405,6 +405,13 @@ class AgentEnsemble : public TestEnsembleIf {
   void overrideConfigFlag(const std::string& fileName);
   void dumpConfigForHwAgent(AgentConfig* agentConf);
 
+  std::map<std::string, int64_t> queryHwAgentFb303RegexCounters(
+      const SwitchID& switchId,
+      const std::string& regex);
+  int64_t queryHwAgentFb303Counter(
+      const SwitchID& switchId,
+      const std::string& key);
+
   void writeConfig(const cfg::SwitchConfig& config);
   void writeConfig(const cfg::AgentConfig& config);
   bool waitForRateOnPort(