[clangd] Ensure that we reply to each call exactly once. NFC (I think!)

[sammccall]

Summary

In debug builds, getting this wrong will trigger asserts.
In production builds, it will send an error reply if none was sent,
and drop redundant replies. (And log).

No tests because this is always a programming error.
(We did have some cases of this, but I fixed them with the new dispatcher).

[ioeric]

clangd/ClangdLSPServer.cpp

182

As discussed offline, we could potentially make this non-copyable to make harder to be duplicated and called more than once. And we could also get rid of the shared State in the functor.

[sammccall]

Make ReplyOnce move-only.

[ioeric]

clangd/ClangdLSPServer.cpp

112

Do we have guarantee that Tracer.Args outlives Reply?

[sammccall]

Clarify Span::Args lifetime.

[sammccall]

clangd/ClangdLSPServer.cpp

112

Yes, assuming the Handler propagates contexts properly.
The args are owned by the context created by the trace::Span.
Added a comment to Span to guarantee this.

clangd/ClangdLSPServer.cpp

else if (auto Handler = Notifications.lookup(Method))
       Handler(std::move(Params));
     else
       log("unhandled notification {0}", Method);
     return true;
   }
 
   bool onCall(StringRef Method, json::Value Params, json::Value ID) override {
     log("<-- {0}({1})", Method, ID);
+    // Calls can be canceled by the client. Add cancellation context.
+    WithContext WithCancel(cancelableRequestContext(ID));
+    ReplyOnce Reply(ID, Method, &Server);
     if (auto Handler = Calls.lookup(Method))
-      Handler(std::move(Params), std::move(ID));
+      Handler(std::move(Params), std::move(Reply));
     else
-      Server.reply(ID, llvm::make_error<LSPError>("method not found",
-                                                  ErrorCode::MethodNotFound));
+      Reply(llvm::make_error<LSPError>("method not found",
     return true;
   }
 
   bool onReply(json::Value ID, Expected<json::Value> Result) override {
     // We ignore replies, just log them.
     if (Result)
       log("<-- reply({0})", ID);
     else
       log("<-- reply({0}) error: {1}", ID, llvm::toString(Result.takeError()));
     return true;
   }
 
   // Bind an LSP method name to a call.
-  template <typename Param, typename Reply>
+  template <typename Param, typename Result>
   void bind(const char *Method,
-            void (ClangdLSPServer::*Handler)(const Param &, Callback<Reply>)) {
+            void (ClangdLSPServer::*Handler)(const Param &, Callback<Result>)) {
     Calls[Method] = [Method, Handler, this](json::Value RawParams,
-                                            json::Value ID) {
+                                            ReplyOnce Reply) {
       Param P;
       if (!fromJSON(RawParams, P)) {
         elog("Failed to decode {0} request.", Method);
-        Server.reply(ID, make_error<LSPError>("failed to decode request",
-                                              ErrorCode::InvalidRequest));
+        Reply(make_error<LSPError>("failed to decode request",
         return;
       }
       trace::Span Tracer(Method);
       SPAN_ATTACH(Tracer, "Params", RawParams);
       auto *Trace = Tracer.Args; // We attach reply from another thread.
-      // Calls can be canceled by the client. Add cancellation context.
-      WithContext WithCancel(cancelableRequestContext(ID));
-      // FIXME: this function should assert it's called exactly once.
-      (Server.*Handler)(P, [this, ID, Trace](llvm::Expected<Reply> Result) {
-        if (Result) {
+      (Server.*Handler)(P, [Trace, Reply](llvm::Expected<Result> Res) {
+        if (Res) {
           if (Trace)
-            (*Trace)["Reply"] = *Result;
-          Server.reply(ID, json::Value(std::move(*Result)));
+            (*Trace)["Reply"] = *Res;
+          Reply(json::Value(std::move(*Res)));
         } else {
-          auto Err = Result.takeError();
+          auto Err = Res.takeError();
           if (Trace)
             (*Trace)["Error"] = llvm::to_string(Err);
-          Server.reply(ID, std::move(Err));
+          Reply(std::move(Err));
         }
       });
     };
   }
 
   // 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](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 {
+    struct State {
+      std::atomic<bool> Replied = {false};
+      json::Value ID;
+      std::string Method;
+      ClangdLSPServer *Server;
+
+      State(const json::Value &ID, StringRef Method, ClangdLSPServer *Server)
+          : ID(ID), Method(Method), Server(Server) {}
+      ~State() {
+        if (!Replied) {
+          elog("No reply to message {0}({1})", Method, ID);
+          assert(false && "must reply to all calls!");
+          Server->reply(ID, make_error<LSPError>("server failed to reply",
+                                                 ErrorCode::InternalError));
+        }
+      }
+    };
+    std::shared_ptr<State> MyState;
+
+  public:
+    ReplyOnce(const json::Value &ID, StringRef Method, ClangdLSPServer *Server)
+        : MyState(std::make_shared<State>(ID, Method, Server)) {}
+
+    void operator()(llvm::Expected<json::Value> Reply) const {
+      if (MyState->Replied.exchange(true)) {
+        elog("Replied twice to message {0}({1})", MyState->Method, MyState->ID);
+        assert(false && "must reply to each call only once!");
+        return;
+      }
+      MyState->Server->reply(MyState->ID, std::move(Reply));
+    }
+  };
+
   llvm::StringMap<std::function<void(json::Value)>> Notifications;
-  llvm::StringMap<std::function<void(json::Value, json::Value)>> Calls;
+  llvm::StringMap<std::function<void(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) {