[Text Index Usage Tracking] Start tracking text index ingress (#44038)

Adds new fields in usage tracking for text index ingress + egress. Ingress tracking is now hooked up in the committer and we include only the indexed field using an `estimated_size()` function.

Text index egress is not tracked yet because it requires some additional changes to pipe through the read size from funrun and we have to decide how we want to bill the size of text index reads.

Follow-up items are to:
- [x] charge for text index size
- [ ] start tracking database egress

GitOrigin-RevId: 5c5df769795e78d79c57f7fedb272bb645bba3e5

Author: jordanhunt22

crates/application/src/function_log.rs

@@ -284,6 +284,8 @@ impl FunctionExecution {
                     storage_write_bytes: self.usage_stats.storage_write_bytes,
                     vector_index_read_bytes: self.usage_stats.vector_index_read_bytes,
                     vector_index_write_bytes: self.usage_stats.vector_index_write_bytes,
+                    text_index_read_bytes: self.usage_stats.text_index_read_bytes,
+                    text_index_write_bytes: self.usage_stats.text_index_write_bytes,
                     memory_used_mb: self.memory_used_mb,
                     return_bytes: self.return_bytes,
                 },

crates/common/src/log_streaming.rs

@@ -67,6 +67,8 @@ pub struct AggregatedFunctionUsageStats {
     pub storage_write_bytes: u64,
     pub vector_index_read_bytes: u64,
     pub vector_index_write_bytes: u64,
+    pub text_index_read_bytes: u64,
+    pub text_index_write_bytes: u64,
     pub memory_used_mb: u64,
     pub return_bytes: Option<u64>,
 }
@@ -974,6 +976,8 @@ mod tests {
                         storage_write_bytes: 0,
                         vector_index_read_bytes: 0,
                         vector_index_write_bytes: 0,
+                        text_index_read_bytes: 0,
+                        text_index_write_bytes: 0,
                         memory_used_mb: 0,
                         return_bytes: Some(64),
                     },

crates/database/src/committer.rs

@@ -93,6 +93,7 @@ use itertools::Itertools;
 use parking_lot::Mutex;
 use prometheus::VMHistogram;
 use rand::Rng;
+use search::TextIndexWriteSize;
 use tokio::sync::{
     mpsc::{
         self,
@@ -757,14 +758,15 @@ impl<RT: Runtime> Committer<RT> {
             .latest_snapshot()
             .unwrap_or_else(|| self.snapshot_manager.read().latest_snapshot());
         for &document_update in ordered_updates.iter() {
-            let (updates, doc_in_vector_index) =
+            let (updates, doc_in_vector_index, text_index_write_size) =
                 latest_pending_snapshot.update(document_update, commit_ts)?;
             index_writes.extend(updates);
             document_writes.push(ValidatedDocumentWrite {
                 commit_ts,
                 id: document_update.id.into(),
                 write: document_update.new_document.clone(),
                 doc_in_vector_index,
+                text_index_write_size,
                 prev_ts: document_update.old_document.as_ref().map(|&(_, ts)| ts),
             });
         }
@@ -1039,6 +1041,7 @@ impl<RT: Runtime> Committer<RT> {
             let ValidatedDocumentWrite {
                 write: document,
                 doc_in_vector_index,
+                text_index_write_size,
                 ..
             } = validated_write;
             if let Some(document) = document {
@@ -1054,28 +1057,35 @@ impl<RT: Runtime> Committer<RT> {
                     .unwrap_or(ComponentPath::root());
                 if let Ok(table_name) = table_mapping.tablet_name(tablet_id) {
                     // Database bandwidth for document writes
-                    if *doc_in_vector_index == DocInVectorIndex::Absent {
-                        usage_tracker.track_database_ingress_size(
+                    usage_tracker.track_database_ingress_size(
+                        component_path.clone().clone(),
+                        table_name.to_string(),
+                        document_write_size as u64,
+                        table_name.is_system(),
+                    );
+                    usage_tracker.track_database_ingress_size_v2(
+                        component_path.clone(),
+                        virtual_system_mapping
+                            .system_to_virtual_table(&table_name)
+                            .unwrap_or(&table_name)
+                            .to_string(),
+                        document_write_size as u64,
+                        table_name.is_system()
+                            && !virtual_system_mapping.has_virtual_table(&table_name),
+                    );
+                    if *doc_in_vector_index == DocInVectorIndex::Present {
+                        usage_tracker.track_vector_ingress_size(
                             component_path.clone(),
                             table_name.to_string(),
                             document_write_size as u64,
                             table_name.is_system(),
                         );
-                        usage_tracker.track_database_ingress_size_v2(
-                            component_path,
-                            virtual_system_mapping
-                                .system_to_virtual_table(&table_name)
-                                .unwrap_or(&table_name)
-                                .to_string(),
-                            document_write_size as u64,
-                            table_name.is_system()
-                                && !virtual_system_mapping.has_virtual_table(&table_name),
-                        );
-                    } else {
-                        usage_tracker.track_vector_ingress_size(
+                    }
+                    if text_index_write_size.0 > 0 {
+                        usage_tracker.track_text_ingress_size(
                             component_path,
                             table_name.to_string(),
-                            document_write_size as u64,
+                            text_index_write_size.0,
                             table_name.is_system(),
                         );
                     }
@@ -1118,6 +1128,7 @@ struct ValidatedDocumentWrite {
     id: InternalDocumentId,
     write: Option<ResolvedDocument>,
     doc_in_vector_index: DocInVectorIndex,
+    text_index_write_size: TextIndexWriteSize,
     prev_ts: Option<Timestamp>,
 }
 

crates/database/src/snapshot_manager.rs

@@ -31,7 +31,10 @@ use indexing::{
     backend_in_memory_indexes::BackendInMemoryIndexes,
     index_registry::IndexRegistry,
 };
-use search::TextIndexManager;
+use search::{
+    TextIndexManager,
+    TextIndexWriteSize,
+};
 use value::{
     ResolvedDocumentId,
     TableMapping,
@@ -220,7 +223,11 @@ impl Snapshot {
         &mut self,
         document_update: &impl DocumentUpdateRef,
         commit_ts: Timestamp,
-    ) -> anyhow::Result<(Vec<DatabaseIndexUpdate>, DocInVectorIndex)> {
+    ) -> anyhow::Result<(
+        Vec<DatabaseIndexUpdate>,
+        DocInVectorIndex,
+        TextIndexWriteSize,
+    )> {
         block_in_place(|| {
             let removal = document_update.old_document();
             let insertion = document_update.new_document();
@@ -268,7 +275,8 @@ impl Snapshot {
                 insertion.cloned(),
             );
 
-            self.text_indexes
+            let text_index_write_size = self
+                .text_indexes
                 .update(
                     &self.index_registry,
                     removal,
@@ -286,7 +294,11 @@ impl Snapshot {
                     WriteTimestamp::Committed(commit_ts),
                 )
                 .context("Vector index update failed")?;
-            Ok((in_memory_index_updates, doc_in_vector_index))
+            Ok((
+                in_memory_index_updates,
+                doc_in_vector_index,
+                text_index_write_size,
+            ))
         })
     }
 

crates/database/src/tests/text_test_utils.rs

@@ -35,6 +35,7 @@ use common::{
     },
     version::MIN_NPM_VERSION_FOR_FUZZY_SEARCH,
 };
+use events::testing::TestUsageEventLogger;
 use futures::try_join;
 use maplit::btreeset;
 use must_let::must_let;
@@ -92,6 +93,7 @@ pub struct TextFixtures {
     pub storage: Arc<dyn Storage>,
     pub db: Database<TestRuntime>,
     pub reader: Arc<dyn PersistenceReader>,
+    pub test_usage_logger: TestUsageEventLogger,
     searcher: Arc<dyn Searcher>,
     segment_term_metadata_fetcher: Arc<dyn SegmentTermMetadataFetcher>,
     writer: TextIndexMetadataWriter<TestRuntime>,
@@ -114,6 +116,7 @@ impl TextFixtures {
             db,
             search_storage,
             searcher,
+            test_usage_logger,
             ..
         } = DbFixtures::new_with_args(
             &rt,
@@ -144,6 +147,7 @@ impl TextFixtures {
             namespace: TableNamespace::test_user(),
             searcher,
             config,
+            test_usage_logger,
         })
     }
 
@@ -193,7 +197,7 @@ impl TextFixtures {
         )
     }
 
-    pub async fn enabled_text_index(&self) -> anyhow::Result<IndexData> {
+    pub async fn enabled_text_index(&self) -> anyhow::Result<TextIndexData> {
         let index_data = self.backfilled_text_index().await?;
         let mut tx = self.db.begin_system().await?;
         IndexModel::new(&mut tx)
@@ -203,7 +207,7 @@ impl TextFixtures {
         Ok(index_data)
     }
 
-    pub async fn backfilled_text_index(&self) -> anyhow::Result<IndexData> {
+    pub async fn backfilled_text_index(&self) -> anyhow::Result<TextIndexData> {
         let index_data = self.insert_backfilling_text_index().await?;
         self.backfill().await?;
 
@@ -249,7 +253,7 @@ impl TextFixtures {
         Ok(ts)
     }
 
-    pub async fn insert_backfilling_text_index(&self) -> anyhow::Result<IndexData> {
+    pub async fn insert_backfilling_text_index(&self) -> anyhow::Result<TextIndexData> {
         let mut tx = self.db.begin_system().await?;
         let index_metadata = backfilling_text_index()?;
         let index_name = &index_metadata.name;
@@ -264,15 +268,17 @@ impl TextFixtures {
         self.db.commit(tx).await?;
 
         let resolved_index_name = TabletIndexName::new(table_id, index_name.descriptor().clone())?;
-        Ok(IndexData {
+        Ok(TextIndexData {
             index_id,
             resolved_index_name,
             index_name: index_name.clone(),
             namespace: self.namespace,
         })
     }
 
-    pub async fn insert_backfilling_text_index_with_document(&self) -> anyhow::Result<IndexData> {
+    pub async fn insert_backfilling_text_index_with_document(
+        &self,
+    ) -> anyhow::Result<TextIndexData> {
         let index_data = self.insert_backfilling_text_index().await?;
         let mut tx = self.db.begin_system().await?;
         add_document(&mut tx, index_data.index_name.table(), "A long text field").await?;
@@ -404,7 +410,7 @@ impl TextFixtures {
 const TABLE_NAME: &str = "table";
 const SEARCH_FIELD: &str = "text";
 
-pub struct IndexData {
+pub struct TextIndexData {
     pub index_id: ResolvedDocumentId,
     pub index_name: IndexName,
     pub resolved_index_name: TabletIndexName,