Make sure source file data is properly tracked.

2 files changed, 21 insertions, 11 deletions

diff --git a/server/TracyWorker.cpp b/server/TracyWorker.cpp
index 1dce5e92..fd04ab8a 100644
--- a/server/TracyWorker.cpp
+++ b/server/TracyWorker.cpp
@@ -3444,8 +3444,9 @@ void Worker::QuerySourceFile( const char* fn, const char* image )
 {
     if( image ) QueryDataTransfer( image, strlen( image ) + 1 );
     QueryDataTransfer( fn, strlen( fn ) + 1 );
-    m_sourceCodeQuery.emplace_back( fn );
-    Query( ServerQuerySourceCode, 0 );
+    m_sourceCodeQuery.emplace( m_nextSourceCodeQuery, fn );
+    Query( ServerQuerySourceCode, m_nextSourceCodeQuery );
+    m_nextSourceCodeQuery++;
 }
 
 void Worker::QueryDataTransfer( const void* ptr, size_t size )
@@ -3503,7 +3504,7 @@ bool Worker::DispatchProcess( const QueueItem& ev, const char*& ptr )
                 m_serverQuerySpaceLeft++;
                 break;
             case QueueType::SourceCode:
-                AddSourceCode( ptr, sz );
+                AddSourceCode( (uint32_t)ev.stringTransfer.ptr, ptr, sz );
                 m_serverQuerySpaceLeft++;
                 break;
             default:
@@ -4161,11 +4162,12 @@ void Worker::AddSymbolCode( uint64_t ptr, const char* data, size_t sz )
 }
 
 
-void Worker::AddSourceCode( const char* data, size_t sz )
+void Worker::AddSourceCode( uint32_t id, const char* data, size_t sz )
 {
-    assert( !m_sourceCodeQuery.empty() );
-    auto file = m_sourceCodeQuery.front();
-    m_sourceCodeQuery.erase( m_sourceCodeQuery.begin() );
+    auto it = m_sourceCodeQuery.find( id );
+    assert( it != m_sourceCodeQuery.end() );
+    auto file = it->second;
+    m_sourceCodeQuery.erase( it );
     if( m_data.sourceFileCache.find( file ) != m_data.sourceFileCache.end() ) return;
     auto src = (char*)m_slab.AllocBig( sz );
     memcpy( src, data, sz );
@@ -4941,8 +4943,7 @@ bool Worker::Process( const QueueItem& ev )
         m_serverQuerySpaceLeft++;
         break;
     case QueueType::AckSourceCodeNotAvailable:
-        assert( !m_sourceCodeQuery.empty() );
-        m_sourceCodeQuery.erase( m_sourceCodeQuery.begin() );
+        ProcessSourceCodeNotAvailable( ev.sourceCodeNotAvailable );
         m_serverQuerySpaceLeft++;
         break;
     case QueueType::AckSymbolCodeNotAvailable:
@@ -7031,6 +7032,13 @@ void Worker::ProcessParamSetup( const QueueParamSetup& ev )
     m_params.push_back( Parameter { ev.idx, StringRef( StringRef::Ptr, ev.name ), bool( ev.isBool ), ev.val } );
 }
 
+void Worker::ProcessSourceCodeNotAvailable( const QueueSourceCodeNotAvailable& ev )
+{
+    auto it = m_sourceCodeQuery.find( ev.id );
+    assert( it != m_sourceCodeQuery.end() );
+    m_sourceCodeQuery.erase( it );
+}
+
 void Worker::ProcessCpuTopology( const QueueCpuTopology& ev )
 {
     auto package = m_data.cpuTopology.find( ev.package );
diff --git a/server/TracyWorker.hpp b/server/TracyWorker.hpp
index 6b324de5..888f6080 100644
--- a/server/TracyWorker.hpp
+++ b/server/TracyWorker.hpp
@@ -744,6 +744,7 @@ private:
     tracy_force_inline void ProcessHwSampleBranchRetired( const QueueHwSample& ev );
     tracy_force_inline void ProcessHwSampleBranchMiss( const QueueHwSample& ev );
     tracy_force_inline void ProcessParamSetup( const QueueParamSetup& ev );
+    tracy_force_inline void ProcessSourceCodeNotAvailable( const QueueSourceCodeNotAvailable& ev );
     tracy_force_inline void ProcessCpuTopology( const QueueCpuTopology& ev );
     tracy_force_inline void ProcessMemNamePayload( const QueueMemNamePayload& ev );
     tracy_force_inline void ProcessFiberEnter( const QueueFiberEnter& ev );
@@ -861,7 +862,7 @@ private:
     void AddExternalThreadName( uint64_t ptr, const char* str, size_t sz );
     void AddFrameImageData( uint64_t ptr, const char* data, size_t sz );
     void AddSymbolCode( uint64_t ptr, const char* data, size_t sz );
-    void AddSourceCode( const char* data, size_t sz );
+    void AddSourceCode( uint32_t id, const char* data, size_t sz );
 
     tracy_force_inline void AddCallstackPayload( uint64_t ptr, const char* data, size_t sz );
     tracy_force_inline void AddCallstackAllocPayload( uint64_t ptr, const char* data, size_t sz );
@@ -1062,7 +1063,8 @@ private:
     size_t m_tmpBufSize = 0;
 
     unordered_flat_map<uint64_t, uint32_t> m_nextCallstack;
-    std::vector<const char*> m_sourceCodeQuery;
+    unordered_flat_map<uint32_t, const char*> m_sourceCodeQuery;
+    uint32_t m_nextSourceCodeQuery = 0;
 };
 
 }