grpc.hpp

Go to the documentation of this file.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License

#ifndef __PROCESS_GRPC_HPP__
#define __PROCESS_GRPC_HPP__

#include <atomic>
#include <chrono>
#include <memory>
#include <thread>
#include <type_traits>

#include <google/protobuf/message.h>

#include <grpc++/grpc++.h>

#include <process/future.hpp>
#include <process/owned.hpp>
#include <process/process.hpp>

#include <stout/duration.hpp>
#include <stout/lambda.hpp>
#include <stout/synchronized.hpp>
#include <stout/try.hpp>


// This file provides libprocess "support" for using gRPC. In
// particular, it defines two wrapper classes: `Channel` (representing a
// connection to a gRPC server) and `client::Runtime`, which integrates
// an event loop waiting for gRPC responses, and provides the `call`
// interface to create an asynchrous gRPC call and return a `Future`.


#define GRPC_RPC(service, rpc) \
  (&service::Stub::Async##rpc)

namespace process {
namespace grpc {

// Forward declarations.
namespace client { class Runtime; }


class Channel
{
public:
  Channel(const std::string& uri,
          const std::shared_ptr<::grpc::ChannelCredentials>& credentials =
            ::grpc::InsecureChannelCredentials())
    : channel(::grpc::CreateChannel(uri, credentials)) {}

private:
  std::shared_ptr<::grpc::Channel> channel;

  friend class client::Runtime;
};


template <typename T>
struct RpcResult
{
  RpcResult(const ::grpc::Status& _status, const T& _response)
    : status(_status), response(_response) {}

  ::grpc::Status status;
  T response;
};

namespace client {

class Runtime
{
public:
  Runtime() : data(new Data()) {}

  template <typename Stub, typename Request, typename Response>
  Future<RpcResult<Response>> call(
      const Channel& channel,
      std::unique_ptr<::grpc::ClientAsyncResponseReader<Response>>(Stub::*rpc)(
          ::grpc::ClientContext*,
          const Request&,
          ::grpc::CompletionQueue*),
      const Request& request)
  {
    static_assert(
        std::is_convertible<Request*, google::protobuf::Message*>::value,
        "Request must be a protobuf message");

    synchronized (data->lock) {
      if (data->terminating) {
        return Failure("Runtime has been terminated.");
      }

      std::shared_ptr<::grpc::ClientContext> context(
          new ::grpc::ClientContext());

      // TODO(chhsiao): Allow the caller to specify a timeout.
      context->set_deadline(
          std::chrono::system_clock::now() + std::chrono::seconds(5));

      // Create a `Promise` and a callback lambda as a tag and invokes
      // an asynchronous gRPC call through the `CompletionQueue`
      // managed by `data`. The `Promise` will be set by the callback
      // upon server response.
      std::shared_ptr<Promise<RpcResult<Response>>> promise(
          new Promise<RpcResult<Response>>);

      promise->future().onDiscard([=] { context->TryCancel(); });

      std::shared_ptr<Response> response(new Response());
      std::shared_ptr<::grpc::Status> status(new ::grpc::Status());

      std::shared_ptr<::grpc::ClientAsyncResponseReader<Response>> reader(
          (Stub(channel.channel).*rpc)(context.get(), request, &data->queue));

      reader->Finish(
          response.get(),
          status.get(),
          new lambda::function<void()>(
              // NOTE: `context` and `reader` need to be held on in
              // order to get updates for the ongoing RPC, and thus
              // are captured here. The lambda itself will later be
              // retrieved and managed in `Data::loop()`.
              [context, reader, response, status, promise]() {
                CHECK(promise->future().isPending());
                if (promise->future().hasDiscard()) {
                  promise->discard();
                  return;
                }

                promise->set(RpcResult<Response>(*status, *response));
              }));

      return promise->future();
    }
  }

  void terminate();

  Future<Nothing> wait();

private:
  struct Data
  {
    Data();
    ~Data();

    void loop();
    void terminate();

    std::unique_ptr<std::thread> looper;
    ::grpc::CompletionQueue queue;
    ProcessBase process;
    std::atomic_flag lock = ATOMIC_FLAG_INIT;
    bool terminating = false;
    Promise<Nothing> terminated;
  };

  std::shared_ptr<Data> data;
};

} // namespace client {

} // namespace grpc {
} // namespace process {

#endif // __PROCESS_GRPC_HPP__