From 18e0a6b5583eec459c5442615bd6bd62ee48708c Mon Sep 17 00:00:00 2001
From: Liquan Pei <liquanpei@gmail.com>
Date: Tue, 16 Jul 2024 12:49:37 -0700
Subject: [PATCH] [ENH] Use instrumented tracing for HNSW provider file read
 (#2525)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- This PR adds tracing for HNSW file copy to local and fixes the cache
initialization from config.
 - New functionality
	 - ...

## Test plan
*How are these changes tested?*

- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
repository](https://github.com/chroma-core/docs)?*
---
 rust/worker/src/index/hnsw_provider.rs | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/rust/worker/src/index/hnsw_provider.rs b/rust/worker/src/index/hnsw_provider.rs
index 0f3ec82554b..5166bd03353 100644
--- a/rust/worker/src/index/hnsw_provider.rs
+++ b/rust/worker/src/index/hnsw_provider.rs
@@ -3,6 +3,7 @@ use super::{
     HnswIndex, HnswIndexConfig, HnswIndexFromSegmentError, Index, IndexConfig,
     IndexConfigFromSegmentError,
 };
+use crate::cache;
 use crate::cache::cache::Cache;
 use crate::config::Configurable;
 use crate::errors::ErrorCodes;
@@ -54,7 +55,7 @@ impl Configurable<(HnswProviderConfig, Storage)> for HnswIndexProvider {
         config: &(HnswProviderConfig, Storage),
     ) -> Result<Self, Box<dyn ChromaError>> {
         let (hnsw_config, storage) = config;
-        let cache = Cache::new(&hnsw_config.hnsw_cache_config);
+        let cache = cache::from_config(&hnsw_config.hnsw_cache_config).await?;
         Ok(Self {
             cache,
             storage: storage.clone(),
@@ -177,7 +178,10 @@ impl HnswIndexProvider {
                     return Err(Box::new(HnswIndexProviderFileError::IOError(e)));
                 }
             };
-            let total_bytes_written = self.copy_stream_to_local_file(reader, file_handle).await?;
+            let total_bytes_written = self
+                .copy_stream_to_local_file(reader, file_handle)
+                .instrument(tracing::info_span!(parent: Span::current(), "hnsw provider file read", file = file))
+                .await?;
             tracing::info!(
                 "Copied {} bytes from storage key: {} to file: {}",
                 total_bytes_written,
@@ -212,6 +216,7 @@ impl HnswIndexProvider {
                     total_bytes_written += chunk.len() as u64;
                 }
                 Err(e) => {
+                    tracing::error!("Failed to copy file: {}", e);
                     return Err(Box::new(HnswIndexProviderFileError::IOError(e)));
                 }
             }