Skip to content

Commit

Permalink
tele JB sync (#41)
Browse files Browse the repository at this point in the history
LSP functions available via HTTP as well

removed unuseful prints

added debug logs; minor

* added debug logs; minor

* implemented suggestions from @Kirill

* fixes for rh and comp_counters

* changes requested from @oleg
  • Loading branch information
valaises authored Dec 19, 2023
1 parent 0363d9e commit 4132002
Show file tree
Hide file tree
Showing 13 changed files with 225 additions and 69 deletions.
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 {
#[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

0 comments on commit 4132002

Please sign in to comment.