[clangd] Add a callback mechanism for handling responses from client.

[hokein]

Summary

The callback will be invoked in clangd when we receive a reply from the client.

This is a prerequisite of implementing a generic mechanism for chainable refactorings (e.g. extract variable and rename), this would allow server to
trigger a new request to the LSP client after receiving a reply from the client.

[ilya-biryukov]

Could you add a bit more context to the description of the change: why do we need the callback and what is the problem with what we have now?
I'm sure this was discussed offline and the change is justified, just wanted to make sure we write it down in the change history.

clang-tools-extra/clangd/ClangdLSPServer.cpp

334

NIT: simplify to

if (CB)
  return CB(std::move(Result));
// rest of the code.

[hokein]

In D65387#1604737, @ilya-biryukov wrote:

Could you add a bit more context to the description of the change: why do we need the callback and what is the problem with what we have now?
I'm sure this was discussed offline and the change is justified, just wanted to make sure we write it down in the change history.

sorry for missing the context here (updated the description). Yeah, this was a prerequisite of implementing a generic mechanism for chainable refactorings based on the offline discussion.

[jkorous]

Hi @hokein,
Do you have any thoughts on how to handle situation when client registers callback but doesn't send a request with registered ID later?

clang-tools-extra/clangd/ClangdLSPServer.cpp

317

Should we care if there's already another callback registered for this StrID?

[sammccall]

clang-tools-extra/clangd/ClangdLSPServer.cpp

150

We want to keep logging errors (not just at verbose level).

If you need a non-destructive way to get the error message for logging, I think it's unfortunately:

if (Result) {
  log();
  handle(move(Result));
} else {
  err = Result.takeError();
  log(err);
  handle(move(err));
}

Note I think our logging function has support for printing both rvalue errors (consuming them) and lvalue errors (not consuming them). But failing that, operator<< doesn't consume the error.

314

we know the underlying ID is an integer (because we generate it above), so why not just key on the integer directly?
In Reply, if the ID isn't an integer then we know there's no callback.

317

I think we should assert on this. NextCallID should never overflow.

317

as Jan points out, this is an unbounded leak if the client never replies to requests.

