Add ServerConnection and outer request handling

Summary:
This adds sort of the main outer loop of the debug-attach command. It
sets up the ServerConnection (that manages the global debugger value) and then
routes messages to/from the inner debugger server and encodes/decodes them as
appropriate.

Reviewed By: krallin

Differential Revision: D43941377

fbshipit-source-id: 9349c2e97b30dee5d99157c283bc2a8290659a04

app/buck2_server/BUCK

@@ -34,6 +34,7 @@ rust_library(
         "fbsource//third-party/rust:chrono",
         "fbsource//third-party/rust:constant_time_eq",
         "fbsource//third-party/rust:crossbeam-channel",
+        "fbsource//third-party/rust:debugserver-types",
         "fbsource//third-party/rust:flate2",
         "fbsource//third-party/rust:futures",
         "fbsource//third-party/rust:inferno",

app/buck2_server/Cargo.toml

@@ -5,6 +5,7 @@ name = "buck2_server"
 version = "0.1.0"
 
 [dependencies]
+debugserver-types = {workspace = true}
 anyhow = { workspace = true }
 async-recursion = { workspace = true }
 async-trait = { workspace = true }

app/buck2_server/src/starlark_debug/mod.rs

@@ -8,15 +8,23 @@
  */
 
 use std::sync::Arc;
+use std::sync::Mutex;
 
 use async_trait::async_trait;
+use buck2_core::fs::project::ProjectRoot;
 use buck2_events::dispatch::EventDispatcher;
 use buck2_interpreter::starlark_debug::StarlarkDebugController;
 use buck2_interpreter::starlark_debug::StarlarkDebuggerHandle;
 use dupe::Dupe;
+use tokio::sync::mpsc;
+
+use crate::starlark_debug::error::StarlarkDebuggerError;
+use crate::starlark_debug::run::ToClientMessage;
+use crate::starlark_debug::server::BuckStarlarkDebuggerServer;
 
 mod error;
 pub mod run;
+mod server;
 
 /// A handle to the debugger server.
 #[derive(Debug, Clone, Dupe)]
@@ -34,10 +42,46 @@ impl StarlarkDebuggerHandle for BuckStarlarkDebuggerHandle {
         unimplemented!("coming soon")
     }
 }
+
+/// We allow only a single debugger to be attached at a time. While it's attached, this will hold the server instance.
+static CURRENT_DEBUGGER: Mutex<Option<Arc<BuckStarlarkDebuggerServer>>> = Mutex::new(None);
+
 /// Used by each command to get a handle to the current debugger. The debugger server will capture the
 /// event dispatcher to send back debugger state snapshots (which indicate that the debugger is attached
 /// and which, if any, threads are paused) while the command is running.
