From da2ce296d881368f2016acff7e43b8c1666802ea Mon Sep 17 00:00:00 2001
From: Adrian Stanea <Adrian.Stanea@analog.com>
Date: Thu, 21 Sep 2023 17:28:04 +0300
Subject: [PATCH] ContextBuilder: solve context close bug using reference count

- added reference count to ContextBuilder to keep track of how many
  external references to the M2k context are present. When the last
  reference is removed, the M2k context can be closed.
- contexts should be closed in the reverse order they were opened:
  contextClose() -> iio_context_destroy() to avoid dangling pointers.

Signed-off-by: Adrian Stanea <Adrian.Stanea@analog.com>
---
 include/libm2k/contextbuilder.hpp |  5 ++
 src/contextbuilder.cpp            | 79 +++++++++++++++++++++++++------
 2 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/include/libm2k/contextbuilder.hpp b/include/libm2k/contextbuilder.hpp
index 0012e53e..5b740684 100644
--- a/include/libm2k/contextbuilder.hpp
+++ b/include/libm2k/contextbuilder.hpp
@@ -113,6 +113,11 @@ class LIBM2K_API ContextBuilder {
 		bool ownsContext = false);
 	static bool m_disable_logging;
 
+	static std::map<std::string, int> reference_count;
+	static void incrementReferenceCount(std::string uri);
+	static void decrementReferenceCount(std::string uri);
+	static bool checkLastReference(std::string uri);
+	static Context* searchInConnectedDevices(std::string uri);
 };
 
 /**
diff --git a/src/contextbuilder.cpp b/src/contextbuilder.cpp
index b9f1cc00..e203a865 100644
--- a/src/contextbuilder.cpp
+++ b/src/contextbuilder.cpp
@@ -52,6 +52,7 @@ std::map<ContextTypes, std::string> ContextBuilder::m_dev_name_map = {
 };
 
 bool ContextBuilder::m_disable_logging = true;
+std::map<std::string, int> ContextBuilder::reference_count = {};
 
 ContextBuilder::ContextBuilder()
 {
@@ -148,19 +149,55 @@ Context* ContextBuilder::buildContext(ContextTypes type, std::string uri,
 	}
 }
 
-Context* ContextBuilder::contextOpen(const char *uri)
+void ContextBuilder::incrementReferenceCount(std::string uri)
 {
-	if (m_disable_logging) {
-		enableLogging(false);
+	if (reference_count.find(uri) != reference_count.end()) {
+		reference_count[uri]++;
+	} else {
+		reference_count.emplace(uri, 1);
 	}
-	LIBM2K_LOG(INFO, "libm2k version: " + getVersion());
+}
+
+void ContextBuilder::decrementReferenceCount(std::string uri)
+{
+	if (reference_count.find(uri) != reference_count.end()) {
+		reference_count[uri]--;
+	}
+}
+
+bool ContextBuilder::checkLastReference(std::string uri)
+{
+	if (reference_count.find(uri) != reference_count.end()) {
+		return (reference_count[uri] == 0);
+	}
+	return false;
+}
 
-	for (Context* dev : s_connectedDevices) {
-		if (dev->getUri() == std::string(uri)) {
+Context* ContextBuilder::searchInConnectedDevices(std::string uri)
+{
+	for (auto dev : s_connectedDevices) {
+		if (dev->getUri() == uri) {
 			return dev;
 		}
 	}
+	return nullptr;
+}
+
 
+Context* ContextBuilder::contextOpen(const char *uri)
+{
+	Context* dev = nullptr;
+	if (m_disable_logging) {
+		enableLogging(false);
+	}
+	LIBM2K_LOG(INFO, "libm2k version: " + getVersion());
+	// use saved device is possible
+	dev = searchInConnectedDevices(uri);
+	if (dev) {
+		incrementReferenceCount(uri);
+		return dev;
+	}
+	// create and save device during first call
 	struct iio_context* ctx = iio_create_context_from_uri(uri);
 	if (!ctx) {
 		return nullptr;
@@ -182,14 +219,16 @@ Context* ContextBuilder::contextOpen(const char *uri)
 
 	ContextTypes dev_type = ContextBuilder::identifyContext(ctx);
 
-	Context* dev = buildContext(dev_type, std::string(uri), ctx, true, true);
+	dev = buildContext(dev_type, std::string(uri), ctx, true, true);
 	s_connectedDevices.push_back(dev);
+	incrementReferenceCount(uri);
 
 	return dev;
 }
 
 Context* ContextBuilder::contextOpen(struct iio_context* ctx, const char* uri)
 {
+	Context* dev = nullptr;
     if (m_disable_logging) {
         enableLogging(false);
     }
@@ -210,21 +249,22 @@ Context* ContextBuilder::contextOpen(struct iio_context* ctx, const char* uri)
 	if (hw_fw_version != nullptr) {
 		LIBM2K_LOG(INFO, "Firmware version: " + std::string(hw_fw_version));
 	}
-
-	for (Context* dev : s_connectedDevices) {
-		if (dev->getUri() == std::string(uri)) {
-			return dev;
-		}
+	// use saved device is possible
+	dev = searchInConnectedDevices(uri);
+	if (dev) {
+		incrementReferenceCount(uri);
+		return dev;
 	}
 
 	if (!ctx) {
 		return nullptr;
 	}
-
+	// create and save device during first call
 	ContextTypes dev_type = ContextBuilder::identifyContext(ctx);
 
-	Context* dev = buildContext(dev_type, std::string(uri), ctx, true);
+	dev = buildContext(dev_type, std::string(uri), ctx, true);
 	s_connectedDevices.push_back(dev);
+	incrementReferenceCount(uri);
 
 	return dev;
 }
@@ -288,6 +328,17 @@ M2k *ContextBuilder::m2kOpen()
 
 void ContextBuilder::contextClose(Context* device, bool deinit)
 {
+	auto uri = device->getUri();
+	if (searchInConnectedDevices(uri) == nullptr) {
+		return;
+	}
+	decrementReferenceCount(uri);
+	bool isLastReference = checkLastReference(uri);
+	if (!isLastReference) {
+		return;
+	}
+
+	reference_count.erase(uri);
 	s_connectedDevices.erase(std::remove(s_connectedDevices.begin(),
 					     s_connectedDevices.end(),
 					     device), s_connectedDevices.end());