<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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/rocket/client/RequestContext.h>

#include <glog/logging.h>

#include <fmt/core.h>
#include <folly/Likely.h>
#include <folly/Range.h>
#include <folly/io/IOBuf.h>
#include <folly/io/IOBufQueue.h>
#include <folly/lang/Assume.h>

#include <thrift/lib/cpp/transport/TTransportException.h>
#include <thrift/lib/cpp2/protocol/CompactProtocol.h>
#include <thrift/lib/cpp2/transport/rocket/RocketException.h>
#include <thrift/lib/cpp2/transport/rocket/client/RequestContextQueue.h>
#include <thrift/lib/cpp2/transport/rocket/client/RocketClientWriteCallback.h>
#include <thrift/lib/cpp2/transport/rocket/framing/Frames.h>
#include <thrift/lib/cpp2/transport/rsocket/gen-cpp2/Config_types.h>

using apache::thrift::transport::TTransportException;

namespace apache {
namespace thrift {
namespace rocket {

folly::Try<void> RequestContext::waitForWriteToComplete() {
  baton_.wait();

  switch (state_) {
    case State::RESPONSE_RECEIVED:
      // Even though this function should only be called for no-response
      // requests, we still use RESPONSE_RECEIVED as the terminal "expected"
      // state for such requests.
      return {};

    case State::REQUEST_ABORTED:
      return folly::Try<void>(
          folly::make_exception_wrapper<TTransportException>(
              TTransportException::NOT_OPEN,
              "Request aborted during client shutdown"));

    case State::WRITE_NOT_SCHEDULED:
    case State::WRITE_SCHEDULED:
    case State::WRITE_SENDING:
    case State::WRITE_SENT:
      LOG(FATAL) << fmt::format(
          "Returned from Baton::wait() with unexpected state {} in {}",
          static_cast<int>(state_),
          __func__);
  }

  folly::assume_unreachable();
}

folly::Try<Payload> RequestContext::waitForResponse(
    std::chrono::milliseconds timeout) {
  // Once the eventual write to the socket succeeds or errors out (via
  // writeSuccess() or writeErr()), a timeout for baton_ will be scheduled via
  // awaitResponseTimeoutHandler_.
  awaitResponseTimeout_ = timeout;
  baton_.wait(awaitResponseTimeoutHandler_);

  switch (state_) {
    case State::WRITE_SENT:
      // writeSuccess() or writeErr() processed this request but a response was
      // not received within the request's allotted timeout. Terminate request
      // with timeout.
      queue_.abortSentRequest(*this);
      return folly::Try<Payload>(
          folly::make_exception_wrapper<TTransportException>(
              TTransportException::TIMED_OUT));

    case State::RESPONSE_RECEIVED:
      return std::move(responsePayload_);

    case State::REQUEST_ABORTED:
      return folly::Try<Payload>(
          folly::make_exception_wrapper<TTransportException>(
              TTransportException::NOT_OPEN,
              "Request aborted during client shutdown"));

    case State::WRITE_NOT_SCHEDULED:
    case State::WRITE_SCHEDULED:
    case State::WRITE_SENDING:
      LOG(FATAL) << fmt::format(
          "Returned from Baton::wait() with unexpected state {} in {}",
          static_cast<int>(state_),
          __func__);
  }

  folly::assume_unreachable();
}

void RequestContext::onPayloadFrame(PayloadFrame&& payloadFrame) {
  DCHECK(!responsePayload_.hasException());
  DCHECK(isRequestResponse());

  // This function is only called on the response payload frame corresponding to
  // a REQUEST_RESPONSE context. Each fragment of the response payload frame
  // should have the next and complete flags.
  DCHECK(payloadFrame.hasNext());
  DCHECK(payloadFrame.hasComplete());

  if (LIKELY(!responsePayload_.hasValue())) {
    responsePayload_.emplace(std::move(payloadFrame.payload()));<commit_msg>thrift: rocket: return underlying exception on client REQUEST_ABORTED<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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/rocket/client/RequestContext.h>

#include <glog/logging.h>

#include <fmt/core.h>
#include <folly/Likely.h>
#include <folly/Range.h>
#include <folly/io/IOBuf.h>
#include <folly/io/IOBufQueue.h>
#include <folly/lang/Assume.h>

#include <thrift/lib/cpp/transport/TTransportException.h>
#include <thrift/lib/cpp2/protocol/CompactProtocol.h>
#include <thrift/lib/cpp2/transport/rocket/RocketException.h>
#include <thrift/lib/cpp2/transport/rocket/client/RequestContextQueue.h>
#include <thrift/lib/cpp2/transport/rocket/client/RocketClientWriteCallback.h>
#include <thrift/lib/cpp2/transport/rocket/framing/Frames.h>
#include <thrift/lib/cpp2/transport/rsocket/gen-cpp2/Config_types.h>

using apache::thrift::transport::TTransportException;

namespace apache {
namespace thrift {
namespace rocket {

namespace {
folly::exception_wrapper requestAbortedException(const folly::Try<Payload>& t) {
  auto reason =
      t.hasException() ? t.exception().what().toStdString() : "unknown reason";
  return folly::make_exception_wrapper<TTransportException>(
      TTransportException::NOT_OPEN,
      fmt::format("Request aborted during client shutdown: {}", reason));
}
} // namespace

folly::Try<void> RequestContext::waitForWriteToComplete() {
  baton_.wait();

  switch (state_) {
    case State::RESPONSE_RECEIVED:
      // Even though this function should only be called for no-response
      // requests, we still use RESPONSE_RECEIVED as the terminal "expected"
      // state for such requests.
      return {};

    case State::REQUEST_ABORTED:
      return folly::Try<void>(requestAbortedException(responsePayload_));

    case State::WRITE_NOT_SCHEDULED:
    case State::WRITE_SCHEDULED:
    case State::WRITE_SENDING:
    case State::WRITE_SENT:
      LOG(FATAL) << fmt::format(
          "Returned from Baton::wait() with unexpected state {} in {}",
          static_cast<int>(state_),
          __func__);
  }

  folly::assume_unreachable();
}

folly::Try<Payload> RequestContext::waitForResponse(
    std::chrono::milliseconds timeout) {
  // Once the eventual write to the socket succeeds or errors out (via
  // writeSuccess() or writeErr()), a timeout for baton_ will be scheduled via
  // awaitResponseTimeoutHandler_.
  awaitResponseTimeout_ = timeout;
  baton_.wait(awaitResponseTimeoutHandler_);

  switch (state_) {
    case State::WRITE_SENT:
      // writeSuccess() or writeErr() processed this request but a response was
      // not received within the request's allotted timeout. Terminate request
      // with timeout.
      queue_.abortSentRequest(*this);
      return folly::Try<Payload>(
          folly::make_exception_wrapper<TTransportException>(
              TTransportException::TIMED_OUT));

    case State::RESPONSE_RECEIVED:
      return std::move(responsePayload_);

    case State::REQUEST_ABORTED:
      return folly::Try<Payload>(requestAbortedException(responsePayload_));

    case State::WRITE_NOT_SCHEDULED:
    case State::WRITE_SCHEDULED:
    case State::WRITE_SENDING:
      LOG(FATAL) << fmt::format(
          "Returned from Baton::wait() with unexpected state {} in {}",
          static_cast<int>(state_),
          __func__);
  }

  folly::assume_unreachable();
}

void RequestContext::onPayloadFrame(PayloadFrame&& payloadFrame) {
  DCHECK(!responsePayload_.hasException());
  DCHECK(isRequestResponse());

  // This function is only called on the response payload frame corresponding to
  // a REQUEST_RESPONSE context. Each fragment of the response payload frame
  // should have the next and complete flags.
  DCHECK(payloadFrame.hasNext());
  DCHECK(payloadFrame.hasComplete());

  if (LIKELY(!responsePayload_.hasValue())) {
    responsePayload_.emplace(std::move(payloadFrame.payload()));
{'fraction_non_alphanumeric': 0.11189688604996008, 'fraction_numerical': 0.0027375384966351088, 'mean_word_length': 4.118505545826036, 'pattern_counts': {'":': 0, '<': 61, '<?xml version=': 0, '>': 53, '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': '1085232', 'n_tokens_mistral': 2573, 'n_tokens_neox': 2565, 'n_words': 651}