Orc/TargetProcess/SimpleRemoteEPCServer.h

0001 //===---- SimpleRemoteEPCServer.h - EPC over abstract channel ---*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 // EPC over simple abstract channel.
0010 //
0011 //===----------------------------------------------------------------------===//
0012
0013 #ifndef LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_SIMPLEREMOTEEPCSERVER_H
0014 #define LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_SIMPLEREMOTEEPCSERVER_H
0015
0016 #include "llvm/ADT/DenseMap.h"
0017 #include "llvm/ADT/FunctionExtras.h"
0018 #include "llvm/Config/llvm-config.h"
0019 #include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
0020 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
0021 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
0022 #include "llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h"
0023 #include "llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorDylibManager.h"
0024 #include "llvm/Support/DynamicLibrary.h"
0025 #include "llvm/Support/Error.h"
0026
0027 #include <condition_variable>
0028 #include <future>
0029 #include <memory>
0030 #include <mutex>
0031
0032 namespace llvm {
0033 namespace orc {
0034
0035 /// A simple EPC server implementation.
0036 class SimpleRemoteEPCServer : public SimpleRemoteEPCTransportClient {
0037 public:
0038   using ReportErrorFunction = unique_function<void(Error)>;
0039
0040   /// Dispatches calls to runWrapper.
0041   class Dispatcher {
0042   public:
0043     virtual ~Dispatcher();
0044     virtual void dispatch(unique_function<void()> Work) = 0;
0045     virtual void shutdown() = 0;
0046   };
0047
0048 #if LLVM_ENABLE_THREADS
0049   class ThreadDispatcher : public Dispatcher {
0050   public:
0051     void dispatch(unique_function<void()> Work) override;
0052     void shutdown() override;
0053
0054   private:
0055     std::mutex DispatchMutex;
0056     bool Running = true;
0057     size_t Outstanding = 0;
0058     std::condition_variable OutstandingCV;
0059   };
0060 #endif
0061
0062   class Setup {
0063     friend class SimpleRemoteEPCServer;
0064
0065   public:
0066     SimpleRemoteEPCServer &server() { return S; }
0067     StringMap<std::vector<char>> &bootstrapMap() { return BootstrapMap; }
0068     template <typename T, typename SPSTagT>
0069     void setBootstrapMapValue(std::string Key, const T &Value) {
0070       std::vector<char> Buffer;
0071       Buffer.resize(shared::SPSArgList<SPSTagT>::size(Value));
0072       shared::SPSOutputBuffer OB(Buffer.data(), Buffer.size());
0073       bool Success = shared::SPSArgList<SPSTagT>::serialize(OB, Value);
0074       (void)Success;
0075       assert(Success && "Bootstrap map value serialization failed");
0076       BootstrapMap[std::move(Key)] = std::move(Buffer);
0077     }
0078     StringMap<ExecutorAddr> &bootstrapSymbols() { return BootstrapSymbols; }
0079     std::vector<std::unique_ptr<ExecutorBootstrapService>> &services() {
0080       return Services;
0081     }
0082     void setDispatcher(std::unique_ptr<Dispatcher> D) { S.D = std::move(D); }
0083     void setErrorReporter(unique_function<void(Error)> ReportError) {
0084       S.ReportError = std::move(ReportError);
0085     }
0086
0087   private:
0088     Setup(SimpleRemoteEPCServer &S) : S(S) {}
0089     SimpleRemoteEPCServer &S;
0090     StringMap<std::vector<char>> BootstrapMap;
0091     StringMap<ExecutorAddr> BootstrapSymbols;
0092     std::vector<std::unique_ptr<ExecutorBootstrapService>> Services;
0093   };
0094
0095   static StringMap<ExecutorAddr> defaultBootstrapSymbols();
0096
0097   template <typename TransportT, typename... TransportTCtorArgTs>
0098   static Expected<std::unique_ptr<SimpleRemoteEPCServer>>
0099   Create(unique_function<Error(Setup &S)> SetupFunction,
0100          TransportTCtorArgTs &&...TransportTCtorArgs) {
0101     auto Server = std::make_unique<SimpleRemoteEPCServer>();
0102     Setup S(*Server);
0103     if (auto Err = SetupFunction(S))
0104       return std::move(Err);
0105
0106     // Set ReportError up-front so that it can be used if construction
0107     // process fails.
0108     if (!Server->ReportError)
0109       Server->ReportError = [](Error Err) {
0110         logAllUnhandledErrors(std::move(Err), errs(), "SimpleRemoteEPCServer ");
0111       };
0112
0113     // Attempt to create transport.
0114     auto T = TransportT::Create(
0115         *Server, std::forward<TransportTCtorArgTs>(TransportTCtorArgs)...);
0116     if (!T)
0117       return T.takeError();
0118     Server->T = std::move(*T);
0119     if (auto Err = Server->T->start())
0120       return std::move(Err);
0121
0122     // If transport creation succeeds then start up services.
0123     Server->Services = std::move(S.services());
0124     Server->Services.push_back(
0125         std::make_unique<rt_bootstrap::SimpleExecutorDylibManager>());
0126     for (auto &Service : Server->Services)
0127       Service->addBootstrapSymbols(S.bootstrapSymbols());
0128
0129     if (auto Err = Server->sendSetupMessage(std::move(S.BootstrapMap),
0130                                             std::move(S.BootstrapSymbols)))
0131       return std::move(Err);
0132     return std::move(Server);
0133   }
0134
0135   /// Set an error reporter for this server.
0136   void setErrorReporter(ReportErrorFunction ReportError) {
0137     this->ReportError = std::move(ReportError);
0138   }
0139
0140   /// Call to handle an incoming message.
0141   ///
0142   /// Returns 'Disconnect' if the message is a 'detach' message from the remote
0143   /// otherwise returns 'Continue'. If the server has moved to an error state,
0144   /// returns an error, which should be reported and treated as a 'Disconnect'.
0145   Expected<HandleMessageAction>
0146   handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr,
0147                 SimpleRemoteEPCArgBytesVector ArgBytes) override;
0148
0149   Error waitForDisconnect();
0150
0151   void handleDisconnect(Error Err) override;
0152
0153 private:
0154   Error sendMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
0155                     ExecutorAddr TagAddr, ArrayRef<char> ArgBytes);
0156
0157   Error sendSetupMessage(StringMap<std::vector<char>> BootstrapMap,
0158                          StringMap<ExecutorAddr> BootstrapSymbols);
0159
0160   Error handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
0161                      SimpleRemoteEPCArgBytesVector ArgBytes);
0162   void handleCallWrapper(uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
0163                          SimpleRemoteEPCArgBytesVector ArgBytes);
0164
0165   shared::WrapperFunctionResult
0166   doJITDispatch(const void *FnTag, const char *ArgData, size_t ArgSize);
0167
0168   static shared::CWrapperFunctionResult jitDispatchEntry(void *DispatchCtx,
0169                                                          const void *FnTag,
0170                                                          const char *ArgData,
0171                                                          size_t ArgSize);
0172
0173   uint64_t getNextSeqNo() { return NextSeqNo++; }
0174   void releaseSeqNo(uint64_t) {}
0175
0176   using PendingJITDispatchResultsMap =
0177       DenseMap<uint64_t, std::promise<shared::WrapperFunctionResult> *>;
0178
0179   std::mutex ServerStateMutex;
0180   std::condition_variable ShutdownCV;
0181   enum { ServerRunning, ServerShuttingDown, ServerShutDown } RunState;
0182   Error ShutdownErr = Error::success();
0183   std::unique_ptr<SimpleRemoteEPCTransport> T;
0184   std::unique_ptr<Dispatcher> D;
0185   std::vector<std::unique_ptr<ExecutorBootstrapService>> Services;
0186   ReportErrorFunction ReportError;
0187
0188   uint64_t NextSeqNo = 0;
0189   PendingJITDispatchResultsMap PendingJITDispatchResults;
0190   std::vector<sys::DynamicLibrary> Dylibs;
0191 };
0192
0193 } // end namespace orc
0194 } // end namespace llvm
0195
0196 #endif // LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_SIMPLEREMOTEEPCSERVER_H