Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

tele JB sync #41

Merged
merged 5 commits into from
Dec 19, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions src/global_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ use crate::telemetry::telemetry_structs;
use crate::vecdb_search::VecdbSearch;
use crate::custom_error::ScratchError;
use hyper::StatusCode;
use tower_lsp::lsp_types::WorkspaceFolder;
use crate::receive_workspace_changes::Document;


#[derive(Debug, StructOpt, Clone)]
Expand Down Expand Up @@ -46,6 +48,10 @@ pub struct Slowdown {
pub requests_in_flight: u64,
}

pub struct LSPBackendDocumentState {
pub document_map: Arc<ARwLock<HashMap<String, Document>>>,
pub workspace_folders: Arc<ARwLock<Option<Vec<WorkspaceFolder>>>>,
}

// #[derive(Debug)]
pub struct GlobalContext {
Expand All @@ -60,6 +66,7 @@ pub struct GlobalContext {
pub telemetry: Arc<StdRwLock<telemetry_structs::Storage>>,
pub vecdb_search: Arc<AMutex<Box<dyn VecdbSearch + Send>>>,
pub ask_shutdown_sender: Arc<Mutex<std::sync::mpsc::Sender<String>>>,
pub lsp_backend_document_state: LSPBackendDocumentState,
}

pub type SharedGlobalContext = Arc<ARwLock<GlobalContext>>;
Expand Down Expand Up @@ -167,6 +174,10 @@ pub async fn create_global_context(
telemetry: Arc::new(StdRwLock::new(telemetry_structs::Storage::new())),
vecdb_search: Arc::new(AMutex::new(Box::new(crate::vecdb_search::VecdbSearchTest::new()))),
ask_shutdown_sender: Arc::new(Mutex::new(ask_shutdown_sender)),
lsp_backend_document_state: LSPBackendDocumentState {
document_map: Arc::new(ARwLock::new(HashMap::new())),
workspace_folders: Arc::new(ARwLock::new(None)),
},
};
(Arc::new(ARwLock::new(cx)), ask_shutdown_receiver, cmdline)
}
6 changes: 6 additions & 0 deletions src/http/routers/v1.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ use crate::http::routers::v1::code_completion::handle_v1_code_completion_web;
use crate::http::routers::v1::graceful_shutdown::handle_v1_graceful_shutdown;
use crate::http::routers::v1::snippet_accepted::handle_v1_snippet_accepted;
use crate::http::routers::v1::telemetry_network::handle_v1_telemetry_network;
use crate::http::routers::v1::lsp_like_handlers::handle_v1_lsp_initialize;
use crate::http::routers::v1::lsp_like_handlers::handle_v1_lsp_did_change;
use crate::http::utils::telemetry_wrapper;

pub mod code_completion;
Expand All @@ -25,6 +27,7 @@ pub mod telemetry_network;
pub mod snippet_accepted;
pub mod caps;
pub mod graceful_shutdown;
mod lsp_like_handlers;

