From e678c7d9ee29cd7de6ab788c866c8a950b843306 Mon Sep 17 00:00:00 2001
From: KCaverly <kyle@zed.dev>
Date: Mon, 11 Sep 2023 10:26:14 -0400
Subject: [PATCH] swap SystemTime for Instant throughout rate_limit_expiry
 tracking

---
 crates/search/src/project_search.rs            |  8 ++++----
 crates/semantic_index/src/embedding.rs         | 18 +++++++++---------
 crates/semantic_index/src/semantic_index.rs    |  2 +-
 .../semantic_index/src/semantic_index_tests.rs |  4 ++--
 4 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/crates/search/src/project_search.rs b/crates/search/src/project_search.rs
index b85d0b9b40..6ca4928803 100644
--- a/crates/search/src/project_search.rs
+++ b/crates/search/src/project_search.rs
@@ -34,7 +34,7 @@ use std::{
     ops::{Not, Range},
     path::PathBuf,
     sync::Arc,
-    time::{Duration, SystemTime},
+    time::{Duration, Instant},
 };
 use util::ResultExt as _;
 use workspace::{
@@ -329,9 +329,9 @@ impl View for ProjectSearchView {
                             Some(format!("Indexing..."))
                         } else {
                             if let Some(rate_limit_expiry) = rate_limit_expiry {
-                                if let Ok(remaining_seconds) =
-                                    rate_limit_expiry.duration_since(SystemTime::now())
-                                {
+                                let remaining_seconds =
+                                    rate_limit_expiry.duration_since(Instant::now());
+                                if remaining_seconds > Duration::from_secs(0) {
                                     Some(format!(
                                         "Remaining files to index (rate limit resets in {}s): {}",
                                         remaining_seconds.as_secs(),
diff --git a/crates/semantic_index/src/embedding.rs b/crates/semantic_index/src/embedding.rs
index 7bac809c97..42d90f0fdb 100644
--- a/crates/semantic_index/src/embedding.rs
+++ b/crates/semantic_index/src/embedding.rs
@@ -16,7 +16,7 @@ use serde::{Deserialize, Serialize};
 use std::env;
 use std::ops::Add;
 use std::sync::Arc;
-use std::time::{Duration, SystemTime};
+use std::time::{Duration, Instant};
 use tiktoken_rs::{cl100k_base, CoreBPE};
 use util::http::{HttpClient, Request};
 
@@ -85,8 +85,8 @@ impl ToSql for Embedding {
 pub struct OpenAIEmbeddings {
     pub client: Arc<dyn HttpClient>,
     pub executor: Arc<Background>,
-    rate_limit_count_rx: watch::Receiver<Option<SystemTime>>,
-    rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<SystemTime>>>>,
+    rate_limit_count_rx: watch::Receiver<Option<Instant>>,
+    rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<Instant>>>>,
 }
 
 #[derive(Serialize)]
@@ -119,14 +119,14 @@ pub trait EmbeddingProvider: Sync + Send {
     async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
     fn max_tokens_per_batch(&self) -> usize;
     fn truncate(&self, span: &str) -> (String, usize);
-    fn rate_limit_expiration(&self) -> Option<SystemTime>;
+    fn rate_limit_expiration(&self) -> Option<Instant>;
 }
 
 pub struct DummyEmbeddings {}
 
 #[async_trait]
 impl EmbeddingProvider for DummyEmbeddings {
-    fn rate_limit_expiration(&self) -> Option<SystemTime> {
+    fn rate_limit_expiration(&self) -> Option<Instant> {
         None
     }
     async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
@@ -174,7 +174,7 @@ impl OpenAIEmbeddings {
         let reset_time = *self.rate_limit_count_tx.lock().borrow();
 
         if let Some(reset_time) = reset_time {
-            if SystemTime::now() >= reset_time {
+            if Instant::now() >= reset_time {
                 *self.rate_limit_count_tx.lock().borrow_mut() = None
             }
         }
@@ -185,7 +185,7 @@ impl OpenAIEmbeddings {
         );
     }
 
-    fn update_reset_time(&self, reset_time: SystemTime) {
+    fn update_reset_time(&self, reset_time: Instant) {
         let original_time = *self.rate_limit_count_tx.lock().borrow();
 
         let updated_time = if let Some(original_time) = original_time {
@@ -232,7 +232,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
         50000
     }
 
-    fn rate_limit_expiration(&self) -> Option<SystemTime> {
+    fn rate_limit_expiration(&self) -> Option<Instant> {
         *self.rate_limit_count_rx.borrow()
     }
     fn truncate(&self, span: &str) -> (String, usize) {
@@ -319,7 +319,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
                     };
 
                     // If we've previously rate limited, increment the duration but not the count
-                    let reset_time = SystemTime::now().add(delay_duration);
+                    let reset_time = Instant::now().add(delay_duration);
                     self.update_reset_time(reset_time);
 
                     log::trace!(
diff --git a/crates/semantic_index/src/semantic_index.rs b/crates/semantic_index/src/semantic_index.rs
index efcc1ba242..115bf5d7a8 100644
--- a/crates/semantic_index/src/semantic_index.rs
+++ b/crates/semantic_index/src/semantic_index.rs
@@ -112,7 +112,7 @@ pub enum SemanticIndexStatus {
     Indexed,
     Indexing {
         remaining_files: usize,
-        rate_limit_expiry: Option<SystemTime>,
+        rate_limit_expiry: Option<Instant>,
     },
 }
 
diff --git a/crates/semantic_index/src/semantic_index_tests.rs b/crates/semantic_index/src/semantic_index_tests.rs
index 4bc95bec62..9035327b2e 100644
--- a/crates/semantic_index/src/semantic_index_tests.rs
+++ b/crates/semantic_index/src/semantic_index_tests.rs
@@ -21,7 +21,7 @@ use std::{
         atomic::{self, AtomicUsize},
         Arc,
     },
-    time::SystemTime,
+    time::{Instant, SystemTime},
 };
 use unindent::Unindent;
 use util::RandomCharIter;
@@ -1275,7 +1275,7 @@ impl EmbeddingProvider for FakeEmbeddingProvider {
         200
     }
 
-    fn rate_limit_expiration(&self) -> Option<SystemTime> {
+    fn rate_limit_expiration(&self) -> Option<Instant> {
         None
     }