<commit_before>/*
 * Copyright 2017-present Facebook, Inc.
 *
 * 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.
 */
#include "thrift/lib/cpp2/transport/rsocket/client/RSClientThriftChannel.h"

#include <proxygen/lib/utils/WheelTimerInstance.h>
#include <thrift/lib/cpp2/protocol/CompactProtocol.h>
#include <thrift/lib/cpp2/transport/rsocket/client/RPCSubscriber.h>

namespace apache {
namespace thrift {

using namespace apache::thrift::transport;
using namespace apache::thrift::detail;
using namespace rsocket;
using namespace yarpl::single;
using proxygen::WheelTimerInstance;

std::unique_ptr<folly::IOBuf> RSClientThriftChannel::serializeMetadata(
    const RequestRpcMetadata& requestMetadata) {
  CompactProtocolWriter writer;
  folly::IOBufQueue queue;
  writer.setOutput(&queue);
  requestMetadata.write(&writer);
  return queue.move();
}

std::unique_ptr<ResponseRpcMetadata> RSClientThriftChannel::deserializeMetadata(
    const folly::IOBuf& buffer) {
  CompactProtocolReader reader;
  auto responseMetadata = std::make_unique<ResponseRpcMetadata>();
  reader.setInput(&buffer);
  responseMetadata->read(&reader);
  return responseMetadata;
}

namespace {
// This default value should match with the
// H2ClientConnection::kDefaultTransactionTimeout's value.
static constexpr std::chrono::milliseconds kDefaultRequestTimeout =
    std::chrono::milliseconds(500);

// Adds a timer that timesout if the observer could not get its onSuccess or
// onError methods being called in a specified time range, which causes onError
// method to be called.
class TimedSingleObserver : public SingleObserverBase<Payload>,
                            public folly::HHWheelTimer::Callback {
 public:
  TimedSingleObserver(
      std::unique_ptr<ThriftClientCallback> callback,
      std::chrono::milliseconds timeout,
      ChannelCounters& counters)
      : timeout_(timeout),
        timer_(timeout, callback->getEventBase()),
        callback_(std::move(callback)),
        counters_(counters) {}

  virtual ~TimedSingleObserver() {
    complete();
  }

  void setDecrementPendingRequestCounter() {
    decrementPendingRequestCounter_ = true;
  }

 protected:
  void onSubscribe(Reference<SingleSubscription> subscription) override {
    auto ref = this->ref_from_this(this);
    SingleObserverBase<Payload>::onSubscribe(std::move(subscription));

    auto evb = callback_->getEventBase();
    evb->runInEventBaseThread([this]() {
      callback_->onThriftRequestSent();
      if (timeout_.count() > 0) {
        timer_.scheduleTimeout(this, timeout_);
      }
    });
  }

  void onSuccess(Payload payload) override {
    auto ref = this->ref_from_this(this);
    auto evb = callback_->getEventBase();
    evb->runInEventBaseThread([ref, payload = std::move(payload)]() mutable {
      if (ref->complete()) {
        ref->callback_->onThriftResponse(
            payload.metadata
                ? RSClientThriftChannel::deserializeMetadata(*payload.metadata)
                : std::make_unique<ResponseRpcMetadata>(),
            std::move(payload.data));
      }
    });
    if (SingleObserverBase<Payload>::subscription()) {
      // TODO: can we get rid of calling the parent functions
      SingleObserverBase<Payload>::onSuccess({});
    }
  }

  void onError(folly::exception_wrapper ew) override {
    auto ref = this->ref_from_this(this);
    auto evb = callback_->getEventBase();
    evb->runInEventBaseThread([ref, ew = std::move(ew)]() mutable {
      if (ref->complete()) {
        ref->callback_->onError(std::move(ew));
        // TODO: Inspect the cases where might we end up in this function.
        // 1- server closes the stream before all the messages are
        // delievered.
        // 2- time outs
      }
    });
    if (SingleObserverBase<Payload>::subscription()) {
      // TODO: can we get rid of calling the parent functions
      SingleObserverBase<Payload>::onError({});
    }
  }

  void timeoutExpired() noexcept override {
    onError(folly::make_exception_wrapper<TTransportException>(
        apache::thrift::transport::TTransportException::TIMED_OUT));
  }

  void callbackCanceled() noexcept override {
    // nothing!
  }

  bool complete() {
    if (alreadySignalled_) {
      return false;
    }
    alreadySignalled_ = true;

    cancelTimeout();
    if (decrementPendingRequestCounter_) {
      counters_.decPendingRequests();
      decrementPendingRequestCounter_ = false;
    }
    return true;
  }

 private:
  std::chrono::milliseconds timeout_;
  // TODO: WheelTimerInstance is part of Proxygen, we need to use another timer.
  WheelTimerInstance timer_;
  std::unique_ptr<ThriftClientCallback> callback_;
  apache::thrift::detail::ChannelCounters& counters_;

  bool alreadySignalled_{false};
  bool decrementPendingRequestCounter_{false};
};
} // namespace

RSClientThriftChannel::RSClientThriftChannel(
    std::shared_ptr<RSocketRequester> rsRequester,
    ChannelCounters& channelCounters)
    : rsRequester_(std::move(rsRequester)), channelCounters_(channelCounters) {}

void RSClientThriftChannel::sendThriftRequest(
    std::unique_ptr<RequestRpcMetadata> metadata,
    std::unique_ptr<folly::IOBuf> payload,
    std::unique_ptr<ThriftClientCallback> callback) noexcept {
  DCHECK(metadata->__isset.kind);

  if (!rsRequester_) {
    if (callback) {
      auto evb = callback->getEventBase();
      evb->runInEventBaseThread([evbCallback = std::move(callback)]() mutable {
        evbCallback->onError(folly::make_exception_wrapper<TTransportException>(
            TTransportException::NOT_OPEN));
      });
    }
    return;<commit_msg>Remove dependency to Proxygen for the HHWheelTimerInstance<commit_after>/*
 * Copyright 2017-present Facebook, Inc.
 *
 * 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.
 */
#include "thrift/lib/cpp2/transport/rsocket/client/RSClientThriftChannel.h"

#include <thrift/lib/cpp2/protocol/CompactProtocol.h>
#include <thrift/lib/cpp2/transport/rsocket/client/RPCSubscriber.h>

namespace apache {
namespace thrift {

using namespace apache::thrift::transport;
using namespace apache::thrift::detail;
using namespace rsocket;
using namespace yarpl::single;

std::unique_ptr<folly::IOBuf> RSClientThriftChannel::serializeMetadata(
    const RequestRpcMetadata& requestMetadata) {
  CompactProtocolWriter writer;
  folly::IOBufQueue queue;
  writer.setOutput(&queue);
  requestMetadata.write(&writer);
  return queue.move();
}

std::unique_ptr<ResponseRpcMetadata> RSClientThriftChannel::deserializeMetadata(
    const folly::IOBuf& buffer) {
  CompactProtocolReader reader;
  auto responseMetadata = std::make_unique<ResponseRpcMetadata>();
  reader.setInput(&buffer);
  responseMetadata->read(&reader);
  return responseMetadata;
}

namespace {

// This default value should match with the
// H2ClientConnection::kDefaultTransactionTimeout's value.
static constexpr std::chrono::milliseconds kDefaultRequestTimeout =
    std::chrono::milliseconds(500);

// Adds a timer that timesout if the observer could not get its onSuccess or
// onError methods being called in a specified time range, which causes onError
// method to be called.
class TimedSingleObserver : public SingleObserver<Payload>,
                            public folly::HHWheelTimer::Callback {
 public:
  TimedSingleObserver(
      std::unique_ptr<ThriftClientCallback> callback,
      std::chrono::milliseconds timeout,
      ChannelCounters& counters)
      : callback_(std::move(callback)),
        counters_(counters),
        timeout_(timeout),
        timer_(callback_->getEventBase()->timer()) {}

  virtual ~TimedSingleObserver() {
    complete();
  }

  void setDecrementPendingRequestCounter() {
    decrementPendingRequestCounter_ = true;
  }

 protected:
  void onSubscribe(Reference<SingleSubscription>) override {
    // Note that we don't need the Subscription object.

    auto evb = callback_->getEventBase();
    evb->runInEventBaseThread([this]() {
      callback_->onThriftRequestSent();
      if (timeout_.count() > 0) {
        timer_.scheduleTimeout(this, timeout_);
      }
    });
  }

  void onSuccess(Payload payload) override {
    auto ref = this->ref_from_this(this);
    auto evb = callback_->getEventBase();
    evb->runInEventBaseThread([ref, payload = std::move(payload)]() mutable {
      if (ref->complete()) {
        ref->callback_->onThriftResponse(
            payload.metadata
                ? RSClientThriftChannel::deserializeMetadata(*payload.metadata)
                : std::make_unique<ResponseRpcMetadata>(),
            std::move(payload.data));
      }
    });
  }

  void onError(folly::exception_wrapper ew) override {
    auto ref = this->ref_from_this(this);
    auto evb = callback_->getEventBase();
    evb->runInEventBaseThread([ref, ew = std::move(ew)]() mutable {
      if (ref->complete()) {
        ref->callback_->onError(std::move(ew));
        // TODO: Inspect the cases where might we end up in this function.
        // 1- server closes the stream before all the messages are
        // delievered.
        // 2- time outs
      }
    });
  }

  void timeoutExpired() noexcept override {
    onError(folly::make_exception_wrapper<TTransportException>(
        apache::thrift::transport::TTransportException::TIMED_OUT));
  }

  void callbackCanceled() noexcept override {
    // nothing!
  }

  bool complete() {
    if (alreadySignalled_) {
      return false;
    }
    alreadySignalled_ = true;

    if (decrementPendingRequestCounter_) {
      counters_.decPendingRequests();
      decrementPendingRequestCounter_ = false;
    }
    return true;
  }

 private:
  std::unique_ptr<ThriftClientCallback> callback_;
  apache::thrift::detail::ChannelCounters& counters_;

  std::chrono::milliseconds timeout_;
  folly::HHWheelTimer& timer_;

  bool alreadySignalled_{false};
  bool decrementPendingRequestCounter_{false};
};
} // namespace

RSClientThriftChannel::RSClientThriftChannel(
    std::shared_ptr<RSocketRequester> rsRequester,
    ChannelCounters& channelCounters)
    : rsRequester_(std::move(rsRequester)), channelCounters_(channelCounters) {}

void RSClientThriftChannel::sendThriftRequest(
    std::unique_ptr<RequestRpcMetadata> metadata,
    std::unique_ptr<folly::IOBuf> payload,
    std::unique_ptr<ThriftClientCallback> callback) noexcept {
  DCHECK(metadata->__isset.kind);

  if (!rsRequester_) {
    if (callback) {
      auto evb = callback->getEventBase();
      evb->runInEventBaseThread([evbCallback = std::move(callback)]() mutable {
        evbCallback->onError(folly::make_exception_wrapper<TTransportException>(
            TTransportException::NOT_OPEN));
      });
    }
    return;
{'fraction_non_alphanumeric': 0.1101571220056667, 'fraction_numerical': 0.0030909247016399073, 'mean_word_length': 4.062146892655368, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 87, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5563707', 'n_tokens_mistral': 3308, 'n_tokens_neox': 3266, 'n_words': 762}