More user id tracking for mus:

components/mus/ws/window_tree_host_impl.cc

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "components/mus/ws/window_tree_host_impl.h"
 
 #include "base/debug/debugger.h"
 #include "base/strings/utf_string_conversions.h"
 #include "components/mus/common/types.h"
 #include "components/mus/ws/client_connection.h"
@@ skipping 86 matching lines @@
       display_manager_(
           DisplayManager::Create(connector, gpu_state, surfaces_state)),
       tree_awaiting_input_ack_(nullptr),
       last_cursor_(0) {
   frame_decoration_values_ = mojom::FrameDecorationValues::New();
   frame_decoration_values_->normal_client_area_insets = mojo::Insets::New();
   frame_decoration_values_->maximized_client_area_insets = mojo::Insets::New();
   frame_decoration_values_->max_title_bar_button_width = 0u;
 
   display_manager_->Init(this);
+
+  connection_manager_->window_manager_factory_registry()->AddObserver(this);
 }
 
 WindowTreeHostImpl::~WindowTreeHostImpl() {
+  connection_manager_->window_manager_factory_registry()->RemoveObserver(this);
+
   DestroyFocusController();
   for (ServerWindow* window : windows_needing_frame_destruction_)
     window->RemoveObserver(this);
 
-  // TODO(sky): this may leave the WindowTreeImpls associated with the
-  // WindowManagerStates still alive. The shutdown ordering is a bit iffy,
-  // figure out what it should be. ConnectionManager::OnConnectionError()
-  // has a check that effects this too.
+  // Destroy any trees, which triggers destroying the WindowManagerState. Copy
+  // off the WindowManagerStates as destruction mutates
+  // |window_manager_state_map_|.
+  std::set<WindowManagerState*> states;
+  for (auto& pair : window_manager_state_map_)
+    states.insert(pair.second.get());
+  for (WindowManagerState* state : states)
+    connection_manager_->DestroyTree(state->tree());
 }
 
 void WindowTreeHostImpl::Init(scoped_ptr<WindowTreeHostConnection> connection) {
   init_called_ = true;
   window_tree_host_connection_ = std::move(connection);
   connection_manager_->AddHost(this);
   InitWindowManagersIfNecessary();
 }
 
 void WindowTreeHostImpl::SetFrameDecorationValues(
@@ skipping 45 matching lines @@
   }
   return nullptr;
 }
 
 WindowManagerState* WindowTreeHostImpl::GetFirstWindowManagerState() {
   return window_manager_state_map_.empty()
              ? nullptr
              : window_manager_state_map_.begin()->second.get();
 }
 
+WindowManagerState* WindowTreeHostImpl::GetWindowManagerStateForUser(
+    UserId user_id) {
+  auto iter = window_manager_state_map_.find(user_id);
+  return iter == window_manager_state_map_.end() ? nullptr : iter->second.get();
+}
+
 void WindowTreeHostImpl::SetCapture(ServerWindow* window,
                                     bool in_nonclient_area) {
   ServerWindow* capture_window = event_dispatcher_.capture_window();
   if (capture_window == window)
     return;
   event_dispatcher_.SetCaptureWindow(window, in_nonclient_area);
 }
 
 mojom::Rotation WindowTreeHostImpl::GetRotation() const {
   return display_manager_->GetRotation();
@@ skipping 175 matching lines @@
                          &WindowTreeHostImpl::OnEventAckTimeout);
 
   tree_awaiting_input_ack_ = connection;
   connection->DispatchInputEvent(target, mojom::Event::From(event));
 }
 
 void WindowTreeHostImpl::CreateWindowManagerStatesFromRegistry() {
   std::vector<WindowManagerFactoryService*> services =
       connection_manager_->window_manager_factory_registry()->GetServices();
   for (WindowManagerFactoryService* service : services) {
-    scoped_ptr<WindowManagerState> wms_ptr(
-        new WindowManagerState(this, service->user_id()));
-    WindowManagerState* wms = wms_ptr.get();
-    window_manager_state_map_[service->user_id()] = std::move(wms_ptr);
-    wms->tree_ = connection_manager_->CreateTreeForWindowManager(
-        this, service->window_manager_factory(), wms->root());
+    if (service->window_manager_factory())
+      CreateWindowManagerStateFromService(service);
   }
 }
 
+void WindowTreeHostImpl::CreateWindowManagerStateFromService(
+    WindowManagerFactoryService* service) {
+  scoped_ptr<WindowManagerState> wms_ptr(
+      new WindowManagerState(this, service->user_id()));
+  WindowManagerState* wms = wms_ptr.get();
+  window_manager_state_map_[service->user_id()] = std::move(wms_ptr);
+  wms->tree_ = connection_manager_->CreateTreeForWindowManager(
+      this, service->window_manager_factory(), wms->root());
+}
+
 void WindowTreeHostImpl::UpdateNativeCursor(int32_t cursor_id) {
   if (cursor_id != last_cursor_) {
     display_manager_->SetCursorById(cursor_id);
     last_cursor_ = cursor_id;
   }
 }
 
 ServerWindow* WindowTreeHostImpl::GetRootWindow() {
   return root_.get();
 }
@@ skipping 201 matching lines @@
   }
 
   DispatchInputEventToWindowImpl(target, in_nonclient_area, event);
 }
 
 void WindowTreeHostImpl::OnWindowDestroyed(ServerWindow* window) {
   windows_needing_frame_destruction_.erase(window);
   window->RemoveObserver(this);
 }
 
+void WindowTreeHostImpl::OnActiveUserIdChanged(UserId id) {
+  // TODO(sky): this likely needs to cancel any pending events and all that.
+}
+
+void WindowTreeHostImpl::OnUserIdAdded(UserId id) {}
+
+void WindowTreeHostImpl::OnUserIdRemoved(UserId id) {
+  if (window_tree_host_connection_)
+    return;
+
+  WindowManagerState* state = GetWindowManagerStateForUser(id);
+  if (!state)
+    return;
+
+  // DestroyTree() calls back to OnWindowTreeConnectionError() and the
+  // WindowManagerState is destroyed (and removed).
+  connection_manager_->DestroyTree(state->tree());
+  DCHECK_EQ(0u, window_manager_state_map_.count(id));
+}
+
+void WindowTreeHostImpl::OnWindowManagerFactorySet(
+    WindowManagerFactoryService* service) {
+  CreateWindowManagerStateFromService(service);
+}
+
 }  // namespace ws
 }  // namespace mus