the most obvious approach to fix this is to bound the size and drop e.g. oldest pending query if we overflow. (we can exploit that IDs are sequential integers, if that's useful)

I mention this because if the size is bounded to say 100 and small in practice, you might end up wanting to use deque or (linked) list

336

elog, this is a protocol error

The message is too low level, it describes the implementation rather than the protocol. Prefer e.g. received a reply with ID {0}, but there was no such call

No need to handle error vs non-error here, it should be logged generally above (whether there was a missing callback or not)

526

While I like keeping the scope small, it's not really possible to evaluate this change without using it at least once.

Can you use the new facility to address this comment?
This would look something like:

• ApplyEdit takes a Callback<ApplyWorkspaceEditResponse>
• callsites look like this

ApplyEdit(*Params.workspaceEdit, bind(std::move(Reply), [](decltype(Reply) Reply, Expected<ApplyWorkspaceEditResponse> Response) {
  if (!Response)
    return Reply(Response.takeError());
  if (!Response->applied)
    return Reply(makeStringError("edits not applied: " + Response->failureReason));
  return Reply("Fix applied");
});

(This is pretty verbose, but seems realistic for what we're actually going to do).

This function is a bit of a complicated monster, so feel free to just handle apply fix, or (better) to factor the code differently so it's reasonable to handle both.

clang-tools-extra/clangd/ClangdLSPServer.h

152

this type is just clangd::Callback<llvm::json::Value>(). (modulo supporting move-only captures, which we want I think)
I'm not sure that's worth abbreviating.

155

the corresponding state for request cancellation lives in MessageHandler, I think this one should too.

Similarly, onReply fits better in MessageHandler than here. call() belongs here though, so I guess call() needs to call something like MessageHandler::bindReply(ID, Callback<T>) or so

157

This API makes the caller of call() responsible for dealing with converting the arbitrary JSON returned by the client into the protocol object, and handling errors (which they'll do inconsistently).

I'd suggest instead a signature like

template <typename Rsp=json::Value>
void call(StringRef Method, json::Value Params, Callback<Rsp> CB=nullptr)

where the template body would wrap CB with conversion and error-handling logic, and then add the "generic" Callback<json::Value> to the map.

you can see bind() - this is similar to how incoming messages are handled.

Making this a template directly forces it to be in the header: to mitigate this you can delegating to a non-template callImpl() as soon as you've erased the type to Callback<llvm::Value>, and/or putting the type-specific call() in MessageHandler (which seems pretty weird to me, so maybe not?)

[hokein]

address comments:

• move the reply callback handling code to MessageHandler;
• prevent memory leakage when LSP clients don't send back the reply;
• use the new machanism for applyTweak and apply fix;
• update all relevant tests;

[hokein]

clang-tools-extra/clangd/ClangdLSPServer.cpp

223

hmm, we will trigger this assertion when clangd exits but not all requests are replied by the client, options:

1. remove the assertion here, but the error message "server failed to reply" returned is misleading, in this case, this is a client failure, not server
2. make sure that all reply callbacks are called, e.g. run the pending reply callback in destructor (returning an LSP error to client)

[sammccall]

clang-tools-extra/clangd/ClangdLSPServer.cpp

223

The latter is pretty, but may be complicated (shutdown protocol is messy and involves receiving multiple messages, what happens if a reply arrives in between?)

I would suggest changing the condition here: if the server's destructor has started then don't attempt to reply. You can have an atomic<bool> in ClangdLSPServer that's set at the top of the destructor

[hokein]

• adjust ReplyOnce assertion logic
• add more tests

[hokein]

some tweaks

[sammccall]

Thanks, this looks a lot better. There are a bunch of details that can be simplified and names that could be clearer, but I think this is the right shape.

clang-tools-extra/clangd/ClangdLSPServer.cpp

161

nit: maybe call this ReplyHandler?

165

this got me thinking - we're using an atomic<int> to assign sequential IDs, but then locking to put them into a map. The reward for using two types of concurrency primitives is our array is almost-but-not-quite sorted.

I think we should demote NextCallID to a regular int, and guard it with the same mutex as ReplyCallbacks. (Just call it CallMutex?) Then bindReply can lock the mutex, pick a call id, add the callback to the map, and return the ID.

I still don't think we need to binary search though. The list is 100 at most and ~1 in practice.

167

nit: a normal loop with if/break seems clearer than find_if here

177

Might just be me, but I think I'd find it slightly to avoid the LogAndRun lambda by overwriting the callback:

// If there was no callback in the table, the client is sending us an unsolicited reply!
if (FoundCB == nullptr)
  Found = [&ID](...) { elog and consume error; }

// inline LogAndRun here

181

elog("more than {0} outstanding LSP calls, forgetting about {1}")

184

nit: we tend to use decltype(Result) to specify the type of Result when we're pseudo-capturing it in a lambda using Bind(). That's not the case here, so please spell out the type.

187

you're returning a string to the callback, I think you meant to return an error?

221

Need a comment about the server being destroyed/unreplied calls case.

258

nit: physically this is a deque, but it functions more like a map and I think it might be easier to refer to that way.

261

nit: I think this part of the comment belongs in the implementation rather than on the variable

268

hmm, maybe just use pair<int, Callback<Value>> here? That way it looks even more like a map

530

I can't parse this name.

Looking at the code... it's doing something like ReplyAfterApplyingEdits

It should take the success message as a parameter, I think - at the moment you're replying "fix applied" even in the tweak case

530

So this method is already a tangled mess, and this patch makes it worse. However, it's *conceptually* exactly the right API, just callbacks and Expected and LSP conspire to be awkward here.
We should refactor this, but this patch isn't the right time. Can you add a //FIXME: this method is tangled and confusing, refactor it or so?

541

nit: edits were not applied (or edits failed to apply)

clang-tools-extra/clangd/ClangdLSPServer.h

54

MessageHandler is a nested class and so has access to Destructing already, no need for this function

136

nit: Destroying
but I think IsBeingDestroyed would be clearer, as destroy is transitive

this needs a comment

153

nit: HandleReply?
lambdas often have names like functions, so this reads like it's a function that replies

168

nit: I know I suggested this name but I think callRaw would be better

clang-tools-extra/clangd/Protocol.h

877 ↗

(On Diff #212796)

needs a default value

= true, I guess

clang-tools-extra/clangd/test/request-reply.test

5 ↗

(On Diff #212796)

nit, extra line

22 ↗

(On Diff #212796)

please use increasing IDs and don't reuse them. (In general, but particularly in this test)

23 ↗

(On Diff #212796)

can you verify that the error comes with the right ID? (and similar below)

[hokein]

address comments.

[hokein]

thanks for the detailed comments!

clang-tools-extra/clangd/test/request-reply.test

22 ↗

(On Diff #212796)

I believe 0 is right here, this is a client reply for the server request 0 (applyEdit).

[sammccall]

clang-tools-extra/clangd/ClangdLSPServer.cpp

175

nit: I think this function could return json::Value, encapsulating the fact that we use integers inside the MessageHandler class.

(Or does something outside the class depend on this?)

198

There's no complicated flow control around this - I'd suggest instead leaving ID uninitialized and letting msan flag this if it goes wrong.

221

You've added a comment that repeats what the code does, please don't do that :-)

Instead, explain the context around this case. e.g.
"There's one legitimate reason to never reply to a request: clangd's request handler sent a call to the client (e.g. applyWorkspaceEdit) and the client never replied. In this case, the ReplyOnce is owned by ClangdServer's reply callback table and is destroyed along with the server. We don't attempt to send a reply in this case, there's little to be gained from doing so."

256

"held" or "to hold"

260

more specific name: MaxReplyCallbacks or even better MaxOutstandingCalls?

clang-tools-extra/clangd/test/fixits-command.test

206 ↗

(On Diff #213304)

this ID is reused

clang-tools-extra/clangd/test/request-reply.test

28 ↗

(On Diff #213304)

ID reuse here (4 was previously used on line 6)

[hokein]

address comments.

[hokein]

clang-tools-extra/clangd/ClangdLSPServer.cpp

175

we just use it for logging, changing it to Value is totally fine.

clang-tools-extra/clangd/test/fixits-command.test

206 ↗

(On Diff #213304)

this ID is not reused, it is for server request applyEdit (see line 168).

clang-tools-extra/clangd/test/request-reply.test

28 ↗

(On Diff #213304)

good catch.

[sammccall]

clang-tools-extra/clangd/test/fixits-command.test

206 ↗

(On Diff #213304)

Doh, sorry :-(

[thakis]

clang-tools-extra/trunk/clangd/test/request-reply.test

6 ↗

(On Diff #213337)

FYI, referring to a file opened as test:///foo.cpp as file:///clangd-test/file.cpp is wrong on Windows. I fixed this in rL374746.

clang-tools-extra/clangd/ClangdLSPServer.h

public:
   ~ClangdLSPServer();
 
   /// Run LSP server loop, communicating with the Transport provided in the
   /// constructor. This method must not be executed more than once.
   ///
   /// \return Whether we shut down cleanly with a 'shutdown' -> 'exit' sequence.
   bool run();
 
 private:
   // Implement DiagnosticsConsumer.
   void onDiagnosticsReady(PathRef File, std::vector<Diag> Diagnostics) override;
   void onFileUpdated(PathRef File, const TUStatus &Status) override;
   void
   onHighlightingsReady(PathRef File,
                        std::vector<HighlightingToken> Highlightings) override;
 
   // LSP methods. Notifications have signature void(const Params&).
private:
   /// Sends a "publishDiagnostics" notification to the LSP client.
   void publishDiagnostics(const URIForFile &File,
                           std::vector<clangd::Diagnostic> Diagnostics);
 
   /// Used to indicate that the 'shutdown' request was received from the
   /// Language Server client.
   bool ShutdownRequestReceived = false;
 
   std::mutex FixItsMutex;
   typedef std::map<clangd::Diagnostic, std::vector<Fix>, LSPDiagnosticCompare>
       DiagnosticToReplacementMap;
   /// Caches FixIts per file and diagnostics
   llvm::StringMap<DiagnosticToReplacementMap> FixItsMap;
 
   // Most code should not deal with Transport directly.
   // MessageHandler deals with incoming messages, use call() etc for outgoing.
   clangd::Transport &Transp;
   class MessageHandler;
   std::unique_ptr<MessageHandler> MsgHandler;
   std::atomic<int> NextCallID = {0};
   std::mutex TranspWriter;
-  void call(StringRef Method, llvm::json::Value Params);
+
+  // A callback that will be invoked when receiving a client reply for a server
+  // request.
+  using ReplyCallback = std::function<void(llvm::Expected<llvm::json::Value>)>;
+  std::mutex ReplyCallbacksMutex;
+  // Request ID => reply Callbacks
+  llvm::StringMap<ReplyCallback> ReplyCallbacks;
+
+  void call(StringRef Method, llvm::json::Value Params, ReplyCallback CB);
+  void onReply(llvm::json::Value ID, llvm::Expected<llvm::json::Value> Result);
   void notify(StringRef Method, llvm::json::Value Params);
 
   const FileSystemProvider &FSProvider;
   /// Options used for code completion
   clangd::CodeCompleteOptions CCOpts;
   /// Options used for diagnostics.
   ClangdDiagnosticOptions DiagOpts;
   /// The supported kinds of the client.
   SymbolKindBitset SupportedSymbolKinds;
   /// The supported completion item kinds of the client.
   CompletionItemKindBitset SupportedCompletionItemKinds;
   /// Whether the client supports CodeAction response objects.
   bool SupportsCodeAction = false;
   /// From capabilities of textDocument/documentSymbol.
   bool SupportsHierarchicalDocumentSymbol = false;
   /// Whether the client supports showing file status.
   bool SupportFileStatus = false;
   /// Which kind of markup should we use in textDocument/hover responses.

clang-tools-extra/clangd/ClangdLSPServer.cpp

else
       Reply(llvm::make_error<LSPError>("method not found",
                                        ErrorCode::MethodNotFound));
     return true;
   }
 
   bool onReply(llvm::json::Value ID,
                llvm::Expected<llvm::json::Value> Result) override {
     WithContext HandlerContext(handlerContext());
-    // We ignore replies, just log them.
-    if (Result)
-      log("<-- reply({0})", ID);
-    else
-      log("<-- reply({0}) error: {1}", ID, llvm::toString(Result.takeError()));
+    // Log the reply.
+    Server.onReply(std::move(ID), std::move(Result));
     return true;
   }
 
   // Bind an LSP method name to a call.
   template <typename Param, typename Result>
   void bind(const char *Method,
             void (ClangdLSPServer::*Handler)(const Param &, Callback<Result>)) {
     Calls[Method] = [Method, Handler, this](llvm::json::Value RawParams,
                                             ReplyOnce Reply) {
       Param P;
       if (fromJSON(RawParams, P)) {
         (Server.*Handler)(P, std::move(Reply));
       } else {
         elog("Failed to decode {0} request.", Method);
         Reply(llvm::make_error<LSPError>("failed to decode request",
                                          ErrorCode::InvalidRequest));
       }
     };
   }
 
   // Bind an LSP method name to a notification.
   template <typename Param>
   void bind(const char *Method,
             void (ClangdLSPServer::*Handler)(const Param &)) {
     Notifications[Method] = [Method, Handler,
                              this](llvm::json::Value RawParams) {
       Param P;
       if (!fromJSON(RawParams, P)) {
         elog("Failed to decode {0} request.", Method);
         return;
       }
       trace::Span Tracer(Method);
       SPAN_ATTACH(Tracer, "Params", RawParams);
       (Server.*Handler)(P);
     };
   }
 
 private:
   // Function object to reply to an LSP call.
   // Each instance must be called exactly once, otherwise:
   //  - the bug is logged, and (in debug mode) an assert will fire
   //  - if there was no reply, an error reply is sent
   //  - if there were multiple replies, only the first is sent
   class ReplyOnce {
     std::atomic<bool> Replied = {false};
     std::chrono::steady_clock::time_point Start;
     llvm::json::Value ID;
     std::string Method;
     ClangdLSPServer *Server; // Null when moved-from.
     llvm::json::Object *TraceArgs;
 
   public:
     ReplyOnce(const llvm::json::Value &ID, llvm::StringRef Method,
               ClangdLSPServer *Server, llvm::json::Object *TraceArgs)
         : Start(std::chrono::steady_clock::now()), ID(ID), Method(Method),
           Server(Server), TraceArgs(TraceArgs) {
       assert(Server);
     }
     ReplyOnce(ReplyOnce &&Other)
         : Replied(Other.Replied.load()), Start(Other.Start),
           ID(std::move(Other.ID)), Method(std::move(Other.Method)),
           Server(Other.Server), TraceArgs(Other.TraceArgs) {
       Other.Server = nullptr;
     }
     ReplyOnce &operator=(ReplyOnce &&) = delete;
     ReplyOnce(const ReplyOnce &) = delete;
     ReplyOnce &operator=(const ReplyOnce &) = delete;
 
     ~ReplyOnce() {
       if (Server && !Replied) {
         elog("No reply to message {0}({1})", Method, ID);
         assert(false && "must reply to all calls!");
         (*this)(llvm::make_error<LSPError>("server failed to reply",
                                            ErrorCode::InternalError));
       }
     }
 
     void operator()(llvm::Expected<llvm::json::Value> Reply) {
       assert(Server && "moved-from!");
       if (Replied.exchange(true)) {
void operator()(llvm::Expected<llvm::json::Value> Reply) {
         std::lock_guard<std::mutex> Lock(Server->TranspWriter);
         Server->Transp.reply(std::move(ID), std::move(Err));
       }
     }
   };
 
   llvm::StringMap<std::function<void(llvm::json::Value)>> Notifications;
   llvm::StringMap<std::function<void(llvm::json::Value, ReplyOnce)>> Calls;
 
   // Method calls may be cancelled by ID, so keep track of their state.
   // This needs a mutex: handlers may finish on a different thread, and that's
   // when we clean up entries in the map.
   mutable std::mutex RequestCancelersMutex;
   llvm::StringMap<std::pair<Canceler, /*Cookie*/ unsigned>> RequestCancelers;
   unsigned NextRequestCookie = 0; // To disambiguate reused IDs, see below.
   void onCancel(const llvm::json::Value &Params) {
     const llvm::json::Value *ID = nullptr;
     if (auto *O = Params.getAsObject())
       ID = O->get("id");
     if (!ID) {
       elog("Bad cancellation request: {0}", Params);
       return;
     }
     auto StrID = llvm::to_string(*ID);
     std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
     auto It = RequestCancelers.find(StrID);
     if (It != RequestCancelers.end())
       It->second.first(); // Invoke the canceler.
   }
return Task.first.derive(llvm::make_scope_exit([this, StrID, Cookie] {
         RequestCancelers.erase(It);
     }));
   }
 
   ClangdLSPServer &Server;
 };
 
 // call(), notify(), and reply() wrap the Transport, adding logging and locking.
-void ClangdLSPServer::call(llvm::StringRef Method, llvm::json::Value Params) {
+void ClangdLSPServer::call(llvm::StringRef Method, llvm::json::Value Params,
+                           ReplyCallback CB) {
   auto ID = NextCallID++;
+  auto StrID = llvm::to_string(ID);
+  {
+    std::lock_guard<std::mutex> Lock(ReplyCallbacksMutex);
+    ReplyCallbacks[StrID] = std::move(CB);
+  }
   log("--> {0}({1})", Method, ID);
-  // We currently don't handle responses, so no need to store ID anywhere.
   std::lock_guard<std::mutex> Lock(TranspWriter);
   Transp.call(Method, std::move(Params), ID);
 }
 
+void ClangdLSPServer::onReply(llvm::json::Value ID,
+                              llvm::Expected<llvm::json::Value> Result) {
+  auto StrID = llvm::to_string(ID);
+  ReplyCallback CB = nullptr;
+  {
+    std::lock_guard<std::mutex> Lock(ReplyCallbacksMutex);
+    CB = ReplyCallbacks.lookup(StrID);
+    ReplyCallbacks.erase(StrID);
+  }
+  // No callback for the client reply, just log them.
+  if (!CB) {
+    if (Result)
+      log("no corresponding callback for the reply ({0})", ID);
+    else
+      log("no corresponding callback for the reply ({0}) error: {1}", ID,
+          llvm::toString(Result.takeError()));
+    return;
+  }
+  CB(std::move(Result));
+}
+
 void ClangdLSPServer::notify(llvm::StringRef Method, llvm::json::Value Params) {
   log("--> {0}", Method);
   std::lock_guard<std::mutex> Lock(TranspWriter);
   Transp.notify(Method, std::move(Params));
 }
 
 void ClangdLSPServer::onInitialize(const InitializeParams &Params,
                                    Callback<llvm::json::Value> Reply) {
void ClangdLSPServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
   Server->onFileEvent(Params);
 }
 
 void ClangdLSPServer::onCommand(const ExecuteCommandParams &Params,
                                 Callback<llvm::json::Value> Reply) {
   auto ApplyEdit = [this](WorkspaceEdit WE) {
     ApplyWorkspaceEditParams Edit;
     Edit.edit = std::move(WE);
     // Ideally, we would wait for the response and if there is no error, we
     // would reply success/failure to the original RPC.
-    call("workspace/applyEdit", Edit);
+    call("workspace/applyEdit", Edit, /*ReplyCallback=*/nullptr);
   };
   if (Params.command == ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND &&
       Params.workspaceEdit) {
     // The flow for "apply-fix" :
     // 1. We publish a diagnostic, including fixits
     // 2. The user clicks on the diagnostic, the editor asks us for code actions
     // 3. We send code actions, with the fixit embedded as context
     // 4. The user selects the fixit, the editor asks us to apply it
     // 5. We unwrap the changes and send them back to the editor
     // 6. The editor applies the changes (applyEdit), and sends us a reply (but
     // we ignore it)
 
     Reply("Fix applied.");
     ApplyEdit(*Params.workspaceEdit);
   } else if (Params.command == ExecuteCommandParams::CLANGD_APPLY_TWEAK &&
              Params.tweakArgs) {
     auto Code = DraftMgr.getDraft(Params.tweakArgs->file.file());
     if (!Code)
       return Reply(llvm::createStringError(
           llvm::inconvertibleErrorCode(),
           "trying to apply a code action for a non-added file"));