From 61cefcfd8a8c81afa15ba0b4d4ec57cebf769ca1 Mon Sep 17 00:00:00 2001
From: AT <manyoso@users.noreply.github.com>
Date: Wed, 15 May 2024 09:01:53 -0500
Subject: [PATCH] Fix destruction and tear down of the embedding thread.
 (#2328)

* Fix destruction and tear down of the embedding thread.

Signed-off-by: Adam Treat <treat.adam@gmail.com>

* Fix order of deletion to prevent use after free.

Signed-off-by: Adam Treat <treat.adam@gmail.com>

---------

Signed-off-by: Adam Treat <treat.adam@gmail.com>
---
 gpt4all-chat/database.cpp | 1 +
 gpt4all-chat/embllm.cpp   | 8 ++++++++
 gpt4all-chat/embllm.h     | 1 +
 3 files changed, 10 insertions(+)

diff --git a/gpt4all-chat/database.cpp b/gpt4all-chat/database.cpp
index 7d3b556d..286de5e4 100644
--- a/gpt4all-chat/database.cpp
+++ b/gpt4all-chat/database.cpp
@@ -552,6 +552,7 @@ Database::~Database()
 {
     m_dbThread.quit();
     m_dbThread.wait();
+    delete m_embLLM;
 }
 
 void Database::scheduleNext(int folder_id, size_t countForFolder)
diff --git a/gpt4all-chat/embllm.cpp b/gpt4all-chat/embllm.cpp
index 913cd84a..d40ad651 100644
--- a/gpt4all-chat/embllm.cpp
+++ b/gpt4all-chat/embllm.cpp
@@ -5,6 +5,7 @@ EmbeddingLLMWorker::EmbeddingLLMWorker()
     : QObject(nullptr)
     , m_networkManager(new QNetworkAccessManager(this))
     , m_model(nullptr)
+    , m_stopGenerating(false)
 {
     moveToThread(&m_workerThread);
     connect(this, &EmbeddingLLMWorker::finished, &m_workerThread, &QThread::quit, Qt::DirectConnection);
@@ -14,6 +15,10 @@ EmbeddingLLMWorker::EmbeddingLLMWorker()
 
 EmbeddingLLMWorker::~EmbeddingLLMWorker()
 {
+    m_stopGenerating = true;
+    m_workerThread.quit();
+    m_workerThread.wait();
+
     if (m_model) {
         delete m_model;
         m_model = nullptr;
@@ -148,6 +153,9 @@ void EmbeddingLLMWorker::requestSyncEmbedding(const QString &text)
 // this function is always called for storage into the database
 void EmbeddingLLMWorker::requestAsyncEmbedding(const QVector<EmbeddingChunk> &chunks)
 {
+    if (m_stopGenerating)
+        return;
+
     if (!hasModel() && !loadModel()) {
         qWarning() << "WARNING: Could not load model for embeddings";
         return;
diff --git a/gpt4all-chat/embllm.h b/gpt4all-chat/embllm.h
index 44ec4ff6..2535a0f5 100644
--- a/gpt4all-chat/embllm.h
+++ b/gpt4all-chat/embllm.h
@@ -58,6 +58,7 @@ private:
     QNetworkAccessManager *m_networkManager;
     std::vector<float> m_lastResponse;
     LLModel *m_model = nullptr;
+    std::atomic<bool> m_stopGenerating;
     QThread m_workerThread;
 };