Record worktree extensions every 5 minutes

Antonio Scandurra and Nathan Sobo created

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

Change summary

crates/collab/src/db.rs        |   8 +-
crates/collab/src/rpc.rs       |  30 +++++-
crates/collab/src/rpc/store.rs |  57 +-------------
crates/project/src/project.rs  |  49 +++++++++++
crates/project/src/worktree.rs |   4 +
crates/rpc/proto/zed.proto     | 140 +++++++++++++++++++----------------
crates/rpc/src/proto.rs        |   2 
7 files changed, 158 insertions(+), 132 deletions(-)

Detailed changes

crates/collab/src/db.rs 🔗

@@ -52,7 +52,7 @@ pub trait Db: Send + Sync {
         &self,
         project_id: ProjectId,
         worktree_id: u64,
-        extensions: HashMap<String, usize>,
+        extensions: HashMap<String, u32>,
     ) -> Result<()>;
 
     /// Get the file counts on the given project keyed by their worktree and extension.
@@ -506,7 +506,7 @@ impl Db for PostgresDb {
         &self,
         project_id: ProjectId,
         worktree_id: u64,
-        extensions: HashMap<String, usize>,
+        extensions: HashMap<String, u32>,
     ) -> Result<()> {
         if extensions.is_empty() {
             return Ok(());
@@ -2255,7 +2255,7 @@ pub mod tests {
         background: Arc<Background>,
         pub users: Mutex<BTreeMap<UserId, User>>,
         pub projects: Mutex<BTreeMap<ProjectId, Project>>,
-        pub worktree_extensions: Mutex<BTreeMap<(ProjectId, u64, String), usize>>,
+        pub worktree_extensions: Mutex<BTreeMap<(ProjectId, u64, String), u32>>,
         pub orgs: Mutex<BTreeMap<OrgId, Org>>,
         pub org_memberships: Mutex<BTreeMap<(OrgId, UserId), bool>>,
         pub channels: Mutex<BTreeMap<ChannelId, Channel>>,
@@ -2442,7 +2442,7 @@ pub mod tests {
             &self,
             project_id: ProjectId,
             worktree_id: u64,
-            extensions: HashMap<String, usize>,
+            extensions: HashMap<String, u32>,
         ) -> Result<()> {
             self.background.simulate_random_delay().await;
             if !self.projects.lock().contains_key(&project_id) {

crates/collab/src/rpc.rs 🔗

@@ -164,6 +164,7 @@ impl Server {
             .add_message_handler(Server::update_project)
             .add_message_handler(Server::register_project_activity)
             .add_request_handler(Server::update_worktree)
+            .add_message_handler(Server::update_worktree_extensions)
             .add_message_handler(Server::start_language_server)
             .add_message_handler(Server::update_language_server)
             .add_message_handler(Server::update_diagnostic_summary)
@@ -996,9 +997,9 @@ impl Server {
     ) -> Result<()> {
         let project_id = ProjectId::from_proto(request.payload.project_id);
         let worktree_id = request.payload.worktree_id;
-        let (connection_ids, metadata_changed, extension_counts) = {
+        let (connection_ids, metadata_changed) = {
             let mut store = self.store_mut().await;
-            let (connection_ids, metadata_changed, extension_counts) = store.update_worktree(
+            let (connection_ids, metadata_changed) = store.update_worktree(
                 request.sender_id,
                 project_id,
                 worktree_id,
@@ -1007,12 +1008,8 @@ impl Server {
                 &request.payload.updated_entries,
                 request.payload.scan_id,
             )?;
-            (connection_ids, metadata_changed, extension_counts.clone())
+            (connection_ids, metadata_changed)
         };
-        self.app_state
-            .db
-            .update_worktree_extensions(project_id, worktree_id, extension_counts)
-            .await?;
 
         broadcast(request.sender_id, connection_ids, |connection_id| {
             self.peer
@@ -1029,6 +1026,25 @@ impl Server {
         Ok(())
     }
 
+    async fn update_worktree_extensions(
+        self: Arc<Server>,
+        request: TypedEnvelope<proto::UpdateWorktreeExtensions>,
+    ) -> Result<()> {
+        let project_id = ProjectId::from_proto(request.payload.project_id);
+        let worktree_id = request.payload.worktree_id;
+        let extensions = request
+            .payload
+            .extensions
+            .into_iter()
+            .zip(request.payload.counts)
+            .collect();
+        self.app_state
+            .db
+            .update_worktree_extensions(project_id, worktree_id, extensions)
+            .await?;
+        Ok(())
+    }
+
     async fn update_diagnostic_summary(
         self: Arc<Server>,
         request: TypedEnvelope<proto::UpdateDiagnosticSummary>,

crates/collab/src/rpc/store.rs 🔗

@@ -3,12 +3,7 @@ use anyhow::{anyhow, Result};
 use collections::{btree_map, hash_map::Entry, BTreeMap, BTreeSet, HashMap, HashSet};
 use rpc::{proto, ConnectionId, Receipt};
 use serde::Serialize;
-use std::{
-    mem,
-    path::{Path, PathBuf},
-    str,
-    time::Duration,
-};
+use std::{mem, path::PathBuf, str, time::Duration};
 use time::OffsetDateTime;
 use tracing::instrument;
 
@@ -59,8 +54,6 @@ pub struct Worktree {
     #[serde(skip)]
     pub entries: BTreeMap<u64, proto::Entry>,
     #[serde(skip)]
-    pub extension_counts: HashMap<String, usize>,
-    #[serde(skip)]
     pub diagnostic_summaries: BTreeMap<PathBuf, proto::DiagnosticSummary>,
     pub scan_id: u64,
 }
@@ -401,7 +394,6 @@ impl Store {
                     for worktree in project.worktrees.values_mut() {
                         worktree.diagnostic_summaries.clear();
                         worktree.entries.clear();
-                        worktree.extension_counts.clear();
                     }
 
                     Ok(Some(UnsharedProject {
@@ -632,7 +624,6 @@ impl Store {
             for worktree in project.worktrees.values_mut() {
                 worktree.diagnostic_summaries.clear();
                 worktree.entries.clear();
-                worktree.extension_counts.clear();
             }
         }
 
@@ -655,7 +646,7 @@ impl Store {
         removed_entries: &[u64],
         updated_entries: &[proto::Entry],
         scan_id: u64,
-    ) -> Result<(Vec<ConnectionId>, bool, HashMap<String, usize>)> {
+    ) -> Result<(Vec<ConnectionId>, bool)> {
         let project = self.write_project(project_id, connection_id)?;
         if !project.online {
             return Err(anyhow!("project is not online"));
@@ -667,45 +658,15 @@ impl Store {
         worktree.root_name = worktree_root_name.to_string();
 
         for entry_id in removed_entries {
-            if let Some(entry) = worktree.entries.remove(&entry_id) {
-                if !entry.is_ignored {
-                    if let Some(extension) = extension_for_entry(&entry) {
-                        if let Some(count) = worktree.extension_counts.get_mut(extension) {
-                            *count = count.saturating_sub(1);
-                        }
-                    }
-                }
-            }
+            worktree.entries.remove(&entry_id);
         }
 
         for entry in updated_entries {
-            if let Some(old_entry) = worktree.entries.insert(entry.id, entry.clone()) {
-                if !old_entry.is_ignored {
-                    if let Some(extension) = extension_for_entry(&old_entry) {
-                        if let Some(count) = worktree.extension_counts.get_mut(extension) {
-                            *count = count.saturating_sub(1);
-                        }
-                    }
-                }
-            }
-
-            if !entry.is_ignored {
-                if let Some(extension) = extension_for_entry(&entry) {
-                    if let Some(count) = worktree.extension_counts.get_mut(extension) {
-                        *count += 1;
-                    } else {
-                        worktree.extension_counts.insert(extension.into(), 1);
-                    }
-                }
-            }
+            worktree.entries.insert(entry.id, entry.clone());
         }
 
         worktree.scan_id = scan_id;
-        Ok((
-            connection_ids,
-            metadata_changed,
-            worktree.extension_counts.clone(),
-        ))
+        Ok((connection_ids, metadata_changed))
     }
 
     pub fn project_connection_ids(
@@ -894,11 +855,3 @@ impl Channel {
         self.connection_ids.iter().copied().collect()
     }
 }
-
-fn extension_for_entry(entry: &proto::Entry) -> Option<&str> {
-    str::from_utf8(&entry.path)
-        .ok()
-        .map(Path::new)
-        .and_then(|p| p.extension())
-        .and_then(|e| e.to_str())
-}

crates/project/src/project.rs 🔗

@@ -52,7 +52,7 @@ use std::{
         atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
         Arc,
     },
-    time::Instant,
+    time::{Duration, Instant},
 };
 use thiserror::Error;
 use util::{post_inc, ResultExt, TryFutureExt as _};
@@ -403,13 +403,22 @@ impl Project {
             });
 
             let (online_tx, online_rx) = watch::channel_with(online);
+            let mut send_extension_counts = None;
             let _maintain_online_status = cx.spawn_weak({
                 let mut online_rx = online_rx.clone();
                 move |this, mut cx| async move {
-                    while online_rx.recv().await.is_some() {
+                    while let Some(online) = online_rx.recv().await {
                         let this = this.upgrade(&cx)?;
+                        if online {
+                            send_extension_counts = Some(
+                                this.update(&mut cx, |this, cx| this.send_extension_counts(cx)),
+                            );
+                        } else {
+                            send_extension_counts.take();
+                        }
+
                         this.update(&mut cx, |this, cx| {
-                            if !this.is_online() {
+                            if !online {
                                 this.unshared(cx);
                             }
                             this.metadata_changed(false, cx)
@@ -463,6 +472,40 @@ impl Project {
         })
     }
 
+    fn send_extension_counts(&self, cx: &mut ModelContext<Self>) -> Task<Option<()>> {
+        cx.spawn_weak(|this, cx| async move {
+            loop {
+                let this = this.upgrade(&cx)?;
+                this.read_with(&cx, |this, cx| {
+                    if let Some(project_id) = this.remote_id() {
+                        for worktree in this.visible_worktrees(cx) {
+                            if let Some(worktree) = worktree.read(cx).as_local() {
+                                let mut extensions = Vec::new();
+                                let mut counts = Vec::new();
+
+                                for (extension, count) in worktree.extension_counts() {
+                                    extensions.push(extension.to_string_lossy().to_string());
+                                    counts.push(*count as u32);
+                                }
+
+                                this.client
+                                    .send(proto::UpdateWorktreeExtensions {
+                                        project_id,
+                                        worktree_id: worktree.id().to_proto(),
+                                        extensions,
+                                        counts,
+                                    })
+                                    .log_err();
+                            }
+                        }
+                    }
+                });
+
+                cx.background().timer(Duration::from_secs(60 * 5)).await;
+            }
+        })
+    }
+
     pub async fn remote(
         remote_id: u64,
         client: Arc<Client>,

crates/project/src/worktree.rs 🔗

@@ -1327,6 +1327,10 @@ impl LocalSnapshot {
         &self.abs_path
     }
 
+    pub fn extension_counts(&self) -> &HashMap<OsString, usize> {
+        &self.extension_counts
+    }
+
     #[cfg(test)]
     pub(crate) fn to_proto(
         &self,

crates/rpc/proto/zed.proto 🔗

@@ -38,72 +38,73 @@ message Envelope {
         UpdateProject update_project = 30;
         RegisterProjectActivity register_project_activity = 31;
         UpdateWorktree update_worktree = 32;
-
-        CreateProjectEntry create_project_entry = 33;
-        RenameProjectEntry rename_project_entry = 34;
-        CopyProjectEntry copy_project_entry = 35;
-        DeleteProjectEntry delete_project_entry = 36;
-        ProjectEntryResponse project_entry_response = 37;
-
-        UpdateDiagnosticSummary update_diagnostic_summary = 38;
-        StartLanguageServer start_language_server = 39;
-        UpdateLanguageServer update_language_server = 40;
-
-        OpenBufferById open_buffer_by_id = 41;
-        OpenBufferByPath open_buffer_by_path = 42;
-        OpenBufferResponse open_buffer_response = 43;
-        UpdateBuffer update_buffer = 44;
-        UpdateBufferFile update_buffer_file = 45;
-        SaveBuffer save_buffer = 46;
-        BufferSaved buffer_saved = 47;
-        BufferReloaded buffer_reloaded = 48;
-        ReloadBuffers reload_buffers = 49;
-        ReloadBuffersResponse reload_buffers_response = 50;
-        FormatBuffers format_buffers = 51;
-        FormatBuffersResponse format_buffers_response = 52;
-        GetCompletions get_completions = 53;
-        GetCompletionsResponse get_completions_response = 54;
-        ApplyCompletionAdditionalEdits apply_completion_additional_edits = 55;
-        ApplyCompletionAdditionalEditsResponse apply_completion_additional_edits_response = 56;
-        GetCodeActions get_code_actions = 57;
-        GetCodeActionsResponse get_code_actions_response = 58;
-        GetHover get_hover = 59;
-        GetHoverResponse get_hover_response = 60;
-        ApplyCodeAction apply_code_action = 61;
-        ApplyCodeActionResponse apply_code_action_response = 62;
-        PrepareRename prepare_rename = 63;
-        PrepareRenameResponse prepare_rename_response = 64;
-        PerformRename perform_rename = 65;
-        PerformRenameResponse perform_rename_response = 66;
-        SearchProject search_project = 67;
-        SearchProjectResponse search_project_response = 68;
-
-        GetChannels get_channels = 69;
-        GetChannelsResponse get_channels_response = 70;
-        JoinChannel join_channel = 71;
-        JoinChannelResponse join_channel_response = 72;
-        LeaveChannel leave_channel = 73;
-        SendChannelMessage send_channel_message = 74;
-        SendChannelMessageResponse send_channel_message_response = 75;
-        ChannelMessageSent channel_message_sent = 76;
-        GetChannelMessages get_channel_messages = 77;
-        GetChannelMessagesResponse get_channel_messages_response = 78;
-
-        UpdateContacts update_contacts = 79;
-        UpdateInviteInfo update_invite_info = 80;
-        ShowContacts show_contacts = 81;
-
-        GetUsers get_users = 82;
-        FuzzySearchUsers fuzzy_search_users = 83;
-        UsersResponse users_response = 84;
-        RequestContact request_contact = 85;
-        RespondToContactRequest respond_to_contact_request = 86;
-        RemoveContact remove_contact = 87;
-
-        Follow follow = 88;
-        FollowResponse follow_response = 89;
-        UpdateFollowers update_followers = 90;
-        Unfollow unfollow = 91;
+        UpdateWorktreeExtensions update_worktree_extensions = 33;
+
+        CreateProjectEntry create_project_entry = 34;
+        RenameProjectEntry rename_project_entry = 35;
+        CopyProjectEntry copy_project_entry = 36;
+        DeleteProjectEntry delete_project_entry = 37;
+        ProjectEntryResponse project_entry_response = 38;
+
+        UpdateDiagnosticSummary update_diagnostic_summary = 39;
+        StartLanguageServer start_language_server = 40;
+        UpdateLanguageServer update_language_server = 41;
+
+        OpenBufferById open_buffer_by_id = 42;
+        OpenBufferByPath open_buffer_by_path = 43;
+        OpenBufferResponse open_buffer_response = 44;
+        UpdateBuffer update_buffer = 45;
+        UpdateBufferFile update_buffer_file = 46;
+        SaveBuffer save_buffer = 47;
+        BufferSaved buffer_saved = 48;
+        BufferReloaded buffer_reloaded = 49;
+        ReloadBuffers reload_buffers = 50;
+        ReloadBuffersResponse reload_buffers_response = 51;
+        FormatBuffers format_buffers = 52;
+        FormatBuffersResponse format_buffers_response = 53;
+        GetCompletions get_completions = 54;
+        GetCompletionsResponse get_completions_response = 55;
+        ApplyCompletionAdditionalEdits apply_completion_additional_edits = 56;
+        ApplyCompletionAdditionalEditsResponse apply_completion_additional_edits_response = 57;
+        GetCodeActions get_code_actions = 58;
+        GetCodeActionsResponse get_code_actions_response = 59;
+        GetHover get_hover = 60;
+        GetHoverResponse get_hover_response = 61;
+        ApplyCodeAction apply_code_action = 62;
+        ApplyCodeActionResponse apply_code_action_response = 63;
+        PrepareRename prepare_rename = 64;
+        PrepareRenameResponse prepare_rename_response = 65;
+        PerformRename perform_rename = 66;
+        PerformRenameResponse perform_rename_response = 67;
+        SearchProject search_project = 68;
+        SearchProjectResponse search_project_response = 69;
+
+        GetChannels get_channels = 70;
+        GetChannelsResponse get_channels_response = 71;
+        JoinChannel join_channel = 72;
+        JoinChannelResponse join_channel_response = 73;
+        LeaveChannel leave_channel = 74;
+        SendChannelMessage send_channel_message = 75;
+        SendChannelMessageResponse send_channel_message_response = 76;
+        ChannelMessageSent channel_message_sent = 77;
+        GetChannelMessages get_channel_messages = 78;
+        GetChannelMessagesResponse get_channel_messages_response = 79;
+
+        UpdateContacts update_contacts = 80;
+        UpdateInviteInfo update_invite_info = 81;
+        ShowContacts show_contacts = 82;
+
+        GetUsers get_users = 83;
+        FuzzySearchUsers fuzzy_search_users = 84;
+        UsersResponse users_response = 85;
+        RequestContact request_contact = 86;
+        RespondToContactRequest respond_to_contact_request = 87;
+        RemoveContact remove_contact = 88;
+
+        Follow follow = 89;
+        FollowResponse follow_response = 90;
+        UpdateFollowers update_followers = 91;
+        Unfollow unfollow = 92;
     }
 }
 
@@ -200,6 +201,13 @@ message UpdateWorktree {
     uint64 scan_id = 6;
 }
 
+message UpdateWorktreeExtensions {
+    uint64 project_id = 1;
+    uint64 worktree_id = 2;
+    repeated string extensions = 3;
+    repeated uint32 counts = 4;
+}
+
 message CreateProjectEntry {
     uint64 project_id = 1;
     uint64 worktree_id = 2;

crates/rpc/src/proto.rs 🔗

@@ -162,6 +162,7 @@ messages!(
     (UpdateLanguageServer, Foreground),
     (UpdateProject, Foreground),
     (UpdateWorktree, Foreground),
+    (UpdateWorktreeExtensions, Background),
 );
 
 request_messages!(
@@ -254,6 +255,7 @@ entity_messages!(
     UpdateLanguageServer,
     UpdateProject,
     UpdateWorktree,
+    UpdateWorktreeExtensions,
 );
 
 entity_messages!(channel_id, ChannelMessageSent);