request_result_rpc_client.h

Go to the documentation of this file.

// Copyright 2016 Google Inc.

#ifndef PAGESPEED_CONTROLLER_REQUEST_RESULT_RPC_CLIENT_H_
#define PAGESPEED_CONTROLLER_REQUEST_RESULT_RPC_CLIENT_H_

#include <memory>

#include "base/logging.h"
#include "pagespeed/controller/controller.grpc.pb.h"
#include "pagespeed/kernel/base/abstract_mutex.h"
#include "pagespeed/kernel/base/function.h"
#include "pagespeed/kernel/base/message_handler.h"
#include "pagespeed/kernel/base/thread_system.h"
#include "pagespeed/kernel/base/thread_annotations.h"
#include "pagespeed/kernel/util/grpc.h"


namespace net_instaweb {

template <typename ReaderWriter>
class RpcHolder {
 public:
  RpcHolder(MessageHandler* handler) : handler_(handler) {}

  Function* CallbackForAsyncCleanup() {
    return MakeFunction(this, &RpcHolder::Finish, &RpcHolder::Error);
  }

  void Finish() {
    rw()->Finish(&status_, MakeFunction(this, &RpcHolder::FinishSucceeded,
                                        &RpcHolder::FinishFailed));
  }

  ::grpc::ClientContext* context() { return &context_; }

  MessageHandler* handler() { return handler_; }

  ReaderWriter* rw() { return rw_.get(); }

  void SetReaderWriter(std::unique_ptr<ReaderWriter> rw) {
    DCHECK(rw_ == nullptr);
    rw_ = std::move(rw);
  }

 private:
  void FinishSucceeded() {
    if (status_.error_code() != ::grpc::StatusCode::OK &&
        status_.error_code() != ::grpc::StatusCode::CANCELLED) {
      MessageType severity =
#ifndef NDEBUG
          (status_.error_code() == ::grpc::StatusCode::ABORTED) ? kFatal :
#endif
                                                                kWarning;
      handler_->Message(severity,
                        "Received error status from CentralController: %d (%s)",
                        status_.error_code(), status_.error_message().c_str());
    }
    delete this;
  }

  void FinishFailed() {
    PS_LOG_WARN(handler_, "RpcHolder Finish failed");
    delete this;
  }

  void Error() {
    PS_LOG_WARN(handler_, "RpcHolder cleanup to CentralController failed");
    Finish(); 
  }

  MessageHandler* handler_;
  ::grpc::ClientContext context_;
  std::unique_ptr<ReaderWriter> rw_;
  ::grpc::Status status_;
};

template <typename RequestT, typename ResponseT, typename CallbackT>
class RequestResultRpcClient {
 public:
  typedef ::grpc::ClientAsyncReaderWriterInterface<RequestT, ResponseT>
      ReaderWriter;

  RequestResultRpcClient(
      ::grpc::CompletionQueue* queue, ThreadSystem* thread_system,
      MessageHandler* handler, CallbackT* callback)
      : mutex_(thread_system->NewMutex()),
        queue_(queue),
        callback_(callback),
        rpc_(new RpcHolder<ReaderWriter>(handler)) {
    CHECK(callback_ != nullptr);
  }

  virtual ~RequestResultRpcClient() {
  }

  void Start(grpc::CentralControllerRpcService::StubInterface* stub) {
    ScopedMutex lock(mutex_.get());
    rpc_->SetReaderWriter(
        StartRpc(stub, rpc_->context(), queue_,
                 MakeFunction(this, &RequestResultRpcClient::BootStrapFinished,
                              &RequestResultRpcClient::StartUpFailed)));
  }

  void SendResultToServer(const RequestT& response) {
    ScopedMutex lock(mutex_.get());
    if (rpc_ == nullptr) {
      return;
    }
    RpcHolder<ReaderWriter>* rpc = rpc_.release();
    rpc->rw()->Write(response, rpc->CallbackForAsyncCleanup());
  }

 private:
  virtual std::unique_ptr<ReaderWriter> StartRpc(
      grpc::CentralControllerRpcService::StubInterface* stub,
      ::grpc::ClientContext* context, ::grpc::CompletionQueue* queue,
      void* tag) = 0;

  virtual void PopulateServerRequest(RequestT* request) = 0;

  void BootStrapFinished() {
    ScopedMutex lock(mutex_.get());
    RequestT req;
    PopulateServerRequest(&req);
    rpc_->rw()->Write(
        req,
        MakeFunction(this, &RequestResultRpcClient::WriteServerRequestComplete,
                     &RequestResultRpcClient::StartUpFailed));
  }

  void WriteServerRequestComplete() {
    ScopedMutex lock(mutex_.get());
    rpc_->rw()->Read(
        &resp_, MakeFunction(
                    this, &RequestResultRpcClient::NotifyClientOfServerDecision,
                    &RequestResultRpcClient::StartUpFailed));
  }

  void NotifyClientOfServerDecision() {
    ScopedMutex lock(mutex_.get());
    DCHECK(rpc_ != nullptr);
    bool ok_to_proceed = resp_.ok_to_proceed();
    resp_.Clear();

    CallbackT* cb = callback_;
    callback_ = nullptr;

    if (ok_to_proceed) {
      lock.Release();
      cb->CallRun();
      return;
    } else {
      rpc_.reset();
      lock.Release();
      cb->CallCancel();
      return;
    }
  }

  void StartUpFailed() LOCKS_EXCLUDED(mutex_) {
    ScopedMutex lock(mutex_.get());
    DCHECK(rpc_ != nullptr);
    if (rpc_ != nullptr) {
      PS_LOG_WARN(rpc_->handler(),
                  "Couldn't get response from CentralController");
      RpcHolder<ReaderWriter>* rpc = rpc_.release();
      rpc->Finish();

      Function* cb = callback_;
      callback_ = nullptr;

      lock.Release();
      cb->CallCancel();
      return;
    }
  }

  std::unique_ptr<AbstractMutex> mutex_;
  ::grpc::CompletionQueue* queue_;
  CallbackT* callback_ GUARDED_BY(mutex_);
  std::unique_ptr<RpcHolder<ReaderWriter>> rpc_ GUARDED_BY(mutex_);
  ResponseT resp_ GUARDED_BY(mutex_);
};

}  

#endif