pub fn make_v1_router() -> Router {
Router::new()
Expand All @@ -35,4 +38,7 @@ pub fn make_v1_router() -> Router {

.route("/caps", telemetry_get!(handle_v1_caps))
.route("/graceful-shutdown", telemetry_get!(handle_v1_graceful_shutdown))

.route("/lsp-initialize", telemetry_post!(handle_v1_lsp_initialize))
.route("/lsp-did-changed", telemetry_post!(handle_v1_lsp_did_change))
}
56 changes: 56 additions & 0 deletions src/http/routers/v1/lsp_like_handlers.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
use axum::Extension;
use axum::response::Result;
use hyper::{Body, Response, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::json;
use url::Url;
use crate::custom_error::ScratchError;
use crate::global_context::SharedGlobalContext;
use crate::receive_workspace_changes;


#[derive(Serialize, Deserialize, Clone)]
struct PostInit {
pub project_roots: Vec<Url>,
}

#[derive(Serialize, Deserialize, Clone)]
struct PostDocument {
pub uri: Url,
pub text: String,
}


pub async fn handle_v1_lsp_initialize(
Extension(global_context): Extension<SharedGlobalContext>,
body_bytes: hyper::body::Bytes,
) -> Result<Response<Body>, ScratchError> {
let post = serde_json::from_slice::<PostInit>(&body_bytes).map_err(|e| {
ScratchError::new(StatusCode::BAD_REQUEST, format!("JSON problem: {}", e))
})?;

Ok(Response::builder()
.status(StatusCode::OK)
.body(Body::from(json!({"success": 1}).to_string()))
.unwrap())
}

pub async fn handle_v1_lsp_did_change(
Extension(global_context): Extension<SharedGlobalContext>,
body_bytes: hyper::body::Bytes,
) -> Result<Response<Body>, ScratchError> {
let post = serde_json::from_slice::<PostDocument>(&body_bytes).map_err(|e| {
ScratchError::new(StatusCode::BAD_REQUEST, format!("JSON problem: {}", e))
})?;

receive_workspace_changes::on_did_change(
global_context.clone(),
&post.uri.to_string(),
&post.text,
).await;

Ok(Response::builder()
.status(StatusCode::OK)
.body(Body::from(json!({"success": 1}).to_string()))
.unwrap())
}
83 changes: 29 additions & 54 deletions src/lsp.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
use std::collections::HashMap;
use std::fmt::Display;
use std::sync::Arc;
use std::time::Instant;

use ropey::Rope;
use serde::{Deserialize, Serialize};
use tokio::net::TcpListener;
use tokio::sync::RwLock as ARwLock;
Expand All @@ -18,6 +16,8 @@ use crate::global_context;
use crate::global_context::CommandLine;
use crate::http::routers::v1::code_completion::handle_v1_code_completion;
use crate::telemetry;
use crate::receive_workspace_changes;


const VERSION: &str = env!("CARGO_PKG_VERSION");

Expand All @@ -34,25 +34,11 @@ impl Display for APIError {
}


#[derive(Debug)]
pub struct Document {
#[allow(dead_code)]
pub language_id: String,
pub text: Rope,
}

impl Document {
fn new(language_id: String, text: Rope) -> Self {
Self { language_id, text }
}
}

// #[derive(Debug)] GlobalContext does not implement Debug
pub struct Backend {
pub gcx: Arc<ARwLock<global_context::GlobalContext>>,
pub client: tower_lsp::Client,
pub document_map: Arc<ARwLock<HashMap<String, Document>>>,
pub workspace_folders: Arc<ARwLock<Option<Vec<WorkspaceFolder>>>>,
}


Expand Down Expand Up @@ -120,11 +106,14 @@ pub struct CompletionRes {

impl Backend {
async fn flat_params_to_code_completion_post(&self, params: &CompletionParams1) -> CodeCompletionPost {
let document_map = self.document_map.read().await;
let document = document_map
.get(params.text_document_position.text_document.uri.as_str())
.unwrap();
let txt = &document.text;
let txt = {
let document_map = self.gcx.read().await.lsp_backend_document_state.document_map.clone();
let document_map = document_map.read().await;
let document = document_map
.get(params.text_document_position.text_document.uri.as_str())
.unwrap();
document.text.clone()
};
CodeCompletionPost {
inputs: CodeCompletionInputs {
sources: HashMap::from([(String::from(&params.text_document_position.text_document.uri.to_string()),
Expand Down Expand Up @@ -182,9 +171,12 @@ impl Backend {
#[tower_lsp::async_trait]
impl LanguageServer for Backend {
async fn initialize(&self, params: InitializeParams) -> Result<InitializeResult> {
*self.workspace_folders.write().await = params.workspace_folders;
info!("LSP client_info {:?}", params.client_info);
info!("LSP workspace_folders {:?}", self.workspace_folders);
{
let mut gcx_locked = self.gcx.write().await;
*gcx_locked.lsp_backend_document_state.workspace_folders.write().await = params.workspace_folders;
info!("LSP workspace_folders {:?}", gcx_locked.lsp_backend_document_state.workspace_folders);
}

let completion_options: CompletionOptions;
completion_options = CompletionOptions {
Expand Down Expand Up @@ -220,35 +212,20 @@ impl LanguageServer for Backend {
// textDocument/didClose

async fn did_open(&self, params: DidOpenTextDocumentParams) {
let rope = ropey::Rope::from_str(&params.text_document.text);
let uri = params.text_document.uri.to_string();
*self
.document_map
.write()
.await
.entry(uri.clone())
.or_insert(Document::new("unknown".to_owned(), Rope::new())) =
Document::new(params.text_document.language_id, rope);
info!("{uri} opened");
receive_workspace_changes::on_did_open(
self.gcx.clone(),
&params.text_document.uri.to_string(),
&params.text_document.text,
&params.text_document.language_id
).await
}

async fn did_change(&self, params: DidChangeTextDocumentParams) {
let t0 = Instant::now();
let rope = ropey::Rope::from_str(&params.content_changes[0].text);
let uri = params.text_document.uri.to_string();
let mut document_map = self.document_map.write().await;
let doc = document_map
.entry(uri.clone())
.or_insert(Document::new("unknown".to_owned(), Rope::new()));
doc.text = rope;
info!("{} changed, save time: {:?}", uri, t0.elapsed());
let t1 = Instant::now();
telemetry::snippets_collection::sources_changed(
receive_workspace_changes::on_did_change(
self.gcx.clone(),
&uri,
&params.content_changes[0].text,
).await;
info!("{} changed, telemetry time: {:?}", uri, t1.elapsed());
&params.text_document.uri.to_string(),
&params.content_changes[0].text
).await
}

async fn did_save(&self, params: DidSaveTextDocumentParams) {
Expand Down Expand Up @@ -281,22 +258,20 @@ impl LanguageServer for Backend {
}
}

fn build_lsp_service(
async fn build_lsp_service(
gcx: Arc<ARwLock<global_context::GlobalContext>>,
) -> (LspService::<Backend>, ClientSocket) {
let (lsp_service, socket) = LspService::build(|client| Backend {
gcx,
client,
document_map: Arc::new(ARwLock::new(HashMap::new())),
workspace_folders: Arc::new(ARwLock::new(None)),
})
.custom_method("refact/getCompletions", Backend::get_completions)
.custom_method("refact/test_if_head_tail_equal_return_added_text", Backend::test_if_head_tail_equal_return_added_text)
.finish();
(lsp_service, socket)
}

pub fn spawn_lsp_task(
pub async fn spawn_lsp_task(
gcx: Arc<ARwLock<global_context::GlobalContext>>,
cmdline: CommandLine
) -> Option<JoinHandle<()>> {
Expand All @@ -313,7 +288,7 @@ pub fn spawn_lsp_task(
Ok((s, addr)) => {
info!("LSP new client connection from {}", addr);
let (read, write) = tokio::io::split(s);
let (lsp_service, socket) = build_lsp_service(gcx_t.clone());
let (lsp_service, socket) = build_lsp_service(gcx_t.clone()).await;
tower_lsp::Server::new(read, write, socket).serve(lsp_service).await;
}
Err(e) => {
Expand All @@ -329,7 +304,7 @@ pub fn spawn_lsp_task(
return Some(tokio::spawn( async move {
let stdin = tokio::io::stdin();
let stdout = tokio::io::stdout();
let (lsp_service, socket) = build_lsp_service(gcx_t.clone());
let (lsp_service, socket) = build_lsp_service(gcx_t.clone()).await;
tower_lsp::Server::new(stdin, stdout, socket).serve(lsp_service).await;
info!("LSP loop exit");
gcx_t.write().await.ask_shutdown_sender.lock().unwrap().send(format!("going-down-because-lsp-exited")).unwrap();
Expand Down
3 changes: 2 additions & 1 deletion src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ mod vecdb_search;
mod lsp;
mod http;
mod background_tasks;
mod receive_workspace_changes;

#[tokio::main]
async fn main() {
Expand Down Expand Up @@ -57,7 +58,7 @@ async fn main() {
}

let mut background_tasks = start_background_tasks(gcx.clone());
let lsp_task = spawn_lsp_task(gcx.clone(), cmdline.clone()); // execution stays inside if stdin-stdout
let lsp_task = spawn_lsp_task(gcx.clone(), cmdline.clone()).await; // execution stays inside if stdin-stdout
if lsp_task.is_some() {
background_tasks.push_back(lsp_task.unwrap())
}
Expand Down
71 changes: 71 additions & 0 deletions src/receive_workspace_changes.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
use std::collections::HashMap;
use std::fmt::Display;
use std::sync::Arc;
use std::time::Instant;
use ropey::Rope;

use tokio::sync::RwLock as ARwLock;
use tracing::{info};

use crate::global_context;
use crate::telemetry;



#[derive(Debug)]
pub struct Document {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do u move Document here?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bc Document is not being used anywhere else but in receive_workspace_changes

#[allow(dead_code)]
pub language_id: String,
pub text: Rope,
}

impl Document {
pub fn new(language_id: String, text: Rope) -> Self {
Self { language_id, text }
}
}


pub async fn on_did_open(
gcx: Arc<ARwLock<global_context::GlobalContext>>,
uri: &String,
text: &String,
language_id: &String,
) {
let t0 = Instant::now();
let gcx_locked = gcx.read().await;
let document_map = &gcx_locked.lsp_backend_document_state.document_map;
let rope = ropey::Rope::from_str(&text);
let mut document_map_locked = document_map.write().await;
*document_map_locked
.entry(uri.clone())
.or_insert(Document::new("unknown".to_owned(), Rope::new())) =
Document::new(language_id.clone(), rope);
info!("{} opened, save time: {:?}", uri, t0.elapsed());
}

pub async fn on_did_change(
gcx: Arc<ARwLock<global_context::GlobalContext>>,
uri: &String,
text: &String,
) {
let t0 = Instant::now();

let gcx_locked = gcx.read().await;
let document_map = &gcx_locked.lsp_backend_document_state.document_map;
let rope = ropey::Rope::from_str(&text);
let mut document_map_locked = document_map.write().await;
let doc = document_map_locked
.entry(uri.clone())
.or_insert(Document::new("unknown".to_owned(), Rope::new()));
doc.text = rope;

info!("{} changed, save time: {:?}", uri, t0.elapsed());
let t1 = Instant::now();
telemetry::snippets_collection::sources_changed(
gcx.clone(),
uri,
text,
).await;
info!("{} changed, telemetry time: {:?}", uri, t1.elapsed());
}
Loading