From a4e7a341c06b47f267bbda6463bceb220c6b174b Mon Sep 17 00:00:00 2001
From: Bartosz Taudul <wolf.pld@gmail.com>
Date: Thu, 1 Aug 2019 23:14:09 +0200
Subject: [PATCH] Proper handling of disconnect request.

---
 client/TracyProfiler.cpp | 32 ++++++++++++++++++++++++++++++++
 client/TracyProfiler.hpp |  1 +
 common/TracyProtocol.hpp |  3 ++-
 server/TracyView.cpp     |  3 ++-
 server/TracyView.hpp     |  1 +
 server/TracyWorker.cpp   |  8 +++++++-
 server/TracyWorker.hpp   |  3 ++-
 7 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/client/TracyProfiler.cpp b/client/TracyProfiler.cpp
index e589d5bb..ca4f7c19 100644
--- a/client/TracyProfiler.cpp
+++ b/client/TracyProfiler.cpp
@@ -2064,6 +2064,9 @@ bool Profiler::HandleServerQuery()
     case ServerQueryFrameName:
         SendString( ptr, (const char*)ptr, QueueType::FrameName );
         break;
+    case ServerQueryDisconnect:
+        HandleDisconnect();
+        return false;
     default:
         assert( false );
         break;
@@ -2072,6 +2075,35 @@ bool Profiler::HandleServerQuery()
     return true;
 }
 
+void Profiler::HandleDisconnect()
+{
+    QueueItem terminate;
+    MemWrite( &terminate.hdr.type, QueueType::Terminate );
+    if( !SendData( (const char*)&terminate, 1 ) ) return;
+    for(;;)
+    {
+        if( m_sock->HasData() )
+        {
+            while( m_sock->HasData() )
+            {
+                if( !HandleServerQuery() ) return;
+            }
+            if( m_bufferOffset != m_bufferStart )
+            {
+                if( !CommitData() ) return;
+            }
+        }
+        else
+        {
+            if( m_bufferOffset != m_bufferStart )
+            {
+                if( !CommitData() ) return;
+            }
+            std::this_thread::sleep_for( std::chrono::milliseconds( 10 ) );
+        }
+    }
+}
+
 void Profiler::CalibrateTimer()
 {
 #ifdef TRACY_HW_TIMER
diff --git a/client/TracyProfiler.hpp b/client/TracyProfiler.hpp
index d61ce7db..1fbdef53 100644
--- a/client/TracyProfiler.hpp
+++ b/client/TracyProfiler.hpp
@@ -487,6 +487,7 @@ private:
     void SendCallstackFrame( uint64_t ptr );
 
     bool HandleServerQuery();
+    void HandleDisconnect();
 
     void CalibrateTimer();
     void CalibrateDelay();
diff --git a/common/TracyProtocol.hpp b/common/TracyProtocol.hpp
index e7f29e94..2a6db75f 100644
--- a/common/TracyProtocol.hpp
+++ b/common/TracyProtocol.hpp
@@ -9,7 +9,7 @@
 namespace tracy
 {
 
-enum : uint32_t { ProtocolVersion = 12 };
+enum : uint32_t { ProtocolVersion = 13 };
 enum : uint32_t { BroadcastVersion = 0 };
 
 using lz4sz_t = uint32_t;
@@ -45,6 +45,7 @@ enum ServerQuery : uint8_t
     ServerQueryPlotName,
     ServerQueryCallstackFrame,
     ServerQueryFrameName,
+    ServerQueryDisconnect
 };
 
 struct ServerQueryPacket
diff --git a/server/TracyView.cpp b/server/TracyView.cpp
index 17e128ac..5ef142c3 100644
--- a/server/TracyView.cpp
+++ b/server/TracyView.cpp
@@ -817,11 +817,12 @@ bool View::DrawConnection()
     const char* stopStr = "Stop";
 #endif
     std::shared_lock<std::shared_mutex> lock( m_worker.GetDataLock() );
-    if( m_worker.IsConnected() )
+    if( !m_disconnectIssued && m_worker.IsConnected() )
     {
         if( ImGui::Button( stopStr ) )
         {
             m_worker.Disconnect();
+            m_disconnectIssued = true;
         }
     }
     else
diff --git a/server/TracyView.hpp b/server/TracyView.hpp
index a2b88527..f9fe45bb 100644
--- a/server/TracyView.hpp
+++ b/server/TracyView.hpp
@@ -271,6 +271,7 @@ private:
     ImGuiTextFilter m_messageFilter;
     bool m_messageFilterWasActive = false;
     int m_visibleMessages = 0;
+    bool m_disconnectIssued = false;
 
     Region m_highlight;
     Region m_highlightZoom;
diff --git a/server/TracyWorker.cpp b/server/TracyWorker.cpp
index e0fbac7e..de242e54 100644
--- a/server/TracyWorker.cpp
+++ b/server/TracyWorker.cpp
@@ -1948,7 +1948,7 @@ void Worker::Exec()
             {
                 continue;
             }
-            if( !m_crashed )
+            if( !m_crashed && !m_disconnect )
             {
                 bool done = true;
                 for( auto& v : m_data.threads )
@@ -4193,6 +4193,12 @@ void Worker::ReadTimelinePre044( FileRead& f, Vector<GpuEvent*>& vec, uint64_t s
     }
 }
 
+void Worker::Disconnect()
+{
+    Query( ServerQueryDisconnect, 0 );
+    m_disconnect = true;
+}
+
 void Worker::Write( FileWrite& f )
 {
     f.Write( FileHeader, sizeof( FileHeader ) );
diff --git a/server/TracyWorker.hpp b/server/TracyWorker.hpp
index b25452f9..4b1e44ef 100644
--- a/server/TracyWorker.hpp
+++ b/server/TracyWorker.hpp
@@ -336,7 +336,7 @@ public:
     bool IsDataStatic() const { return !m_thread.joinable(); }
     bool IsBackgroundDone() const { return m_backgroundDone.load( std::memory_order_relaxed ); }
     void Shutdown() { m_shutdown.store( true, std::memory_order_relaxed ); }
-    void Disconnect() { Shutdown(); }   // TODO: Needs proper implementation.
+    void Disconnect();
 
     void Write( FileWrite& f );
     int GetTraceVersion() const { return m_traceVersion; }
@@ -498,6 +498,7 @@ private:
     std::string m_hostInfo;
     bool m_terminate = false;
     bool m_crashed = false;
+    bool m_disconnect = false;
     LZ4_streamDecode_t* m_stream;
     char* m_buffer;
     int m_bufferOffset;