From 75d2944fb61be58fba48ef52b95069ebe8fb33b2 Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Tue, 10 Oct 2023 21:35:20 -0700
Subject: [PATCH] feat: support loading the source code index whenever it's
 ready in file system (#530)

* feat: support loading index whenever it's ready

* fix test
---
 crates/tabby/src/serve/completions.rs |  4 +-
 crates/tabby/src/serve/mod.rs         | 26 ++++--------
 crates/tabby/src/serve/search.rs      | 60 +++++++++++++++++++++++----
 3 files changed, 62 insertions(+), 28 deletions(-)

diff --git a/crates/tabby/src/serve/completions.rs b/crates/tabby/src/serve/completions.rs
index e0ab8be..5c55fde 100644
--- a/crates/tabby/src/serve/completions.rs
+++ b/crates/tabby/src/serve/completions.rs
@@ -128,12 +128,12 @@ pub struct CompletionState {
 impl CompletionState {
     pub fn new(
         engine: Arc<Box<dyn TextGeneration>>,
-        index_server: Option<Arc<IndexServer>>,
+        index_server: Arc<IndexServer>,
         prompt_template: Option<String>,
     ) -> Self {
         Self {
             engine,
-            prompt_builder: prompt::PromptBuilder::new(prompt_template, index_server),
+            prompt_builder: prompt::PromptBuilder::new(prompt_template, Some(index_server)),
         }
     }
 }
diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs
index c15e852..4cf03fb 100644
--- a/crates/tabby/src/serve/mod.rs
+++ b/crates/tabby/src/serve/mod.rs
@@ -22,7 +22,7 @@ use tabby_common::{
 use tabby_download::Downloader;
 use tokio::time::sleep;
 use tower_http::{cors::CorsLayer, timeout::TimeoutLayer};
-use tracing::{debug, info, warn};
+use tracing::{info, warn};
 use utoipa::{openapi::ServerBuilder, OpenApi};
 use utoipa_swagger_ui::SwaggerUi;
 
@@ -165,14 +165,7 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 }
 
 fn api_router(args: &ServeArgs) -> Router {
-    let index_server = match IndexServer::load() {
-        Ok(index_server) => Some(Arc::new(index_server)),
-        Err(err) => {
-            debug!("Load index failed due to `{}`", err);
-            None
-        }
-    };
-
+    let index_server = Arc::new(IndexServer::new());
     let completion_state = {
         let (
             engine,
@@ -235,15 +228,12 @@ fn api_router(args: &ServeArgs) -> Router {
         })
     }
 
-    if let Some(index_server) = index_server {
-        info!("Index is ready, enabling /v1beta/search API route");
-        routers.push({
-            Router::new().route(
-                "/v1beta/search",
-                routing::get(search::search).with_state(index_server),
-            )
-        })
-    }
+    routers.push({
+        Router::new().route(
+            "/v1beta/search",
+            routing::get(search::search).with_state(index_server),
+        )
+    });
 
     let mut root = Router::new();
     for router in routers {
diff --git a/crates/tabby/src/serve/search.rs b/crates/tabby/src/serve/search.rs
index 9e97fd1..87a37d7 100644
--- a/crates/tabby/src/serve/search.rs
+++ b/crates/tabby/src/serve/search.rs
@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use anyhow::Result;
 use axum::{
@@ -14,7 +14,8 @@ use tantivy::{
     schema::Field,
     DocAddress, Document, Index, IndexReader,
 };
-use tracing::instrument;
+use tokio::{sync::OnceCell, task, time::sleep};
+use tracing::{debug, instrument, log::info};
 use utoipa::{IntoParams, ToSchema};
 
 #[derive(Deserialize, IntoParams)]
@@ -60,9 +61,9 @@ pub struct HitDocument {
     tag = "v1beta",
     responses(
         (status = 200, description = "Success" , body = SearchResponse, content_type = "application/json"),
-        (status = 405, description = "When code search is not enabled, the endpoint will returns 405 Method Not Allowed"),
-    )
-)]
+        (status = 501, description = "When code search is not enabled, the endpoint will returns 501 Not Implemented"),
+        )
+    )]
 #[instrument(skip(state, query))]
 pub async fn search(
     State(state): State<Arc<IndexServer>>,
@@ -73,13 +74,13 @@ pub async fn search(
         query.limit.unwrap_or(20),
         query.offset.unwrap_or(0),
     ) else {
-        return Err(StatusCode::INTERNAL_SERVER_ERROR);
+        return Err(StatusCode::NOT_IMPLEMENTED);
     };
 
     Ok(Json(serp))
 }
 
-pub struct IndexServer {
+struct IndexServerImpl {
     reader: IndexReader,
     query_parser: QueryParser,
 
@@ -91,7 +92,7 @@ pub struct IndexServer {
     field_name: Field,
 }
 
-impl IndexServer {
+impl IndexServerImpl {
     pub fn load() -> Result<Self> {
         let index = Index::open_in_dir(path::index_dir())?;
         index.register_tokenizer();
@@ -162,3 +163,46 @@ fn get_field(doc: &Document, field: Field) -> String {
         .unwrap()
         .to_owned()
 }
+
+static IMPL: OnceCell<IndexServerImpl> = OnceCell::const_new();
+
+pub struct IndexServer {}
+
+impl IndexServer {
+    pub fn new() -> Self {
+        task::spawn(IMPL.get_or_init(|| async {
+            task::spawn(IndexServer::worker())
+                .await
+                .expect("Failed to create IndexServerImpl")
+        }));
+        Self {}
+    }
+
+    fn get_cell(&self) -> Option<&IndexServerImpl> {
+        IMPL.get()
+    }
+
+    async fn worker() -> IndexServerImpl {
+        loop {
+            match IndexServerImpl::load() {
+                Ok(index_server) => {
+                    info!("Index is ready, enabling server...");
+                    return index_server;
+                }
+                Err(err) => {
+                    debug!("Source code index is not ready `{}`", err);
+                }
+            };
+
+            sleep(Duration::from_secs(60)).await;
+        }
+    }
+
+    pub fn search(&self, q: &str, limit: usize, offset: usize) -> tantivy::Result<SearchResponse> {
+        if let Some(imp) = self.get_cell() {
+            imp.search(q, limit, offset)
+        } else {
+            Err(tantivy::TantivyError::InternalError("Not Ready".to_owned()))
+        }
+    }
+}