From 729a6da92c9d12b85e9ffcaf7f4da387e378b826 Mon Sep 17 00:00:00 2001
From: Max Ng <maxgg@meta.com>
Date: Wed, 4 Dec 2024 22:58:26 -0800
Subject: [PATCH] Fix MirrorOnDropReport removal error

Summary: Fix TAM object removal error due to dangling references from switch and port.

Reviewed By: nivinl

Differential Revision: D66789586

fbshipit-source-id: 598d54f3a8375f96dbbebbeb34fb57c43de5e394
---
 fboss/agent/hw/sai/switch/SaiTamManager.h     |  2 +
 .../hw/sai/switch/npu/bcm/SaiTamManager.cpp   | 40 ++++++++++++++-----
 .../hw/sai/switch/npu/tajo/SaiTamManager.cpp  |  2 +
 .../agent/hw/sai/switch/oss/SaiTamManager.cpp |  2 +
 4 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/fboss/agent/hw/sai/switch/SaiTamManager.h b/fboss/agent/hw/sai/switch/SaiTamManager.h
index 6c478c77fb625..5700f5275e6e9 100644
--- a/fboss/agent/hw/sai/switch/SaiTamManager.h
+++ b/fboss/agent/hw/sai/switch/SaiTamManager.h
@@ -63,6 +63,8 @@ class SaiTamManager {
   }
 
  private:
+  void updateTamObjectOnSwitchAndPort(PortID portId);
+
   SaiStore* saiStore_;
   SaiManagerTable* managerTable_;
   SaiPlatform* platform_;
diff --git a/fboss/agent/hw/sai/switch/npu/bcm/SaiTamManager.cpp b/fboss/agent/hw/sai/switch/npu/bcm/SaiTamManager.cpp
index 6a2ac58b71dba..69deaa8e83d0f 100644
--- a/fboss/agent/hw/sai/switch/npu/bcm/SaiTamManager.cpp
+++ b/fboss/agent/hw/sai/switch/npu/bcm/SaiTamManager.cpp
@@ -107,6 +107,8 @@ SaiTamManager::SaiTamManager(
 void SaiTamManager::addMirrorOnDropReport(
     const std::shared_ptr<MirrorOnDropReport>& report) {
 #if defined(BRCM_SAI_SDK_DNX_GTE_11_0)
+  XLOG(INFO) << "Creating MirrorOnDropReport " << report->getID();
+
   // Create report
   auto& reportStore = saiStore_->get<SaiTamReportTraits>();
   auto reportTraits =
@@ -180,8 +182,8 @@ void SaiTamManager::addMirrorOnDropReport(
     auto event = eventStore.setObject(eventTraits, eventTraits);
     events.push_back(event);
     eventIds.push_back(event->adapterKey());
-    +XLOG(INFO) << "Created event ID " << eventId << " with tam event "
-                << event->adapterKey() << " and aging group " << agingGroupKey;
+    XLOG(INFO) << "Created event ID " << eventId << " with tam event "
+               << event->adapterKey() << " and aging group " << agingGroupKey;
   }
 
   // Create tam
@@ -193,15 +195,6 @@ void SaiTamManager::addMirrorOnDropReport(
   auto& tamStore = saiStore_->get<SaiTamTraits>();
   auto tam = tamStore.setObject(tamTraits, tamTraits);
 
-  // Associate TAM with port
-  XLOG(INFO) << "Associating TAM object with port "
-             << report->getMirrorPortId();
-  managerTable_->portManager().setTamObject(
-      PortID(report->getMirrorPortId()), {tam->adapterKey()});
-
-  // Associate TAM with switch
-  managerTable_->switchManager().setTamObject({tam->adapterKey()});
-
   auto tamHandle = std::make_unique<SaiTamHandle>();
   tamHandle->report = reportObj;
   tamHandle->action = action;
@@ -212,12 +205,24 @@ void SaiTamManager::addMirrorOnDropReport(
   tamHandle->tam = tam;
   tamHandle->portId = report->getMirrorPortId();
   tamHandles_.emplace(report->getID(), std::move(tamHandle));
+
+  // Associate TAM with port
+  XLOG(INFO) << "Associating TAM object " << tam->adapterKey()
+             << " with switch and port " << report->getMirrorPortId();
+  updateTamObjectOnSwitchAndPort(report->getMirrorPortId());
 #endif
 }
 
 void SaiTamManager::removeMirrorOnDropReport(
     const std::shared_ptr<MirrorOnDropReport>& report) {
+  XLOG(INFO) << "Removing MirrorOnDropReport " << report->getID();
+  auto handle = std::move(tamHandles_[report->getID()]);
   tamHandles_.erase(report->getID());
+
+  // Unbind the TAM object from port and switch before letting it destruct.
+  XLOG(INFO) << "Unassociating TAM object " << handle->tam->adapterKey()
+             << " from switch and port " << report->getMirrorPortId();
+  updateTamObjectOnSwitchAndPort(handle->portId);
 }
 
 void SaiTamManager::changeMirrorOnDropReport(
@@ -235,4 +240,17 @@ std::vector<PortID> SaiTamManager::getAllMirrorOnDropPortIds() {
   return portIds;
 }
 
+void SaiTamManager::updateTamObjectOnSwitchAndPort(PortID portId) {
+  std::vector<sai_object_id_t> portTamIds;
+  std::vector<sai_object_id_t> switchTamIds;
+  for (const auto& [_, tamHandle] : tamHandles_) {
+    if (tamHandle->portId == portId) {
+      portTamIds.push_back(tamHandle->tam->adapterKey());
+    }
+    switchTamIds.push_back(tamHandle->tam->adapterKey());
+  }
+  managerTable_->portManager().setTamObject(portId, portTamIds);
+  managerTable_->switchManager().setTamObject(switchTamIds);
+}
+
 } // namespace facebook::fboss
diff --git a/fboss/agent/hw/sai/switch/npu/tajo/SaiTamManager.cpp b/fboss/agent/hw/sai/switch/npu/tajo/SaiTamManager.cpp
index e29b9cf94ff89..014174549a887 100644
--- a/fboss/agent/hw/sai/switch/npu/tajo/SaiTamManager.cpp
+++ b/fboss/agent/hw/sai/switch/npu/tajo/SaiTamManager.cpp
@@ -111,4 +111,6 @@ std::vector<PortID> SaiTamManager::getAllMirrorOnDropPortIds() {
   return {};
 }
 
+void SaiTamManager::updateTamObjectOnSwitchAndPort(PortID /* portId */) {}
+
 } // namespace facebook::fboss
diff --git a/fboss/agent/hw/sai/switch/oss/SaiTamManager.cpp b/fboss/agent/hw/sai/switch/oss/SaiTamManager.cpp
index d3a69732df63a..ace1820c47b97 100644
--- a/fboss/agent/hw/sai/switch/oss/SaiTamManager.cpp
+++ b/fboss/agent/hw/sai/switch/oss/SaiTamManager.cpp
@@ -23,4 +23,6 @@ std::vector<PortID> SaiTamManager::getAllMirrorOnDropPortIds() {
   return {};
 }
 
+void SaiTamManager::updateTamObjectOnSwitchAndPort(PortID /* portId */) {}
+
 } // namespace facebook::fboss