unittests/ExecutionEngine/Orc/QueueChannel.h

   1 //===----------------------- Queue.h - RPC Queue ------------------*-c++-*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #ifndef LLVM_UNITTESTS_EXECUTIONENGINE_ORC_QUEUECHANNEL_H
  10 #define LLVM_UNITTESTS_EXECUTIONENGINE_ORC_QUEUECHANNEL_H
  11
  12 #include "llvm/ExecutionEngine/Orc/RawByteChannel.h"
  13 #include "llvm/Support/Error.h"
  14
  15 #include <condition_variable>
  16 #include <queue>
  17
  18 namespace llvm {
  19
  20 class QueueChannelError : public ErrorInfo<QueueChannelError> {
  21 public:
  22   static char ID;
  23 };
  24
  25 class QueueChannelClosedError
  26     : public ErrorInfo<QueueChannelClosedError, QueueChannelError> {
  27 public:
  28   static char ID;
  29   std::error_code convertToErrorCode() const override {
  30     return inconvertibleErrorCode();
  31   }
  32
  33   void log(raw_ostream &OS) const override {
  34     OS << "Queue closed";
  35   }
  36 };
  37
  38 class Queue : public std::queue<char> {
  39 public:
  40   using ErrorInjector = std::function<Error()>;
  41
  42   Queue()
  43     : ReadError([]() { return Error::success(); }),
  44       WriteError([]() { return Error::success(); }) {}
  45
  46   Queue(const Queue&) = delete;
  47   Queue& operator=(const Queue&) = delete;
  48   Queue(Queue&&) = delete;
  49   Queue& operator=(Queue&&) = delete;
  50
  51   std::mutex &getMutex() { return M; }
  52   std::condition_variable &getCondVar() { return CV; }
  53   Error checkReadError() { return ReadError(); }
  54   Error checkWriteError() { return WriteError(); }
  55   void setReadError(ErrorInjector NewReadError) {
  56     {
  57       std::lock_guard<std::mutex> Lock(M);
  58       ReadError = std::move(NewReadError);
  59     }
  60     CV.notify_one();
  61   }
  62   void setWriteError(ErrorInjector NewWriteError) {
  63     std::lock_guard<std::mutex> Lock(M);
  64     WriteError = std::move(NewWriteError);
  65   }
  66 private:
  67   std::mutex M;
  68   std::condition_variable CV;
  69   std::function<Error()> ReadError, WriteError;
  70 };
  71
  72 class QueueChannel : public orc::rpc::RawByteChannel {
  73 public:
  74   QueueChannel(std::shared_ptr<Queue> InQueue,
  75                std::shared_ptr<Queue> OutQueue)
  76       : InQueue(InQueue), OutQueue(OutQueue) {}
  77
  78   QueueChannel(const QueueChannel&) = delete;
  79   QueueChannel& operator=(const QueueChannel&) = delete;
  80   QueueChannel(QueueChannel&&) = delete;
  81   QueueChannel& operator=(QueueChannel&&) = delete;
  82
  83   Error readBytes(char *Dst, unsigned Size) override {
  84     std::unique_lock<std::mutex> Lock(InQueue->getMutex());
  85     while (Size) {
  86       {
  87         Error Err = InQueue->checkReadError();
  88         while (!Err && InQueue->empty()) {
  89           InQueue->getCondVar().wait(Lock);
  90           Err = InQueue->checkReadError();
  91         }
  92         if (Err)
  93           return Err;
  94       }
  95       *Dst++ = InQueue->front();
  96       --Size;
  97       ++NumRead;
  98       InQueue->pop();
  99     }
 100     return Error::success();
 101   }
 102
 103   Error appendBytes(const char *Src, unsigned Size) override {
 104     std::unique_lock<std::mutex> Lock(OutQueue->getMutex());
 105     while (Size--) {
 106       if (Error Err = OutQueue->checkWriteError())
 107         return Err;
 108       OutQueue->push(*Src++);
 109       ++NumWritten;
 110     }
 111     OutQueue->getCondVar().notify_one();
 112     return Error::success();
 113   }
 114
 115   Error send() override { return Error::success(); }
 116
 117   void close() {
 118     auto ChannelClosed = []() { return make_error<QueueChannelClosedError>(); };
 119     InQueue->setReadError(ChannelClosed);
 120     InQueue->setWriteError(ChannelClosed);
 121     OutQueue->setReadError(ChannelClosed);
 122     OutQueue->setWriteError(ChannelClosed);
 123   }
 124
 125   uint64_t NumWritten = 0;
 126   uint64_t NumRead = 0;
 127
 128 private:
 129
 130   std::shared_ptr<Queue> InQueue;
 131   std::shared_ptr<Queue> OutQueue;
 132 };
 133
 134 inline std::pair<std::unique_ptr<QueueChannel>, std::unique_ptr<QueueChannel>>
 135 createPairedQueueChannels() {
 136   auto Q1 = std::make_shared<Queue>();
 137   auto Q2 = std::make_shared<Queue>();
 138   auto C1 = std::make_unique<QueueChannel>(Q1, Q2);
 139   auto C2 = std::make_unique<QueueChannel>(Q2, Q1);
 140   return std::make_pair(std::move(C1), std::move(C2));
 141 }
 142
 143 }
 144
 145 #endif