-pub fn create_debugger_handle(_events: EventDispatcher) -> Option<BuckStarlarkDebuggerHandle> {
-    // We don't yet support a way to set the global debugger, so return no handle.
-    None
+pub fn create_debugger_handle(events: EventDispatcher) -> Option<BuckStarlarkDebuggerHandle> {
+    CURRENT_DEBUGGER
+        .lock()
+        .unwrap()
+        .as_ref()
+        .and_then(|v| v.new_handle(events))
+}
+
+/// Manages setting/unsetting the CURRENT_DEBUGGER while the starlark debug-attach command is running.
+struct ServerConnection(Arc<BuckStarlarkDebuggerServer>);
+
+impl ServerConnection {
+    fn new(
+        to_client_send: mpsc::UnboundedSender<ToClientMessage>,
+        project_root: ProjectRoot,
+    ) -> anyhow::Result<Self> {
+        let mut locked = CURRENT_DEBUGGER.lock().unwrap();
+        if locked.is_some() {
+            return Err(StarlarkDebuggerError::DebuggerAlreadyAttached.into());
+        }
+
+        let server = Arc::new(BuckStarlarkDebuggerServer::new(
+            to_client_send,
+            project_root,
+        ));
+        *locked = Some(server.dupe());
+        Ok(Self(server))
+    }
+}
+
+impl Drop for ServerConnection {
+    fn drop(&mut self) {
+        *CURRENT_DEBUGGER.lock().unwrap() = None;
+        let _ignored = self.0.detach();
+    }
 }

app/buck2_server/src/starlark_debug/run.rs

@@ -11,9 +11,36 @@ use buck2_events::dispatch::span_async;
 use buck2_server_ctx::command_end::command_end;
 use buck2_server_ctx::ctx::ServerCommandContextTrait;
 use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
+use debugserver_types as dap;
+use tokio::select;
+use tokio::sync::mpsc;
+use tokio_stream::StreamExt;
+use tracing::debug;
 
+use crate::starlark_debug::error::StarlarkDebuggerInternalError;
+use crate::starlark_debug::ServerConnection;
 use crate::streaming_request_handler::StreamingRequestHandler;
 
+/// Messages from the debugger server to its client (the cli `buck2 starlark debug-attach` which
+/// then forwards them along through its stdout).
+#[derive(Debug)]
+#[allow(unused)] // temporary
+pub(crate) enum ToClientMessage {
+    Event(dap::Event),
+    Response(dap::Response),
+    Shutdown(anyhow::Result<()>),
+}
+
+impl ToClientMessage {
+    fn pretty_string(&self) -> anyhow::Result<String> {
+        match self {
+            ToClientMessage::Event(ev) => Ok(serde_json::to_string_pretty(&ev)?),
+            ToClientMessage::Response(resp) => Ok(serde_json::to_string_pretty(&resp)?),
+            ToClientMessage::Shutdown(res) => Ok(format!("{:?}", res)),
+        }
+    }
+}
+
 /// Wraps `run_dap_server` with a command start/end span.
 pub(crate) async fn run_dap_server_command(
     ctx: Box<dyn ServerCommandContextTrait>,
@@ -38,9 +65,94 @@ pub(crate) async fn run_dap_server_command(
 }
 
 async fn run_dap_server(
-    _ctx: Box<dyn ServerCommandContextTrait>,
-    _partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::DapMessage>,
-    _req: StreamingRequestHandler<buck2_cli_proto::DapRequest>,
+    ctx: Box<dyn ServerCommandContextTrait>,
+    mut partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::DapMessage>,
+    mut req: StreamingRequestHandler<buck2_cli_proto::DapRequest>,
 ) -> anyhow::Result<buck2_cli_proto::DapResponse> {
-    unimplemented!("coming soon")
+    let (to_client_send, mut to_client_recv) = mpsc::unbounded_channel();
+    let server_connection = ServerConnection::new(to_client_send, ctx.project_root().clone())?;
+
+    let mut seq = 0;
+
+    // We can get requests from the client (e.g. vscode) and we can get responses/events from the server
+    // and mostly we just forward them along to the other side.
+    // If an error is encountered or the server sends the Shutdown message, we'll break out, disconnect
+    // from the server and return the final response.
+    let response = loop {
+        select! {
+            request = req.next() => {
+                let request = match request {
+                    Some(v) => v.map_err(|e| anyhow::anyhow!("debugserver req error: {}", e))?,
+                    None => {
+                        // client disconnected.
+                        break buck2_cli_proto::DapResponse {};
+                    }
+                };
+
+
+                let debugserver_req: dap::Request = serde_json::from_str(&request.dap_json)?;
+                debug!("received request {}", &serde_json::to_string_pretty(&debugserver_req)?);
+                server_connection.0.send_request(debugserver_req)?;
+            }
+            message = to_client_recv.recv() => {
+                match message {
+                    Some(v) => {
+                        if let Some(v) = handle_outgoing_message(&mut seq, &mut partial_result_dispatcher, v)? {
+                            break v;
+                        }
+                    }
+                    None => return Err(StarlarkDebuggerInternalError::UnexpectedDebuggerShutdown.into())
+                }
+            }
+        }
+    };
+
+    debug!("returning debugserver response");
+    Ok(response)
+}
+
+/// Converts the ToClientMessage to either a partial result (and dispatches it) and/or the final
+/// response (for a Shutdown message).
+fn handle_outgoing_message(
+    seq: &mut u32,
+    partial_result_dispatcher: &mut PartialResultDispatcher<buck2_cli_proto::DapMessage>,
+    message: ToClientMessage,
+) -> anyhow::Result<Option<buck2_cli_proto::DapResponse>> {
+    debug!("sending message {}", &message.pretty_string()?);
+
+    let this_seq = *seq as i64;
+    *seq += 1;
+
+    let (dap_json, response) = match message {
+        ToClientMessage::Event(mut ev) => {
+            ev.seq = this_seq;
+            (serde_json::to_vec(&ev)?, None)
+        }
+        ToClientMessage::Response(mut resp) => {
+            resp.seq = this_seq;
+            (serde_json::to_vec(&resp)?, None)
+        }
+        ToClientMessage::Shutdown(res) => {
+            let exit_code = match res {
+                Ok(..) => 0,
+                Err(e) => {
+                    debug!("server shutdown with error {:?}", &e);
+                    1
+                }
+            };
+
+            (
+                serde_json::to_vec(&dap::ExitedEvent {
+                    type_: "event".to_owned(),
+                    seq: this_seq,
+                    event: "exited".to_owned(),
+                    body: dap::ExitedEventBody { exit_code },
+                })?,
+                Some(buck2_cli_proto::DapResponse {}),
+            )
+        }
+    };
+
+    partial_result_dispatcher.emit(buck2_cli_proto::DapMessage { dap_json });
+    Ok(response)
 }

app/buck2_server/src/starlark_debug/server.rs

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under both the MIT license found in the
+ * LICENSE-MIT file in the root directory of this source tree and the Apache
+ * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+ * of this source tree.
+ */
+
+use std::sync::Arc;
+
+use buck2_core::fs::project::ProjectRoot;
+use buck2_events::dispatch::EventDispatcher;
+use debugserver_types as dap;
+use tokio::sync::mpsc;
+
+use crate::starlark_debug::run::ToClientMessage;
+use crate::starlark_debug::BuckStarlarkDebuggerHandle;
+
+/// The buck starlark debugger server. Most of the work is managed by the single-threaded server state.
+///
+/// There will be several references to the BuckStarlarkDebuggerServer instance and it will forward messages
+/// along to the state.
+#[derive(Debug)]
+pub(crate) struct BuckStarlarkDebuggerServer {}
+
+impl BuckStarlarkDebuggerServer {
+    pub(crate) fn new(
+        _to_client: mpsc::UnboundedSender<ToClientMessage>,
+        _project_root: ProjectRoot,
+    ) -> Self {
+        Self {}
+    }
+    pub(crate) fn new_handle(
+        self: &Arc<Self>,
+        _events: EventDispatcher,
+    ) -> Option<BuckStarlarkDebuggerHandle> {
+        unimplemented!("coming soon")
+    }
+
+    /// Called to forward along requests from the DAP client.
+    pub(crate) fn send_request(&self, _req: dap::Request) -> anyhow::Result<()> {
+        unimplemented!("coming soon")
+    }
+
+    /// Called when the DAP client has disconnected.
+    pub(crate) fn detach(&self) -> anyhow::Result<()> {
+        Ok(())
+    }
+}