Fix randomized test failure caused by unsharing while guest was joining

Co-Authored-By: Nathan Sobo <nathan@zed.dev>

crates/collab/src/rpc.rs

@@ -336,9 +336,11 @@ impl Server {
     }
 
     async fn sign_out(self: &mut Arc<Self>, connection_id: ConnectionId) -> Result<()> {
-        println!("Signing out {:?}", connection_id);
         self.peer.disconnect(connection_id);
-        let removed_connection = self.store_mut().await.remove_connection(connection_id)?;
+
+        let removed_user_id = {
+            let mut store = self.store_mut().await;
+            let removed_connection = store.remove_connection(connection_id)?;
 
             for (project_id, project) in removed_connection.hosted_projects {
                 broadcast(connection_id, project.guests.keys().copied(), |conn_id| {
@@ -347,8 +349,17 @@ impl Server {
                 });
             }
 
-        for (project_id, peer_ids) in removed_connection.guest_project_ids {
+            for project_id in removed_connection.guest_project_ids {
-            broadcast(connection_id, peer_ids, |conn_id| {
+                if let Some(project) = store.project(project_id).trace_err() {
+                    if project.guests.is_empty() {
+                        self.peer
+                            .send(
+                                project.host_connection_id,
+                                proto::ProjectUnshared { project_id },
+                            )
+                            .trace_err();
+                    } else {
+                        broadcast(connection_id, project.connection_ids(), |conn_id| {
                             self.peer.send(
                                 conn_id,
                                 proto::RemoveProjectCollaborator {
@@ -358,9 +369,13 @@ impl Server {
                             )
                         });
                     }
+                }
+            }
 
-        self.update_user_contacts(removed_connection.user_id)
+            removed_connection.user_id
-            .await?;
+        };
+
+        self.update_user_contacts(removed_user_id).await?;
 
         Ok(())
     }
@@ -829,7 +844,6 @@ impl Server {
         request: TypedEnvelope<proto::UpdateBuffer>,
         response: Response<proto::UpdateBuffer>,
     ) -> Result<()> {
-        println!("{:?}: update buffer", request.sender_id);
         let receiver_ids = self
             .store()
             .await

crates/collab/src/rpc/store.rs

@@ -51,7 +51,7 @@ pub type ReplicaId = u16;
 pub struct RemovedConnectionState {
     pub user_id: UserId,
     pub hosted_projects: HashMap<u64, Project>,
-    pub guest_project_ids: HashMap<u64, Vec<ConnectionId>>,
+    pub guest_project_ids: HashSet<u64>,
     pub contact_ids: HashSet<UserId>,
 }
 
@@ -134,10 +134,8 @@ impl Store {
         for project_id in connection.projects.clone() {
             if let Ok(project) = self.unregister_project(project_id, connection_id) {
                 result.hosted_projects.insert(project_id, project);
-            } else if let Ok(project) = self.leave_project(connection_id, project_id) {
+            } else if self.leave_project(connection_id, project_id).is_ok() {
-                result
+                result.guest_project_ids.insert(project_id);
-                    .guest_project_ids
-                    .insert(project_id, project.connection_ids);
             }
         }
 

crates/project/src/project.rs

@@ -259,6 +259,7 @@ impl Project {
         client.add_model_message_handler(Self::handle_register_worktree);
         client.add_model_message_handler(Self::handle_unregister_worktree);
         client.add_model_message_handler(Self::handle_unregister_project);
+        client.add_model_message_handler(Self::handle_project_unshared);
         client.add_model_message_handler(Self::handle_update_buffer_file);
         client.add_model_message_handler(Self::handle_update_buffer);
         client.add_model_message_handler(Self::handle_update_diagnostic_summary);
@@ -557,7 +558,7 @@ impl Project {
     }
 
     fn unregister(&mut self, cx: &mut ModelContext<Self>) {
-        self.unshare(cx);
+        self.unshared(cx);
         for worktree in &self.worktrees {
             if let Some(worktree) = worktree.upgrade(cx) {
                 worktree.update(cx, |worktree, _| {
@@ -817,12 +818,7 @@ impl Project {
         }
     }
 
-    fn share(
+    fn share(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
-        &mut self,
-        requester_user_id: UserId,
-        requester_peer_id: PeerId,
-        cx: &mut ModelContext<Self>,
-    ) -> Task<Result<()>> {
         let project_id;
         if let ProjectClientState::Local {
             remote_id_rx,
@@ -883,7 +879,7 @@ impl Project {
         })
     }
 
-    fn unshare(&mut self, cx: &mut ModelContext<Self>) {
+    fn unshared(&mut self, cx: &mut ModelContext<Self>) {
         if let ProjectClientState::Local { is_shared, .. } = &mut self.client_state {
             if !*is_shared {
                 return;
@@ -930,10 +926,7 @@ impl Project {
             let client = self.client.clone();
             cx.foreground()
                 .spawn(async move {
-                    println!("SHARING because {} wanted to join!!!", requester_id);
                     share.await?;
-                    println!("DONE SHARING!!!!!");
-
                     client.send(proto::RespondToJoinProjectRequest {
                         requester_id,
                         project_id,
@@ -3802,6 +3795,16 @@ impl Project {
         Ok(())
     }
 
+    async fn handle_project_unshared(
+        this: ModelHandle<Self>,
+        _: TypedEnvelope<proto::ProjectUnshared>,
+        _: Arc<Client>,
+        mut cx: AsyncAppContext,
+    ) -> Result<()> {
+        this.update(&mut cx, |this, cx| this.unshared(cx));
+        Ok(())
+    }
+
     async fn handle_add_collaborator(
         this: ModelHandle<Self>,
         mut envelope: TypedEnvelope<proto::AddProjectCollaborator>,
@@ -3843,19 +3846,7 @@ impl Project {
                     buffer.update(cx, |buffer, cx| buffer.remove_peer(replica_id, cx));
                 }
             }
-            println!(
+
-                "{} observed {:?} leaving",
-                this.user_store
-                    .read(cx)
-                    .current_user()
-                    .unwrap()
-                    .github_login,
-                peer_id
-            );
-            if this.collaborators.is_empty() {
-                println!("UNSHARING!!!!");
-                this.unshare(cx);
-            }
             cx.emit(Event::CollaboratorLeft(peer_id));
             cx.notify();
             Ok(())

crates/rpc/proto/zed.proto

@@ -21,20 +21,21 @@ message Envelope {
         LeaveProject leave_project = 15;
         AddProjectCollaborator add_project_collaborator = 16;
         RemoveProjectCollaborator remove_project_collaborator = 17;
+        ProjectUnshared project_unshared = 18;
 
-        GetDefinition get_definition = 18;
+        GetDefinition get_definition = 19;
-        GetDefinitionResponse get_definition_response = 19;
+        GetDefinitionResponse get_definition_response = 20;
-        GetReferences get_references = 20;
+        GetReferences get_references = 21;
-        GetReferencesResponse get_references_response = 21;
+        GetReferencesResponse get_references_response = 22;
-        GetDocumentHighlights get_document_highlights = 22;
+        GetDocumentHighlights get_document_highlights = 23;
-        GetDocumentHighlightsResponse get_document_highlights_response = 23;
+        GetDocumentHighlightsResponse get_document_highlights_response = 24;
-        GetProjectSymbols get_project_symbols = 24;
+        GetProjectSymbols get_project_symbols = 25;
-        GetProjectSymbolsResponse get_project_symbols_response = 25;
+        GetProjectSymbolsResponse get_project_symbols_response = 26;
-        OpenBufferForSymbol open_buffer_for_symbol = 26;
+        OpenBufferForSymbol open_buffer_for_symbol = 27;
-        OpenBufferForSymbolResponse open_buffer_for_symbol_response = 27;
+        OpenBufferForSymbolResponse open_buffer_for_symbol_response = 28;
 
-        RegisterWorktree register_worktree = 28;
+        RegisterWorktree register_worktree = 29;
-        UnregisterWorktree unregister_worktree = 29;
+        UnregisterWorktree unregister_worktree = 30;
         UpdateWorktree update_worktree = 31;
 
         CreateProjectEntry create_project_entry = 32;
@@ -213,6 +214,10 @@ message RemoveProjectCollaborator {
     uint32 peer_id = 2;
 }
 
+message ProjectUnshared {
+    uint64 project_id = 1;
+}
+
 message GetDefinition {
      uint64 project_id = 1;
      uint64 buffer_id = 2;

crates/rpc/src/proto.rs

@@ -128,6 +128,7 @@ messages!(
     (ProjectEntryResponse, Foreground),
     (RegisterProjectResponse, Foreground),
     (Ping, Foreground),
+    (ProjectUnshared, Foreground),
     (RegisterProject, Foreground),
     (RegisterWorktree, Foreground),
     (ReloadBuffers, Foreground),
@@ -225,6 +226,7 @@ entity_messages!(
     OpenBufferForSymbol,
     PerformRename,
     PrepareRename,
+    ProjectUnshared,
     ReloadBuffers,
     RemoveProjectCollaborator,
     RequestJoinProject,