Use absl::Mutex in TrackManager (#4685)

* Use absl::Mutex in TrackManager

TrackManager is a class that is not as easy as it should be to navigate. The main complexity lies in the fact that there are two threads that are accessing the class. The capture thread is inserting timers and therefore creating the tracks when needed and the main thread is reading the tracks to sort then and render the visible ones (a detailed picture can be found in http://go/stadia-orbit-track-creation#heading=h.vpgt3chmpmpb).

We had issues in the class because of multi-threading (and currently it seems Daniel found a new one), so in this PR I'm cleaning this class by replacing the recursive_mutex with an absl::Mutex. I don't see how this could solve the crash found by Daniel but at least the class might be easier to understand when looking for the issue.

We were using a recursive_mutex because some public methods were calling another ones where both of then needs the mutex to be locked. But as we don't have recursion, this can be solved by splitting some methods.

In this PR we are:
- Guarding all the containers that needs to be locked (ABSL_GUARDED_BY)
- Locking all the public methods that are accessing the containers.
- Not locking the private methods, as they are only called by the class itself.
- Manually assuring that the locked methods don't call another public method that lock the mutex.

I'm not sure about if the last point is the best thing to do but I feel that this is an improvement, do you have another suggestion?

Test: Load a capture with all kind of tracks. Also take a capture.

* Fixup: Reader Lock and Shared_Locks_Required

* Fixup: WriterMutexLock

Author: karupayun

src/OrbitGl/TrackManager.cpp

@@ -61,7 +61,11 @@ TrackManager::TrackManager(TrackContainer* track_container, TimelineInfoInterfac
 }
 
 std::vector<Track*> TrackManager::GetAllTracks() const {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::ReaderMutexLock lock(&mutex_);
+  return GetAllTracksInternal();
+}
+
+std::vector<Track*> TrackManager::GetAllTracksInternal() const {
   std::vector<Track*> tracks;
   for (const auto& track : all_tracks_) {
     tracks.push_back(track.get());
@@ -74,7 +78,7 @@ std::vector<Track*> TrackManager::GetAllTracks() const {
 }
 
 std::vector<FrameTrack*> TrackManager::GetFrameTracks() const {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::ReaderMutexLock lock(&mutex_);
   std::vector<FrameTrack*> tracks;
   for (const auto& [unused_id, track] : frame_tracks_) {
     tracks.push_back(track.get());
@@ -83,7 +87,7 @@ std::vector<FrameTrack*> TrackManager::GetFrameTracks() const {
 }
 
 void TrackManager::SortTracks() {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   // Gather all tracks regardless of the process in sorted order
   std::vector<Track*> all_processes_sorted_tracks;
 
@@ -124,7 +128,7 @@ void TrackManager::SortTracks() {
   if (absl::GetFlag(FLAGS_enable_tracepoint_feature)) {
     if (app_ != nullptr && app_->HasCaptureData()) {
       ThreadTrack* tracepoints_track =
-          GetOrCreateThreadTrack(orbit_base::kAllThreadsOfAllProcessesTid);
+          GetOrCreateThreadTrackInternal(orbit_base::kAllThreadsOfAllProcessesTid);
       if (!tracepoints_track->IsEmpty()) {
         all_processes_sorted_tracks.push_back(tracepoints_track);
       }
@@ -133,7 +137,7 @@ void TrackManager::SortTracks() {
 
   // Process track.
   if (app_ != nullptr && app_->HasCaptureData()) {
-    ThreadTrack* process_track = GetOrCreateThreadTrack(orbit_base::kAllProcessThreadsTid);
+    ThreadTrack* process_track = GetOrCreateThreadTrackInternal(orbit_base::kAllProcessThreadsTid);
     if (!process_track->IsEmpty()) {
       all_processes_sorted_tracks.push_back(process_track);
     }
@@ -360,7 +364,7 @@ void TrackManager::AddFrameTrack(const std::shared_ptr<FrameTrack>& frame_track)
 }
 
 void TrackManager::RemoveFrameTrack(uint64_t function_id) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   deleted_tracks_.push_back(frame_tracks_[function_id]);
   frame_tracks_.erase(function_id);
 
@@ -439,7 +443,7 @@ Track* TrackManager::GetOrCreateTrackFromTimerInfo(const TimerInfo& timer_info)
 }
 
 SchedulerTrack* TrackManager::GetOrCreateSchedulerTrack() {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   if (scheduler_track_ == nullptr) {
     auto [unused, timer_data] = capture_data_->CreateTimerData();
     scheduler_track_ =
@@ -451,7 +455,11 @@ SchedulerTrack* TrackManager::GetOrCreateSchedulerTrack() {
 }
 
 ThreadTrack* TrackManager::GetOrCreateThreadTrack(uint32_t tid) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
+  return GetOrCreateThreadTrackInternal(tid);
+}
+
+ThreadTrack* TrackManager::GetOrCreateThreadTrackInternal(uint32_t tid) {
   std::shared_ptr<ThreadTrack> track = thread_tracks_[tid];
   if (track == nullptr) {
     auto* thread_track_data_provider = capture_data_->GetThreadTrackDataProvider();
@@ -466,7 +474,7 @@ ThreadTrack* TrackManager::GetOrCreateThreadTrack(uint32_t tid) {
 }
 
 std::optional<ThreadTrack*> TrackManager::GetThreadTrack(uint32_t tid) const {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::ReaderMutexLock lock(&mutex_);
   if (!thread_tracks_.contains(tid)) {
     return std::nullopt;
   }
@@ -475,7 +483,7 @@ std::optional<ThreadTrack*> TrackManager::GetThreadTrack(uint32_t tid) const {
 }
 
 GpuTrack* TrackManager::GetOrCreateGpuTrack(uint64_t timeline_hash) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   std::string timeline =
       app_->GetStringManager()->Get(timeline_hash).value_or(std::to_string(timeline_hash));
   std::shared_ptr<GpuTrack> track = gpu_tracks_[timeline];
@@ -492,7 +500,7 @@ GpuTrack* TrackManager::GetOrCreateGpuTrack(uint64_t timeline_hash) {
 }
 
 VariableTrack* TrackManager::GetOrCreateVariableTrack(std::string_view name) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
 
   auto existing_track = variable_tracks_.find(name);
   if (existing_track != variable_tracks_.end()) return existing_track->second.get();
@@ -505,7 +513,7 @@ VariableTrack* TrackManager::GetOrCreateVariableTrack(std::string_view name) {
 }
 
 AsyncTrack* TrackManager::GetOrCreateAsyncTrack(std::string_view name) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
 
   auto existing_track = async_tracks_.find(name);
   if (existing_track != async_tracks_.end()) return existing_track->second.get();
@@ -520,7 +528,7 @@ AsyncTrack* TrackManager::GetOrCreateAsyncTrack(std::string_view name) {
 }
 
 FrameTrack* TrackManager::GetOrCreateFrameTrack(uint64_t function_id) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   if (auto track_it = frame_tracks_.find(function_id); track_it != frame_tracks_.end()) {
     return track_it->second.get();
   }
@@ -541,49 +549,64 @@ FrameTrack* TrackManager::GetOrCreateFrameTrack(uint64_t function_id) {
   return track.get();
 }
 
+SystemMemoryTrack* TrackManager::GetSystemMemoryTrack() const {
+  absl::ReaderMutexLock lock(&mutex_);
+  return system_memory_track_.get();
+}
+
+CGroupAndProcessMemoryTrack* TrackManager::GetCGroupAndProcessMemoryTrack() const {
+  absl::ReaderMutexLock lock(&mutex_);
+  return cgroup_and_process_memory_track_.get();
+}
+
+PageFaultsTrack* TrackManager::GetPageFaultsTrack() const {
+  absl::ReaderMutexLock lock(&mutex_);
+  return page_faults_track_.get();
+}
+
 SystemMemoryTrack* TrackManager::CreateAndGetSystemMemoryTrack() {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   if (system_memory_track_ == nullptr) {
     system_memory_track_ = std::make_shared<SystemMemoryTrack>(
         track_container_, timeline_info_, viewport_, layout_, module_manager_, capture_data_);
     AddTrack(system_memory_track_);
   }
 
-  return GetSystemMemoryTrack();
+  return system_memory_track_.get();
 }
 
 CGroupAndProcessMemoryTrack* TrackManager::CreateAndGetCGroupAndProcessMemoryTrack(
     std::string_view cgroup_name) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   if (cgroup_and_process_memory_track_ == nullptr) {
     cgroup_and_process_memory_track_ = std::make_shared<CGroupAndProcessMemoryTrack>(
         track_container_, timeline_info_, viewport_, layout_, std::string{cgroup_name},
         module_manager_, capture_data_);
     AddTrack(cgroup_and_process_memory_track_);
   }
 
-  return GetCGroupAndProcessMemoryTrack();
+  return cgroup_and_process_memory_track_.get();
 }
 
 PageFaultsTrack* TrackManager::CreateAndGetPageFaultsTrack(std::string_view cgroup_name,
                                                            uint64_t memory_sampling_period_ms) {
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  absl::WriterMutexLock lock(&mutex_);
   if (page_faults_track_ == nullptr) {
     page_faults_track_ = std::make_shared<PageFaultsTrack>(
         track_container_, timeline_info_, viewport_, layout_, std::string{cgroup_name},
         memory_sampling_period_ms, module_manager_, capture_data_);
     AddTrack(page_faults_track_);
   }
-  return GetPageFaultsTrack();
+  return page_faults_track_.get();
 }
 
 // TODO(b/177200020): Move to TrackContainer after assuring to have only one thread in the UI.
 std::pair<uint64_t, uint64_t> TrackManager::GetTracksMinMaxTimestamps() const {
   uint64_t min_time = std::numeric_limits<uint64_t>::max();
   uint64_t max_time = std::numeric_limits<uint64_t>::min();
 
-  std::lock_guard<std::recursive_mutex> lock(mutex_);
-  for (auto& track : GetAllTracks()) {
+  absl::ReaderMutexLock lock(&mutex_);
+  for (auto& track : GetAllTracksInternal()) {
     if (!track->IsEmpty()) {
       min_time = std::min(min_time, track->GetMinTime());
       max_time = std::max(max_time, track->GetMaxTime());

src/OrbitGl/include/OrbitGl/TrackManager.h

@@ -76,16 +76,12 @@ class TrackManager {
   VariableTrack* GetOrCreateVariableTrack(std::string_view name);
   AsyncTrack* GetOrCreateAsyncTrack(std::string_view name);
   FrameTrack* GetOrCreateFrameTrack(uint64_t function_id);
-  [[nodiscard]] SystemMemoryTrack* GetSystemMemoryTrack() const {
-    return system_memory_track_.get();
-  }
+  [[nodiscard]] SystemMemoryTrack* GetSystemMemoryTrack() const;
   [[nodiscard]] SystemMemoryTrack* CreateAndGetSystemMemoryTrack();
-  [[nodiscard]] CGroupAndProcessMemoryTrack* GetCGroupAndProcessMemoryTrack() const {
-    return cgroup_and_process_memory_track_.get();
-  }
+  [[nodiscard]] CGroupAndProcessMemoryTrack* GetCGroupAndProcessMemoryTrack() const;
   [[nodiscard]] CGroupAndProcessMemoryTrack* CreateAndGetCGroupAndProcessMemoryTrack(
       std::string_view cgroup_name);
-  PageFaultsTrack* GetPageFaultsTrack() const { return page_faults_track_.get(); }
+  PageFaultsTrack* GetPageFaultsTrack() const;
   PageFaultsTrack* CreateAndGetPageFaultsTrack(std::string_view cgroup_name,
                                                uint64_t memory_sampling_period_ms);
 
@@ -103,31 +99,40 @@ class TrackManager {
   [[nodiscard]] int FindMovingTrackIndex();
   void UpdateMovingTrackPositionInVisibleTracks();
   void SortTracks();
-  [[nodiscard]] std::vector<ThreadTrack*> GetSortedThreadTracks();
+  [[nodiscard]] std::vector<ThreadTrack*> GetSortedThreadTracks()
+      ABSL_SHARED_LOCKS_REQUIRED(mutex_);
+  [[nodiscard]] std::vector<Track*> GetAllTracksInternal() const ABSL_SHARED_LOCKS_REQUIRED(mutex_);
+  ThreadTrack* GetOrCreateThreadTrackInternal(uint32_t tid) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
   // Filter tracks that are already sorted in sorted_tracks_.
   void UpdateVisibleTrackList();
   void DeletePendingTracks();
 
-  void AddTrack(const std::shared_ptr<Track>& track);
+  void AddTrack(const std::shared_ptr<Track>& track) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
   void AddFrameTrack(const std::shared_ptr<FrameTrack>& frame_track);
 
-  // TODO(b/174655559): Use absl's mutex here.
-  mutable std::recursive_mutex mutex_;
-
-  std::vector<std::shared_ptr<Track>> all_tracks_;
-  absl::flat_hash_map<uint32_t, std::shared_ptr<ThreadTrack>> thread_tracks_;
-  std::map<std::string, std::shared_ptr<AsyncTrack>, std::less<>> async_tracks_;
-  std::map<std::string, std::shared_ptr<VariableTrack>, std::less<>> variable_tracks_;
+  // This mutex is needed because two different threads will access the following containers. The
+  // capture thread will insert tracks and the main thread will read the tracks from the containers
+  // to generate the sorted list of visible tracks.
+  mutable absl::Mutex mutex_;
+
+  std::vector<std::shared_ptr<Track>> all_tracks_ ABSL_GUARDED_BY(mutex_);
+  absl::flat_hash_map<uint32_t, std::shared_ptr<ThreadTrack>> thread_tracks_
+      ABSL_GUARDED_BY(mutex_);
+  std::map<std::string, std::shared_ptr<AsyncTrack>, std::less<>> async_tracks_
+      ABSL_GUARDED_BY(mutex_);
+  std::map<std::string, std::shared_ptr<VariableTrack>, std::less<>> variable_tracks_
+      ABSL_GUARDED_BY(mutex_);
   // Mapping from timeline to GPU tracks. Timeline name is used for stable ordering. In particular
   // we want the marker tracks next to their queue track. E.g. "gfx" and "gfx_markers" should appear
   // next to each other.
-  std::map<std::string, std::shared_ptr<GpuTrack>> gpu_tracks_;
+  std::map<std::string, std::shared_ptr<GpuTrack>> gpu_tracks_ ABSL_GUARDED_BY(mutex_);
   // Mapping from function id to frame tracks.
-  std::map<uint64_t, std::shared_ptr<FrameTrack>> frame_tracks_;
-  std::shared_ptr<SchedulerTrack> scheduler_track_;
-  std::shared_ptr<SystemMemoryTrack> system_memory_track_;
-  std::shared_ptr<CGroupAndProcessMemoryTrack> cgroup_and_process_memory_track_;
-  std::shared_ptr<PageFaultsTrack> page_faults_track_;
+  std::map<uint64_t, std::shared_ptr<FrameTrack>> frame_tracks_ ABSL_GUARDED_BY(mutex_);
+  std::shared_ptr<SchedulerTrack> scheduler_track_ ABSL_GUARDED_BY(mutex_);
+  std::shared_ptr<SystemMemoryTrack> system_memory_track_ ABSL_GUARDED_BY(mutex_);
+  std::shared_ptr<CGroupAndProcessMemoryTrack> cgroup_and_process_memory_track_
+      ABSL_GUARDED_BY(mutex_);
+  std::shared_ptr<PageFaultsTrack> page_faults_track_ ABSL_GUARDED_BY(mutex_);
 
   // This intermediatly stores tracks that have been deleted from one of the track vectors above
   // so they can safely be removed from the list of sorted and visible tracks.