From 046e2bae9d18da5dd0122e04dd161a8a45e1bb81 Mon Sep 17 00:00:00 2001
From: Paul Hollinsky <phollinsky@intrepidcs.com>
Date: Tue, 1 Sep 2020 15:52:41 -0400
Subject: [PATCH] Ensure proper closure of the heartbeat thread on reopen

---
 device/device.cpp              | 39 +++++++++++++++++++++++++---------
 include/icsneo/device/device.h |  9 +-------
 2 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/device/device.cpp b/device/device.cpp
index 0ca2ca0..d28f2fa 100644
--- a/device/device.cpp
+++ b/device/device.cpp
@@ -69,6 +69,12 @@ bool Device::SerialStringIsNumeric(const std::string& serial) {
 	return isdigit(serial[0]) && isdigit(serial[1]);
 }
 
+Device::~Device() {
+	if(isMessagePollingEnabled())
+		disableMessagePolling();
+	close();
+}
+
 uint16_t Device::getTimestampResolution() const {
 	return com->decoder->timestampResolution;
 }
@@ -205,14 +211,24 @@ bool Device::open() {
 		handleInternalMessage(message);
 	}));
 
-	std::atomic<bool> receivedMessage{false};
-	messageReceivedCallbackID = com->addMessageCallback(MessageCallback(filter, [&](std::shared_ptr<Message> message) {
-		receivedMessage = true;
-	}));
-
-	heartbeatThread = std::thread([&]() {
+	heartbeatThread = std::thread([this]() {
 		EventManager::GetInstance().downgradeErrorsOnCurrentThread();
-		while(true) {
+
+		MessageFilter filter;
+		filter.includeInternalInAny = true;
+		std::atomic<bool> receivedMessage{false};
+		auto messageReceivedCallbackID = com->addMessageCallback(MessageCallback(filter, [&receivedMessage](std::shared_ptr<Message> message) {
+			receivedMessage = true;
+		}));
+
+		// Give the device time to get situated
+		auto i = 150;
+		while(!stopHeartbeatThread && i != 0) {
+			std::this_thread::sleep_for(std::chrono::milliseconds(50));
+			i--;
+		}
+
+		while(!stopHeartbeatThread) {
 			// Wait for 110ms for a possible heartbeat
 			std::this_thread::sleep_for(std::chrono::milliseconds(110));
 			if(!receivedMessage) {
@@ -229,6 +245,8 @@ bool Device::open() {
 			}
 			receivedMessage = false;
 		}
+
+		com->removeMessageCallback(messageReceivedCallbackID);
 	});
 
 	forEachExtension([](const std::shared_ptr<DeviceExtension>& ext) { ext->onDeviceOpen(); return true; });
@@ -249,11 +267,12 @@ bool Device::close() {
 	if(internalHandlerCallbackID)
 		com->removeMessageCallback(internalHandlerCallbackID);
 
-	if(messageReceivedCallbackID)
-		com->removeMessageCallback(messageReceivedCallbackID);
-
 	internalHandlerCallbackID = 0;
 
+	if(heartbeatThread.joinable())
+		heartbeatThread.join();
+	stopHeartbeatThread = false;
+
 	forEachExtension([](const std::shared_ptr<DeviceExtension>& ext) { ext->onDeviceClose(); return true; });
 	return com->close();
 }
diff --git a/include/icsneo/device/device.h b/include/icsneo/device/device.h
index 417b099..c379a82 100644
--- a/include/icsneo/device/device.h
+++ b/include/icsneo/device/device.h
@@ -27,13 +27,7 @@ namespace icsneo {
 
 class Device {
 public:
-	virtual ~Device() {
-		if(isMessagePollingEnabled())
-			disableMessagePolling();
-		close();
-		if(heartbeatThread.joinable())
-			heartbeatThread.join();
-	}
+	virtual ~Device();
 
 	static std::string SerialNumToString(uint32_t serial);
 	static uint32_t SerialStringToNum(const std::string& serial);
@@ -103,7 +97,6 @@ protected:
 	bool online = false;
 	int messagePollingCallbackID = 0;
 	int internalHandlerCallbackID = 0;
-	int messageReceivedCallbackID = 0;
 	device_eventhandler_t report;
 
 	// START Initialization Functions