From 02de5f710e2579f2079a28944e90b48b34d91cac Mon Sep 17 00:00:00 2001
From: Jade <spacey-sooty@proton.me>
Date: Fri, 2 May 2025 13:07:45 +0800
Subject: [PATCH] [cscore] Fix memory leak in usbviewer example (#7922)

Signed-off-by: Jade Turner <spacey-sooty@proton.me>
---
 cscore/examples/usbviewer/usbviewer.cpp | 64 +++++++++++++------------
 1 file changed, 34 insertions(+), 30 deletions(-)

diff --git a/cscore/examples/usbviewer/usbviewer.cpp b/cscore/examples/usbviewer/usbviewer.cpp
index cfb9aea928..0cb31dfb3f 100644
--- a/cscore/examples/usbviewer/usbviewer.cpp
+++ b/cscore/examples/usbviewer/usbviewer.cpp
@@ -3,13 +3,17 @@
 // the WPILib BSD license file in the root directory of this project.
 
 #include <atomic>
+#include <memory>
 #include <thread>
+#include <utility>
 #include <vector>
 
 #include <imgui.h>
 #include <imgui_internal.h>
 #include <opencv2/core/core.hpp>
+#include <opencv2/core/mat.hpp>
 #include <opencv2/imgproc.hpp>
+#include <wpi/mutex.h>
 #include <wpi/print.h>
 #include <wpi/spinlock.h>
 #include <wpigui.h>
@@ -20,10 +24,10 @@
 namespace gui = wpi::gui;
 
 int main() {
-  std::atomic<cv::Mat*> latestFrame{nullptr};
-  std::vector<cv::Mat*> sharedFreeList;
-  wpi::spinlock sharedFreeListMutex;
-  std::vector<cv::Mat*> sourceFreeList;
+  wpi::spinlock latestFrameMutex;
+  std::unique_ptr<cv::Mat> latestFrame;
+  wpi::mutex freeListMutex;
+  std::vector<std::unique_ptr<cv::Mat>> freeList;
   std::atomic<bool> stopCamera{false};
 
   cs::UsbCamera camera{"usbcam", 0};
@@ -41,36 +45,31 @@ int main() {
         continue;
       }
 
-      // get or create a mat, prefer sourceFreeList over sharedFreeList
-      cv::Mat* out;
-      if (!sourceFreeList.empty()) {
-        out = sourceFreeList.back();
-        sourceFreeList.pop_back();
-      } else {
-        {
-          std::scoped_lock lock(sharedFreeListMutex);
-          for (auto mat : sharedFreeList) {
-            sourceFreeList.emplace_back(mat);
-          }
-          sharedFreeList.clear();
-        }
-        if (!sourceFreeList.empty()) {
-          out = sourceFreeList.back();
-          sourceFreeList.pop_back();
+      // get or create a mat
+      std::unique_ptr<cv::Mat> out;
+      {
+        std::scoped_lock lock{freeListMutex};
+        if (!freeList.empty()) {
+          out = std::move(freeList.back());
+          freeList.pop_back();
         } else {
-          out = new cv::Mat;
+          out = std::make_unique<cv::Mat>();
         }
       }
 
       // convert to RGBA
       cv::cvtColor(frame, *out, cv::COLOR_BGR2RGBA);
 
-      // make available
-      auto prev = latestFrame.exchange(out);
+      {
+        // make available
+        std::scoped_lock lock{latestFrameMutex};
+        latestFrame.swap(out);
+      }
 
-      // put prev on free list
-      if (prev) {
-        sourceFreeList.emplace_back(prev);
+      // put the previous frame on free list
+      if (out) {
+        std::scoped_lock lock{freeListMutex};
+        freeList.emplace_back(std::move(out));
       }
     }
   });
@@ -79,7 +78,11 @@ int main() {
   gui::Initialize("Hello World", 1024, 768);
   gui::Texture tex;
   gui::AddEarlyExecute([&] {
-    auto frame = latestFrame.exchange(nullptr);
+    std::unique_ptr<cv::Mat> frame;
+    {
+      std::scoped_lock lock{latestFrameMutex};
+      latestFrame.swap(frame);
+    }
     if (frame) {
       // create or update texture
       if (!tex || frame->cols != tex.GetWidth() ||
@@ -89,9 +92,10 @@ int main() {
       } else {
         tex.Update(frame->data);
       }
-      // put back on shared freelist
-      std::scoped_lock lock(sharedFreeListMutex);
-      sharedFreeList.emplace_back(frame);
+      {
+        std::scoped_lock lock{freeListMutex};
+        freeList.emplace_back(std::move(frame));
+      }
     }
 
     ImGui::SetNextWindowSize(ImVec2(640, 480), ImGuiCond_FirstUseEver);